@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f1178f3a06

Revert DSv4 eager workspace reuse (#52836)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>

Woosuk Kwon committed Aug 18, 2026 at 17:02 UTC f1178f3a06fa30a0cc282376924210cedad08c44
15 files changed +30 -354
csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+6 -24
@@ -942,10 +942,9 @@ static void launchFullCacheKernel(
942 // ────────────────────────────────────────────────────────────────────────────
943 // Torch op wrapper
944 // ────────────────────────────────────────────────────────────────────────────
945 -void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
945 +torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
946 torch::stable::Tensor const& q_in, // [N, num_heads_q, 512] bf16
947 torch::stable::Tensor const& kv, // [N, 512] bf16 (read-only)
948 - torch::stable::Tensor& q_out, // [N, q_head_padded, 512]
948 torch::stable::Tensor& k_cache, // [num_blocks, block_bytes] uint8
949 torch::stable::Tensor const& slot_mapping, // [N] int64
950 torch::stable::Tensor const& position_ids, // [N] int64
@@ -971,16 +970,8 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
970 STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
971 STD_TORCH_CHECK(q_in.scalar_type() == kv.scalar_type(),
972 "q_in and kv dtype must match");
974 - STD_TORCH_CHECK(q_out.device() == q_in.device() && q_out.is_contiguous(),
975 - "q_out must be contiguous and on the same device as q_in");
976 - STD_TORCH_CHECK(q_out.scalar_type() == q_in.scalar_type(),
977 - "q_out dtype must match q_in");
973 STD_TORCH_CHECK(q_head_padded >= q_in.size(1),
974 "q_head_padded must be >= q_in.size(1) (num_heads_q)");
980 - STD_TORCH_CHECK(q_out.dim() == 3 && q_out.size(0) == q_in.size(0) &&
981 - q_out.size(1) == q_head_padded &&
982 - q_out.size(2) == q_in.size(2),
983 - "q_out shape [N, q_head_padded, 512]");
975 STD_TORCH_CHECK(k_cache.scalar_type() == torch::headeronly::ScalarType::Byte,
976 "k_cache must be uint8");
977 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(
999 q_in.get_device_index());
1000 const cudaStream_t stream = get_current_cuda_stream(q_in.get_device_index());
1001
1002 + // Allocate the padded q output. The kernel writes every element (live
1003 + // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe.
1004 + auto q_out = torch::stable::new_empty(
1005 + q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type());
1006 +
1007 VLLM_STABLE_DISPATCH_HALF_TYPES(
1008 q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
1009 using qkv_scalar_t = scalar_t;
@@ -1024,20 +1020,6 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
1020 num_heads_q_padded, cache_block_size_i, kv_block_stride,
1021 stream);
1022 });
1027 -}
1028 -
1029 -torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
1030 - torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv,
1031 - torch::stable::Tensor& k_cache,
1032 - torch::stable::Tensor const& slot_mapping,
1033 - torch::stable::Tensor const& position_ids,
1034 - torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
1035 - double eps, int64_t cache_block_size) {
1036 - auto q_out = torch::stable::new_empty(
1037 - q_in, {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.scalar_type());
1038 - fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
1039 - q_in, kv, q_out, k_cache, slot_mapping, position_ids, cos_sin_cache,
1040 - q_head_padded, eps, cache_block_size);
1023 return q_out;
1024 }
1025
csrc/libtorch_stable/ops.h
-8
@@ -269,14 +269,6 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
269 torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
270 double eps, int64_t cache_block_size);
271
272 -void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
273 - torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv,
274 - torch::stable::Tensor& q_out, torch::stable::Tensor& k_cache,
275 - torch::stable::Tensor const& slot_mapping,
276 - torch::stable::Tensor const& position_ids,
277 - torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
278 - double eps, int64_t cache_block_size);
279 -
272 void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
273 torch::stable::Tensor& q, torch::stable::Tensor const& kv,
274 torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
csrc/libtorch_stable/torch_bindings.cpp
-7
@@ -433,11 +433,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
433 "Tensor q_in, Tensor kv, Tensor! k_cache, "
434 "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
435 "int q_head_padded, float eps, int cache_block_size) -> Tensor");
436 - ops.def(
437 - "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out("
438 - "Tensor q_in, Tensor kv, Tensor! q_out, Tensor! k_cache, "
439 - "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
440 - "int q_head_padded, float eps, int cache_block_size) -> ()");
436
437 // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate
438 // FP8 tensor, and KV into a contiguous 512-wide token-strided cache.
@@ -773,8 +768,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
768 ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
769 ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert",
770 TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert));
776 - ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out",
777 - TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out));
771 ops.impl(
772 "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert",
773 TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert));
tests/kernels/test_compressor_kv_cache.py
-18
@@ -20,7 +20,6 @@ import torch
20
21 from vllm import _custom_ops as ops
22 from vllm.models.deepseek_v4.common.ops import (
23 - compute_global_topk_indices_and_lens,
23 dequantize_and_gather_k_cache,
24 quantize_and_insert_k_cache,
25 )
@@ -113,23 +112,6 @@ def test_get_dspark_swa_index_width(
112 assert get_dspark_swa_index_width(window_size, num_speculative_tokens) == expected
113
114
116 -def test_compute_global_topk_reuses_output_buffers():
117 - device = "cuda"
118 - topk_indices = torch.tensor(
119 - [[0, 3, -1], [1, 2, -1]], dtype=torch.int32, device=device
120 - )
121 - token_to_req = torch.tensor([0, 1], dtype=torch.int32, device=device)
122 - block_table = torch.tensor([[5, 7], [11, 13]], dtype=torch.int32, device=device)
123 - is_valid = torch.tensor([True, False], device=device)
124 - args = (topk_indices, token_to_req, block_table, 2, is_valid)
125 - expected = compute_global_topk_indices_and_lens(*args)
126 - outputs = tuple(torch.empty_like(tensor) for tensor in expected)
127 - actual = compute_global_topk_indices_and_lens(*args, output_buffers=outputs)
128 - for result, output, reference in zip(actual, outputs, expected):
129 - assert result.data_ptr() == output.data_ptr()
130 - torch.testing.assert_close(result, reference)
131 -
132 -
115 def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
116 """PyTorch reference for UE8M0 FP8 quantization (per-block, power-of-2 scale).
117
tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+2 -12
@@ -257,18 +257,8 @@ def test_q_path_matches_reference(num_tokens: int, n_heads: int, padded_heads: i
257 num_blocks, bs, HEAD_BYTES, dtype=torch.uint8, device=device
258 ).view(num_blocks, -1)
259 slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device)
260 - q_out = torch.empty(num_tokens, padded_heads, HEAD_DIM, dtype=dtype, device=device)
261 - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
262 - q,
263 - kv,
264 - q_out,
265 - k_cache,
266 - slot_mapping,
267 - positions,
268 - cos_sin_cache,
269 - padded_heads,
270 - eps,
271 - bs,
260 + q_out = _call_fused(
261 + q, padded_heads, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
262 )
263
264 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
-26
@@ -150,23 +150,6 @@ def test_fused_indexer_q_rope_quant_matches_unfused(
150 q_quant_ref, weights_ref = _reference(
151 positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
152 )
153 - output_buffers: tuple[torch.Tensor, ...] | None = None
154 - OUTPUT_BUFFER_TEST_NUM_TOKENS = 7
155 - if num_tokens == OUTPUT_BUFFER_TEST_NUM_TOKENS and cache_dtype == torch.float32:
156 - if use_fp4:
157 - q_ref, q_scale_ref = q_quant_ref
158 - output_buffers = (
159 - torch.empty_like(q_ref),
160 - torch.empty_like(q_scale_ref)
161 - .view(torch.uint8)
162 - .reshape(num_tokens, N_HEAD, -1),
163 - torch.empty_like(weights_ref),
164 - )
165 - else:
166 - output_buffers = (
167 - torch.empty_like(q_quant_ref),
168 - torch.empty_like(weights_ref),
169 - )
153 # use_cutedsl=False: force the triton path even when cutedsl is installed
154 # by patching the dispatcher's has_cutedsl() binding to return False.
155 cutedsl_patch = (
@@ -186,17 +169,8 @@ def test_fused_indexer_q_rope_quant_matches_unfused(
169 softmax_scale,
170 head_scale,
171 use_fp4,
189 - output_buffers=output_buffers,
172 )
173
192 - if output_buffers is not None:
193 - if use_fp4:
194 - assert q_quant_fused[0].data_ptr() == output_buffers[0].data_ptr()
195 - assert q_quant_fused[1].data_ptr() == output_buffers[1].data_ptr()
196 - else:
197 - assert q_quant_fused.data_ptr() == output_buffers[0].data_ptr()
198 - assert weights_fused.data_ptr() == output_buffers[-1].data_ptr()
199 -
174 if use_fp4:
175 q_quant_ref, q_scale_ref = q_quant_ref
176 q_quant_fused, q_scale_fused = q_quant_fused
vllm/models/deepseek_v4/attention.py
+2 -35
@@ -29,7 +29,6 @@ from vllm.models.deepseek_v4.common.ops import (
29 from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE
30
31 if TYPE_CHECKING:
32 - from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool
32 from vllm.v1.attention.backends.mla.sparse_swa import (
33 DeepseekSparseSWAMetadata,
34 )
@@ -185,7 +184,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
184 prefix: str,
185 topk_indices_buffer: torch.Tensor | None = None,
186 aux_stream_list: list[torch.cuda.Stream] | None = None,
188 - eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None,
187 ) -> None:
188 super().__init__()
189 config = vllm_config.model_config.hf_config
@@ -274,7 +272,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
272 )
273 self.indexer_rotary_emb = self.rotary_emb
274 self.topk_indices_buffer = topk_indices_buffer
277 - self.eager_scratch_pool = eager_scratch_pool
275
276 self.indexer = None
277 if self.compress_ratio == 4:
@@ -296,7 +293,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
293 compress_ratio=self.compress_ratio,
294 prefix=f"{prefix}.indexer",
295 aux_stream=indexer_aux_stream,
299 - eager_scratch_pool=eager_scratch_pool,
296 )
297
298 self._prepare_and_attn_fn = self._prepare_and_attn
@@ -355,7 +351,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
351 rotate=True,
352 prefix=f"{prefix}.compressor",
353 k_cache_prefix=self.prefix,
358 - eager_scratch_pool=eager_scratch_pool,
354 )
355
356 def forward(
@@ -637,24 +632,10 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
632 if cache_dtype == torch.uint8:
633 # fp8_ds_mla UE8M0 paged path. Horizontally fused:
634 # Q side: per-head RMSNorm (no weight) + GPT-J RoPE, zero-filling
640 - # the padding head slots.
635 + # the padding head slots; the kernel allocates and returns
636 + # the padded q tensor.
637 # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert.
638 swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
643 - if self.eager_scratch_pool is not None:
644 - q_out = self.eager_scratch_pool.q_out(q.shape[0])
645 - torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(
646 - q,
647 - kv,
648 - q_out,
649 - swa_kv_cache_2d,
650 - swa_metadata.slot_mapping,
651 - positions,
652 - cos_sin_cache,
653 - self.padded_heads,
654 - self.eps,
655 - swa_metadata.block_size,
656 - )
657 - return q_out
639 return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
640 q,
641 kv,
@@ -703,13 +684,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
684 )
685 return q_fp8
686
706 - def _global_topk_output_buffers(
707 - self, topk_indices: torch.Tensor
708 - ) -> tuple[torch.Tensor, torch.Tensor] | None:
709 - if self.compress_ratio != 4 or self.eager_scratch_pool is None:
710 - return None
711 - return self.eager_scratch_pool.global_topk_outputs(topk_indices)
712 -
687 def get_attn_backend(self) -> type[AttentionBackend]:
688 return self.backend_cls
689
@@ -792,7 +766,6 @@ class DeepseekV4Indexer(nn.Module):
766 compress_ratio: int = 1,
767 prefix: str = "",
768 aux_stream: torch.cuda.Stream | None = None,
795 - eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None,
769 ):
770 super().__init__()
771 self.vllm_config = vllm_config
@@ -805,7 +778,6 @@ class DeepseekV4Indexer(nn.Module):
778 self.rope_dim = config.qk_rope_head_dim # 64
779 self.q_lora_rank = q_lora_rank # 1536
780 self.compress_ratio = compress_ratio
808 - self.eager_scratch_pool = eager_scratch_pool
781 self.use_fp4_kv = dsa_indexer_uses_fp4(vllm_config)
782 logger.info_once(
783 "Using %s indexer cache for Lightning Indexer.",
@@ -869,7 +841,6 @@ class DeepseekV4Indexer(nn.Module):
841 prefix=f"{prefix}.compressor",
842 k_cache_prefix=self.k_cache.prefix,
843 use_fp4_cache=self.use_fp4_kv,
872 - eager_scratch_pool=eager_scratch_pool,
844 )
845
846 self.indexer_op = SparseAttnIndexer(
@@ -933,9 +904,6 @@ class DeepseekV4Indexer(nn.Module):
904 # ReplicatedLinear returns (output, bias); bias is None.
905 q, _ = self.wq_b(qr)
906 q = q.view(-1, self.n_head, self.head_dim)
936 - outputs = None
937 - if self.eager_scratch_pool is not None and self.use_fp4_kv:
938 - outputs = self.eager_scratch_pool.indexer_q_outputs(q.shape[0])
907 return fused_indexer_q_rope_quant(
908 positions,
909 q,
@@ -944,7 +912,6 @@ class DeepseekV4Indexer(nn.Module):
912 self.softmax_scale,
913 self.n_head**-0.5,
914 use_fp4=self.use_fp4_kv,
947 - output_buffers=outputs,
915 )
916
917 # compressor returns None and writes K to the indexer KV cache; the
vllm/models/deepseek_v4/common/ops/cache_utils.py
+2 -10
@@ -438,7 +438,6 @@ def compute_global_topk_indices_and_lens(
438 block_table: torch.Tensor,
439 block_size: int,
440 is_valid_token: torch.Tensor,
441 - output_buffers: tuple[torch.Tensor, torch.Tensor] | None = None,
441 ) -> tuple[torch.Tensor, torch.Tensor]:
442 """Map local topk indices to global KV cache slots and count valid entries.
443
@@ -448,15 +447,8 @@ def compute_global_topk_indices_and_lens(
447 3. Masking padding tokens to length 0
448 """
449 num_tokens = topk_indices.shape[0]
451 - if output_buffers is None:
452 - global_topk_indices = torch.empty_like(topk_indices)
453 - topk_lens = torch.empty(
454 - num_tokens, dtype=torch.int32, device=topk_indices.device
455 - )
456 - else:
457 - global_topk_indices, topk_lens = output_buffers
458 - assert global_topk_indices.shape == topk_indices.shape
459 - assert topk_lens.shape == (num_tokens,)
450 + global_topk_indices = torch.empty_like(topk_indices)
451 + topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)
452 _compute_global_topk_indices_and_lens_kernel[(num_tokens,)](
453 global_topk_indices,
454 global_topk_indices.stride(0),
vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
+12 -30
@@ -295,7 +295,6 @@ def fused_indexer_q_rope_quant(
295 index_weights_softmax_scale: float,
296 index_weights_head_scale: float,
297 use_fp4: bool = False,
298 - output_buffers: tuple[torch.Tensor, ...] | None = None,
298 ) -> tuple[
299 torch.Tensor | tuple[torch.Tensor, torch.Tensor],
300 torch.Tensor,
@@ -333,13 +332,7 @@ def fused_indexer_q_rope_quant(
332 num_index_q_heads = index_q.shape[1]
333 index_q_head_dim = index_q.shape[2]
334
336 - if output_buffers is None:
337 - index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
338 - else:
339 - expected_num_buffers = 3 if use_fp4 else 2
340 - assert len(output_buffers) == expected_num_buffers
341 - index_weights_out = output_buffers[-1]
342 - assert index_weights_out.shape == index_weights.shape
335 + index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
336
337 if use_fp4:
338 assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, (
@@ -347,23 +340,16 @@ def fused_indexer_q_rope_quant(
340 f"size {MXFP4_BLOCK_SIZE}"
341 )
342 num_scale_blocks = index_q_head_dim // MXFP4_BLOCK_SIZE
350 - packed_shape = (num_tokens, num_index_q_heads, index_q_head_dim // 2)
351 - scale_shape = (num_tokens, num_index_q_heads, num_scale_blocks)
352 - if output_buffers is None:
353 - index_q_packed = torch.empty(
354 - packed_shape,
355 - dtype=torch.uint8,
356 - device=index_q.device,
357 - )
358 - index_q_scale = torch.empty(
359 - scale_shape,
360 - dtype=torch.uint8,
361 - device=index_q.device,
362 - )
363 - else:
364 - index_q_packed, index_q_scale, _ = output_buffers
365 - assert index_q_packed.shape == packed_shape
366 - assert index_q_scale.shape == scale_shape
343 + index_q_packed = torch.empty(
344 + (num_tokens, num_index_q_heads, index_q_head_dim // 2),
345 + dtype=torch.uint8,
346 + device=index_q.device,
347 + )
348 + index_q_scale = torch.empty(
349 + (num_tokens, num_index_q_heads, num_scale_blocks),
350 + dtype=torch.uint8,
351 + device=index_q.device,
352 + )
353 if has_cutedsl():
354 # lazily import, otherwise some tests fail due to CUDA driver init failure.
355 from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
@@ -432,11 +418,7 @@ def fused_indexer_q_rope_quant(
418 fp8_dtype = current_platform.fp8_dtype()
419 use_fnuz = fp8_dtype == torch.float8_e4m3fnuz
420 fp8_max = 224.0 if use_fnuz else 448.0
435 - if output_buffers is None:
436 - index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype)
437 - else:
438 - index_q_fp8, _ = output_buffers
439 - assert index_q_fp8.shape == index_q.shape
421 + index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype)
422 if has_cutedsl():
423 # lazily import, otherwise some tests fail due to CUDA driver init failure.
424 from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
vllm/models/deepseek_v4/compressor.py
+1 -10
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import dataclass
5 -from typing import TYPE_CHECKING, Any, ClassVar, cast
5 +from typing import Any, ClassVar, cast
6
7 import torch
8 from torch import nn
@@ -35,9 +35,6 @@ from vllm.v1.kv_cache_interface import (
35 SlidingWindowMLASpec,
36 )
37
38 -if TYPE_CHECKING:
39 - from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool
40 -
38
39 def _prefer_two_stage_compressor() -> bool:
40 # Platforms that favor the triton variant of two-stage compressor split.
@@ -229,7 +226,6 @@ class DeepseekCompressor(nn.Module):
226 prefix: str = "",
227 k_cache_prefix="",
228 use_fp4_cache: bool = False,
232 - eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None,
229 ):
230 super().__init__()
231 self.compress_ratio = compress_ratio
@@ -239,7 +235,6 @@ class DeepseekCompressor(nn.Module):
235 self.prefix = prefix
236 self.k_cache_prefix = k_cache_prefix
237 self.use_fp4_cache = use_fp4_cache
242 - self.eager_scratch_pool = eager_scratch_pool
238
239 config = vllm_config.model_config.hf_config
240 self.rope_head_dim = config.qk_rope_head_dim
@@ -433,10 +428,6 @@ class DeepseekCompressor(nn.Module):
428 store_full_fp8=store_full_fp8,
429 fp8_scale=fp8_scale,
430 )
436 - if not self.overlap and self.eager_scratch_pool is not None:
437 - extra_kwargs["compress_scratch"] = (
438 - self.eager_scratch_pool.compressor_scratch(num_actual)
439 - )
431 elif self._use_two_stage_fused_compressor:
432 # head=512 cr>=128 (no overlap): two-pass split compressor on the
433 # prefill suffix, single-pass on the decode prefix.
vllm/models/deepseek_v4/eager_scratch.py deleted
-137
@@ -1,137 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from math import prod
5 -
6 -import torch
7 -
8 -from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE
9 -from vllm.utils.math_utils import round_up
10 -
11 -
12 -class DeepseekV4EagerScratchPool:
13 - """Model-wide outputs and scratch used inside the attention eager break."""
14 -
15 - _ALIGNMENT = 256
16 -
17 - def __init__(
18 - self,
19 - max_num_tokens: int,
20 - padded_q_heads: int,
21 - q_head_dim: int,
22 - index_q_heads: int,
23 - index_q_head_dim: int,
24 - index_topk: int,
25 - device: torch.device | str,
26 - ) -> None:
27 - self.max_num_tokens = max_num_tokens
28 - self.index_topk = index_topk
29 - self._q = torch.empty(
30 - (max_num_tokens, padded_q_heads, q_head_dim),
31 - dtype=torch.bfloat16,
32 - device=device,
33 - )
34 -
35 - fp4_specs = (
36 - ((max_num_tokens, index_q_heads, index_q_head_dim // 2), torch.uint8),
37 - (
38 - (
39 - max_num_tokens,
40 - index_q_heads,
41 - index_q_head_dim // MXFP4_BLOCK_SIZE,
42 - ),
43 - torch.uint8,
44 - ),
45 - ((max_num_tokens, index_q_heads), torch.float32),
46 - )
47 - global_specs = (
48 - ((max_num_tokens, index_topk), torch.int32),
49 - ((max_num_tokens,), torch.int32),
50 - )
51 - compressor_specs = (((max_num_tokens, q_head_dim), torch.float32),)
52 - # FP4 indexer is C4 only, global mapping after FP4 indexer
53 - # compressor scratch is C128 only
54 - # so here we use max instead of sum
55 - aux_bytes = max(
56 - self._packed_size(specs)
57 - for specs in (fp4_specs, global_specs, compressor_specs)
58 - )
59 - storage = torch.empty(aux_bytes, dtype=torch.uint8, device=device)
60 -
61 - self._q_outputs: dict[int, torch.Tensor] = {}
62 - fp4_values, fp4_scales, fp4_weights = self._views(storage, fp4_specs)
63 - self._fp4_template = (fp4_values, fp4_scales, fp4_weights)
64 - self._fp4_outputs: dict[
65 - int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]
66 - ] = {}
67 - global_indices, global_lens = self._views(storage, global_specs)
68 - self._global_template = (global_indices, global_lens)
69 - self._global_outputs: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
70 - self._compressor_template = self._views(storage, compressor_specs)[0]
71 - self._compressor_outputs: dict[int, torch.Tensor] = {}
72 - self._storage = storage
73 -
74 - @classmethod
75 - def _packed_size(
76 - cls, specs: tuple[tuple[tuple[int, ...], torch.dtype], ...]
77 - ) -> int:
78 - offset = 0
79 - for shape, dtype in specs:
80 - offset = round_up(offset, cls._ALIGNMENT) + prod(shape) * dtype.itemsize
81 - return round_up(offset, cls._ALIGNMENT)
82 -
83 - @classmethod
84 - def _views(
85 - cls,
86 - storage: torch.Tensor,
87 - specs: tuple[tuple[tuple[int, ...], torch.dtype], ...],
88 - ) -> list[torch.Tensor]:
89 - offset = 0
90 - views = []
91 - for shape, dtype in specs:
92 - offset = round_up(offset, cls._ALIGNMENT)
93 - num_bytes = prod(shape) * dtype.itemsize
94 - views.append(storage[offset : offset + num_bytes].view(dtype).view(shape))
95 - offset += num_bytes
96 - return views
97 -
98 - def q_out(self, num_tokens: int) -> torch.Tensor:
99 - output = self._q_outputs.get(num_tokens)
100 - if output is None:
101 - output = self._q[:num_tokens]
102 - self._q_outputs[num_tokens] = output
103 - return output
104 -
105 - def compressor_scratch(self, num_tokens: int) -> torch.Tensor:
106 - output = self._compressor_outputs.get(num_tokens)
107 - if output is None:
108 - output = self._compressor_template[:num_tokens]
109 - self._compressor_outputs[num_tokens] = output
110 - return output
111 -
112 - def indexer_q_outputs(
113 - self,
114 - num_tokens: int,
115 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
116 - output = self._fp4_outputs.get(num_tokens)
117 - if output is None:
118 - values, scales, weights = self._fp4_template
119 - output = (
120 - values[:num_tokens],
121 - scales[:num_tokens],
122 - weights[:num_tokens],
123 - )
124 - self._fp4_outputs[num_tokens] = output
125 - return output
126 -
127 - def global_topk_outputs(
128 - self, topk_indices: torch.Tensor
129 - ) -> tuple[torch.Tensor, torch.Tensor]:
130 - num_tokens, topk = topk_indices.shape
131 - assert topk == self.index_topk
132 - output = self._global_outputs.get(num_tokens)
133 - if output is None:
134 - indices, lens = self._global_template
135 - output = (indices[:num_tokens], lens[:num_tokens])
136 - self._global_outputs[num_tokens] = output
137 - return output
vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
-4
@@ -763,9 +763,6 @@ class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
763 attn_metadata.block_table[:num_decodes],
764 block_size,
765 is_valid,
766 - output_buffers=self._global_topk_output_buffers(
767 - self.topk_indices_buffer[:num_decode_tokens]
768 - ),
766 )
767 )
768 extra_sparse_indices = global_indices.view(num_decode_tokens, 1, -1)
@@ -855,7 +852,6 @@ class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
852 attn_metadata.block_table,
853 block_size,
854 swa_metadata.is_valid_token[prefill_token_slice],
858 - output_buffers=self._global_topk_output_buffers(local_topk_indices),
855 )
856 )
857
vllm/models/deepseek_v4/nvidia/flashmla.py
-3
@@ -197,9 +197,6 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
197 attn_metadata.block_table[:num_decodes],
198 block_size,
199 is_valid,
200 - output_buffers=self._global_topk_output_buffers(
201 - self.topk_indices_buffer[:num_decode_tokens]
202 - ),
200 )
201 topk_indices = global_indices.view(num_decode_tokens, 1, -1)
202 else:
vllm/models/deepseek_v4/nvidia/model.py
-20
@@ -72,7 +72,6 @@ from vllm.models.common.ops.sequence_parallel import (
72 sp_shard,
73 )
74 from vllm.models.deepseek_v4.attention import DeepseekV4Attention
75 -from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool
75 from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import (
76 DeepseekV4FlashInferMLAAttention,
77 DeepseekV4FlashInferSM120Attention,
@@ -815,7 +814,6 @@ class DeepseekV4DecoderLayer(nn.Module):
814 prefix,
815 topk_indices_buffer: torch.Tensor | None = None,
816 aux_stream_list: list[torch.cuda.Stream] | None = None,
818 - eager_scratch_pool: DeepseekV4EagerScratchPool | None = None,
817 ):
818 super().__init__()
819
@@ -829,7 +827,6 @@ class DeepseekV4DecoderLayer(nn.Module):
827 prefix=f"{prefix}.attn",
828 topk_indices_buffer=topk_indices_buffer,
829 aux_stream_list=aux_stream_list,
832 - eager_scratch_pool=eager_scratch_pool,
830 )
831 if self.use_sequence_parallel:
832 self.attn.wo_b.reduce_results = False
@@ -1017,22 +1014,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1014 # (compressor kv_score, indexer.weights_proj, indexer.compressor
1015 # kv_score). fused_wqa_wkv stays on the default stream.
1016 aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
1020 - padded_heads = _select_dsv4_attn_cls(vllm_config).get_padded_num_q_heads(
1021 - config.num_attention_heads // get_tensor_model_parallel_world_size()
1022 - )
1023 - self.eager_scratch_pool: DeepseekV4EagerScratchPool | None = None
1024 - if not vllm_config.parallel_config.use_ubatching:
1025 - # TODO: support dbo if needed
1026 - # this requires the buffer to have ubatch dim
1027 - self.eager_scratch_pool = DeepseekV4EagerScratchPool(
1028 - vllm_config.scheduler_config.max_num_batched_tokens,
1029 - padded_heads,
1030 - config.head_dim,
1031 - config.index_n_heads,
1032 - config.index_head_dim,
1033 - config.index_topk,
1034 - current_platform.device_type,
1035 - )
1017
1018 # Reserved topk indices buffer for all Indexer layers to reuse.
1019 self.topk_indices_buffer = torch.empty(
@@ -1058,7 +1039,6 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1039 prefix=prefix,
1040 topk_indices_buffer=self.topk_indices_buffer,
1041 aux_stream_list=aux_stream_list,
1061 - eager_scratch_pool=self.eager_scratch_pool,
1042 ),
1043 prefix=f"{prefix}.layers",
1044 )
vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py
+5 -10
@@ -2097,7 +2097,6 @@ def compress_norm_rope_store_cutedsl(
2097 store_full_kv: bool = False,
2098 store_full_fp8: bool = False,
2099 fp8_scale: torch.Tensor | None = None,
2100 - compress_scratch: torch.Tensor | None = None,
2100 ) -> None:
2101 if compress_ratio == 4:
2102 # For C4A, the single fused kernel is faster than the two-kernel version.
@@ -2130,15 +2129,11 @@ def compress_norm_rope_store_cutedsl(
2129 )
2130 else:
2131 # For C128, the two-kernel version is faster than the single fused kernel.
2133 - if compress_scratch is None:
2134 - compressed_kv = torch.empty(
2135 - (num_actual, head_dim),
2136 - dtype=torch.float32,
2137 - device=state_cache.device,
2138 - )
2139 - else:
2140 - assert compress_scratch.shape == (num_actual, head_dim)
2141 - compressed_kv = compress_scratch
2132 + compressed_kv = torch.empty(
2133 + (num_actual, head_dim),
2134 + dtype=torch.float32,
2135 + device=state_cache.device,
2136 + )
2137 split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
2138 state_cache,
2139 token_to_req_indices,