[atom-vllm]: enable DSpark speculative decoding for DeepSeek-V4-Flash-0731

atom/model_ops/embed_head.py

@@ -2,7 +2,6 @@
 # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
 
 import torch
-import torch.nn.functional as F
 import triton
 import triton.language as tl
 from aiter.dist.communication_op import tensor_model_parallel_all_gather
@@ -174,7 +173,10 @@ def forward(self, x: torch.Tensor):
             )
             y = get_tp_group().all_reduce(y, ca_fp8_quant=False)
         else:
-            y = F.embedding(x, self.weight)
+            # Not F.embedding: async scheduling can pass the -1 spec
+            # placeholder, which reads below the table and GPU-faults.
+            # replicated_embedding returns zeros out of range.
+            y = replicated_embedding(x, self.weight)
         return y
         # if self.tp_size > 1:
         #     mask = torch.logical_and(x >= self.vocab_start_idx, x < self.vocab_end_idx)
@@ -222,6 +224,23 @@ def forward(self, x: torch.Tensor):
         return replicated_embedding(x, self.weight)
 
 
+class _UnquantizedHeadMethod:
+    """``QuantizeMethodBase``-shaped shim so vLLM's ``LogitsProcessor`` can drive
+    an ATOM head (EAGLE3/DSpark share the target's ``lm_head`` with the draft).
+
+    ``apply`` is the LOCAL shard GEMM only: vLLM's ``_apply_head`` does its own
+    TP gather, and ``ParallelLMHead.forward`` would double-gather.
+    """
+
+    @staticmethod
+    def apply(
+        layer: nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return tgemm.mm(x, layer.weight, layer.bias if bias is None else bias)
+
+
 class ParallelLMHead(VocabParallelEmbedding):
 
     def __init__(
@@ -232,6 +251,8 @@ def __init__(
         **kwargs,
     ):
         super().__init__(num_embeddings, embedding_dim)
+        # Plain object, so nn.Module keeps it in __dict__, not as a submodule.
+        self.quant_method = _UnquantizedHeadMethod()
         if bias:
             self.bias = atom_parameter(
                 torch.empty(self.num_embeddings_per_partition),

atom/models/deepseek_v4.py

@@ -126,6 +126,12 @@
 
 logger = logging.getLogger(__name__)
 
+# COMPILE BOUNDARY: `DeepseekV4Model.forward` is the only traced entry point.
+# `aux_hidden_state_layers` is read inside it, so Dynamo bakes in the branch and
+# the ids; from `--level 2` up the custom dispatcher replays code object 0
+# without evaluating guards, so mutating it after the first forward is silently
+# ignored. Set it at load time (`set_aux_hidden_state_layers`) and never after.
+
 # ---------------------------------------------------------------------------
 # Classical KV cache scatter / gather helpers (PR3-pre2c-B).
 #
@@ -4554,17 +4560,28 @@ def __init__(
         self.hc_head_base = atom_parameter(torch.empty(hc_mult, dtype=torch.float32))
         self.hc_head_scale = atom_parameter(torch.empty(1, dtype=torch.float32))
 
+        # 1-based layer ids whose residual is also returned, for EAGLE3 /
+        # DSpark / DFlash drafts. Set via `set_aux_hidden_state_layers`.
+        self.aux_hidden_state_layers: tuple[int, ...] = ()
+
     def forward(
         self,
         input_ids: torch.Tensor,  # [num_tokens] int  flat ragged-batch token ids
         positions: torch.Tensor,  # [num_tokens] int  abs positions (required)
     ) -> torch.Tensor:  # [num_tokens, hc, dim]  pre-hc_head residual stream
         """Forward over `num_tokens` flat ragged-batch tokens.
 
+        TRACED ENTRY POINT — see the COMPILE BOUNDARY note at the top of the file.
+
         Returns the mHC residual stack `[num_tokens, hc, dim]` BEFORE hc_head
         reduction — `hc_head + RMSNorm + LM head` are all deferred to
         `compute_logits`. Returning the hc-shaped residual lets the (future)
         MTP draft consume it without re-expanding from a dim-reduced state.
+
+        With `aux_hidden_state_layers` set, returns `(h, aux_hidden_states)`
+        instead; each aux entry is that layer's residual averaged over hc
+        (`[num_tokens, dim]`), matching vLLM's DSv4 target, which is what the
+        DSpark drafter was trained against.
         """
         assert input_ids.dim() == 1, f"input_ids must be 1D, got {input_ids.shape}"
         # PCP note: under PCP, `input_ids`/`positions` arrive already round-robin-
@@ -4578,12 +4595,38 @@ def forward(
         h = h.unsqueeze(-2).repeat(1, self.hc_mult, 1)
         hc_state = HCState(residual=h, post_mix=None, comb_mix=None, x_prev=None)
 
-        for layer in self.layers:
+        aux_layers = self.aux_hidden_state_layers
+        if not aux_layers:
+            for layer in self.layers:
+                hc_state = layer(hc_state, positions)
+            h = self.layers[-1].hc_post(
+                hc_state.x_prev, hc_state.residual, hc_state.post_mix, hc_state.comb_mix
+            )
+            return h
+
+        # Aux path. `hc_post` is the deferred mHC reconstruction: it turns the
+        # layer's `x_prev` + residual + mixing tensors into that layer's
+        # post-layer `[num_tokens, hc, dim]` residual, which is what aux wants.
+        aux_hidden_states: list[torch.Tensor] = []
+        final_aux_recon: torch.Tensor | None = None
+        for idx, layer in enumerate(self.layers):
             hc_state = layer(hc_state, positions)
-        h = self.layers[-1].hc_post(
-            hc_state.x_prev, hc_state.residual, hc_state.post_mix, hc_state.comb_mix
-        )
-        return h
+            if (idx + 1) in aux_layers:
+                final_aux_recon = layer.hc_post(
+                    hc_state.x_prev,
+                    hc_state.residual,
+                    hc_state.post_mix,
+                    hc_state.comb_mix,
+                )
+                aux_hidden_states.append(final_aux_recon.mean(dim=1))
+        if final_aux_recon is not None and len(self.layers) in aux_layers:
+            # Last layer was an aux layer; reuse its reconstruction.
+            h = final_aux_recon
+        else:
+            h = self.layers[-1].hc_post(
+                hc_state.x_prev, hc_state.residual, hc_state.post_mix, hc_state.comb_mix
+            )
+        return h, aux_hidden_states
 
 
 class DeepseekV4ForCausalLM(nn.Module):
@@ -4782,6 +4825,11 @@ def forward(
         else:
             ctx.context.input_ids = input_ids
         h = self.model(input_ids, positions)
+        # Aux layers make the inner model return `(h, aux_list)`. Unpack so the
+        # PCP fixups below keep operating on a single tensor; re-attach at exit.
+        aux_hidden_states = None
+        if isinstance(h, tuple):
+            h, aux_hidden_states = h
 
         # ----- PCP: all-gather shards, restore original order, drop pad -----
         if use_pcp:
@@ -4794,8 +4842,36 @@ def forward(
                 h = pcp_allgather_rerange(h, pcp_size)
                 if pad > 0:
                     h = h[:n_global]
+                if aux_hidden_states is not None:
+                    aux_hidden_states = [
+                        (
+                            pcp_allgather_rerange(a, pcp_size)[:n_global]
+                            if pad > 0
+                            else pcp_allgather_rerange(a, pcp_size)
+                        )
+                        for a in aux_hidden_states
+                    ]
+        if aux_hidden_states is not None:
+            return h, aux_hidden_states
         return h
 
+    # ----- EAGLE3 / DSpark / DFlash auxiliary hidden states -----
+    # ATOM's server-mode convention. The plugin's
+    # `_enable_eagle3_target_interface` bridges it onto vLLM's `SupportsEagle3`.
+
+    def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
+        """`layers` are 1-based: id `i` means "after `self.layers[i-1]`"."""
+        self.model.aux_hidden_state_layers = tuple(int(i) for i in layers)
+
+    def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
+        # Checkpoint ids are 0-based; +1 converts to the 1-based convention.
+        for attr in ("dspark_target_layer_ids", "target_layer_ids"):
+            ids = getattr(self.hf_config, attr, None)
+            if ids:
+                return tuple(int(i) + 1 for i in ids)
+        num_layers = len(self.model.layers)
+        return (2, num_layers // 2, num_layers - 3)
+
     def compute_logits(
         self,
         hidden_states: torch.Tensor,  # [num_tokens, hc, dim]  pre-hc_head residual