neuron: sampling flags per request batch, NUMA-local CPU binding, sync/op-site tracing

- fused_sample: _host_flags/_greedy_verify flags persist while the batch's requests are unchanged
  (each device read drained the queue and waited for the whole forward). Tiny golden 9/9; full model
  step time unchanged (660 ms), driver dflash launch 71 -> 10.5 ms host.
- bind_cpu on Neuron: the device's NUMA node (DEEPINFER_NEURON_CPU_BIND=node|core|legacy); the generic
  split gave each rank one CPU and ranks 0-15 the wrong node. Neutral on TPOT (228.2 ms).
- profiled followers exit through sys.exit so the runtime writes their trace.
- op_census: trace_syncs / census_sites (DEEPINFER_NEURON_TRACE_SYNCS=1|2|3).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py

@@ -127,20 +127,26 @@ def _plain_logprobs_and_greedy(logits, temperatures):  # noqa: ANN001, ANN202
 
 _COMPILED: dict = {}
 _FLAGS: dict = {}
+_FLAGS_EPOCH: list = [None]
 
 
-def reset_host_flags() -> None:
-    _FLAGS.clear()
+def reset_host_flags(requests: tuple | None = None) -> None:
+    """Start a step. The flags are per-request sampling parameters, so they hold while the step's requests are those
+    of the previous step (``requests``: their ids): only a changed batch reads them again (a device read drains the
+    queue, and the first read of a step waits for the whole forward)."""
+    if requests is None or requests != _FLAGS_EPOCH[0]:
+        _FLAGS.clear()
+    _FLAGS_EPOCH[0] = requests
 
 
 def _host_flags(temperatures, top_p, top_k, vocab: int) -> tuple[bool, bool]:  # noqa: ANN001
     """(any row filters by top-p / top-k, every row is greedy) for the step's sampling parameters. Reading them is a
     device sync, and a DSpark step samples the same parameter tensors six times (verification, then each draft),
-    so the flags are kept per tensor until the step ends (reset_host_flags: a later step's tensors may reuse the
-    address)."""
+    so the flags are kept until the batch changes (reset_host_flags)."""
     from deepinfer.sampling_params import SAMPLING_EPS
 
-    key = tuple((t.data_ptr(), t.shape[0]) for t in (temperatures, top_p, top_k))
+    # By shape, not address: a step may hand over the same parameters in a fresh buffer.
+    key = ("filter", tuple(t.shape[0] for t in (temperatures, top_p, top_k)), vocab)
     flags = _FLAGS.get(key)
     if flags is None:
         t_h, p_h, k_h = temperatures.cpu(), top_p.cpu(), top_k.cpu()
@@ -193,7 +199,7 @@ def _greedy_verify(
     Sampled (non-greedy) verification is not ported (its int64 RNG reject draw)."""
     from deepinfer.sampling_params import SAMPLING_EPS
 
-    key = ("greedy", temperatures.data_ptr(), temperatures.shape[0])
+    key = ("greedy", temperatures.shape[0])
     if key not in _FLAGS:
         _FLAGS[key] = bool((temperatures.cpu() <= SAMPLING_EPS).all())
     if not _FLAGS[key]:

deepinfer/neuron_ops/op_census.py

@@ -4,6 +4,7 @@ separate NEFF launch, so the census shows what still runs eagerly between and ar
 from __future__ import annotations
 
 import contextlib
+import os
 from collections.abc import Iterator
 
 import torch
@@ -33,3 +34,68 @@ def census(label: str, top: int = 40) -> Iterator[None]:
     graphs = [e for e in prof.key_averages() if "Compiled" in e.key or "neuron" in e.key.lower()]
     lines += [f"{e.count:6d}  {e.cpu_time_total / 1e3:9.2f} ms  {e.key}" for e in graphs[:10]]
     logger.info(f"Neuron op census ({label}): {total} aten ops outside bookkeeping\n" + "\n".join(lines))
