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 <noreply@anthropic.com>

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

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()

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

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.
 

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