[DSv4 Perf] Skip empty c128 kernel launch, around 2x kernel performance improvement. (#48957)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 22, 2026 at 10:55 UTC
37e370fe936fbee062b7a4bd502375794d859b5f
2 files changed
+53
-2
tests/kernels/test_compressor_kv_cache.py
+21
@@ -13,6 +13,7 @@ These tests cover:
13
"""
14
15
import math
16
+from types import SimpleNamespace
17
18
import pytest
19
import torch
@@ -27,6 +28,7 @@ from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
28
_fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
29
_launch_two_stage_sparse_attn_compressor,
30
)
31
+from vllm.models.deepseek_v4.compressor import _get_c128_boundary
32
from vllm.platforms import current_platform
33
34
from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
@@ -58,6 +60,25 @@ def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
60
return x_fp8, scales
61
62
63
+@pytest.mark.parametrize(
64
+ ("starts", "query_start_loc", "expected"),
65
+ [
66
+ ([0], [0, 127], False),
67
+ ([0], [0, 128], True),
68
+ ([127], [0, 1], True),
69
+ ([128], [0, 127], False),
70
+ ([1, 255], [0, 1, 2], True),
71
+ (None, [0, 1], None),
72
+ ],
73
+)
74
+def test_get_c128_boundary(starts, query_start_loc, expected):
75
+ metadata = SimpleNamespace(
76
+ _num_computed_tokens_cpu=None if starts is None else torch.tensor(starts),
77
+ query_start_loc_cpu=torch.tensor(query_start_loc),
78
+ )
79
+ assert _get_c128_boundary(metadata) is expected
80
+
81
+
82
# ── Test A: DeepseekV4 Attention path ──────────────────────────────────────────────
83
84
vllm/models/deepseek_v4/compressor.py
+32
-2
@@ -7,7 +7,7 @@ from typing import Any, ClassVar, cast
7
import torch
8
from torch import nn
9
10
-from vllm.config import VllmConfig, get_current_vllm_config
10
+from vllm.config import CUDAGraphMode, VllmConfig, get_current_vllm_config
11
from vllm.forward_context import get_forward_context
12
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
13
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -42,6 +42,19 @@ def _prefer_two_stage_compressor() -> bool:
42
return current_platform.is_rocm()
43
44
45
+def _get_c128_boundary(metadata: CommonAttentionMetadata) -> bool | None:
46
+ starts = metadata._num_computed_tokens_cpu
47
+ if starts is None:
48
+ return None
49
+
50
+ starts_list = starts.tolist()
51
+ query_start_loc = metadata.query_start_loc_cpu.tolist()
52
+ return any(
53
+ start % 128 + query_start_loc[i + 1] - query_start_loc[i] >= 128
54
+ for i, start in enumerate(starts_list)
55
+ )
56
+
57
+
58
class CompressorBackend(AttentionBackend):
59
def __init__(self):
60
super().__init__()
@@ -90,6 +103,7 @@ class CompressorMetadata:
103
104
token_to_req_indices: torch.Tensor | None = None # [num_tokens]
105
num_decode_tokens: int | None = None
106
+ c128_boundary: bool | None = None
107
108
109
class CompressorMetadataBuilder(AttentionMetadataBuilder):
@@ -127,6 +141,11 @@ class CompressorMetadataBuilder(AttentionMetadataBuilder):
141
block_size=self.block_size,
142
token_to_req_indices=token_to_req_indices,
143
num_decode_tokens=num_decode_tokens,
144
+ c128_boundary=(
145
+ _get_c128_boundary(common_attn_metadata)
146
+ if self.block_size == 8
147
+ else None
148
+ ),
149
)
150
151
@@ -317,7 +336,8 @@ class DeepseekCompressor(nn.Module):
336
)
337
338
# Get the metadata and handle dummy profiling run.
320
- attn_metadata = get_forward_context().attn_metadata
339
+ forward_context = get_forward_context()
340
+ attn_metadata = forward_context.attn_metadata
341
if not isinstance(attn_metadata, dict):
342
return
343
@@ -359,6 +379,16 @@ class DeepseekCompressor(nn.Module):
379
pdl_kwargs=pdl_kwargs,
380
)
381
382
+ # full graph cannot branch on per-step CPU metadata after capture
383
+ if (
384
+ current_platform.is_cuda()
385
+ and self.head_dim == 512
386
+ and self.compress_ratio == 128
387
+ and forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL
388
+ and state_metadata.c128_boundary is False
389
+ ):
390
+ return
391
+
392
# Fused: compress → RMSNorm → RoPE → FP8 quant → KV cache write.
393
# RoPE requirements (kernel applies forward GPT-J style rotation):
394
# - is_neox_style=False (interleaved pairs, NOT split-half)