neuron: compiled decode verification behind DEEPINFER_NEURON_COMPILE_VERIFY (default off)

Sampler._verification runs as one compiled region for greedy batches without structured output: the native
sampling structs enter as Python mirrors (neuron_ops/sampling_mirror.py), the flags are read once per batch
eagerly (prime_host_flags), and greedy mirrors carry no seeds / positions. Untested on device yet.

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

deepinfer/environ.py

@@ -144,6 +144,8 @@ class envs:
     DEEPINFER_NEURON_NUM_KV_BLOCKS = 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)
+    # The decode verification (greedy batch, no structured output) as one compiled region (issue #29).
+    DEEPINFER_NEURON_COMPILE_VERIFY = EnvVar.bool(default=False)
     # Neuron only: compile the decode step with fullgraph=True (raise on the first graph break).
     DEEPINFER_NEURON_FULLGRAPH = EnvVar.bool()
 

deepinfer/model_executor/layers/sampler.py

@@ -5,13 +5,14 @@
 
 from __future__ import annotations
 
+import functools
 from typing import TYPE_CHECKING
 
 import torch
 import torch.nn as nn
 from deep_select.interface import topk
 
-from deepinfer import envs
+from deepinfer import PLATFORM, envs
 from deepinfer.logger import init_logger
 from deepinfer.model_executor.input_metadata import InputMetadata
 from deepinfer.model_executor.layers.mtp_info_pool import MTPInfoPool
@@ -56,6 +57,30 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+_COMPILED_VERIFY: list = []
+
+
+def _compiled_verification():  # noqa: ANN202
+    """``Sampler._verification`` (undecorated) as a compiled Neuron region: eagerly it is ~300 small ops a decode step,
+    each its own NEFF launch with ~0.5 ms of device idle around it (issue #29)."""
+    if not _COMPILED_VERIFY:
+        from deepinfer.neuron_ops import neuron_compile  # noqa: F401  (registers the "deepinfer_neuron" backend)
+        from deepinfer.neuron_ops import regions
+        from deepinfer.neuron_ops.transfer_to_neuron import hide_cuda
+
+        regions.raise_dynamo_limits()
+        compiled = torch.compile(
+            Sampler._verification.__wrapped__, backend=envs.DEEPINFER_NEURON_COMPILE_BACKEND, fullgraph=False, dynamic=False
+        )
+
+        def run(*args, **kwargs):  # noqa: ANN002, ANN003, ANN202
+            with cpu_timer.record("verification"), hide_cuda():
+                return compiled(*args, **kwargs)
+
+        _COMPILED_VERIFY.append(run)
+    return _COMPILED_VERIFY[0]
+
+
 class Sampler(nn.Module):
     """从模型 hidden states 采样下一个 token，支持 MTP（Multi-Token Prediction）投机采样。
 
@@ -212,6 +237,16 @@ class Sampler(nn.Module):
         else:
             # Decode sampling.
             speculative_inputs = driver_inputs.speculative
+            verify, verify_driver, verify_decode = self._verification, driver_inputs, decode_inputs
+            if PLATFORM.is_neuron and envs.DEEPINFER_NEURON_COMPILE_VERIFY and so_indices is None and not self.force_topk:
+                from deepinfer.model_executor.layers.tilelang_ops.neuron.fused_sample import prime_host_flags
+                from deepinfer.neuron_ops.sampling_mirror import mirror_decode, mirror_driver
+
+                params = driver_inputs.kernel.logits
+                _, all_greedy = prime_host_flags(params.temperatures, params.top_ps, params.top_ks, logits.shape[1])
+                if all_greedy:
+                    verify = functools.partial(_compiled_verification(), self)
+                    verify_driver, verify_decode = mirror_driver(driver_inputs, greedy=True), mirror_decode(decode_inputs)
             logits = Movable(logits)
             (
                 sampled_tokens,
@@ -222,10 +257,10 @@ class Sampler(nn.Module):
                 num_verified_tokens,
                 batch_confidence_error,
                 emitted_stop_token,
-            ) = self._verification(
+            ) = verify(
                 logits,
-                driver_inputs,
-                decode_inputs,
+                verify_driver,
+                verify_decode,
                 so_indices,
                 so_bitmasks,
                 req_cnt,

deepinfer/model_executor/layers/tilelang_ops/fused_sample.py

@@ -63,7 +63,7 @@ def fused_sample(
         ``max(target_prob_i - exp(draft_logprob_i - DRAFT_PROB_PENALTY), 0)``.
     """
     assert logits.stride(1) == 1