+
+
+@contextlib.contextmanager
+def trace_syncs(label: str) -> Iterator[None]:
+    """Log every device -> host read in the block (each drains the Neuron execution queue) with its caller."""
+    import collections
+    import traceback
+
+    hits: collections.Counter = collections.Counter()
+    saved = {}
+
+    def site() -> str:
+        for frame in reversed(traceback.extract_stack()[:-2]):
+            if "/deepinfer/" in frame.filename and "op_census.py" not in frame.filename:
+                return f"{frame.filename.split('/deepinfer/')[-1]}:{frame.lineno} {frame.name}"
+        return "?"
+
+    def wrap(name: str, check=None):  # noqa: ANN001, ANN202
+        orig = getattr(torch.Tensor, name)
+        saved[name] = orig
+
+        def fn(self, *a, **k):  # noqa: ANN001, ANN002, ANN003, ANN202
+            if self.device.type == "neuron" and (check is None or check(*a, **k)):
+                key = f"{name:8s} {site()}"
+                if key not in hits and os.environ.get("DEEPINFER_NEURON_TRACE_SYNCS") == "2":
+                    logger.info(f"host sync {key}:\n" + "".join(traceback.format_stack()[-14:-1]))
+                hits[key] += 1
+            return orig(self, *a, **k)
+
+        setattr(torch.Tensor, name, fn)
+
+    def to_cpu(*a, **k) -> bool:  # noqa: ANN002, ANN003
+        dev = k.get("device", a[0] if a else None)
+        return isinstance(dev, (str, torch.device)) and str(dev).startswith("cpu")
+
+    for name in ("item", "cpu", "tolist", "__bool__", "numpy"):
+        wrap(name)
+    wrap("to", to_cpu)
+    try:
+        yield
+    finally:
+        for name, orig in saved.items():
+            setattr(torch.Tensor, name, orig)
+        logger.info(f"Neuron host syncs ({label}): {sum(hits.values())}\n" + "\n".join(f"{n:4d}  {k}" for k, n in hits.most_common()))
+
+
+@contextlib.contextmanager
+def census_sites(label: str, top: int = 45) -> Iterator[None]:
+    """Count the ops run in the block (outermost aten op of each dispatch) by their innermost DeepInfer call site:
+    each eager op on a NeuronCore is its own NEFF launch."""
+    import collections
+
+    from torch.profiler import ProfilerActivity, profile
+
+    with profile(activities=[ProfilerActivity.CPU], with_stack=True) as prof:
+        yield
+    hits: collections.Counter = collections.Counter()
+    events = [e for e in prof.events() if e.name.startswith("aten::") and e.name not in _FREE]
+    inner = {id(e) for e in events for c in e.cpu_children if c.name.startswith("aten::")}
+    for e in events:
+        if e.cpu_parent is not None and e.cpu_parent.name.startswith("aten::"):
+            continue  # count each dispatch once, at its outermost aten op
+        where = next((f for f in (e.stack or []) if "/deepinfer/" in f and "op_census.py" not in f), "?")
+        hits[where.split("/deepinfer/")[-1]] += 1
+    logger.info(f"Neuron op sites ({label}): {sum(hits.values())} ops\n" + "\n".join(f"{n:5d}  {k}" for k, n in hits.most_common(top)))

deepinfer/utils.py

@@ -140,16 +140,59 @@ def allow_cuda_sync() -> AbstractContextManager[None]:
     return _sync_debug_arbiter.scope("allow")
 
 
+def _neuron_bind_cpus(local_rank: int) -> list[int] | None:
+    """CPUs for the rank on logical NeuronCore ``local_rank`` (of CUDA_VISIBLE_DEVICES), on its device's NUMA node.
+
+    The generic split (physical cores // ranks) gave each of 64 ranks one logical CPU, numbered by rank: the step's
+    dispatch thread and the runtime's threads shared it, and devices 0-3 (NUMA 1 on trn2.48xlarge) got CPUs of NUMA 0.
+    DEEPINFER_NEURON_CPU_BIND: ``node`` (default) the device's whole NUMA node, ``core`` one physical core and its
+    hyperthread on that node, ``legacy`` the generic split. None falls back to the generic split."""
+    mode = os.environ.get("DEEPINFER_NEURON_CPU_BIND", "node")
+    if mode == "legacy":
+        return None
+    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):
+        logger.warning("neuron-ls unavailable: generic CPU binding")
+        return None
+    visible = os.environ.get("CUDA_VISIBLE_DEVICES")
+    core = int(visible.split(",")[local_rank]) if visible else local_rank
+    device = next((d for d in devices if core in d["neuroncore_ids"]), None)
+    if device is None:
+        return None
+
+    def cpus(spec: str) -> list[int]:
+        out = []
+        for part in spec.split(","):
+            lo, _, hi = part.partition("-")
+            out.extend(range(int(lo), int(hi or lo) + 1))
+        return out
+
+    node_cpus = cpus(device["cpu_affinity"])
+    if mode != "core":
+        return node_cpus
+    # One physical core per NeuronCore of the node, in core order, plus its hyperthread sibling.
+    node_cores = sorted(c for d in devices if d["numa_node"] == device["numa_node"] for c in d["neuroncore_ids"])
+    physical = sorted({min(cpus(open(f"/sys/devices/system/cpu/cpu{c}/topology/thread_siblings_list").read().strip())) for c in node_cpus})
+    first = physical[node_cores.index(core) % len(physical)]
+    return cpus(open(f"/sys/devices/system/cpu/cpu{first}/topology/thread_siblings_list").read().strip())
+
+
 def bind_cpu(local_rank: int) -> None:
     from deepinfer import PLATFORM
 
     if PLATFORM.is_neuron and envs.DEEPINFER_NEURON_DEVICE == "cpu":
         # Neuron CPU mode computes on the host: a 1-core share (96 cores / 64 NeuronCores) would serialize it.
         return
