[ROCm][CI] Fix Kimi K3 KDA on ROCm (#50262)

Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>

.buildkite/test-amd.yaml

@@ -1576,10 +1576,14 @@ steps:
   - vllm/third_party/flash_linear_attention/ops/kda.py
   - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
   - vllm/third_party/flash_linear_attention/ops/l2norm.py
-  - tests/kernels/test_kda.py
+  - vllm/models/kimi_k3/nvidia/kda.py
+  - vllm/models/kimi_k3/nvidia/kda_metadata.py
+  - vllm/models/kimi_k3/nvidia/ops/third_party/kda/
+  - tests/models/kimi_k3/test_kda.py
+  - tests/models/kimi_k3/test_kda_metadata.py
   - vllm/platforms/rocm.py
   commands:
-  - pytest -v -s kernels/test_kda.py
+  - pytest -v -s models/kimi_k3/test_kda.py models/kimi_k3/test_kda_metadata.py
 
 - label: Kernels Mamba Test # TBD
   timeout_in_minutes: 180

csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

@@ -942,10 +942,9 @@ static void launchFullCacheKernel(
 // ────────────────────────────────────────────────────────────────────────────
 // Torch op wrapper
 // ────────────────────────────────────────────────────────────────────────────
-void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
+torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
     torch::stable::Tensor const& q_in,           // [N, num_heads_q, 512] bf16
     torch::stable::Tensor const& kv,             // [N, 512] bf16 (read-only)
-    torch::stable::Tensor& q_out,                // [N, q_head_padded, 512]
     torch::stable::Tensor& k_cache,              // [num_blocks, block_bytes] uint8
     torch::stable::Tensor const& slot_mapping,   // [N] int64
     torch::stable::Tensor const& position_ids,   // [N] int64
@@ -971,16 +970,8 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
   STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
   STD_TORCH_CHECK(q_in.scalar_type() == kv.scalar_type(),
                   "q_in and kv dtype must match");
-  STD_TORCH_CHECK(q_out.device() == q_in.device() && q_out.is_contiguous(),
-                  "q_out must be contiguous and on the same device as q_in");
-  STD_TORCH_CHECK(q_out.scalar_type() == q_in.scalar_type(),
-                  "q_out dtype must match q_in");
   STD_TORCH_CHECK(q_head_padded >= q_in.size(1),
                   "q_head_padded must be >= q_in.size(1) (num_heads_q)");
-  STD_TORCH_CHECK(q_out.dim() == 3 && q_out.size(0) == q_in.size(0) &&
-                      q_out.size(1) == q_head_padded &&
-                      q_out.size(2) == q_in.size(2),
-                  "q_out shape [N, q_head_padded, 512]");
   STD_TORCH_CHECK(k_cache.scalar_type() == torch::headeronly::ScalarType::Byte,
                   "k_cache must be uint8");
   STD_TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
@@ -1008,6 +999,11 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
       q_in.get_device_index());
   const cudaStream_t stream = get_current_cuda_stream(q_in.get_device_index());
 
+  // Allocate the padded q output.  The kernel writes every element (live
+  // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe.
+  auto q_out = torch::stable::new_empty(
+      q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type());
+
   VLLM_STABLE_DISPATCH_HALF_TYPES(
       q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
         using qkv_scalar_t = scalar_t;
@@ -1024,20 +1020,6 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
                 num_heads_q_padded, cache_block_size_i, kv_block_stride,
                 stream);
       });
-}
-
-torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv,
-    torch::stable::Tensor& k_cache,
-    torch::stable::Tensor const& slot_mapping,
-    torch::stable::Tensor const& position_ids,
-    torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
-    double eps, int64_t cache_block_size) {
-  auto q_out = torch::stable::new_empty(
-      q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type());
-  fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
-      q_in, kv, q_out, k_cache, slot_mapping, position_ids, cos_sin_cache,
-      q_head_padded, eps, cache_block_size);
   return q_out;
 }
 

csrc/libtorch_stable/ops.h

@@ -269,14 +269,6 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
     torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
     double eps, int64_t cache_block_size);
 
