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

[DSv4 Perf] Optimize workspace reuse for eager break, 3.9% E2E TTFT improvement. (#49236)

Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Wentao Ye committed Jul 31, 2026 at 11:45 UTC df71917cf17c97d81f48122e6baa7c13d184ff90

15 files changed +354 -30

csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

+24 -6

@@ -942,9 +942,10 @@ static void launchFullCacheKernel(

942 // ────────────────────────────────────────────────────────────────────────────

943 // Torch op wrapper

944 // ────────────────────────────────────────────────────────────────────────────

945 -torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(

945 +void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(

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]

949 torch::stable::Tensor& k_cache, // [num_blocks, block_bytes] uint8

950 torch::stable::Tensor const& slot_mapping, // [N] int64

951 torch::stable::Tensor const& position_ids, // [N] int64

@@ -970,8 +971,16 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(

971 STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");

972 STD_TORCH_CHECK(q_in.scalar_type() == kv.scalar_type(),

973 "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");

978 STD_TORCH_CHECK(q_head_padded >= q_in.size(1),

979 "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]");

984 STD_TORCH_CHECK(k_cache.scalar_type() == torch::headeronly::ScalarType::Byte,

985 "k_cache must be uint8");

986 STD_TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,

@@ -999,11 +1008,6 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(

1008 q_in.get_device_index());

1009 const cudaStream_t stream = get_current_cuda_stream(q_in.get_device_index());

1010

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 -

1011 VLLM_STABLE_DISPATCH_HALF_TYPES(

1012 q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {

1013 using qkv_scalar_t = scalar_t;

@@ -1020,6 +1024,20 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(

1024 num_heads_q_padded, cache_block_size_i, kv_block_stride,

1025 stream);

1026 });

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);

1041 return q_out;

1042 }

1043

csrc/libtorch_stable/ops.h

@@ -269,6 +269,14 @@ 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 +

280 void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(

281 torch::stable::Tensor& q, torch::stable::Tensor const& kv,

282 torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,

csrc/libtorch_stable/torch_bindings.cpp

@@ -433,6 +433,11 @@ 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) -> ()");

441

442 // FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate

443 // FP8 tensor, and KV into a contiguous 512-wide token-strided cache.

@@ -751,6 +756,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {

756 ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));

757 ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert",

758 TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert));

759 + ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out",

760 + TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out));

761 ops.impl(

762 "fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert",

763 TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert));

tests/kernels/test_compressor_kv_cache.py

+18

@@ -20,6 +20,7 @@ import torch

21 from vllm import _custom_ops as ops

22 from vllm.models.deepseek_v4.common.ops import (

23 + compute_global_topk_indices_and_lens,

24 dequantize_and_gather_k_cache,

25 quantize_and_insert_k_cache,

26 )

@@ -34,6 +35,23 @@ from vllm.platforms import current_platform

35 from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4

38 +def test_compute_global_topk_reuses_output_buffers():

39 + device = "cuda"

40 + topk_indices = torch.tensor(

41 + [[0, 3, -1], [1, 2, -1]], dtype=torch.int32, device=device

42 + )

43 + token_to_req = torch.tensor([0, 1], dtype=torch.int32, device=device)

44 + block_table = torch.tensor([[5, 7], [11, 13]], dtype=torch.int32, device=device)

45 + is_valid = torch.tensor([True, False], device=device)

46 + args = (topk_indices, token_to_req, block_table, 2, is_valid)

47 + expected = compute_global_topk_indices_and_lens(*args)

48 + outputs = tuple(torch.empty_like(tensor) for tensor in expected)

49 + actual = compute_global_topk_indices_and_lens(*args, output_buffers=outputs)

50 + for result, output, reference in zip(actual, outputs, expected):

51 + assert result.data_ptr() == output.data_ptr()

52 + torch.testing.assert_close(result, reference)

53 +

54 +

55 def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):

