Neuron: compiled log-softmax + greedy argmax in the sampler; eager op census (#27)

The op census of a decode step (DEEPINFER_NEURON_REGION_TIMING) found 85 eager
aten ops in the forward outside the region graphs but 1905 in sampling (the
top-p/top-k bisection and the int64 philox draw), ~0.6 s of NEFF launches.

- The log-softmax / top-p / top-k part runs as one compiled graph: 1.4 ms
  instead of 46.7 ms on one core for the [1, 4096] force_topk row, logprobs
  within 1e-9.
- Greedy rows (temperature clamped to SAMPLING_EPS) have one-hot weights, so
  the two-level draw returns their argmax; an all-greedy batch takes the argmax
  from the graph and skips the draw. The philox draw stays eager for sampled
  rows: neuronx-cc graphs reject 64-bit constants outside the int32 range
  (NCC_ESFH001), an int32 philox is #3.

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

deepinfer/environ.py

@@ -122,6 +122,8 @@ class envs:
     # (issue #26, the vllm-neuron#71 layout). 0 = eager prefill. Needs max_num_batched_tokens <= bucket and
     # DEEPINFER_MAX_BACKTRACKING_TOKENS >= bucket (the SWA / compressor rings must hold a whole chunk).
     DEEPINFER_NEURON_PREFILL_BUCKET = EnvVar.int(default=0)
+    # Neuron only: run the sampler's process_sampling as one compiled graph instead of ~1900 eager ops a step.
+    DEEPINFER_NEURON_COMPILE_SAMPLER = EnvVar.bool(default=True)
     # Neuron only: compile the decode step with fullgraph=True (raise on the first graph break).
     DEEPINFER_NEURON_FULLGRAPH = EnvVar.bool()
 

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

@@ -105,6 +105,47 @@ def compute_logprobs(
     return torch.where(filter_active.unsqueeze(1), filtered, plain)
 
 
+def _sample_eager(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_size: int):  # noqa: ANN001, ANN202
+    logprobs = compute_logprobs(logits, temperatures, top_p, top_k)
+    # Exponentiate the logprobs so the weights equal exp(logprobs_out) bit for bit.
+    return logprobs, two_level_sample(torch.exp(logprobs), seed, sample_positions, 0, vocab_size)
+
+
+def _logprobs_and_greedy(logits, temperatures, top_p, top_k):  # noqa: ANN001, ANN202
+    logprobs = compute_logprobs(logits, temperatures, top_p, top_k)
+    return logprobs, logprobs.argmax(dim=1).to(torch.int32)
+
+
+_COMPILED = None
+
+
+def _sample(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_size: int):  # noqa: ANN001, ANN202
+    """On a NeuronCore the log-softmax / top-p / top-k part is one compiled graph (issue #27): eagerly the sampler
+    is ~1900 small ops a decode step, each its own NEFF launch (~0.6 s).
+
+    Greedy rows (temperature clamped to SAMPLING_EPS) have one-hot weights, so the two-level draw returns their argmax; a batch of only
+    greedy rows takes the argmax and skips the draw. The philox draw stays eager: it is int64 arithmetic, and
+    neuronx-cc graphs reject 64-bit constants outside the int32 range (NCC_ESFH001).
+    """
+    global _COMPILED
+    from deepinfer import envs
+
+    if logits.device.type != "neuron" or torch.compiler.is_compiling() or not envs.DEEPINFER_NEURON_COMPILE_SAMPLER:
+        return _sample_eager(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_size)
+    from deepinfer.neuron_ops import neuron_compile  # noqa: F401  (registers the "deepinfer_neuron" backend)
+    from deepinfer.neuron_ops.transfer_to_neuron import hide_cuda
+
+    if _COMPILED is None:
+        _COMPILED = torch.compile(_logprobs_and_greedy, backend=envs.DEEPINFER_NEURON_COMPILE_BACKEND, dynamic=False)
+    with hide_cuda():
+        logprobs, greedy = _COMPILED(logits, temperatures, top_p, top_k)
+    from deepinfer.sampling_params import SAMPLING_EPS
+
+    if bool((temperatures.cpu() <= SAMPLING_EPS).all()):  # host sync on the [rows] temperatures
+        return logprobs, greedy
+    return logprobs, two_level_sample(torch.exp(logprobs), seed, sample_positions, 0, vocab_size)
+
+
 def fused_sample_tl(
     vocab_size: int,
     do_sample: bool,
@@ -135,11 +176,13 @@ def fused_sample_tl(
         lse_workspace: torch.Tensor | None,
     ) -> None:
         assert logits.shape[1] == vocab_size and logits.stride(1) == 1
-        logprobs.copy_(compute_logprobs(logits, temperatures, top_p, top_k))
-        if do_sample:
-            seed = seq_seeds.view(torch.int64) if seq_seeds.dtype == torch.uint64 else seq_seeds.to(torch.int64)
-            seed = seed + int(seed_offset)  # uint64 wraparound == int64 wraparound on the same bits
-            # Exponentiate the written logprobs so the weights equal exp(logprobs_out) bit for bit.
-            candidate_token_ids.copy_(two_level_sample(torch.exp(logprobs), seed, sample_positions, 0, vocab_size))
+        if not do_sample:
+            logprobs.copy_(compute_logprobs(logits, temperatures, top_p, top_k))
+            return
+        seed = seq_seeds.view(torch.int64) if seq_seeds.dtype == torch.uint64 else seq_seeds.to(torch.int64)
+        seed = seed + int(seed_offset)  # uint64 wraparound == int64 wraparound on the same bits
+        out_logprobs, tokens = _sample(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_size)
+        logprobs.copy_(out_logprobs)
+        candidate_token_ids.copy_(tokens)
 
     return _fused_sample

deepinfer/neuron_ops/op_census.py

@@ -0,0 +1,35 @@
+"""Count the eager ops of one step on Neuron (issue #27): every aten op outside a compiled region graph is a
+separate NEFF launch, so the census shows what still runs eagerly between and around the region graphs."""
+
+from __future__ import annotations
+
+import contextlib
+from collections.abc import Iterator
+
+import torch
+
+from deepinfer.logger import init_logger
+
+logger = init_logger(__name__)
+
+#: Bookkeeping ops that launch nothing on the device.
+_FREE = ("aten::view", "aten::_reshape_alias", "aten::as_strided", "aten::alias", "aten::detach", "aten::lift_fresh",
+         "aten::t", "aten::unsqueeze", "aten::squeeze", "aten::expand", "aten::permute", "aten::transpose",
+         "aten::slice", "aten::select", "aten::reshape", "aten::empty", "aten::empty_strided", "aten::resolve_conj",
+         "aten::resolve_neg", "aten::_unsafe_view", "aten::split", "aten::chunk", "aten::unbind", "aten::narrow")
+
+
+@contextlib.contextmanager
+def census(label: str, top: int = 40) -> Iterator[None]:
+    """Log the ops run inside the block, by count (the outermost aten op of each dispatch)."""
+    from torch.profiler import ProfilerActivity, profile
+
+    with profile(activities=[ProfilerActivity.CPU]) as prof:
+        yield
+    rows = [e for e in prof.key_averages() if e.key.startswith("aten::") and e.key not in _FREE]
+    rows.sort(key=lambda e: -e.count)
+    total = sum(e.count for e in rows)
+    lines = [f"{e.count:6d}  {e.cpu_time_total / 1e3:9.2f} ms  {e.key}" for e in rows[:top]]
+    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))

deepinfer/worker/model_runner.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import contextlib
 import copy
 import time
 from dataclasses import dataclass
@@ -380,7 +381,10 @@ class ModelRunner:
             kwargs["kv_caches"] = kv_caches
 
         profiler = _maybe_profile_sample(self)
+        from deepinfer.neuron_ops.op_census import census
+
         with (
+            census("sample") if profiler is not None else contextlib.nullcontext(),
             cpu_timer.record("launch_sample"),
             cuda_timer.record("device_sample", scale=1),
             forbid_cuda_sync(),

deepinfer/worker/neuron_graph_runner.py

@@ -8,6 +8,7 @@ outside the graph.
 
 from __future__ import annotations
 
+import contextlib
 import dataclasses
 import math
 import time
@@ -200,7 +201,9 @@ class NeuronGraphRunner(CudaGraphRunner):
             profiler = cProfile.Profile()
             profiler.enable()
         started = time.perf_counter()
-        with hide_cuda():
+        from deepinfer.neuron_ops.op_census import census
+
+        with hide_cuda(), census(f"forward bs={batch_size}") if profiler is not None else contextlib.nullcontext():
             out = self.compiled[batch_size](input_metadata_device)
         torch.neuron.synchronize()
         total = time.perf_counter() - started