-void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
-    torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv,
-    torch::stable::Tensor& q_out, torch::stable::Tensor& k_cache,
-    torch::stable::Tensor const& slot_mapping,
-    torch::stable::Tensor const& position_ids,
-    torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
-    double eps, int64_t cache_block_size);
-
 void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
     torch::stable::Tensor& q, torch::stable::Tensor const& kv,
     torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,

csrc/libtorch_stable/torch_bindings.cpp

@@ -433,11 +433,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "Tensor q_in, Tensor kv, Tensor! k_cache, "
       "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
       "int q_head_padded, float eps, int cache_block_size) -> Tensor");
-  ops.def(
-      "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out("
-      "Tensor q_in, Tensor kv, Tensor! q_out, Tensor! k_cache, "
-      "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
-      "int q_head_padded, float eps, int cache_block_size) -> ()");
 
   // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate
   // FP8 tensor, and KV into a contiguous 512-wide token-strided cache.
@@ -756,8 +751,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
   ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert",
            TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert));
-  ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out",
-           TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out));
   ops.impl(
       "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert",
       TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert));

tests/kernels/test_compressor_kv_cache.py

@@ -20,7 +20,6 @@ import torch
 
 from vllm import _custom_ops as ops
 from vllm.models.deepseek_v4.common.ops import (
-    compute_global_topk_indices_and_lens,
     dequantize_and_gather_k_cache,
     quantize_and_insert_k_cache,
 )
@@ -35,23 +34,6 @@ from vllm.platforms import current_platform
 from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
 
 