-    if isinstance(sampling_inputs, SamplingKernelInputs):
+    if isinstance(sampling_inputs, SamplingKernelInputs) or hasattr(sampling_inputs, "seq_seeds"):  # or its mirror
         do_sample = True
         transform_inputs = sampling_inputs.logits
         seq_seeds = sampling_inputs.seq_seeds
@@ -187,7 +187,7 @@ def fused_sample_ref(
     The reference may differ numerically from the TileLang implementation.
     Its in-place update of ``logits`` is safe for the caller in sampler_utils.py.
     """
-    if isinstance(sampling_inputs, SamplingKernelInputs):
+    if isinstance(sampling_inputs, SamplingKernelInputs) or hasattr(sampling_inputs, "seq_seeds"):  # or its mirror
         do_sample = True
         transform_inputs = sampling_inputs.logits
         seq_seeds = sampling_inputs.seq_seeds

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

@@ -147,7 +147,7 @@ def _host_flags(temperatures, top_p, top_k, vocab: int) -> tuple[bool, bool]:  #
 
     # 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)
+    flags = _FLAGS.get(key, _FLAGS.get("batch"))
     if flags is None:
         t_h, p_h, k_h = temperatures.cpu(), top_p.cpu(), top_k.cpu()
         flags = (bool(((p_h < 1.0) | ((k_h >= 0) & (k_h < vocab))).any()), bool((t_h <= SAMPLING_EPS).all()))
@@ -155,6 +155,14 @@ def _host_flags(temperatures, top_p, top_k, vocab: int) -> tuple[bool, bool]:  #
     return flags
 
 
+def prime_host_flags(temperatures, top_p, top_k, vocab: int) -> tuple[bool, bool]:  # noqa: ANN001
+    """Read the batch's flags eagerly (once per batch) so a compiled region finds them (``_host_flags`` falls back to
+    them for its expanded per-row parameters: expansion changes neither any-filter nor all-greedy)."""
+    if "batch" not in _FLAGS:
+        _FLAGS["batch"] = _host_flags(temperatures, top_p, top_k, vocab)
+    return _FLAGS["batch"]
+
+
 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).
@@ -166,7 +174,15 @@ def _sample(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_si
     global _COMPILED
     from deepinfer import envs
 
-    if logits.device.type != "neuron" or torch.compiler.is_compiling() or not envs.DEEPINFER_NEURON_COMPILE_SAMPLER:
+    if torch.compiler.is_compiling() and logits.device.type == "neuron":
+        # Inside the compiled verification (DEEPINFER_NEURON_COMPILE_VERIFY): the same choice as below, traced.
+        filtering, all_greedy = _host_flags(temperatures, top_p, top_k, logits.shape[1])
+        fn = _logprobs_and_greedy if filtering else _plain_logprobs_and_greedy
+        logprobs, greedy = fn(logits, temperatures, top_p, top_k) if filtering else fn(logits, temperatures)
+        if all_greedy:
+            return logprobs, greedy
+        return logprobs, two_level_sample(torch.exp(logprobs), seed, sample_positions, 0, vocab_size)
+    if logits.device.type != "neuron" 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
@@ -200,9 +216,10 @@ def _greedy_verify(
     from deepinfer.sampling_params import SAMPLING_EPS
 
     key = ("greedy", temperatures.shape[0])
-    if key not in _FLAGS:
-        _FLAGS[key] = bool((temperatures.cpu() <= SAMPLING_EPS).all())
-    if not _FLAGS[key]:
+    greedy = _FLAGS.get(key, _FLAGS["batch"][1] if "batch" in _FLAGS else None)
+    if greedy is None:
+        greedy = _FLAGS[key] = bool((temperatures.cpu() <= SAMPLING_EPS).all())
+    if not greedy:
         unsupported("fused_sample_tl(do_reject_sampling=True) for non-greedy rows")
     rows = read_rows.long()
     draft = draft_ids.index_select(0, rows)
@@ -244,8 +261,11 @@ def fused_sample_tl(
         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
+        if seq_seeds is None:  # a greedy batch's mirror (neuron_ops/sampling_mirror.py): no draw
+            seed = None
+        else:
+            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)

deepinfer/model_executor/layers/tilelang_ops/sampler.py

@@ -86,6 +86,18 @@ def expand_sampling_params(
         SamplingKernelInputs with [total_q] tensors, preserving each field's dtype.
     """
     total_q = row_to_req.shape[0]
