Neuron: sampling flags read once a step; runtime profiling on the driver only; one-lookup embedding

- fused_sample._host_flags: the top-p/top-k and greedy flags of the step's
  sampling parameters are read from the device once a step (a DSpark step
  samples the same tensors six times); reset_host_flags at each sample step.
- NEURON_RT_INSPECT_* stays on the driver (64 ranks would overwrite one
  output directory).
- VocabParallelEmbedding.ids_in_range: one F.embedding instead of the
  zero-filling gather (~5 eager NEFF launches).
- bench33 (512 in, TP=64, nextn=5): TPOT 272.6 -> 268.2 ms.

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

deepinfer/engine/llm_engine.py

@@ -394,15 +394,18 @@ class LLMEngine:
             )
             # Neuron: the follower claims core cp_rank of CUDA_VISIBLE_DEVICES (deepinfer/platform.py); the spawned
             # process copies this environment at start.
-            saved = os.environ.get("DEEPINFER_NEURON_LOCAL_RANK")
-            os.environ["DEEPINFER_NEURON_LOCAL_RANK"] = str(cp_rank)
+            # Runtime profiling (NEURON_RT_INSPECT_*) stays on the driver: 64 ranks would overwrite one output dir.
+            overrides = {"DEEPINFER_NEURON_LOCAL_RANK": str(cp_rank), "NEURON_RT_INSPECT_ENABLE": "0"}
+            saved = {k: os.environ.get(k) for k in overrides}
+            os.environ.update(overrides)
             try:
                 worker.start()
             finally:
-                if saved is None:
-                    os.environ.pop("DEEPINFER_NEURON_LOCAL_RANK", None)
-                else:
-                    os.environ["DEEPINFER_NEURON_LOCAL_RANK"] = saved
+                for k, v in saved.items():
+                    if v is None:
+                        os.environ.pop(k, None)
+                    else:
+                        os.environ[k] = v
         logger.info("_init_workers: non-driver workers started")
 
         driver_kwargs = {

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

@@ -126,6 +126,27 @@ def _plain_logprobs_and_greedy(logits, temperatures):  # noqa: ANN001, ANN202
 
 
 _COMPILED: dict = {}
+_FLAGS: dict = {}
+
+
+def reset_host_flags() -> None:
+    _FLAGS.clear()
+
+
+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)."""
+    from deepinfer.sampling_params import SAMPLING_EPS
+
+    key = tuple((t.data_ptr(), t.shape[0]) for t in (temperatures, top_p, top_k))
+    flags = _FLAGS.get(key)
+    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()))
+        _FLAGS[key] = flags
+    return flags
 
 
 def _sample(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_size: int):  # noqa: ANN001, ANN202
@@ -151,17 +172,15 @@ def _sample(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_si
             _COMPILED[fn] = torch.compile(fn, backend=envs.DEEPINFER_NEURON_COMPILE_BACKEND, dynamic=False)
         return _COMPILED[fn]
 
-    # Host copies of the [rows] parameters (they come from the host; the sync waits for the logits anyway).
-    t_h, p_h, k_h = temperatures.cpu(), top_p.cpu(), top_k.cpu()
     # The top-p / top-k bisection is 32 rounds over the vocabulary (~0.5 s at 129280 even compiled): only rows that
     # filter need it, and the default parameters filter nothing.
-    filtering = bool(((p_h < 1.0) | ((k_h >= 0) & (k_h < logits.shape[1]))).any())
+    filtering, all_greedy = _host_flags(temperatures, top_p, top_k, logits.shape[1])
     with hide_cuda():
         if filtering:
             logprobs, greedy = compiled(_logprobs_and_greedy)(logits, temperatures, top_p, top_k)
         else:
             logprobs, greedy = compiled(_plain_logprobs_and_greedy)(logits, temperatures)
-    if bool((t_h <= SAMPLING_EPS).all()):
+    if all_greedy:
         return logprobs, greedy
     return logprobs, two_level_sample(torch.exp(logprobs), seed, sample_positions, 0, vocab_size)
 
@@ -174,7 +193,10 @@ def _greedy_verify(
     Sampled (non-greedy) verification is not ported (its int64 RNG reject draw)."""
     from deepinfer.sampling_params import SAMPLING_EPS
 
-    if not bool((temperatures.cpu() <= SAMPLING_EPS).all()):
+    key = ("greedy", temperatures.data_ptr(), temperatures.shape[0])
+    if key not in _FLAGS:
+        _FLAGS[key] = bool((temperatures.cpu() <= SAMPLING_EPS).all())
+    if not _FLAGS[key]:
         unsupported("fused_sample_tl(do_reject_sampling=True) for non-greedy rows")
     rows = read_rows.long()
     draft = draft_ids.index_select(0, rows)

deepinfer/model_executor/layers/vocab_parallel_embedding.py

@@ -3,6 +3,7 @@
 
 import torch
 
+from deepinfer import PLATFORM
 from deepinfer.model_executor.layers.tilelang_ops.gather import gather
 from deepinfer.model_executor.utils import reconstruct_param
 
@@ -23,6 +24,10 @@ class VocabParallelEmbedding(torch.nn.Module):
         if self.weight.device != input_.device:
             # Host-resident table (Neuron, see DeepseekForCausalLM.load_model): move the looked-up rows only.
             return gather(self.weight, input_.cpu()).to(input_.device)
+        if PLATFORM.is_neuron and getattr(self, "ids_in_range", False) and not torch.compiler.is_compiling():
+            # Eagerly the zero-filling gather is ~5 NEFF launches; callers whose ids are always in range (DSpark's
+            # Markov embedding of sampled draft tokens) take one embedding lookup.
+            return torch.nn.functional.embedding(input_.long().clamp(0, self.weight.shape[0] - 1), self.weight)
         output = gather(self.weight, input_)
         return output
 

deepinfer/worker/model_runner.py

@@ -444,6 +444,10 @@ class ModelRunner:
         if self.nextn:
             kwargs["kv_caches"] = kv_caches
 
+        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)
         profiler = _maybe_profile_sample(self)
         from deepinfer.neuron_ops.op_census import census