-def test_compute_global_topk_reuses_output_buffers():
-    device = "cuda"
-    topk_indices = torch.tensor(
-        [[0, 3, -1], [1, 2, -1]], dtype=torch.int32, device=device
-    )
-    token_to_req = torch.tensor([0, 1], dtype=torch.int32, device=device)
-    block_table = torch.tensor([[5, 7], [11, 13]], dtype=torch.int32, device=device)
-    is_valid = torch.tensor([True, False], device=device)
-    args = (topk_indices, token_to_req, block_table, 2, is_valid)
-    expected = compute_global_topk_indices_and_lens(*args)
-    outputs = tuple(torch.empty_like(tensor) for tensor in expected)
-    actual = compute_global_topk_indices_and_lens(*args, output_buffers=outputs)
-    for result, output, reference in zip(actual, outputs, expected):
-        assert result.data_ptr() == output.data_ptr()
-        torch.testing.assert_close(result, reference)
-
-
 def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
     """PyTorch reference for UE8M0 FP8 quantization (per-block, power-of-2 scale).
 

tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py

@@ -257,18 +257,8 @@ def test_q_path_matches_reference(num_tokens: int, n_heads: int, padded_heads: i
         num_blocks, bs, HEAD_BYTES, dtype=torch.uint8, device=device
     ).view(num_blocks, -1)
     slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device)
-    q_out = torch.empty(num_tokens, padded_heads, HEAD_DIM, dtype=dtype, device=device)
-    torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
-        q,
-        kv,
-        q_out,
-        k_cache,
-        slot_mapping,
-        positions,
-        cos_sin_cache,
-        padded_heads,
-        eps,
-        bs,
+    q_out = _call_fused(
+        q, padded_heads, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
     )
 
     torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=1e-2, atol=1e-2)

tests/kernels/test_fused_indexer_q_rope_quant.py

@@ -150,23 +150,6 @@ def test_fused_indexer_q_rope_quant_matches_unfused(
     q_quant_ref, weights_ref = _reference(
         positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
     )
-    output_buffers: tuple[torch.Tensor, ...] | None = None
-    OUTPUT_BUFFER_TEST_NUM_TOKENS = 7
-    if num_tokens == OUTPUT_BUFFER_TEST_NUM_TOKENS and cache_dtype == torch.float32:
-        if use_fp4:
-            q_ref, q_scale_ref = q_quant_ref
-            output_buffers = (
-                torch.empty_like(q_ref),
-                torch.empty_like(q_scale_ref)
-                .view(torch.uint8)
-                .reshape(num_tokens, N_HEAD, -1),
-                torch.empty_like(weights_ref),
-            )
-        else:
-            output_buffers = (
-                torch.empty_like(q_quant_ref),
-                torch.empty_like(weights_ref),
-            )
     # use_cutedsl=False: force the triton path even when cutedsl is installed
     # by patching the dispatcher's has_cutedsl() binding to return False.
     cutedsl_patch = (
@@ -186,17 +169,8 @@ def test_fused_indexer_q_rope_quant_matches_unfused(
             softmax_scale,
             head_scale,
             use_fp4,
-            output_buffers=output_buffers,
         )
 
-    if output_buffers is not None:
-        if use_fp4:
-            assert q_quant_fused[0].data_ptr() == output_buffers[0].data_ptr()
-            assert q_quant_fused[1].data_ptr() == output_buffers[1].data_ptr()
-        else:
-            assert q_quant_fused.data_ptr() == output_buffers[0].data_ptr()
-        assert weights_fused.data_ptr() == output_buffers[-1].data_ptr()
-
     if use_fp4:
         q_quant_ref, q_scale_ref = q_quant_ref
         q_quant_fused, q_scale_fused = q_quant_fused

tests/v1/sample/test_head_dtype.py

@@ -14,6 +14,7 @@ import torch
 from vllm import LLM, SamplingParams
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
     UnquantizedEmbeddingMethod,
 )
 
@@ -28,6 +29,7 @@ class _FakeLmHead:
         self.weight = weight
         self.quant_method = object() if quantized else UnquantizedEmbeddingMethod()
         self.shard_indices = shard_indices
+        self.tp_size = 1
 
 
 def _build_processor(vocab_size: int) -> LogitsProcessor:
@@ -135,11 +137,59 @@ def test_fp32_head_rejects_quantized_lm_head(default_vllm_config):
         lp._get_logits(torch.randn(4, 16, dtype=torch.bfloat16), lm_head, None)
 
 
+def test_replicated_lm_head_skips_tp_communication_and_preserves_processing(
+    default_vllm_config,
+):
+    from unittest import mock
+
+    vocab_size, hidden_size = 12, 8
+    soft_cap, scale = 2.0, 0.5
+    lp = LogitsProcessor(
+        vocab_size,
+        soft_cap=soft_cap,
+        scale=scale,
+    )
+    lp.head_dtype = torch.float32
+
+    hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16)
+    weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16)
+    world_size_getter = (
+        "vllm.model_executor.layers.vocab_parallel_embedding."
+        "get_tensor_model_parallel_world_size"
+    )
+    with mock.patch(world_size_getter, return_value=2):
+        lm_head = ParallelLMHead(
+            vocab_size,
+            hidden_size,
+            params_dtype=torch.bfloat16,
+            disable_tp=True,
+        )
+    lm_head.weight_loader(lm_head.weight, weight)
+    assert lm_head.tp_size == 1
+
+    with mock.patch.object(lp, "_gather_logits") as gather_mock:
+        logits = lp(lm_head, hidden_states)
+
+    gather_mock.assert_not_called()
+
+    expected = torch.nn.functional.linear(hidden_states.float(), weight.float())
+    expected = torch.tanh(expected / soft_cap) * soft_cap * scale
+    torch.testing.assert_close(logits, expected)
+
+    all_gather_path = (
+        "vllm.model_executor.layers.logits_processor.tensor_model_parallel_all_gather"
+    )
+    with mock.patch(all_gather_path) as all_gather:
+        top = lp.get_top_tokens(lm_head, hidden_states)
+
+    all_gather.assert_not_called()
+    assert torch.equal(top, expected.argmax(dim=-1))
+
+
 def test_get_top_tokens_honors_head_dtype(default_vllm_config):
     # The spec-decode local-argmax path (get_top_tokens) must run the lm_head
     # in head_dtype too, not just _get_logits.
     import types
-    from unittest import mock
 
     vocab_size, hidden_size = 64, 16
     lp = _build_processor(vocab_size)
@@ -154,13 +204,7 @@ def test_get_top_tokens_honors_head_dtype(default_vllm_config):
         ),
     )
 
-    with mock.patch(
-        "vllm.model_executor.layers.logits_processor."
-        "get_tensor_model_parallel_world_size",
-        return_value=1,
-    ):
-        top = lp.get_top_tokens(lm_head, hidden_states, None)
-
+    top = lp.get_top_tokens(lm_head, hidden_states, None)
     expected = torch.nn.functional.linear(hidden_states.float(), weight.float()).argmax(
         dim=-1
     )

vllm/model_executor/layers/logits_processor.py

@@ -7,7 +7,6 @@ import torch.nn.functional as F
 
 from vllm.config import get_current_vllm_config
 from vllm.distributed import (
-    get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_gather,
     tensor_model_parallel_gather,
 )
@@ -145,7 +144,8 @@ class LogitsProcessor(PluggableLayer):
         logits = self._apply_head(lm_head, hidden_states, embedding_bias)
 
         # Gather logits for TP
-        logits = self._gather_logits(logits)
+        if lm_head.tp_size > 1:
+            logits = self._gather_logits(logits)
 
         # Remove paddings in vocab (if any).
         if logits is not None:
@@ -169,7 +169,7 @@ class LogitsProcessor(PluggableLayer):
                 "The local argmax reduction optimization is not supported for "
                 "non-positive logit scaling factors."
             )
-        tp_size = get_tensor_model_parallel_world_size()
+        tp_size = lm_head.tp_size
 
         logits = self._apply_head(lm_head, hidden_states, embedding_bias)
         if self.soft_cap is not None:

vllm/model_executor/layers/vocab_parallel_embedding.py

@@ -232,6 +232,7 @@ class VocabParallelEmbedding(PluggableLayer):
         padding_size: padding size for the vocabulary.
         quant_config: quant config for the layer
         prefix: full name of the layer in the state dict
+        disable_tp: If true, tensor parallelism will be disabled for this layer.
     """  # noqa: E501
 
     # --8<-- [end:vocab_parallel_embedding]
