A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 18:15 UTC
[DSv4] Refactor compressor & Fix ROCm compatibility (#43710)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
May 26, 2026 at 19:56 UTC
adaa5e455ad89525e8afa69b0633ec9c4b720f7e
9 files changed
+364
-239
vllm/models/deepseek_v4/common/ops/__init__.py
+2
@@ -10,6 +10,7 @@ from .cache_utils import (
10
from .fused_indexer_q import MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant
11
from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant
12
from .fused_qk_rmsnorm import fused_q_kv_rmsnorm
13
+from .save_partial_states import save_partial_states
14
15
__all__ = [
16
"MXFP4_BLOCK_SIZE",
@@ -20,4 +21,5 @@ __all__ = [
21
"fused_inv_rope_fp8_quant",
22
"fused_q_kv_rmsnorm",
23
"quantize_and_insert_k_cache",
24
+ "save_partial_states",
25
]
vllm/models/deepseek_v4/common/ops/cache_utils.py
+1
-1
@@ -366,7 +366,7 @@ def dequantize_and_gather_k_cache(
366
) -> None:
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.dequant_gather_k_cutedsl import (
369
+ from vllm.models.deepseek_v4.nvidia.ops import (
370
dequantize_and_gather_k_cache_cutedsl,
371
)
372
vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py
+78
-34
@@ -11,12 +11,6 @@ Three specialized kernels:
11
- _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn:
12
head=128, MXFP4 (block=32), 4 ue8m0 bytes
13
14
-Additional cutedsl kernels:
15
- - _compress_kv_sparse_attn_cutedsl / _norm_rope_insert_sparse_attn_cutedsl:
16
- CuTe DSL split kernels for C128
17
- - _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl:
18
- CuTe DSL fused kernels for C4
19
-
14
RoPE is register-based via tl.reshape -> tl.split -> tl.interleave (or the
15
even/odd halves are consumed directly for MXFP4, no interleave needed).
16
FP8 UE8M0 quant uses tl.reshape to tile [N_QUANT_BLOCKS, QUANT_BLOCK] for
@@ -25,43 +19,93 @@ even/odd halves, producing (N_QUANT_BLOCKS, MXFP4_BLOCK/2) packed nibbles
19
and N_QUANT_BLOCKS ue8m0 bytes.
20
"""
21
28
-from functools import cache
22
+from typing import Any
23
+
24
+import torch
25
26
from vllm.triton_utils import tl, triton
27
28
from .fused_indexer_q import _fp32x2_to_fp4x2
29
30
35
-@cache
36
-def _get_sparse_attn_cutedsl_impls():
37
- from .sparse_attn_compress_cutedsl import (
38
- _compress_kv_sparse_attn_cutedsl,
39
- _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl,
40
- _norm_rope_insert_sparse_attn_cutedsl,
41
- )
42
-
43
- return (
44
- _compress_kv_sparse_attn_cutedsl,
45
- _norm_rope_insert_sparse_attn_cutedsl,
46
- _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl,
31
+def compress_norm_rope_store_triton(
32
+ state_cache: torch.Tensor,
33
+ num_actual: int,
34
+ token_to_req_indices: torch.Tensor,
35
+ positions: torch.Tensor,
36
+ slot_mapping: torch.Tensor,
37
+ block_table: torch.Tensor,
38
+ block_size: int,
39
+ state_width: int,
40
+ cos_sin_cache: torch.Tensor,
41
+ kv_cache: torch.Tensor,
42
+ k_cache_metadata: Any,
43
+ pdl_kwargs: dict,
44
+ head_dim: int,
45
+ rope_head_dim: int,
46
+ compress_ratio: int,
47
+ overlap: bool,
48
+ use_fp4_cache: bool,
49
+ rms_norm_weight: torch.Tensor,
50
+ rms_norm_eps: float,
51
+ quant_block: int,
52
+ token_stride: int,
53
+ scale_dim: int,
54
+) -> None:
55
+ """Shared triton launcher for the fused compress+norm+RoPE+insert path.
56
+
57
+ Picks one of the three kernels in this module based on ``head_dim`` and
58
+ ``use_fp4_cache``. Identical launch signature for all three.
59
+ """
60
+ if head_dim == 512:
61
+ kernel = _fused_kv_compress_norm_rope_insert_sparse_attn
62
+ num_warps = 4
63
+ elif use_fp4_cache:
64
+ kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
65
+ num_warps = 1
66
+ else:
67
+ kernel = _fused_kv_compress_norm_rope_insert_indexer_attn
68
+ num_warps = 1
69
+
70
+ kernel[(num_actual,)](
71
+ # state cache
72
+ state_cache,
73
+ state_cache.stride(0),
74
+ state_cache.stride(1),
75
+ # metadata
76
+ token_to_req_indices,
77
+ positions,
78
+ slot_mapping,
79
+ block_table,
80
+ block_table.stride(0),
81
+ block_size,
82
+ # RMSNorm
83
+ rms_norm_weight,
84
+ rms_norm_eps,
85
+ # RoPE
86
+ cos_sin_cache,
87
+ cos_sin_cache.stride(0),
88
+ # KV cache
89
+ kv_cache,
90
+ k_cache_metadata.slot_mapping,
91
+ kv_cache.shape[1], # paged KV cache block size (tokens per block)
92
+ # constexprs
93
+ HEAD_SIZE=head_dim,
94
+ TRITON_BLOCK_SIZE=triton.next_power_of_2(head_dim),
95
+ STATE_WIDTH=state_width,
96
+ COMPRESS_RATIO=compress_ratio,
97
+ OVERLAP=overlap,
98
+ ROPE_HEAD_DIM=rope_head_dim,
99
+ FP8_MAX=448.0,
100
+ QUANT_BLOCK=quant_block,
101
+ TOKEN_STRIDE=token_stride,
102
+ SCALE_DIM=scale_dim,
103
+ KV_BLOCK_STRIDE=kv_cache.stride(0),
104
+ num_warps=num_warps,
105
+ **pdl_kwargs,
106
)
107
108
50
-def _compress_kv_sparse_attn_cutedsl(*args, **kwargs):
51
- """CuTe DSL sparse-attention compress wrapper."""
52
- return _get_sparse_attn_cutedsl_impls()[0](*args, **kwargs)
53
-
54
-
55
-def _norm_rope_insert_sparse_attn_cutedsl(*args, **kwargs):
56
- """CuTe DSL RMSNorm/RoPE/FP8-store wrapper."""
57
- return _get_sparse_attn_cutedsl_impls()[1](*args, **kwargs)
58
-
59
-
60
-def _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(*args, **kwargs):
61
- """CuTe DSL fused C4 sparse-attention compressor wrapper."""
62
- return _get_sparse_attn_cutedsl_impls()[2](*args, **kwargs)
63
-
64
-
109
# =============================================================================
110
# DeepseekV4 Attention path (head=512, nope=448 FP8 + rope=64 bf16)
111
# =============================================================================
vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
+2
-2
@@ -346,7 +346,7 @@ def fused_indexer_q_rope_quant(
346
)
347
if has_cutedsl():
348
# lazily import, otherwise some tests fail due to CUDA driver init failure.
349
- from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
349
+ from vllm.models.deepseek_v4.nvidia.ops import (
350
fused_indexer_q_rope_quant_mxfp4_cutedsl,
351
)
352
@@ -400,7 +400,7 @@ def fused_indexer_q_rope_quant(
400
index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
401
if has_cutedsl():
402
# lazily import, otherwise some tests fail due to CUDA driver init failure.
403
- from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import (
403
+ from vllm.models.deepseek_v4.nvidia.ops import (
404
fused_indexer_q_rope_quant_fp8_cutedsl,
405
)
406
vllm/models/deepseek_v4/common/ops/save_partial_states.py
new
+101
@@ -0,0 +1,101 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm.triton_utils import tl, triton
7
+
8
+
9
+def save_partial_states(
10
+ kv: torch.Tensor,
11
+ score: torch.Tensor,
12
+ ape: torch.Tensor,
13
+ positions: torch.Tensor,
14
+ state_cache: torch.Tensor,
15
+ slot_mapping: torch.Tensor,
16
+ block_size: int,
17
+ state_width: int,
18
+ compress_ratio: int,
19
+ pdl_kwargs: dict | None = None,
20
+) -> None:
21
+ """Write packed [kv, score+ape] partial states into the compressor cache.
22
+
23
+ One program per token; pads (slot_id == -1) are skipped.
24
+ """
25
+ num_actual = slot_mapping.shape[0]
26
+ head_size = kv.shape[-1]
27
+ _save_partial_states_kernel[(num_actual,)](
28
+ kv,
29
+ kv.stride(0),
30
+ score,
31
+ score.stride(0),
32
+ ape,
33
+ ape.stride(0),
34
+ positions,
35
+ state_cache,
36
+ state_cache.stride(0),
37
+ state_cache.stride(1),
38
+ slot_mapping,
39
+ block_size,
40
+ HEAD_SIZE=head_size,
41
+ TRITON_BLOCK_SIZE=triton.next_power_of_2(head_size),
42
+ STATE_WIDTH=state_width,
43
+ COMPRESS_RATIO=compress_ratio,
44
+ **(pdl_kwargs or {}),
45
+ )
46
+
47
+
48
+@triton.jit
49
+def _save_partial_states_kernel(
50
+ kv_ptr,
51
+ kv_stride,
52
+ score_ptr,
53
+ score_stride,
54
+ ape_ptr,
55
+ ape_stride,
56
+ positions_ptr,
57
+ state_cache_ptr,
58
+ state_cache_stride0,
59
+ state_cache_stride1,
60
+ slot_mapping_ptr,
61
+ block_size,
62
+ HEAD_SIZE: tl.constexpr,
63
+ TRITON_BLOCK_SIZE: tl.constexpr,
64
+ # state_cache last dim packs [kv_state, score_state], each STATE_WIDTH wide.
65
+ STATE_WIDTH: tl.constexpr,
66
+ COMPRESS_RATIO: tl.constexpr,
67
+):
68
+ token_idx = tl.program_id(0)
69
+ slot_id = tl.load(slot_mapping_ptr + token_idx)
70
+
71
+ # Skip padded / invalid tokens (slot_id == -1 is the PAD sentinel used
72
+ # by vLLM). During CUDA graph replay the batch may contain padding
73
+ # tokens whose slot_mapping is -1; writing to kv_state[-1] would be an
74
+ # illegal memory access.
75
+ if slot_id < 0:
76
+ return
77
+
78
+ block_idx = slot_id // block_size
79
+ pos_in_block = slot_id % block_size
80
+ base_ptr = (
81
+ state_cache_ptr
82
+ + block_idx * state_cache_stride0
83
+ + pos_in_block * state_cache_stride1
84
+ )
85
+
86
+ block = tl.arange(0, TRITON_BLOCK_SIZE)
87
+ mask = block < HEAD_SIZE
88
+
89
+ kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask)
90
+ tl.store(base_ptr + block, kv, mask=mask)
91
+
92
+ # Fused: score += ape[position % compress_ratio]
93
+ position = tl.load(positions_ptr + token_idx)
94
+ ape_row = position % COMPRESS_RATIO
95
+ ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask)
96
+ score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask)
97
+ tl.store(
98
+ base_ptr + STATE_WIDTH + block,
99
+ score + ape,
100
+ mask=mask,
101
+ )
vllm/models/deepseek_v4/compressor.py
+68
-198
@@ -13,15 +13,13 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
13
from vllm.model_executor.layers.layernorm import RMSNorm
14
from vllm.model_executor.layers.linear import MergedColumnParallelLinear
15
from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
16
- _compress_kv_sparse_attn_cutedsl,
17
- _fused_kv_compress_norm_rope_insert_indexer_attn,
18
- _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
19
- _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl,
20
- _norm_rope_insert_sparse_attn_cutedsl,
16
+ compress_norm_rope_store_triton,
17
)
18
from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE
19
+from vllm.models.deepseek_v4.common.ops.save_partial_states import (
20
+ save_partial_states,
21
+)
22
from vllm.platforms import current_platform
24
-from vllm.triton_utils import tl, triton
23
from vllm.v1.attention.backend import (
24
AttentionBackend,
25
AttentionCGSupport,
@@ -173,6 +171,16 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
171
172
173
class DeepseekCompressor(nn.Module):
174
+ """DeepSeek V4 KV/score compressor.
175
+
176
+ Owns the linear / norm / state-cache / ape state and the shared forward
177
+ prologue (kv/score split, save_partial_states launch). The
178
+ compress → norm → RoPE → store step is dispatched to a triton kernel
179
+ (``compress_norm_rope_store_triton``) by default, except for the NVIDIA
180
+ head_dim=128 indexer path which uses the cutedsl kernel
181
+ (``compress_norm_rope_store_cutedsl``) for better performance.
182
+ """
183
+
184
def __init__(
185
self,
186
vllm_config: VllmConfig,
@@ -242,32 +250,18 @@ class DeepseekCompressor(nn.Module):
250
assert not use_fp4_cache, (
251
"MXFP4 cache is only supported for indexer (head=128)"
252
)
245
- self._use_cutedsl_sparse_compressor = True
246
- self._use_cutedsl_fused_sparse_compressor = self.compress_ratio == 4
247
- self._compress_kernel = _compress_kv_sparse_attn_cutedsl
248
- self._norm_rope_store_kernel = _norm_rope_insert_sparse_attn_cutedsl
249
- self._fused_sparse_kernel = (
250
- _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl
251
- )
253
self._quant_block = 64
254
self._token_stride = self.nope_head_dim + self.rope_head_dim * 2
255
self._scale_dim = self.nope_head_dim // 64 + 1 # 7 real + 1 pad
255
- self._num_warps = 4
256
elif self.head_dim == 128:
257
- self._use_cutedsl_sparse_compressor = False
257
if use_fp4_cache:
259
- self._fused_kernel = (
260
- _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
261
- )
258
self._quant_block = MXFP4_BLOCK_SIZE
259
self._token_stride = self.head_dim // 2
260
self._scale_dim = self.head_dim // MXFP4_BLOCK_SIZE
261
else:
266
- self._fused_kernel = _fused_kv_compress_norm_rope_insert_indexer_attn
262
self._quant_block = 128
263
self._token_stride = self.head_dim
264
self._scale_dim = 4 # single float32 scale
270
- self._num_warps = 1
265
else:
266
raise ValueError(
267
f"Unsupported head_dim for fused quant+cache: {self.head_dim}"
@@ -312,29 +306,22 @@ class DeepseekCompressor(nn.Module):
306
)
307
308
# Store the KV and score (with fused APE addition) in the state.
315
- # NOTE: PDL is disabled — both this kernel and _fused_kernel below
316
- # depend on preceding kernel outputs (kv/score from the cublas GEMM;
317
- # state_cache from this kernel) but neither emits/waits on PDL grid
318
- # dependency primitives, so launch_pdl=True caused a read-after-write
319
- # race and non-deterministic output.
320
- _save_partial_states_kernel[(num_actual,)](
321
- kv,
322
- kv.stride(0),
323
- score,
324
- score.stride(0),
325
- self.ape,
326
- self.ape.stride(0),
327
- positions,
328
- state_cache,
329
- state_cache.stride(0),
330
- state_cache.stride(1),
331
- slot_mapping,
332
- block_size,
333
- HEAD_SIZE=kv.shape[-1],
334
- TRITON_BLOCK_SIZE=triton.next_power_of_2(kv.shape[-1]),
335
- STATE_WIDTH=state_width,
336
- COMPRESS_RATIO=self.compress_ratio,
337
- **pdl_kwargs,
309
+ # NOTE: PDL is disabled — both this kernel and the compress kernels
310
+ # below depend on preceding kernel outputs (kv/score from the cublas
311
+ # GEMM; state_cache from this kernel) but neither emits/waits on PDL
312
+ # grid dependency primitives, so launch_pdl=True caused a
313
+ # read-after-write race and non-deterministic output.
314
+ save_partial_states(
315
+ kv=kv,
316
+ score=score,
317
+ ape=self.ape,
318
+ positions=positions,
319
+ state_cache=state_cache,
320
+ slot_mapping=slot_mapping,
321
+ block_size=block_size,
322
+ state_width=state_width,
323
+ compress_ratio=self.compress_ratio,
324
+ pdl_kwargs=pdl_kwargs,
325
)
326
327
# Fused: compress → RMSNorm → RoPE → FP8 quant → KV cache write.
@@ -348,161 +335,44 @@ class DeepseekCompressor(nn.Module):
335
k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix])
336
kv_cache = self._static_forward_context[self.k_cache_prefix].kv_cache
337
351
- if self._use_cutedsl_sparse_compressor:
352
- if self._use_cutedsl_fused_sparse_compressor:
353
- self._fused_sparse_kernel(
354
- state_cache,
355
- token_to_req_indices,
356
- positions,
357
- slot_mapping,
358
- block_table,
359
- block_size,
360
- self.norm.weight,
361
- self.rms_norm_eps,
362
- cos_sin_cache,
363
- kv_cache,
364
- k_cache_metadata.slot_mapping,
365
- kv_cache.shape[1], # paged KV cache block size
366
- kv_cache.stride(0),
367
- head_size=self.head_dim,
368
- state_width=state_width,
369
- rope_head_dim=self.rope_head_dim,
370
- fp8_max=448.0,
371
- quant_block=self._quant_block,
372
- token_stride=self._token_stride,
373
- scale_dim=self._scale_dim,
374
- compress_ratio=self.compress_ratio,
375
- overlap=self.overlap,
376
- )
338
+ if current_platform.is_cuda():
339
+ # NVIDIA GPUs.
340
+ if self.head_dim == 512:
341
+ from .nvidia.ops import compress_norm_rope_store_cutedsl
342
+
343
+ # Main compressor path.
344
+ # Use a cutedsl kernel for better performance.
345
+ compress_norm_rope_store_fn = compress_norm_rope_store_cutedsl
346
else:
378
- compressed_kv = torch.empty(
379
- (num_actual, self.head_dim),
380
- dtype=torch.float32,
381
- device=state_cache.device,
382
- )
383
- self._compress_kernel(
384
- state_cache,
385
- token_to_req_indices,
386
- positions,
387
- slot_mapping,
388
- block_table,
389
- block_size,
390
- compressed_kv,
391
- head_size=self.head_dim,
392
- state_width=state_width,
393
- compress_ratio=self.compress_ratio,
394
- overlap=self.overlap,
395
- )
396
- self._norm_rope_store_kernel(
397
- compressed_kv,
398
- positions,
399
- slot_mapping,
400
- self.norm.weight,
401
- self.rms_norm_eps,
402
- cos_sin_cache,
403
- kv_cache,
404
- k_cache_metadata.slot_mapping,
405
- kv_cache.shape[1], # paged KV cache block size
406
- kv_cache.stride(0),
407
- head_size=self.head_dim,
408
- rope_head_dim=self.rope_head_dim,
409
- fp8_max=448.0,
410
- quant_block=self._quant_block,
411
- token_stride=self._token_stride,
412
- scale_dim=self._scale_dim,
413
- compress_ratio=self.compress_ratio,
414
- )
347
+ # Indexer path (head_dim == 128).
348
+ # Use a triton kernel.
349
+ compress_norm_rope_store_fn = compress_norm_rope_store_triton
350
else:
416
- self._fused_kernel[(num_actual,)](
417
- # state cache
418
- state_cache,
419
- state_cache.stride(0),
420
- state_cache.stride(1),
421
- # metadata
422
- token_to_req_indices,
423
- positions,
424
- slot_mapping,
425
- block_table,
426
- block_table.stride(0),
427
- block_size,
428
- # RMSNorm
429
- self.norm.weight,
430
- self.rms_norm_eps,
431
- # RoPE
432
- cos_sin_cache,
433
- cos_sin_cache.stride(0),
434
- # KV cache
435
- kv_cache,
436
- k_cache_metadata.slot_mapping,
437
- kv_cache.shape[1], # paged KV cache block size (tokens per block)
438
- # constexprs
439
- HEAD_SIZE=self.head_dim,
440
- TRITON_BLOCK_SIZE=triton.next_power_of_2(self.head_dim),
441
- STATE_WIDTH=state_width,
442
- COMPRESS_RATIO=self.compress_ratio,
443
- OVERLAP=self.overlap,
444
- ROPE_HEAD_DIM=self.rope_head_dim,
445
- FP8_MAX=448.0,
446
- QUANT_BLOCK=self._quant_block,
447
- TOKEN_STRIDE=self._token_stride,
448
- SCALE_DIM=self._scale_dim,
449
- KV_BLOCK_STRIDE=kv_cache.stride(0),
450
- num_warps=self._num_warps,
451
- **pdl_kwargs,
452
- )
453
-
351
+ # AMD GPUs.
352
+ # Always use a triton kernel.
353
+ compress_norm_rope_store_fn = compress_norm_rope_store_triton
354
455
-@triton.jit
456
-def _save_partial_states_kernel(
457
- kv_ptr,
458
- kv_stride,
459
- score_ptr,
460
- score_stride,
461
- ape_ptr,
462
- ape_stride,
463
- positions_ptr,
464
- state_cache_ptr,
465
- state_cache_stride0,
466
- state_cache_stride1,
467
- slot_mapping_ptr,
468
- block_size,
469
- HEAD_SIZE: tl.constexpr,
470
- TRITON_BLOCK_SIZE: tl.constexpr,
471
- # state_cache last dim packs [kv_state, score_state], each STATE_WIDTH wide.
472
- STATE_WIDTH: tl.constexpr,
473
- COMPRESS_RATIO: tl.constexpr,
474
-):
475
- token_idx = tl.program_id(0)
476
- slot_id = tl.load(slot_mapping_ptr + token_idx)
477
-
478
- # Skip padded / invalid tokens (slot_id == -1 is the PAD sentinel used
479
- # by vLLM). During CUDA graph replay the batch may contain padding
480
- # tokens whose slot_mapping is -1; writing to kv_state[-1] would be an
481
- # illegal memory access.
482
- if slot_id < 0:
483
- return
484
-
485
- block_idx = slot_id // block_size
486
- pos_in_block = slot_id % block_size
487
- base_ptr = (
488
- state_cache_ptr
489
- + block_idx * state_cache_stride0
490
- + pos_in_block * state_cache_stride1
491
- )
492
-
493
- block = tl.arange(0, TRITON_BLOCK_SIZE)
494
- mask = block < HEAD_SIZE
495
-
496
- kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask)
497
- tl.store(base_ptr + block, kv, mask=mask)
498
-
499
- # Fused: score += ape[position % compress_ratio]
500
- position = tl.load(positions_ptr + token_idx)
501
- ape_row = position % COMPRESS_RATIO
502
- ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask)
503
- score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask)
504
- tl.store(
505
- base_ptr + STATE_WIDTH + block,
506
- score + ape,
507
- mask=mask,
508
- )
355
+ compress_norm_rope_store_fn(
356
+ state_cache=state_cache,
357
+ num_actual=num_actual,
358
+ token_to_req_indices=token_to_req_indices,
359
+ positions=positions,
360
+ slot_mapping=slot_mapping,
361
+ block_table=block_table,
362
+ block_size=block_size,
363
+ state_width=state_width,
364
+ cos_sin_cache=cos_sin_cache,
365
+ kv_cache=kv_cache,
366
+ k_cache_metadata=k_cache_metadata,
367
+ pdl_kwargs=pdl_kwargs,
368
+ head_dim=self.head_dim,
369
+ rope_head_dim=self.rope_head_dim,
370
+ compress_ratio=self.compress_ratio,
371
+ overlap=self.overlap,
372
+ use_fp4_cache=self.use_fp4_cache,
373
+ rms_norm_weight=self.norm.weight,
374
+ rms_norm_eps=self.rms_norm_eps,
375
+ quant_block=self._quant_block,
376
+ token_stride=self._token_stride,
377
+ scale_dim=self._scale_dim,
378
+ )
vllm/models/deepseek_v4/nvidia/model.py
+1
-1
@@ -59,7 +59,7 @@ from vllm.models.deepseek_v4.attention import (
59
DeepseekV4MLAModules,
60
DeepseekV4MultiHeadLatentAttentionWrapper,
61
)
62
-from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
62
+from vllm.models.deepseek_v4.nvidia.ops import prepare_megamoe_inputs
63
from vllm.platforms import current_platform
64
from vllm.sequence import IntermediateTensors
65
from vllm.utils.torch_utils import direct_register_custom_op
vllm/models/deepseek_v4/nvidia/ops/__init__.py
+16
@@ -6,3 +6,19 @@ These modules import ``cutlass``/``cutedsl`` at module top level, so they must
6
not be imported on non-CUDA platforms. Callers should gate on
7
``vllm.utils.import_utils.has_cutedsl()`` before importing from here.
8
"""
9
+
10
+from .dequant_gather_k_cutedsl import dequantize_and_gather_k_cache_cutedsl
11
+from .fused_indexer_q_cutedsl import (
12
+ fused_indexer_q_rope_quant_fp8_cutedsl,
13
+ fused_indexer_q_rope_quant_mxfp4_cutedsl,
14
+)
15
+from .prepare_megamoe import prepare_megamoe_inputs
16
+from .sparse_attn_compress_cutedsl import compress_norm_rope_store_cutedsl
17
+
18
+__all__ = [
19
+ "compress_norm_rope_store_cutedsl",
20
+ "dequantize_and_gather_k_cache_cutedsl",
21
+ "fused_indexer_q_rope_quant_fp8_cutedsl",
22
+ "fused_indexer_q_rope_quant_mxfp4_cutedsl",
23
+ "prepare_megamoe_inputs",
24
+]
vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py
renamed
+95
-3
@@ -8,6 +8,7 @@ The public wrappers provide the C4 fused and C128 split kernels.
8
from __future__ import annotations
9
10
from functools import cache
11
+from typing import Any
12
13
import cutlass
14
import cutlass.cute as cute
@@ -1086,7 +1087,7 @@ class SparseAttnNormRopeStoreKernel:
1087
)
1088
1089
1089
-def _compress_kv_sparse_attn_cutedsl(
1090
+def compress_kv_sparse_attn_cutedsl(
1091
state_cache: torch.Tensor,
1092
token_to_req_indices: torch.Tensor,
1093
positions: torch.Tensor,
@@ -1118,7 +1119,7 @@ def _compress_kv_sparse_attn_cutedsl(
1119
)
1120
1121
1121
-def _norm_rope_insert_sparse_attn_cutedsl(
1122
+def norm_rope_insert_sparse_attn_cutedsl(
1123
compressed_kv: torch.Tensor,
1124
positions: torch.Tensor,
1125
slot_mapping: torch.Tensor,
@@ -1174,7 +1175,7 @@ def _norm_rope_insert_sparse_attn_cutedsl(
1175
)
1176
1177
1177
-def _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
1178
+def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
1179
state_cache: torch.Tensor,
1180
token_to_req_indices: torch.Tensor,
1181
positions: torch.Tensor,
@@ -1238,3 +1239,94 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
1239
kv_slot_mapping,
1240
kv_cache_block_size,
1241
)
1242
+
1243
+
1244
+def compress_norm_rope_store_cutedsl(
1245
+ state_cache: torch.Tensor,
1246
+ num_actual: int,
1247
+ token_to_req_indices: torch.Tensor,
1248
+ positions: torch.Tensor,
1249
+ slot_mapping: torch.Tensor,
1250
+ block_table: torch.Tensor,
1251
+ block_size: int,
1252
+ state_width: int,
1253
+ cos_sin_cache: torch.Tensor,
1254
+ kv_cache: torch.Tensor,
1255
+ k_cache_metadata: Any,
1256
+ pdl_kwargs: dict,
1257
+ head_dim: int,
1258
+ rope_head_dim: int,
1259
+ compress_ratio: int,
1260
+ overlap: bool,
1261
+ use_fp4_cache: bool,
1262
+ rms_norm_weight: torch.Tensor,
1263
+ rms_norm_eps: float,
1264
+ quant_block: int,
1265
+ token_stride: int,
1266
+ scale_dim: int,
1267
+) -> None:
1268
+ if compress_ratio == 4:
1269
+ # For C4A, the single fused kernel is faster than the two-kernel version.
1270
+ fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
1271
+ state_cache,
1272
+ token_to_req_indices,
1273
+ positions,
1274
+ slot_mapping,
1275
+ block_table,
1276
+ block_size,
1277
+ rms_norm_weight,
1278
+ rms_norm_eps,
1279
+ cos_sin_cache,
1280
+ kv_cache,
1281
+ k_cache_metadata.slot_mapping,
1282
+ kv_cache.shape[1], # paged KV cache block size
1283
+ kv_cache.stride(0),
1284
+ head_size=head_dim,
1285
+ state_width=state_width,
1286
+ rope_head_dim=rope_head_dim,
1287
+ fp8_max=448.0,
1288
+ quant_block=quant_block,
1289
+ token_stride=token_stride,
1290
+ scale_dim=scale_dim,
1291
+ compress_ratio=compress_ratio,
1292
+ overlap=overlap,
1293
+ )
1294
+ else:
1295
+ # For C128, the two-kernel version is faster than the single fused kernel.
1296
+ compressed_kv = torch.empty(
1297
+ (num_actual, head_dim),
1298
+ dtype=torch.float32,
1299
+ device=state_cache.device,
1300
+ )
1301
+ compress_kv_sparse_attn_cutedsl(
1302
+ state_cache,
1303
+ token_to_req_indices,
1304
+ positions,
1305
+ slot_mapping,
1306
+ block_table,
1307
+ block_size,
1308
+ compressed_kv,
1309
+ head_size=head_dim,
1310
+ state_width=state_width,
1311
+ compress_ratio=compress_ratio,
1312
+ overlap=overlap,
1313
+ )
1314
+ norm_rope_insert_sparse_attn_cutedsl(
1315
+ compressed_kv,
1316
+ positions,
1317
+ slot_mapping,
1318
+ rms_norm_weight,
1319
+ rms_norm_eps,
1320
+ cos_sin_cache,
1321
+ kv_cache,
1322
+ k_cache_metadata.slot_mapping,
1323
+ kv_cache.shape[1], # paged KV cache block size
1324
+ kv_cache.stride(0),
1325
+ head_size=head_dim,
1326
+ rope_head_dim=rope_head_dim,
1327
+ fp8_max=448.0,
1328
+ quant_block=quant_block,
1329
+ token_stride=token_stride,
1330
+ scale_dim=scale_dim,
1331
+ compress_ratio=compress_ratio,
1332
+ )