56 """PyTorch reference for UE8M0 FP8 quantization (per-block, power-of-2 scale).

tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py

+12 -2

@@ -257,8 +257,18 @@ 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 = _call_fused(

261 - q, padded_heads, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs

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,

272 )

273

274 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,6 +150,23 @@ 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 + )

170 # use_cutedsl=False: force the triton path even when cutedsl is installed

171 # by patching the dispatcher's has_cutedsl() binding to return False.

172 cutedsl_patch = (

@@ -169,8 +186,17 @@ def test_fused_indexer_q_rope_quant_matches_unfused(

186 softmax_scale,

187 head_scale,

188 use_fp4,

189 + output_buffers=output_buffers,

190 )

191

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 +

200 if use_fp4:

201 q_quant_ref, q_scale_ref = q_quant_ref

202 q_quant_fused, q_scale_fused = q_quant_fused

vllm/models/deepseek_v4/attention.py

+35 -2

@@ -29,6 +29,7 @@ from vllm.models.deepseek_v4.common.ops import (

29 from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE

31 if TYPE_CHECKING:

32 + from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool

33 from vllm.v1.attention.backends.mla.sparse_swa import (

34 DeepseekSparseSWAMetadata,

35 )

@@ -181,6 +182,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

182 prefix: str,

183 topk_indices_buffer: torch.Tensor | None = None,

184 aux_stream_list: list[torch.cuda.Stream] | None = None,

185 + eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None,

186 ) -> None:

187 super().__init__()

188 config = vllm_config.model_config.hf_config

@@ -269,6 +271,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

271 )

272 self.indexer_rotary_emb = self.rotary_emb

273 self.topk_indices_buffer = topk_indices_buffer

274 + self.eager_scratch_pool = eager_scratch_pool

275

276 self.indexer = None

277 if self.compress_ratio == 4:

@@ -290,6 +293,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

293 compress_ratio=self.compress_ratio,

294 prefix=f"{prefix}.indexer",

295 aux_stream=indexer_aux_stream,

296 + eager_scratch_pool=eager_scratch_pool,

297 )

298

299 # Will be None on ROCm for now.

@@ -340,6 +344,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

344 rotate=True,

345 prefix=f"{prefix}.compressor",

346 k_cache_prefix=self.prefix,

347 + eager_scratch_pool=eager_scratch_pool,

348 )

349

350 def forward(

@@ -571,10 +576,24 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

576 if cache_dtype == torch.uint8:

577 # fp8_ds_mla UE8M0 paged path. Horizontally fused:

578 # Q side: per-head RMSNorm (no weight) + GPT-J RoPE, zero-filling

574 - # the padding head slots; the kernel allocates and returns

575 - # the padded q tensor.

579 + # the padding head slots.

580 # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert.

581 swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)

582 + if self.eager_scratch_pool is not None:

583 + q_out = self.eager_scratch_pool.q_out(q.shape[0])

584 + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert_out(

585 + q,

586 + kv,

587 + q_out,

588 + swa_kv_cache_2d,

589 + swa_metadata.slot_mapping,

590 + positions,

591 + cos_sin_cache,

592 + self.padded_heads,

593 + self.eps,

594 + swa_metadata.block_size,

595 + )

596 + return q_out

597 return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(

598 q,

599 kv,

@@ -623,6 +642,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

642 )

643 return q_fp8

644

645 + def _global_topk_output_buffers(

646 + self, topk_indices: torch.Tensor

647 + ) -> tuple[torch.Tensor, torch.Tensor] | None:

648 + if self.compress_ratio != 4 or self.eager_scratch_pool is None:

649 + return None

650 + return self.eager_scratch_pool.global_topk_outputs(topk_indices)

651 +

652 def get_attn_backend(self) -> type[AttentionBackend]:

653 return self.backend_cls

654

@@ -702,6 +728,7 @@ class DeepseekV4Indexer(nn.Module):

728 compress_ratio: int = 1,

729 prefix: str = "",

730 aux_stream: torch.cuda.Stream | None = None,

731 + eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None,

732 ):

733 super().__init__()

734 self.vllm_config = vllm_config

@@ -714,6 +741,7 @@ class DeepseekV4Indexer(nn.Module):

741 self.rope_dim = config.qk_rope_head_dim # 64

742 self.q_lora_rank = q_lora_rank # 1536

743 self.compress_ratio = compress_ratio

744 + self.eager_scratch_pool = eager_scratch_pool

745 self.use_fp4_kv = self.vllm_config.attention_config.use_fp4_indexer_cache

746 logger.info_once(

747 "Using %s indexer cache for Lightning Indexer.",

@@ -777,6 +805,7 @@ class DeepseekV4Indexer(nn.Module):

805 prefix=f"{prefix}.compressor",

806 k_cache_prefix=self.k_cache.prefix,

807 use_fp4_cache=self.use_fp4_kv,

808 + eager_scratch_pool=eager_scratch_pool,

809 )

810

811 self.indexer_op = SparseAttnIndexer(

@@ -837,6 +866,9 @@ class DeepseekV4Indexer(nn.Module):

866 # ReplicatedLinear returns (output, bias); bias is None.

867 q, _ = self.wq_b(qr)

868 q = q.view(-1, self.n_head, self.head_dim)

869 + outputs = None

870 + if self.eager_scratch_pool is not None and self.use_fp4_kv:

871 + outputs = self.eager_scratch_pool.indexer_q_outputs(q.shape[0])

872 return fused_indexer_q_rope_quant(

873 positions,

874 q,

@@ -845,6 +877,7 @@ class DeepseekV4Indexer(nn.Module):

877 self.softmax_scale,

878 self.n_head**-0.5,

879 use_fp4=self.use_fp4_kv,

880 + output_buffers=outputs,

881 )

882

883 # compressor returns None and writes K to the indexer KV cache; the

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

+10 -2

@@ -438,6 +438,7 @@ 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,

442 ) -> tuple[torch.Tensor, torch.Tensor]:

443 """Map local topk indices to global KV cache slots and count valid entries.