@@ -245,12 +246,19 @@ class VocabParallelEmbedding(PluggableLayer):
         padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
+        *,
+        disable_tp: bool = False,
     ):
         super().__init__()
 
         # Keep the input dimensions.
-        tp_rank = get_tensor_model_parallel_rank()
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.disable_tp = disable_tp
+        if disable_tp:
+            tp_rank, self.tp_size = 0, 1
+        else:
+            tp_rank = get_tensor_model_parallel_rank()
+            self.tp_size = get_tensor_model_parallel_world_size()
+        self.tp_rank = tp_rank
         self.num_embeddings = num_embeddings
         self.padding_size = padding_size
         self.org_vocab_size = org_num_embeddings or num_embeddings
@@ -323,6 +331,13 @@ class VocabParallelEmbedding(PluggableLayer):
             params_dtype=params_dtype,
             weight_loader=self.weight_loader,
         )
+        self.update_param_tp_status()
+
+    def update_param_tp_status(self):
+        for param in self.parameters():
+            if isinstance(param, BasevLLMParameter):
+                param.tp_rank = self.tp_rank
+                param.tp_size = self.tp_size
 
     @classmethod
     def _get_indices(
@@ -487,9 +502,9 @@ class VocabParallelEmbedding(PluggableLayer):
         # Mask the output embedding.
         if self.tp_size > 1:
             output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
-        # Reduce across all the model parallel GPUs.
-        output = tensor_model_parallel_all_reduce(output_parallel)
-        return output
+            # Reduce across all the model parallel GPUs.
+            return tensor_model_parallel_all_reduce(output_parallel)
+        return output_parallel
 
     def extra_repr(self) -> str:
         s = f"num_embeddings={self.num_embeddings_per_partition}"
@@ -516,6 +531,7 @@ class ParallelLMHead(VocabParallelEmbedding):
         params_dtype: type of the parameters.
         org_num_embeddings: original vocabulary size (without LoRA).
         padding_size: padding size for the vocabulary.
+        disable_tp: If true, tensor parallelism will be disabled for this layer.
     """
 
     # --8<-- [end:parallel_lm_head]
@@ -530,6 +546,8 @@ class ParallelLMHead(VocabParallelEmbedding):
         padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
+        *,
+        disable_tp: bool = False,
     ):
         super().__init__(
             num_embeddings,
@@ -539,6 +557,7 @@ class ParallelLMHead(VocabParallelEmbedding):
             padding_size,
             quant_config,
             prefix,
+            disable_tp=disable_tp,
         )
         self.quant_config = quant_config
         if bias:

vllm/model_executor/models/qwen3_dspark.py

@@ -24,7 +24,6 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
-    VocabParallelEmbedding,
 )
 
 from .qwen3_dflash import DFlashQwen3ForCausalLM, DFlashQwen3Model
@@ -40,6 +39,10 @@ class DSparkMarkovHead(nn.Module):
     ``vocab_size``); ``markov_w2`` projects it to a draft-vocab bias
     (``draft_vocab_size``) added to the base draft logits. The two sizes
     coincide for full-vocab drafts.
+
+    Both weights are replicated because the head runs sequentially for every
+    draft position. Sharding them would add an all-reduce and a full-vocab
+    gather to each position.
     """
 
     def __init__(
@@ -50,19 +53,24 @@ class DSparkMarkovHead(nn.Module):
         prefix: str,
     ) -> None:
         super().__init__()
-        # TODO(ben): profile for which (if any) it makes sense to replicate or TP-shard
-        self.markov_w1 = VocabParallelEmbedding(
-            vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w1")
-        )
+        self.markov_w1 = nn.Embedding(vocab_size, markov_rank)
         self.markov_w2 = ParallelLMHead(
-            draft_vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w2")
+            draft_vocab_size,
+            markov_rank,
+            bias=False,
+            prefix=maybe_prefix(prefix, "markov_w2"),
+            disable_tp=True,
         )
 
     def embed(self, token_ids: torch.Tensor) -> torch.Tensor:
         """r-dim Markov embedding of ``token_ids`` ([B] -> [B, r])."""
         return self.markov_w1(token_ids)
 
-    def bias(self, markov_embed: torch.Tensor, logits_processor) -> torch.Tensor:
+    def bias(
+        self,
+        markov_embed: torch.Tensor,
+        logits_processor: LogitsProcessor,
+    ) -> torch.Tensor:
         """Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V])."""
         return logits_processor(self.markov_w2, markov_embed)
 

vllm/models/deepseek_v4/attention.py

@@ -29,7 +29,6 @@ from vllm.models.deepseek_v4.common.ops import (
 from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE
 
 if TYPE_CHECKING:
-    from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool
     from vllm.v1.attention.backends.mla.sparse_swa import (
         DeepseekSparseSWAMetadata,
     )
@@ -182,7 +181,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         prefix: str,
         topk_indices_buffer: torch.Tensor | None = None,
         aux_stream_list: list[torch.cuda.Stream] | None = None,
-        eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None,
     ) -> None:
         super().__init__()
         config = vllm_config.model_config.hf_config
@@ -271,7 +269,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         )
         self.indexer_rotary_emb = self.rotary_emb
         self.topk_indices_buffer = topk_indices_buffer