+    if PLATFORM.is_neuron and not isinstance(params, SamplingKernelInputs):
+        # The compiled Neuron verification's mirror (neuron_ops/sampling_mirror.py): plain gathers, and no seeds or
+        # positions for a greedy batch.
+        from deepinfer.neuron_ops.sampling_mirror import KernelMirror, LogitsMirror
+
+        req = row_to_req.to(torch.int64)
+        lg = params.logits
+        positions = None
+        if params.sample_positions is not None:
+            positions = params.sample_positions.to(torch.int32)[req] + local_step.to(torch.int32) * DRAFT_POS_HASH_MUL
+        seeds = None if params.seq_seeds is None else params.seq_seeds.view(torch.int64)[req].view(torch.uint64)
+        return KernelMirror(LogitsMirror(lg.temperatures[req], lg.top_ps[req], lg.top_ks[req]), seeds, positions)
 
     temperatures_q = torch.empty(total_q, dtype=torch.float32, device=params.logits.temperatures.device)
     top_ps_q = torch.empty_like(temperatures_q)

deepinfer/neuron_ops/sampling_mirror.py

@@ -0,0 +1,74 @@
+"""Python mirrors of the native sampling-input structs, for the compiled Neuron verification (issue #29).
+
+Dynamo cannot read the nanobind structs' fields (each access breaks the graph) nor construct them, so the driver's
+decode inputs enter the compiled ``Sampler._verification`` as these dataclasses: same field names, plain tensors.
+A greedy batch carries no seeds or sample positions (``None``): only the RNG draw reads them, and the expansion of
+the uint64 seeds is int64 arithmetic neuronx-cc graphs reject.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class LogitsMirror:
+    temperatures: torch.Tensor
+    top_ps: torch.Tensor
+    top_ks: torch.Tensor
+
+
+@dataclass
+class KernelMirror:
+    logits: LogitsMirror
+    seq_seeds: torch.Tensor | None
+    sample_positions: torch.Tensor | None
+
+
+@dataclass
+class SpeculativeMirror:
+    mtp_info_read_rows: torch.Tensor
+    mtp_info_write_slots: torch.Tensor
+    nextn_tokens: torch.Tensor
+    stop_token_ids: torch.Tensor
+
+
+@dataclass
+class DriverMirror:
+    kernel: KernelMirror
+    speculative: SpeculativeMirror | None
+    golden_output: torch.Tensor | None
+    max_top_logprobs: int
+
+
+@dataclass
+class DecodeMirror:
+    mtp_first_token_positions: torch.Tensor
+    mtp_seqlens_q: torch.Tensor
+
+
+def mirror_kernel(kernel, greedy: bool) -> KernelMirror:  # noqa: ANN001
+    logits = kernel.logits
+    return KernelMirror(
+        LogitsMirror(logits.temperatures, logits.top_ps, logits.top_ks),
+        None if greedy else kernel.seq_seeds,
+        None if greedy else kernel.sample_positions,
+    )
+
+
+def mirror_driver(driver, greedy: bool) -> DriverMirror:  # noqa: ANN001
+    spec = driver.speculative
+    return DriverMirror(
+        kernel=mirror_kernel(driver.kernel, greedy),
+        speculative=None
+        if spec is None
+        else SpeculativeMirror(spec.mtp_info_read_rows, spec.mtp_info_write_slots, spec.nextn_tokens, spec.stop_token_ids),
+        golden_output=driver.golden_output,
+        max_top_logprobs=driver.max_top_logprobs,
+    )
+
+
+def mirror_decode(decode) -> DecodeMirror:  # noqa: ANN001
+    return DecodeMirror(decode.mtp_first_token_positions, decode.mtp_seqlens_q)