From ee9e3aaed963b4d7d4ba2527baa8668fe788bd12 Mon Sep 17 00:00:00 2001 From: whn09 Date: Sun, 25 Oct 2026 06:06:09 +0800 Subject: [PATCH] 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 --- deepinfer/environ.py | 5 +- .../layers/vocab_parallel_embedding.py | 13 +++-- deepinfer/model_executor/models/mtp/dflash.py | 53 +++++++++++++++++++ 3 files changed, 65 insertions(+), 6 deletions(-) diff --git a/deepinfer/environ.py b/deepinfer/environ.py index b880f7c..63da639 100644 --- a/deepinfer/environ.py +++ b/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). diff --git a/deepinfer/model_executor/layers/vocab_parallel_embedding.py b/deepinfer/model_executor/layers/vocab_parallel_embedding.py index 61ca149..92f5540 100644 --- a/deepinfer/model_executor/layers/vocab_parallel_embedding.py +++ b/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 diff --git a/deepinfer/model_executor/models/mtp/dflash.py b/deepinfer/model_executor/models/mtp/dflash.py index 32465ca..a4b3f79 100644 --- a/deepinfer/model_executor/models/mtp/dflash.py +++ b/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,