-        self.eager_scratch_pool = eager_scratch_pool
 
         self.indexer = None
         if self.compress_ratio == 4:
@@ -293,7 +290,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
                 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.
@@ -344,7 +340,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
                 rotate=True,
                 prefix=f"{prefix}.compressor",
                 k_cache_prefix=self.prefix,
-                eager_scratch_pool=eager_scratch_pool,
             )
 
     def forward(
@@ -573,24 +568,10 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         if cache_dtype == torch.uint8:
             # fp8_ds_mla UE8M0 paged path. Horizontally fused:
             #   Q side:  per-head RMSNorm (no weight) + GPT-J RoPE, zero-filling
-            #            the padding head slots.
+            #            the padding head slots; the kernel allocates and returns
+            #            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.eager_scratch_pool is not None:
-                q_out = self.eager_scratch_pool.q_out(q.shape[0])
-                torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
-                    q,
-                    kv,
-                    q_out,
-                    swa_kv_cache_2d,
-                    swa_metadata.slot_mapping,
-                    positions,
-                    cos_sin_cache,
-                    self.padded_heads,
-                    self.eps,
-                    swa_metadata.block_size,
-                )
-                return q_out
             return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
                 q,
                 kv,
@@ -639,13 +620,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         )
         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
 
@@ -725,7 +699,6 @@ class DeepseekV4Indexer(nn.Module):
         compress_ratio: int = 1,
         prefix: str = "",
         aux_stream: torch.cuda.Stream | None = None,
