From 7f558a3775f677afeed7828352f4f882b627f802 Mon Sep 17 00:00:00 2001 From: whn09 Date: Thu, 29 Oct 2026 12:10:31 +0800 Subject: [PATCH] neuron: workers bind memory to their own core's NUMA node; prompt chunks run the attention half as one graph; glibc trims freed heap early (TTFT 796 -> 720 ms on the bench prompt, jitter +-25 -> +-3 ms) - NUMA: each worker sees one core, so which_numa(local_rank) failed on all but local rank 0 and every worker inherited its MPOL_BIND to node 1: all host memory landed on one node and a run could be OOM-killed. - DEEPINFER_NEURON_WHOLE_ATTN_MIN_ROWS=512: no cuts inside the attention half of a >=512-row step, only before the MoE (its collectives are another group kind). - MALLOC_TRIM_THRESHOLD_/MALLOC_TOP_PAD_ defaults: compile-time host memory per rank stops growing ~15 GB. (MALLOC_ARENA_MAX=2 bounds it too but costs ~50 ms TTFT.) Co-Authored-By: Claude Opus 5.5 --- .../model_executor/models/decoder_layer.py | 2 +- deepinfer/neuron_ops/regions.py | 10 +++++++++- deepinfer/platform.py | 6 ++++++ deepinfer/utils.py | 13 +++++++++++++ deepinfer/worker/base.py | 18 ++++++++++++++++-- 5 files changed, 45 insertions(+), 4 deletions(-) diff --git a/deepinfer/model_executor/models/decoder_layer.py b/deepinfer/model_executor/models/decoder_layer.py index 786ea25..2f2aa99 100644 --- a/deepinfer/model_executor/models/decoder_layer.py +++ b/deepinfer/model_executor/models/decoder_layer.py @@ -175,7 +175,7 @@ class DeepseekDecoderLayer(HCLayerMixin): # The static MoE's NKI path reads only the bf16 hidden: drop the ffn norm's fp8 copy before the split, # where it would be a graph output (and its encode would run) for nothing. hidden_states = hidden_states_bf16 - neuron_regions.split(coarse=True) + neuron_regions.split(coarse=True, keep=True) capacity = envs.DEEPINFER_NEURON_MOE_CAPACITY if capacity and hidden_states_bf16.shape[0] == envs.DEEPINFER_NEURON_PREFILL_BUCKET: # A prefill chunk: the experts over their routed rows only; the overflow check is a graph break, and an diff --git a/deepinfer/neuron_ops/regions.py b/deepinfer/neuron_ops/regions.py index 5ef0897..aa6ecf9 100644 --- a/deepinfer/neuron_ops/regions.py +++ b/deepinfer/neuron_ops/regions.py @@ -146,7 +146,13 @@ def set_step_rows(rows: int) -> None: _STEP_ROWS[0] = rows -def split(coarse: bool = False) -> None: +#: DEEPINFER_NEURON_WHOLE_ATTN_MIN_ROWS: steps of at least this many rows (prompt chunks) run the attention half of a +#: layer as one graph, cutting only at the ``keep`` boundaries (before the MoE: its collectives are of another group +#: kind, and both kinds in one graph hang the runtime). 0: off. A 512-row chunk: TTFT 849 -> 720 ms (issue #29). +_WHOLE_ATTN_MIN_ROWS = int(os.environ.get("DEEPINFER_NEURON_WHOLE_ATTN_MIN_ROWS", "512")) + + +def split(coarse: bool = False, keep: bool = False) -> None: """Cut the current region graph here (only while compiling; free in eager). Always for prefill chunks under DEEPINFER_NEURON_REGION_SPLIT_PREFILL: neuronx-cc runs the split 128-row attention ~2x faster (62 vs 122 ms a layer, issue #26), while decode's 1-row graphs are faster whole. Multi-row steps up to _SPLIT_MIN_ROWS rows cut @@ -154,6 +160,8 @@ def split(coarse: bool = False) -> None: if not torch.compiler.is_compiling(): return rows = _STEP_ROWS[0] + if _WHOLE_ATTN_MIN_ROWS and rows >= _WHOLE_ATTN_MIN_ROWS and not keep: + return multi_row = envs.DEEPINFER_NEURON_REGION_SPLIT_PREFILL and rows > 1 if envs.DEEPINFER_NEURON_REGION_SPLIT or (multi_row and (coarse or rows > _SPLIT_MIN_ROWS)): _boundary() diff --git a/deepinfer/platform.py b/deepinfer/platform.py index 6d32273..f35d960 100644 --- a/deepinfer/platform.py +++ b/deepinfer/platform.py @@ -254,6 +254,12 @@ class PlatformNeuron(Platform): # with 64 ranks on one host the pollers crowd the cores at every step boundary, the main threads get # preempted, and the slowest rank paces every collective. Off: decode step 580 -> 522 ms (issue #29). os.environ.setdefault("NEURON_RT_XU_WORKER_ENABLE_LOW_LATENCY_COMPLETION", "0") + # glibc returns freed heap to the OS early: 64 ranks' compile-time host memory otherwise grows ~15 GB each + # and a prompt chunk's larger graphs ran a NUMA node out (kernel OOM kill). Set before the workers spawn, + # which read them at startup. MALLOC_ARENA_MAX=2 (vllm-neuron#71's setting) also bounds it but costs + # ~50 ms of TTFT (issue #29). + os.environ.setdefault("MALLOC_TRIM_THRESHOLD_", "131072") + os.environ.setdefault("MALLOC_TOP_PAD_", "131072") # torch-neuronx numbers a process's torch device by LOCAL_RANK out of LOCAL_WORLD_SIZE when the first # "neuron" process group is created (distributed/backend.py, else it derives them from the hosts of all # ranks). With one core per process every rank must be device 0 of a 1-device world: `cuda:N` maps to diff --git a/deepinfer/utils.py b/deepinfer/utils.py index 45d37ec..019d61b 100644 --- a/deepinfer/utils.py +++ b/deepinfer/utils.py @@ -140,6 +140,19 @@ def allow_cuda_sync() -> AbstractContextManager[None]: return _sync_debug_arbiter.scope("allow") +def neuron_numa_node(core: int) -> int | None: + """NUMA node of the device holding physical NeuronCore ``core`` (neuron-ls), None when unknown.""" + import json + import subprocess + + try: + devices = json.loads(subprocess.run(["neuron-ls", "--json-output"], capture_output=True, check=True, timeout=30).stdout) + except (OSError, subprocess.SubprocessError, ValueError): + return None + device = next((d for d in devices if core in d["neuroncore_ids"]), None) + return None if device is None else int(device["numa_node"]) + + def _neuron_bind_cpus(local_rank: int) -> list[int] | None: """CPUs for the rank on logical NeuronCore ``local_rank`` (the physical core id), on its device's NUMA node. diff --git a/deepinfer/worker/base.py b/deepinfer/worker/base.py index 734e5e8..58cfdfb 100644 --- a/deepinfer/worker/base.py +++ b/deepinfer/worker/base.py @@ -192,7 +192,7 @@ class WorkerBase: def init_model(self) -> None: configure_worker_environment() - _bind_numa(self.local_rank) + _bind_numa(self.local_rank, self.gpu_id) self.device = torch.device(f"cuda:{self.local_rank}") torch.cuda.set_device(self.device) configure_worker_kernels() @@ -360,8 +360,22 @@ def _init_distributed_environment( logger.info("After init_parallel_groups") -def _bind_numa(local_rank: int) -> None: +def _bind_numa(local_rank: int, gpu_id: int | None = None) -> None: # Bind process and memory affinity using the device NUMA topology. + if PLATFORM.is_neuron and gpu_id is not None: + # Neuron: each worker sees one core, so which_numa(local_rank) fails on all but local rank 0, and the workers + # inherited that rank's MPOL_BIND to its node: every rank's host memory (weights in loading, Engram tables) + # landed on NUMA node 1 while node 0 stayed free, and node 1 ran out (kernel OOM kill, issue #29). Bind each + # worker to the node of its own core instead. + from deep_cluster.topology import bind_numa + + from deepinfer.utils import neuron_numa_node + + node = neuron_numa_node(gpu_id) + if node is not None and bind_numa(node) == 0: + return + logger.warning(f"Neuron NUMA binding failed for core {gpu_id}") + return try: from deep_cluster.topology import bind_numa, which_numa