Neuron DSpark: draft sampling as one region graph for greedy rows; token embedding on device by default

- DFlash._greedy_drafts: post norm, LM head, Markov bias, confidence head and
  the argmax of the 5 draft steps as one region graph when every row is greedy
  with no top-p/top-k filter.
- DEEPINFER_NEURON_HOST_EMBED defaults off: in torch-neuronx every host<->device
  copy drains the stream, and the host lookup (D2H ids, H2D rows) left the cores
  idle for 50-100 ms gaps between the draft graphs (runtime trace bench36).
- Eager device embedding lookups are one F.embedding with a validity mask (no
  host scalar operand).
- bench34 / bench37 (512 in, 128 out, TP=64, nextn=5): TPOT 268.2 -> 264.0 ->
  234.7 ms, step 1025 -> 907 ms, TTFT 5.44 -> 5.01 s.

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

deepinfer/environ.py

@@ -120,8 +120,9 @@ class envs:
     DEEPINFER_NEURON_REGION_SPLIT_PREFILL = EnvVar.bool(default=True)
     # Neuron only: synchronize around every region graph call and log per-kind times each decode step (profiling).
     DEEPINFER_NEURON_REGION_TIMING = EnvVar.bool()
-    # Keep the token embedding (and the text model's unused vision tower) in host memory on Neuron.
-    DEEPINFER_NEURON_HOST_EMBED = EnvVar.bool(default=True)
+    # Keep the token embedding (and the text model's unused vision tower) in host memory on Neuron. Off by default:
+    # a host lookup is a device->host and a host->device copy, and in torch-neuronx each drains the device stream.
+    DEEPINFER_NEURON_HOST_EMBED = EnvVar.bool(default=False)
     # Neuron only, regions mode: run prompt chunks through the decode region graphs padded to this many tokens
     # (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/model_executor/layers/vocab_parallel_embedding.py

@@ -24,10 +24,15 @@ 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)
+        if PLATFORM.is_neuron and not torch.compiler.is_compiling():
+            # Eagerly the zero-filling gather is ~5 NEFF launches, and its scalar operand is a host->device copy,
+            # which drains the stream. One lookup instead; out-of-range ids still read zeros unless the caller's
+            # ids are always valid (ids_in_range: DSpark's Markov embedding of sampled draft tokens).
+            n = self.weight.shape[0]
+            rows = torch.nn.functional.embedding(input_.long().clamp(0, n - 1), self.weight)
+            if getattr(self, "ids_in_range", False):
+                return rows
+            return rows * ((input_ >= 0) & (input_ < n)).unsqueeze(-1).to(rows.dtype)
         output = gather(self.weight, input_)
         return output
 

deepinfer/model_executor/models/mtp/dflash.py

@@ -169,6 +169,7 @@ class DFlash(nn.Module):
                     self.markov_rank,
                     weight=modules[-1].shared_head.markov_w1.weight,
                 )
+                self.markov_w1.ids_in_range = True  # looked up with sampled draft tokens
                 self.markov_w2 = ReplicatedLinear(
                     self.markov_rank,
                     cfg.vocab_size,
@@ -297,6 +298,19 @@ class DFlash(nn.Module):
         req_cnt = driver_input.golden_tokens.shape[0]
         sampling_input = driver_input.sampling
 
+        greedy = self._neuron_greedy_draft_params(hidden_states_before_norm, sampling_input)
+        if greedy is not None:
+            # Neuron, all rows greedy: the whole draft loop is one region graph (eagerly it is ~400 small NEFF
+            # launches a step, ~0.3 s of host dispatch with the cores idle).
+            next_n_token_ids, next_n_token_logprobs, confidence_rates = neuron_regions.entry(self, "_greedy_drafts")(
+                hidden_states_before_norm, driver_input.golden_tokens, greedy
+            )
+            prediction = MTPPredictInfo(
+                token_ids=next_n_token_ids, confidence_rates=confidence_rates if self.enable_confidence else None
+            )
+            self.mtp_info_pool.store(sampling_input.mtp_info_write_slots, prediction, next_n_token_logprobs, None)
+            return prediction
+
         hidden_states = self.post_norm(hidden_states_before_norm)
         logits_next = get_logits(hidden_states, self.lm_head.weight, self.vocab_size)
         sample_positions_table = (
@@ -356,6 +370,45 @@ class DFlash(nn.Module):
         )
         return prediction
 
+    def _neuron_greedy_draft_params(self, hidden_states: torch.Tensor, sampling_input) -> torch.Tensor | None:  # noqa: ANN001
+        """The temperatures when the draft loop can run as one Neuron graph: region graphs on, no top-k space, every
+        row greedy with no top-p / top-k filter (the flags are host-cached for the step). Otherwise None."""
+        if not (PLATFORM.is_neuron and hidden_states.device.type == "neuron") or self.force_topk:
+            return None
+        if envs.DEEPINFER_ENFORCE_EAGER or envs.DEEPINFER_NEURON_GRAPH_MODE != "regions":
+            return None
+        from deepinfer.model_executor.layers.tilelang_ops.neuron.fused_sample import _host_flags
+
+        params = sampling_input.kernel_inputs.logits
+        filtering, all_greedy = _host_flags(params.temperatures, params.top_ps, params.top_ks, self.vocab_size)
+        return params.temperatures if all_greedy and not filtering else None
+
+    def _greedy_drafts(
+        self, hidden_states_before_norm: torch.Tensor, golden_tokens: torch.Tensor, temperatures: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """sample_drafts for greedy rows as one graph: (token ids [req, n], logprobs [req, n, vocab], confidence
+        [req, n]). Same arithmetic as the eager loop (greedy rows sample their argmax)."""
+        from deepinfer.model_executor.layers.tilelang_ops.neuron.fused_sample import _plain_logprobs_and_greedy
+
+        n = self.max_MTP_value
+        hidden_states = self.post_norm(hidden_states_before_norm)
+        logits_next = get_logits(hidden_states, self.lm_head.weight, self.vocab_size)
+        token = golden_tokens
+        ids, logprobs, confidence = [], [], []
+        for step_i in range(n):
+            logits = logits_next[step_i::n]
+            w1 = None
+            if self.enable_markov:
+                w1 = self.markov_w1(token)
+                logits = self.markov_w2(w1, out_acc=logits.contiguous())
+            if self.enable_confidence:
+                confidence.append(self.confidence_head(hidden_states_before_norm[step_i::n], w1, step_i=step_i))
+            lp, token = _plain_logprobs_and_greedy(logits, temperatures)
+            ids.append(token)
+            logprobs.append(lp)
+        conf = torch.stack(confidence, dim=1) if confidence else torch.zeros_like(torch.stack(ids, dim=1), dtype=torch.float32)
+        return torch.stack(ids, dim=1), torch.stack(logprobs, dim=1), conf
+
     def run_drafts(
         self,
         main_input_metadata: InputMetadataV4,