-        eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None,
     ):
         super().__init__()
         self.vllm_config = vllm_config
@@ -738,7 +711,6 @@ class DeepseekV4Indexer(nn.Module):
         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.vllm_config.attention_config.use_fp4_indexer_cache
         logger.info_once(
             "Using %s indexer cache for Lightning Indexer.",
@@ -802,7 +774,6 @@ class DeepseekV4Indexer(nn.Module):
             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(
@@ -863,9 +834,6 @@ class DeepseekV4Indexer(nn.Module):
             # 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,
@@ -874,7 +842,6 @@ class DeepseekV4Indexer(nn.Module):
                 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

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

@@ -438,7 +438,6 @@ 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.
 
@@ -448,15 +447,8 @@ def compute_global_topk_indices_and_lens(
     3. Masking padding tokens to length 0
     """
     num_tokens = topk_indices.shape[0]
-    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,)
+    global_topk_indices = torch.empty_like(topk_indices)
+    topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)
     _compute_global_topk_indices_and_lens_kernel[(num_tokens,)](
         global_topk_indices,
         global_topk_indices.stride(0),

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

@@ -295,7 +295,6 @@ 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,
@@ -333,13 +332,7 @@ def fused_indexer_q_rope_quant(
     num_index_q_heads = index_q.shape[1]
     index_q_head_dim = index_q.shape[2]
 
-    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
+    index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
 
     if use_fp4:
         assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, (
@@ -347,23 +340,16 @@ def fused_indexer_q_rope_quant(
             f"size {MXFP4_BLOCK_SIZE}"
         )
         num_scale_blocks = index_q_head_dim // MXFP4_BLOCK_SIZE
-        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
+        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,
+        )
         if has_cutedsl():
             # lazily import, otherwise some tests fail due to CUDA driver init failure.
             from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
@@ -432,11 +418,7 @@ 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
-    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
+    index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype)
     if has_cutedsl():
         # lazily import, otherwise some tests fail due to CUDA driver init failure.
         from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (

vllm/models/deepseek_v4/compressor.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, ClassVar, cast
+from typing import Any, ClassVar, cast
 
 import torch
 from torch import nn
@@ -35,9 +35,6 @@ from vllm.v1.kv_cache_interface import (
     SlidingWindowMLASpec,
 )
 
-if TYPE_CHECKING:
-    from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool
-
 
 def _prefer_two_stage_compressor() -> bool:
     # Platforms that favor the triton variant of two-stage compressor split.
@@ -229,7 +226,6 @@ class DeepseekCompressor(nn.Module):
         prefix: str = "",
         k_cache_prefix="",
         use_fp4_cache: bool = False,
-        eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None,
     ):
         super().__init__()
         self.compress_ratio = compress_ratio
@@ -239,7 +235,6 @@ class DeepseekCompressor(nn.Module):
         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 = vllm_config.model_config.hf_config
         self.rope_head_dim = config.qk_rope_head_dim
@@ -433,10 +428,6 @@ class DeepseekCompressor(nn.Module):
                 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.

vllm/models/deepseek_v4/eager_scratch.py

@@ -1,137 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from math import prod
-
-import torch
-
-from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE
-from vllm.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,
-        padded_q_heads: int,
-        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
-        self._q = torch.empty(
-            (max_num_tokens, padded_q_heads, q_head_dim),
-            dtype=torch.bfloat16,
-            device=device,
-        )
-
-        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

vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py

@@ -748,9 +748,6 @@ class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
                         attn_metadata.block_table[:num_decodes],
                         block_size,
                         is_valid,
-                        output_buffers=self._global_topk_output_buffers(
-                            self.topk_indices_buffer[:num_decode_tokens]
-                        ),
                     )
                 )
                 extra_sparse_indices = global_indices.view(num_decode_tokens, 1, -1)
@@ -840,7 +837,6 @@ class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
                     attn_metadata.block_table,
                     block_size,
                     swa_metadata.is_valid_token[prefill_token_slice],
-                    output_buffers=self._global_topk_output_buffers(local_topk_indices),
                 )
             )
 

vllm/models/deepseek_v4/nvidia/flashmla.py

@@ -170,9 +170,6 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
                     attn_metadata.block_table[:num_decodes],
                     block_size,
                     is_valid,
-                    output_buffers=self._global_topk_output_buffers(
-                        self.topk_indices_buffer[:num_decode_tokens]
-                    ),
                 )
                 topk_indices = global_indices.view(num_decode_tokens, 1, -1)
             else:

vllm/models/deepseek_v4/nvidia/model.py

@@ -66,7 +66,6 @@ from vllm.model_executor.models.utils import (
 )
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.models.deepseek_v4.attention import DeepseekV4Attention
-from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool
 from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import (
     DeepseekV4FlashInferMLAAttention,
     DeepseekV4FlashInferSM120Attention,
@@ -799,7 +798,6 @@ class DeepseekV4DecoderLayer(nn.Module):
         prefix,
         topk_indices_buffer: torch.Tensor | None = None,
         aux_stream_list: list[torch.cuda.Stream] | None = None,
-        eager_scratch_pool: DeepseekV4EagerScratchPool | None = None,
     ):
         super().__init__()
 
@@ -812,7 +810,6 @@ class DeepseekV4DecoderLayer(nn.Module):
             prefix=f"{prefix}.attn",
             topk_indices_buffer=topk_indices_buffer,
             aux_stream_list=aux_stream_list,
-            eager_scratch_pool=eager_scratch_pool,
         )
         self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")
 
@@ -989,22 +986,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
         # (compressor kv_score, indexer.weights_proj, indexer.compressor
         # kv_score). fused_wqa_wkv stays on the default stream.
         aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
-        padded_heads = _select_dsv4_attn_cls(vllm_config).get_padded_num_q_heads(
-            config.num_attention_heads // get_tensor_model_parallel_world_size()
-        )
-        self.eager_scratch_pool: DeepseekV4EagerScratchPool | None = None
-        if not vllm_config.parallel_config.use_ubatching:
-            # TODO: support dbo if needed
-            # this requires the buffer to have ubatch dim
-            self.eager_scratch_pool = DeepseekV4EagerScratchPool(
-                vllm_config.scheduler_config.max_num_batched_tokens,
-                padded_heads,
-                config.head_dim,
-                config.index_n_heads,
-                config.index_head_dim,
-                config.index_topk,
-                current_platform.device_type,
-            )
 
         # Reserved topk indices buffer for all Indexer layers to reuse.
         self.topk_indices_buffer = torch.empty(
@@ -1030,7 +1011,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
                 prefix=prefix,
                 topk_indices_buffer=self.topk_indices_buffer,
                 aux_stream_list=aux_stream_list,
-                eager_scratch_pool=self.eager_scratch_pool,
             ),
             prefix=f"{prefix}.layers",
         )

vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py

@@ -2097,7 +2097,6 @@ def compress_norm_rope_store_cutedsl(
     store_full_kv: bool = False,
     store_full_fp8: bool = False,
     fp8_scale: torch.Tensor | None = None,
-    compress_scratch: torch.Tensor | None = None,
 ) -> None:
     if compress_ratio == 4:
         # For C4A, the single fused kernel is faster than the two-kernel version.
@@ -2130,15 +2129,11 @@ def compress_norm_rope_store_cutedsl(
         )
     else:
         # For C128, the two-kernel version is faster than the single fused kernel.
-        if compress_scratch is None:
-            compressed_kv = torch.empty(
-                (num_actual, head_dim),
-                dtype=torch.float32,
-                device=state_cache.device,
-            )
-        else:
-            assert compress_scratch.shape == (num_actual, head_dim)
-            compressed_kv = compress_scratch
+        compressed_kv = torch.empty(
+            (num_actual, head_dim),
+            dtype=torch.float32,
+            device=state_cache.device,
+        )
         split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
             state_cache,
             token_to_req_indices,

vllm/models/kimi_k3/nvidia/kda.py

@@ -164,6 +164,8 @@ def is_flashkda_supported(
     dtype: torch.dtype,
     lower_bound: float | None,
 ) -> bool:
+    if not current_platform.is_cuda():
+        return False
     capability = current_platform.get_device_capability()
     return (
         capability is not None