444

@@ -447,8 +448,15 @@ def compute_global_topk_indices_and_lens(

448 3. Masking padding tokens to length 0

449 """

450 num_tokens = topk_indices.shape[0]

450 - global_topk_indices = torch.empty_like(topk_indices)

451 - topk_lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)

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,)

460 _compute_global_topk_indices_and_lens_kernel[(num_tokens,)](

461 global_topk_indices,

462 global_topk_indices.stride(0),

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

+30 -12

@@ -295,6 +295,7 @@ 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,

299 ) -> tuple[

300 torch.Tensor | tuple[torch.Tensor, torch.Tensor],

301 torch.Tensor,

@@ -332,7 +333,13 @@ def fused_indexer_q_rope_quant(

333 num_index_q_heads = index_q.shape[1]

334 index_q_head_dim = index_q.shape[2]

335

335 - index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)

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

343

344 if use_fp4:

345 assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, (

@@ -340,16 +347,23 @@ def fused_indexer_q_rope_quant(

347 f"size {MXFP4_BLOCK_SIZE}"

348 )

349 num_scale_blocks = index_q_head_dim // MXFP4_BLOCK_SIZE

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 - )

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

367 if has_cutedsl():

368 # lazily import, otherwise some tests fail due to CUDA driver init failure.

369 from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (

@@ -418,7 +432,11 @@ def fused_indexer_q_rope_quant(

432 fp8_dtype = current_platform.fp8_dtype()

433 use_fnuz = fp8_dtype == torch.float8_e4m3fnuz

434 fp8_max = 224.0 if use_fnuz else 448.0

421 - index_q_fp8 = torch.empty_like(index_q, dtype=fp8_dtype)

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

440 if has_cutedsl():

441 # lazily import, otherwise some tests fail due to CUDA driver init failure.

442 from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (

vllm/models/deepseek_v4/compressor.py

+10 -1

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from dataclasses import dataclass

5 -from typing import Any, ClassVar, cast

5 +from typing import TYPE_CHECKING, Any, ClassVar, cast

7 import torch

8 from torch import nn

@@ -35,6 +35,9 @@ from vllm.v1.kv_cache_interface import (

35 SlidingWindowMLASpec,

36 )

38 +if TYPE_CHECKING:

39 + from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool

40 +

42 def _prefer_two_stage_compressor() -> bool:

43 # Platforms that favor the triton variant of two-stage compressor split.

@@ -226,6 +229,7 @@ class DeepseekCompressor(nn.Module):

229 prefix: str = "",

230 k_cache_prefix="",

231 use_fp4_cache: bool = False,

232 + eager_scratch_pool: "DeepseekV4EagerScratchPool | None" = None,

233 ):

234 super().__init__()

235 self.compress_ratio = compress_ratio

@@ -235,6 +239,7 @@ class DeepseekCompressor(nn.Module):

239 self.prefix = prefix

240 self.k_cache_prefix = k_cache_prefix

241 self.use_fp4_cache = use_fp4_cache

242 + self.eager_scratch_pool = eager_scratch_pool

243

244 config = vllm_config.model_config.hf_config

245 self.rope_head_dim = config.qk_rope_head_dim

@@ -428,6 +433,10 @@ class DeepseekCompressor(nn.Module):

433 store_full_fp8=store_full_fp8,

434 fp8_scale=fp8_scale,

435 )

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 + )

440 elif self._use_two_stage_fused_compressor:

441 # head=512 cr>=128 (no overlap): two-pass split compressor on the

442 # prefill suffix, single-pass on the decode prefix.

vllm/models/deepseek_v4/eager_scratch.py new

+137

@@ -0,0 +1,137 @@

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

@@ -745,6 +745,9 @@ class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):

745 attn_metadata.block_table[:num_decodes],

746 block_size,

747 is_valid,

748 + output_buffers=self._global_topk_output_buffers(

749 + self.topk_indices_buffer[:num_decode_tokens]

750 + ),

751 )

752 )

753 extra_sparse_indices = global_indices.view(num_decode_tokens, 1, -1)

@@ -834,6 +837,7 @@ class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):

837 attn_metadata.block_table,

838 block_size,

839 swa_metadata.is_valid_token[prefill_token_slice],

840 + output_buffers=self._global_topk_output_buffers(local_topk_indices),

841 )

842 )

843

vllm/models/deepseek_v4/nvidia/flashmla.py

@@ -176,6 +176,9 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):

176 attn_metadata.block_table[:num_decodes],

177 block_size,

178 is_valid,

179 + output_buffers=self._global_topk_output_buffers(

180 + self.topk_indices_buffer[:num_decode_tokens]

181 + ),

182 )

183 topk_indices = global_indices.view(num_decode_tokens, 1, -1)

184 else:

vllm/models/deepseek_v4/nvidia/model.py

+20

@@ -66,6 +66,7 @@ from vllm.model_executor.models.utils import (

66 )

67 from vllm.model_executor.utils import set_weight_attrs

68 from vllm.models.deepseek_v4.attention import DeepseekV4Attention

69 +from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool

70 from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import (

71 DeepseekV4FlashInferMLAAttention,

72 DeepseekV4FlashInferSM120Attention,

@@ -798,6 +799,7 @@ class DeepseekV4DecoderLayer(nn.Module):

799 prefix,

800 topk_indices_buffer: torch.Tensor | None = None,

801 aux_stream_list: list[torch.cuda.Stream] | None = None,

802 + eager_scratch_pool: DeepseekV4EagerScratchPool | None = None,

803 ):

804 super().__init__()

805

@@ -810,6 +812,7 @@ class DeepseekV4DecoderLayer(nn.Module):

812 prefix=f"{prefix}.attn",

813 topk_indices_buffer=topk_indices_buffer,

814 aux_stream_list=aux_stream_list,

815 + eager_scratch_pool=eager_scratch_pool,

816 )

817 self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")

818

@@ -986,6 +989,22 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

989 # (compressor kv_score, indexer.weights_proj, indexer.compressor

990 # kv_score). fused_wqa_wkv stays on the default stream.

991 aux_stream_list = [torch.cuda.Stream() for _ in range(3)]

992 + padded_heads = _select_dsv4_attn_cls(vllm_config).get_padded_num_q_heads(

993 + config.num_attention_heads // get_tensor_model_parallel_world_size()

994 + )

995 + self.eager_scratch_pool: DeepseekV4EagerScratchPool | None = None

996 + if not vllm_config.parallel_config.use_ubatching:

997 + # TODO: support dbo if needed

998 + # this requires the buffer to have ubatch dim

999 + self.eager_scratch_pool = DeepseekV4EagerScratchPool(

1000 + vllm_config.scheduler_config.max_num_batched_tokens,

1001 + padded_heads,

1002 + config.head_dim,

1003 + config.index_n_heads,

1004 + config.index_head_dim,

1005 + config.index_topk,

1006 + current_platform.device_type,

1007 + )

1008

1009 # Reserved topk indices buffer for all Indexer layers to reuse.

1010 self.topk_indices_buffer = torch.empty(

@@ -1011,6 +1030,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

1030 prefix=prefix,

1031 topk_indices_buffer=self.topk_indices_buffer,

1032 aux_stream_list=aux_stream_list,

1033 + eager_scratch_pool=self.eager_scratch_pool,

1034 ),

1035 prefix=f"{prefix}.layers",

1036 )

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

+10 -5

@@ -2097,6 +2097,7 @@ 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,

2101 ) -> None:

2102 if compress_ratio == 4:

2103 # For C4A, the single fused kernel is faster than the two-kernel version.

@@ -2129,11 +2130,15 @@ def compress_norm_rope_store_cutedsl(

2130 )

2131 else:

2132 # For C128, the two-kernel version is faster than the single fused kernel.

2132 - compressed_kv = torch.empty(

2133 - (num_actual, head_dim),

2134 - dtype=torch.float32,

2135 - device=state_cache.device,

2136 - )

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

2142 split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(

2143 state_cache,

2144 token_to_req_indices,