fix(spec-decode): initialize DeepSeek V4 DSpark top-k buffer

aphrodite/models/deepseek_v4/attention.py

@@ -29,6 +29,7 @@
 from aphrodite.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE
 
 if TYPE_CHECKING:
+    from aphrodite.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool
     from aphrodite.v1.attention.backends.mla.sparse_swa import (
         DeepseekSparseSWAMetadata,
     )
@@ -181,6 +182,7 @@ def __init__(
         topk_indices_buffer: torch.Tensor | None = None,
         q_workspace: torch.Tensor | None = None,
         aux_stream_list: list[torch.cuda.Stream] | None = None,
+        eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None,
     ) -> None:
         super().__init__()
         config = aphrodite_config.model_config.hf_config
@@ -270,6 +272,7 @@ def __init__(
         self.indexer_rotary_emb = self.rotary_emb
         self.topk_indices_buffer = topk_indices_buffer
         self.q_workspace = q_workspace
+        self.eager_scratch_pool = eager_scratch_pool
 
         self.indexer = None
         if self.compress_ratio == 4:
@@ -289,6 +292,7 @@ def __init__(
                 compress_ratio=self.compress_ratio,
                 prefix=f"{prefix}.indexer",
                 aux_stream=indexer_aux_stream,
+                eager_scratch_pool=eager_scratch_pool,
             )
 
         # Will be None on ROCm for now.
@@ -337,6 +341,7 @@ def __init__(
                 rotate=True,
                 prefix=f"{prefix}.compressor",
                 k_cache_prefix=self.prefix,
+                eager_scratch_pool=eager_scratch_pool,
             )
 
     def forward(
@@ -565,7 +570,9 @@ def _fused_qnorm_rope_kv_insert(
             #            the padded q tensor.
             #   KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert.
             swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
-            if self.q_workspace is None:
+            if self.eager_scratch_pool is not None:
+                q_out = self.eager_scratch_pool.q_out(q.shape[0])
+            elif self.q_workspace is None:
                 q_out = torch.empty(
                     (q.shape[0], self.padded_heads, self.head_dim),
                     dtype=q.dtype,
@@ -622,6 +629,11 @@ def _fused_qnorm_rope_kv_insert(
         )
         return q_fp8
 
+    def _global_topk_output_buffers(self, topk_indices: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor] | None:
+        if self.compress_ratio != 4 or self.eager_scratch_pool is None:
+            return None
+        return self.eager_scratch_pool.global_topk_outputs(topk_indices)
+
     def get_attn_backend(self) -> type[AttentionBackend]:
         return self.backend_cls
 
@@ -699,6 +711,7 @@ def __init__(
         compress_ratio: int = 1,
         prefix: str = "",
         aux_stream: torch.cuda.Stream | None = None,
+        eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None,
     ):
         super().__init__()
         self.aphrodite_config = aphrodite_config
@@ -711,6 +724,7 @@ def __init__(
         self.rope_dim = config.qk_rope_head_dim  # 64
         self.q_lora_rank = q_lora_rank  # 1536
         self.compress_ratio = compress_ratio
+        self.eager_scratch_pool = eager_scratch_pool
         self.use_fp4_kv = self.aphrodite_config.attention_config.use_fp4_indexer_cache
         logger.info_once(
             "Using %s indexer cache for Lightning Indexer.",
@@ -768,6 +782,7 @@ def __init__(
             prefix=f"{prefix}.compressor",
             k_cache_prefix=self.k_cache.prefix,
             use_fp4_cache=self.use_fp4_kv,
+            eager_scratch_pool=eager_scratch_pool,
         )
 
         self.indexer_op = SparseAttnIndexer(
@@ -822,6 +837,9 @@ def wq_b_and_q_quant():
             # ReplicatedLinear returns (output, bias); bias is None.
             q, _ = self.wq_b(qr)
             q = q.view(-1, self.n_head, self.head_dim)
+            outputs = None
+            if self.eager_scratch_pool is not None and self.use_fp4_kv:
+                outputs = self.eager_scratch_pool.indexer_q_outputs(q.shape[0])
             return fused_indexer_q_rope_quant(
                 positions,
                 q,
@@ -830,6 +848,7 @@ def wq_b_and_q_quant():
                 self.softmax_scale,
                 self.n_head**-0.5,
                 use_fp4=self.use_fp4_kv,
+                output_buffers=outputs,
             )
 
         # compressor returns None and writes K to the indexer KV cache; the

aphrodite/models/deepseek_v4/common/ops/cache_utils.py

@@ -431,6 +431,7 @@ def compute_global_topk_indices_and_lens(
     block_table: torch.Tensor,
     block_size: int,
     is_valid_token: torch.Tensor,
+    output_buffers: tuple[torch.Tensor, torch.Tensor] | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Map local topk indices to global KV cache slots and count valid entries.
 
@@ -440,8 +441,13 @@ def compute_global_topk_indices_and_lens(
     3. Masking padding tokens to length 0
     """
     num_tokens = topk_indices.shape[0]
-    global_topk_indices = torch.empty_like(topk_indices)
-    topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)
+    if output_buffers is None:
+        global_topk_indices = torch.empty_like(topk_indices)
+        topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)
+    else:
+        global_topk_indices, topk_lens = output_buffers
+        assert global_topk_indices.shape == topk_indices.shape
+        assert topk_lens.shape == (num_tokens,)
     _compute_global_topk_indices_and_lens_kernel[(num_tokens,)](
         global_topk_indices,
         global_topk_indices.stride(0),

aphrodite/models/deepseek_v4/common/ops/fused_indexer_q.py

@@ -276,6 +276,7 @@ def fused_indexer_q_rope_quant(
     index_weights_softmax_scale: float,
     index_weights_head_scale: float,
     use_fp4: bool = False,
+    output_buffers: tuple[torch.Tensor, ...] | None = None,
 ) -> tuple[
     torch.Tensor | tuple[torch.Tensor, torch.Tensor],
     torch.Tensor,
@@ -313,23 +314,36 @@ def fused_indexer_q_rope_quant(
     num_index_q_heads = index_q.shape[1]
     index_q_head_dim = index_q.shape[2]
 
-    index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
+    if output_buffers is None:
+        index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
+    else:
+        expected_num_buffers = 3 if use_fp4 else 2
+        assert len(output_buffers) == expected_num_buffers
+        index_weights_out = output_buffers[-1]
+        assert index_weights_out.shape == index_weights.shape
 
     if use_fp4:
         assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, (
             f"head_dim={index_q_head_dim} must be a multiple of MXFP4 block size {MXFP4_BLOCK_SIZE}"
         )
         num_scale_blocks = index_q_head_dim // MXFP4_BLOCK_SIZE
-        index_q_packed = torch.empty(
-            (num_tokens, num_index_q_heads, index_q_head_dim // 2),
-            dtype=torch.uint8,
-            device=index_q.device,
-        )
-        index_q_scale = torch.empty(
-            (num_tokens, num_index_q_heads, num_scale_blocks),
-            dtype=torch.uint8,
-            device=index_q.device,
-        )
+        packed_shape = (num_tokens, num_index_q_heads, index_q_head_dim // 2)
+        scale_shape = (num_tokens, num_index_q_heads, num_scale_blocks)
+        if output_buffers is None:
+            index_q_packed = torch.empty(
+                packed_shape,
+                dtype=torch.uint8,
+                device=index_q.device,
+            )
+            index_q_scale = torch.empty(
+                scale_shape,
+                dtype=torch.uint8,
+                device=index_q.device,
+            )
+        else:
+            index_q_packed, index_q_scale, _ = output_buffers
+        assert index_q_packed.shape == packed_shape
+        assert index_q_scale.shape == scale_shape
         if has_cutedsl():
             # lazily import, otherwise some tests fail due to CUDA driver init failure.
             from aphrodite.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
@@ -398,7 +412,11 @@ def fused_indexer_q_rope_quant(
     fp8_dtype = current_platform.fp8_dtype()
     use_fnuz = fp8_dtype == torch.float8_e4m3fnuz
     fp8_max = 224.0 if use_fnuz else 448.0
-    index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype)
+    if output_buffers is None:
+        index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype)
+    else:
+        index_q_fp8, _ = output_buffers
+        assert index_q_fp8.shape == index_q.shape
     if has_cutedsl():
         # lazily import, otherwise some tests fail due to CUDA driver init failure.
         from aphrodite.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (

aphrodite/models/deepseek_v4/compressor.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from dataclasses import dataclass
-from typing import Any, ClassVar, cast
+from typing import TYPE_CHECKING, Any, ClassVar, cast
 
 import torch
 from torch import nn
@@ -39,6 +39,9 @@
     SlidingWindowMLASpec,
 )
 
+if TYPE_CHECKING:
+    from aphrodite.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool
+
 
 def _prefer_two_stage_compressor() -> bool:
     # Platforms that favor the triton variant of two-stage compressor split.
@@ -219,6 +222,7 @@ def __init__(
         prefix: str = "",
         k_cache_prefix="",
         use_fp4_cache: bool = False,
+        eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None,
     ):
         super().__init__()
         self.compress_ratio = compress_ratio
@@ -228,6 +232,7 @@ def __init__(
         self.prefix = prefix
         self.k_cache_prefix = k_cache_prefix
         self.use_fp4_cache = use_fp4_cache
+        self.eager_scratch_pool = eager_scratch_pool
 
         config = aphrodite_config.model_config.hf_config
         self.rope_head_dim = config.qk_rope_head_dim
@@ -399,6 +404,8 @@ def forward(
                 store_full_fp8=store_full_fp8,
                 fp8_scale=fp8_scale,
             )
+            if not self.overlap and self.eager_scratch_pool is not None:
+                extra_kwargs["compress_scratch"] = self.eager_scratch_pool.compressor_scratch(num_actual)
         elif self._use_two_stage_fused_compressor:
             # head=512 cr>=128 (no overlap): two-pass split compressor on the
             # prefill suffix, single-pass on the decode prefix.

aphrodite/models/deepseek_v4/eager_scratch.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from math import prod
+
+import torch
+
+from aphrodite.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE
+from aphrodite.utils.math_utils import round_up
+
+
+class DeepseekV4EagerScratchPool:
+    """Model-wide outputs and scratch used inside the attention eager break."""
+
+    _ALIGNMENT = 256
+
+    def __init__(
+        self,
+        max_num_tokens: int,
+        q_workspace: torch.Tensor,
+        q_head_dim: int,
+        index_q_heads: int,
+        index_q_head_dim: int,
+        index_topk: int,
+        device: torch.device | str,
+    ) -> None:
+        self.max_num_tokens = max_num_tokens
+        self.index_topk = index_topk
+        assert q_workspace.shape[0] >= max_num_tokens
+        assert q_workspace.shape[-1] == q_head_dim
+        self._q = q_workspace
+
+        fp4_specs = (
+            ((max_num_tokens, index_q_heads, index_q_head_dim // 2), torch.uint8),
+            (
+                (
+                    max_num_tokens,
+                    index_q_heads,
+                    index_q_head_dim // MXFP4_BLOCK_SIZE,
+                ),
+                torch.uint8,
+            ),
+            ((max_num_tokens, index_q_heads), torch.float32),
+        )
+        global_specs = (
+            ((max_num_tokens, index_topk), torch.int32),
+            ((max_num_tokens,), torch.int32),
+        )
+        compressor_specs = (((max_num_tokens, q_head_dim), torch.float32),)
+        # FP4 indexer is C4 only, global mapping after FP4 indexer
+        # compressor scratch is C128 only
+        # so here we use max instead of sum
+        aux_bytes = max(self._packed_size(specs) for specs in (fp4_specs, global_specs, compressor_specs))
+        storage = torch.empty(aux_bytes, dtype=torch.uint8, device=device)
+
+        self._q_outputs: dict[int, torch.Tensor] = {}
+        fp4_values, fp4_scales, fp4_weights = self._views(storage, fp4_specs)
+        self._fp4_template = (fp4_values, fp4_scales, fp4_weights)
+        self._fp4_outputs: dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}
+        global_indices, global_lens = self._views(storage, global_specs)
+        self._global_template = (global_indices, global_lens)
+        self._global_outputs: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
+        self._compressor_template = self._views(storage, compressor_specs)[0]
+        self._compressor_outputs: dict[int, torch.Tensor] = {}
+        self._storage = storage
+
+    @classmethod
+    def _packed_size(cls, specs: tuple[tuple[tuple[int, ...], torch.dtype], ...]) -> int:
+        offset = 0
+        for shape, dtype in specs:
+            offset = round_up(offset, cls._ALIGNMENT) + prod(shape) * dtype.itemsize
+        return round_up(offset, cls._ALIGNMENT)
+
+    @classmethod
+    def _views(
+        cls,
+        storage: torch.Tensor,
+        specs: tuple[tuple[tuple[int, ...], torch.dtype], ...],
+    ) -> list[torch.Tensor]:
+        offset = 0
+        views = []
+        for shape, dtype in specs:
+            offset = round_up(offset, cls._ALIGNMENT)
+            num_bytes = prod(shape) * dtype.itemsize
+            views.append(storage[offset : offset + num_bytes].view(dtype).view(shape))
+            offset += num_bytes
+        return views
+
+    def q_out(self, num_tokens: int) -> torch.Tensor:
+        output = self._q_outputs.get(num_tokens)
+        if output is None:
+            output = self._q[:num_tokens]
+            self._q_outputs[num_tokens] = output
+        return output
+
+    def compressor_scratch(self, num_tokens: int) -> torch.Tensor:
+        output = self._compressor_outputs.get(num_tokens)
+        if output is None:
+            output = self._compressor_template[:num_tokens]
+            self._compressor_outputs[num_tokens] = output
+        return output
+
+    def indexer_q_outputs(
+        self,
+        num_tokens: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        output = self._fp4_outputs.get(num_tokens)
+        if output is None:
+            values, scales, weights = self._fp4_template
+            output = (
+                values[:num_tokens],
+                scales[:num_tokens],
+                weights[:num_tokens],
+            )
+            self._fp4_outputs[num_tokens] = output
+        return output
+
+    def global_topk_outputs(self, topk_indices: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        num_tokens, topk = topk_indices.shape
+        assert topk == self.index_topk
+        output = self._global_outputs.get(num_tokens)
+        if output is None:
+            indices, lens = self._global_template
+            output = (indices[:num_tokens], lens[:num_tokens])
+            self._global_outputs[num_tokens] = output
+        return output

aphrodite/models/deepseek_v4/nvidia/dspark.py

@@ -86,12 +86,19 @@ def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = "") -> No
         )
         self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
+        self.topk_indices_buffer = torch.empty(
+            aphrodite_config.scheduler_config.max_num_batched_tokens,
+            config.index_topk,
+            dtype=torch.int32,
+        )
+
         current_aphrodite_config = get_current_aphrodite_config()
         self.layers = nn.ModuleList(
             [
                 DeepseekV4DecoderLayer(
                     current_aphrodite_config,
                     prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"),
+                    topk_indices_buffer=self.topk_indices_buffer,
                 )
                 for i in range(self.num_dspark_layers)
             ]