-    phy_cpus = psutil.cpu_count(logical=False)
-    ngpus = get_actual_num_gpus()
-    num_cpus_per_gpu = phy_cpus // ngpus
-    bind_cpus = list(range(local_rank * num_cpus_per_gpu, (local_rank + 1) * num_cpus_per_gpu))
+    bind_cpus = _neuron_bind_cpus(local_rank) if PLATFORM.is_neuron else None
+    if bind_cpus is None:
+        phy_cpus = psutil.cpu_count(logical=False)
+        ngpus = get_actual_num_gpus()
+        num_cpus_per_gpu = phy_cpus // ngpus
+        bind_cpus = list(range(local_rank * num_cpus_per_gpu, (local_rank + 1) * num_cpus_per_gpu))
     logger.info(f"bind_cpu gpu_id: {local_rank}, cpus: {bind_cpus}.")
     import deep_cluster.topology
 

deepinfer/worker/base.py

@@ -10,6 +10,7 @@ import functools
 import gc
 import math
 import os
+import sys
 import signal
 import time
 from typing import Any
@@ -148,6 +149,10 @@ class WorkerBase:
         dist.destroy_process_group()
 
         if not self.is_driver_worker:
+            if os.environ.get("NEURON_RT_INSPECT_ENABLE") == "1":
+                # A profiled Neuron follower (DEEPINFER_NEURON_PROFILE_RANKS): the runtime writes its trace in the
+                # teardown os._exit() skips.
+                sys.exit(0)
             os._exit(0)
 
     def maybe_gc_collect(self) -> None:

deepinfer/worker/model_runner.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import contextlib
+import os
 import copy
 import time
 from dataclasses import dataclass
@@ -447,11 +448,20 @@ class ModelRunner:
         if PLATFORM.is_neuron:
             from deepinfer.model_executor.layers.tilelang_ops.neuron.fused_sample import reset_host_flags
 
-            reset_host_flags()  # sampling-parameter flags are read once a step (see _host_flags)
+            # Sampling-parameter flags are read once per batch of requests (see _host_flags).
+            reset_host_flags(tuple(seq.request_id for seq in seq_metadata_list) if seq_metadata_list else None)
         profiler = _maybe_profile_sample(self)
         from deepinfer.neuron_ops.op_census import census
 
+        from deepinfer.neuron_ops.op_census import census_sites, trace_syncs
+
+        self._trace_sample_calls = getattr(self, "_trace_sample_calls", 0) + 1
+        tracing = PLATFORM.is_neuron and os.environ.get("DEEPINFER_NEURON_TRACE_SYNCS", "0") != "0" and self._trace_sample_calls in (6, 7)
         with (
+            trace_syncs(f"sample call {self._trace_sample_calls}") if tracing else contextlib.nullcontext(),
+            census_sites(f"sample call {self._trace_sample_calls}")
+            if tracing and os.environ.get("DEEPINFER_NEURON_TRACE_SYNCS") == "3"
+            else contextlib.nullcontext(),
             census("sample") if profiler is not None else contextlib.nullcontext(),
             cpu_timer.record("launch_sample"),
             cuda_timer.record("device_sample", scale=1),