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 3, 22:45 UTC
[DSv4] Improved dequant gather K cache kernel (#42236)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Thien Tran committed
May 11, 2026 at 21:41 UTC
724ed2fc352bb388aea7269ca5c6690637c30ce7
5 files changed
+658
-100
tests/kernels/test_compressor_kv_cache.py
+141
-7
@@ -3,11 +3,12 @@
3
"""
4
Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant.
5
6
-Four test functions cover five paths:
6
+These tests cover:
7
A) DeepseekV4 Attention: head_dim=512 (448 FP8 nope + 64 bf16 rope), quant_block=64
8
- B) Indexer: head_dim=128 (all FP8), quant_block=128
9
- C) DeepseekV4 Attention magnitude range: correctness across small/large values
10
- D) Indexer fused Triton kernel: compress+norm+rope+quant+insert
8
+ B) Fused dequant+gather K cache
9
+ C) Indexer: head_dim=128 (all FP8), quant_block=128
10
+ D) DeepseekV4 Attention magnitude range: correctness across small/large values
11
+ E) Indexer fused Triton kernel: compress+norm+rope+quant+insert
12
"""
13
14
import math
@@ -134,7 +135,140 @@ def test_deepseek_v4_attention_quant_cache_roundtrip(num_tokens: int, block_size
135
)
136
137
137
-# ── Test B: Indexer path ────────────────────────────────────────────────────
138
+# ── Test B: Fused dequant+gather K cache ────────────────────────────────────
139
+
140
+
141
+def _dequantize_and_gather_k_cache_reference(
142
+ out: torch.Tensor,
143
+ k_cache: torch.Tensor,
144
+ seq_lens: torch.Tensor,
145
+ gather_lens: torch.Tensor | None,
146
+ block_table: torch.Tensor,
147
+ block_size: int,
148
+ offset: int,
149
+) -> None:
150
+ fp8_dim = 448
151
+ bf16_dim = 64
152
+ scale_dim = 8
153
+ quant_block = 64
154
+ token_data_size = fp8_dim + bf16_dim * 2
155
+
156
+ for req_id in range(seq_lens.shape[0]):
157
+ seq_len = seq_lens[req_id].item()
158
+ gather_len = gather_lens[req_id].item() if gather_lens is not None else seq_len
159
+ start_pos = seq_len - gather_len
160
+
161
+ for i in range(gather_len):
162
+ pos = start_pos + i
163
+ pos_in_block = pos % block_size
164
+ block_idx = block_table[req_id, pos // block_size].item()
165
+ cache_block = k_cache[block_idx].view(-1)
166
+
167
+ token_data_start = pos_in_block * token_data_size
168
+ fp8_bytes = cache_block[token_data_start : token_data_start + fp8_dim]
169
+ fp8_vals = fp8_bytes.view(torch.float8_e4m3fn).float()
170
+
171
+ scale_start = block_size * token_data_size + pos_in_block * scale_dim
172
+ encoded_scales = cache_block[scale_start : scale_start + scale_dim]
173
+ scales = torch.exp2(encoded_scales[:7].float() - 127.0)
174
+ dequant = fp8_vals * scales.repeat_interleave(quant_block)
175
+
176
+ bf16_start = token_data_start + fp8_dim
177
+ bf16_bytes = cache_block[bf16_start : bf16_start + bf16_dim * 2]
178
+ bf16_tail = bf16_bytes.view(torch.bfloat16)
179
+
180
+ out[req_id, offset + i, :fp8_dim] = dequant
181
+ out[req_id, offset + i, fp8_dim:] = bf16_tail
182
+
183
+
184
+@pytest.mark.parametrize(
185
+ ("seq_lens_host", "gather_lens_host", "offset"),
186
+ [
187
+ ([9, 23, 7], None, 0),
188
+ ([19, 8, 257], [6, 8, 129], 5),
189
+ ],
190
+)
191
+def test_dequantize_and_gather_k_cache(
192
+ seq_lens_host: list[int],
193
+ gather_lens_host: list[int] | None,
194
+ offset: int,
195
+):
196
+ block_size = 64
197
+ head_dim = 512
198
+ nope_dim = 448
199
+ scale_dim = 8
200
+ head_bytes = nope_dim + (head_dim - nope_dim) * 2 + scale_dim
201
+ device = "cuda"
202
+ num_reqs = len(seq_lens_host)
203
+ num_tokens = sum(seq_lens_host)
204
+ max_gather_len = max(gather_lens_host or seq_lens_host)
205
+ max_blocks_per_seq = math.ceil(max(seq_lens_host) / block_size)
206
+ num_blocks = sum(math.ceil(seq_len / block_size) for seq_len in seq_lens_host)
207
+
208
+ compressed_kv = torch.randn(
209
+ num_tokens, head_dim, dtype=torch.bfloat16, device=device
210
+ )
211
+
212
+ # Randomize physical pages so the test covers block-table translation.
213
+ # Keep padded block-table entries invalid to catch accidental reads.
214
+ physical_blocks = torch.randperm(num_blocks, device=device)
215
+ block_table = torch.full(
216
+ (num_reqs, max_blocks_per_seq), int(-1e6), dtype=torch.int32, device=device
217
+ )
218
+ start = 0
219
+ for req_id, seq_len in enumerate(seq_lens_host):
220
+ num_req_blocks = math.ceil(seq_len / block_size)
221
+ req_blocks = physical_blocks[start : start + num_req_blocks]
222
+ block_table[req_id, :num_req_blocks] = req_blocks
223
+ start += num_req_blocks
224
+
225
+ # Build slot_mapping for quantize_and_insert_k_cache.
226
+ slot_mapping = torch.empty(num_tokens, dtype=torch.int64, device=device)
227
+ start = 0
228
+ for req_id, seq_len in enumerate(seq_lens_host):
229
+ logical_pos = torch.arange(seq_len, dtype=torch.int64, device=device)
230
+ block_idx = block_table[req_id, logical_pos // block_size].to(torch.int64)
231
+ token_slots = block_idx * block_size + logical_pos % block_size
232
+ slot_mapping[start : start + seq_len] = token_slots
233
+ start += seq_len
234
+
235
+ # Insert compressed K into the paged cache layout used by the gather op.
236
+ k_cache = torch.empty(
237
+ num_blocks, block_size, head_bytes, dtype=torch.uint8, device=device
238
+ )
239
+ k_cache_2d = k_cache.view(num_blocks, -1)
240
+ quantize_and_insert_k_cache(compressed_kv, k_cache_2d, slot_mapping, block_size)
241
+
242
+ out_shape = (num_reqs, offset + max_gather_len + 3, head_dim)
243
+ ref_out = torch.empty(out_shape, dtype=torch.bfloat16, device=device)
244
+ actual_out = torch.empty_like(ref_out)
245
+ seq_lens = torch.tensor(seq_lens_host, dtype=torch.int32, device=device)
246
+ gather_lens = (
247
+ torch.tensor(gather_lens_host, dtype=torch.int32, device=device)
248
+ if gather_lens_host is not None
249
+ else None
250
+ )
251
+
252
+ # Compare production gather against a PyTorch reference for valid output rows.
253
+ _dequantize_and_gather_k_cache_reference(
254
+ ref_out, k_cache, seq_lens, gather_lens, block_table, block_size, offset
255
+ )
256
+ dequantize_and_gather_k_cache(
257
+ actual_out, k_cache, seq_lens, gather_lens, block_table, block_size, offset
258
+ )
259
+ torch.accelerator.synchronize()
260
+
261
+ # only check non-padded content
262
+ for req_id, seq_len in enumerate(seq_lens_host):
263
+ gather_len = (
264
+ gather_lens_host[req_id] if gather_lens_host is not None else seq_len
265
+ )
266
+ actual = actual_out[req_id, offset : offset + gather_len]
267
+ expected = ref_out[req_id, offset : offset + gather_len]
268
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
269
+
270
+
271
+# ── Test C: Indexer path ────────────────────────────────────────────────────
272
273
274
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 17])
@@ -254,7 +388,7 @@ def test_indexer_gather_accepts_upper_bound_output():
388
assert torch.all(dst_scale[valid_tokens:] == sentinel)
389
390
257
-# ── Test C: DeepseekV4 attention with values at different magnitudes ───────────
391
+# ── Test D: DeepseekV4 attention with values at different magnitudes ───────────
392
393
394
def test_deepseek_v4_quant_magnitude_range():
@@ -316,7 +450,7 @@ def test_deepseek_v4_quant_magnitude_range():
450
)
451
452
319
-# ── Test D: Indexer fused K-cache insert (Triton kernels) ────────────────────
453
+# ── Test E: Indexer fused K-cache insert (Triton kernels) ────────────────────
454
#
455
# Both kernels share the same Triton signature; use_fp4 selects between them.
456
# Full pipeline: state-cache gather → softmax-weighted compress → RMSNorm →
vllm/v1/attention/ops/deepseek_v4_ops/cache_utils.py
+30
-1
@@ -17,6 +17,7 @@ preparation.
17
import torch
18
19
from vllm.triton_utils import tl, triton
20
+from vllm.utils.import_utils import has_cutedsl
21
22
23
@triton.jit
@@ -303,7 +304,7 @@ def _dequantize_and_gather_k_kernel(
304
tl.store(output_row_ptr + bf16_output_offset + chunk_offsets, bf16_vals)
305
306
306
-def dequantize_and_gather_k_cache(
307
+def dequantize_and_gather_k_cache_triton(
308
# [num_reqs, max_num_tokens, head_size]
309
out: torch.Tensor,
310
# [num_blocks, block_size, head_bytes]
@@ -349,6 +350,34 @@ def dequantize_and_gather_k_cache(
350
)
351
352
353
+def dequantize_and_gather_k_cache(
354
+ # [num_reqs, max_num_tokens, head_size]
355
+ out: torch.Tensor,
356
+ # [num_blocks, block_size, head_bytes]
357
+ k_cache: torch.Tensor,
358
+ # [num_reqs]
359
+ seq_lens: torch.Tensor,
360
+ # [num_reqs]
361
+ gather_lens: torch.Tensor | None,
362
+ # [num_reqs, max_blocks_per_seq]
363
+ block_table: torch.Tensor,
364
+ block_size: int,
365
+ offset: int,
366
+) -> None:
367
+ if has_cutedsl():
368
+ # lazily import, otherwise some tests fail due to CUDA driver init failure.
369
+ from .dequant_gather_k_cutedsl import dequantize_and_gather_k_cache_cutedsl
370
+
371
+ dequantize_and_gather_k_cache_cutedsl(
372
+ out, k_cache, seq_lens, gather_lens, block_table, block_size, offset
373
+ )
374
+ return
375
+
376
+ dequantize_and_gather_k_cache_triton(
377
+ out, k_cache, seq_lens, gather_lens, block_table, block_size, offset
378
+ )
379
+
380
+
381
def compute_global_topk_indices_and_lens(
382
topk_indices: torch.Tensor,
383
token_to_req_indices: torch.Tensor,
vllm/v1/attention/ops/deepseek_v4_ops/cutedsl_utils.py
new
+145
@@ -0,0 +1,145 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import cutlass
5
+import cutlass.cute as cute
6
+from cutlass import Float32, Uint32
7
+from cutlass._mlir import ir
8
+from cutlass._mlir.dialects import llvm, vector
9
+from cutlass.cutlass_dsl import T, dsl_user_op
10
+
11
+
12
+@dsl_user_op
13
+def _recast_val(x, dtype, *, loc=None, ip=None):
14
+ return dtype(llvm.bitcast(dtype.mlir_type, x.ir_value(loc=loc, ip=ip)))
15
+
16
+
17
+@dsl_user_op
18
+def _fp32x2_to_bf16x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32:
19
+ out = llvm.inline_asm(
20
+ T.i32(),
21
+ [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
22
+ "cvt.rn.bf16x2.f32 $0, $2, $1;",
23
+ "=r,f,f",
24
+ has_side_effects=False,
25
+ is_align_stack=False,
26
+ )
27
+ return Uint32(out)
28
+
29
+
30
+@dsl_user_op
31
+def _bf16x2_to_fp32(data: Uint32, *, loc=None, ip=None) -> tuple[Float32, Float32]:
32
+ out = llvm.inline_asm(
33
+ llvm.StructType.get_literal([T.f32(), T.f32()]),
34
+ [data.ir_value(loc=loc, ip=ip)],
35
+ "shl.b32 $0, $2, 16;\n\tand.b32 $1, $2, 0xFFFF0000;\n",
36
+ "=f,=f,r",
37
+ has_side_effects=False,
38
+ is_align_stack=False,
39
+ )
40
+ return (
41
+ Float32(llvm.extractvalue(T.f32(), out, [0], loc=loc, ip=ip)),
42
+ Float32(llvm.extractvalue(T.f32(), out, [1], loc=loc, ip=ip)),
43
+ )
44
+
45
+
46
+@dsl_user_op
47
+def _bf16x2_abs(a: Uint32, *, loc=None, ip=None) -> Uint32:
48
+ out = llvm.inline_asm(
49
+ T.i32(),
50
+ [a.ir_value(loc=loc, ip=ip)],
51
+ "abs.bf16x2 $0, $1;",
52
+ "=r,r",
53
+ has_side_effects=False,
54
+ is_align_stack=False,
55
+ )
56
+ return Uint32(out)
57
+
58
+
59
+@dsl_user_op
60
+def _bf16x2_max(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
61
+ out = llvm.inline_asm(
62
+ T.i32(),
63
+ [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
64
+ "max.bf16x2 $0, $1, $2;",
65
+ "=r,r,r",
66
+ has_side_effects=False,
67
+ is_align_stack=False,
68
+ )
69
+ return Uint32(out)
70
+
71
+
72
+@dsl_user_op
73
+def _bf16x2_mul(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
74
+ out = llvm.inline_asm(
75
+ T.i32(),
76
+ [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
77
+ "mul.rn.bf16x2 $0, $1, $2;",
78
+ "=r,r,r",
79
+ has_side_effects=False,
80
+ is_align_stack=False,
81
+ )
82
+ return Uint32(out)
83
+
84
+
85
+@dsl_user_op
86
+def _fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
87
+ # there is only fp8->fp16 conversion, hence we need to go
88
+ # round trip through fp16.
89
+ out = llvm.inline_asm(
90
+ llvm.StructType.get_literal([T.i32()] * 2),
91
+ [x.ir_value(loc=loc, ip=ip)],
92
+ "{\n\t"
93
+ ".reg .b16 x0, x1;\n\t"
94
+ ".reg .b16 t00, t01, t10, t11;\n\t"
95
+ "mov.b32 {x0, x1}, $2;\n\t"
96
+ "cvt.rn.f16x2.e4m3x2 $0, x0;\n\t"
97
+ "cvt.rn.f16x2.e4m3x2 $1, x1;\n\t"
98
+ "mov.b32 {t00, t01}, $0;\n\t"
99
+ "mov.b32 {t10, t11}, $1;\n\t"
100
+ "cvt.rn.bf16.f16 t00, t00;\n\t"
101
+ "cvt.rn.bf16.f16 t01, t01;\n\t"
102
+ "cvt.rn.bf16.f16 t10, t10;\n\t"
103
+ "cvt.rn.bf16.f16 t11, t11;\n\t"
104
+ "mov.b32 $0, {t00, t01};\n\t"
105
+ "mov.b32 $1, {t10, t11};\n\t"
106
+ "}\n",
107
+ "=r,=r,r",
108
+ has_side_effects=False,
109
+ is_align_stack=False,
110
+ )
111
+ vec = vector.from_elements(
112
+ ir.VectorType.get([2], T.i32(), loc=loc),
113
+ [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(2)],
114
+ loc=loc,
115
+ ip=ip,
116
+ )
117
+ return cute.TensorSSA(vec, 2, Uint32)
118
+
119
+
120
+@dsl_user_op
121
+def _fp32x8_to_fp4x8(
122
+ vals: cute.Tensor,
123
+ offset: cutlass.Constexpr[int],
124
+ *,
125
+ loc=None,
126
+ ip=None,
127
+) -> Uint32:
128
+ # Pack eight scaled FP32 values into four E2M1x2 bytes, returned as one b32.
129
+ assert vals.element_type is Float32
130
+ out = llvm.inline_asm(
131
+ T.i32(),
132
+ [vals[offset + i].ir_value(loc=loc, ip=ip) for i in range(8)],
133
+ "{\n\t"
134
+ ".reg .b8 x0, x1, x2, x3;\n\t"
135
+ "cvt.rn.satfinite.e2m1x2.f32 x0, $2, $1;\n\t"
136
+ "cvt.rn.satfinite.e2m1x2.f32 x1, $4, $3;\n\t"
137
+ "cvt.rn.satfinite.e2m1x2.f32 x2, $6, $5;\n\t"
138
+ "cvt.rn.satfinite.e2m1x2.f32 x3, $8, $7;\n\t"
139
+ "mov.b32 $0, {x0, x1, x2, x3};\n\t"
140
+ "}\n",
141
+ "=r,f,f,f,f,f,f,f,f",
142
+ has_side_effects=False,
143
+ is_align_stack=False,
144
+ )
145
+ return Uint32(out)
vllm/v1/attention/ops/deepseek_v4_ops/dequant_gather_k_cutedsl.py
new
+334
@@ -0,0 +1,334 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from functools import cache
5
+
6
+import cutlass
7
+import cutlass.cute as cute
8
+import torch
9
+from cuda.bindings.driver import CUstream
10
+from cutlass import BFloat16, Int32, Uint8, Uint32
11
+from cutlass.cute.nvgpu import cpasync
12
+from quack.compile_utils import make_fake_tensor
13
+
14
+from vllm.v1.attention.ops.deepseek_v4_ops.cutedsl_utils import (
15
+ _bf16x2_mul,
16
+ _fp8x4_to_bf16x4,
17
+)
18
+
19
+
20
+def dequantize_and_gather_k_cache_cutedsl(
21
+ out: torch.Tensor,
22
+ k_cache: torch.Tensor,
23
+ seq_lens: torch.Tensor,
24
+ gather_lens: torch.Tensor | None,
25
+ block_table: torch.Tensor,
26
+ block_size: int,
27
+ offset: int,
28
+) -> None:
29
+ DequantGatherKCacheKernel.compile(
30
+ block_size=block_size,
31
+ has_gather_lens=gather_lens is not None,
32
+ )(out, k_cache, seq_lens, gather_lens, block_table, offset)
33
+
34
+
35
+class DequantGatherKCacheKernel:
36
+ # Hard-coded for DSv4.
37
+ head_dim = 512
38
+ group_size = 64 # 1 scale per 64 elems
39
+
40
+ def __init__(self, fp8_dim: int = 448, block_size: int = 64):
41
+ self.fp8_dim = fp8_dim
42
+ self.bf16_dim = self.head_dim - fp8_dim
43
+ self.data_dim = fp8_dim + self.bf16_dim * 2
44
+ self.block_size = block_size
45
+
46
+ self.num_warps = 4
47
+ self.tb_size = self.num_warps * 32
48
+ self.num_stages = 4
49
+
50
+ @cute.jit
51
+ def __call__(
52
+ self,
53
+ out: cute.Tensor,
54
+ k_cache: cute.Tensor,
55
+ seq_lens: cute.Tensor,
56
+ gather_lens: cute.Tensor | None,
57
+ block_table: cute.Tensor,
58
+ offset: Int32,
59
+ stream: CUstream,
60
+ ):
61
+ # Split k_cache into k_data and k_scale. Each [block_size, head_bytes]
62
+ # block is actually a concat of
63
+ # [block_size, fp8_dim + bf16_dim * 2] and [block_size, 8].
64
+ k_data = cute.make_tensor(
65
+ k_cache.iterator,
66
+ layout=cute.make_layout(
67
+ (k_cache.shape[0], self.block_size, self.data_dim),
68
+ stride=(k_cache.stride[0], self.data_dim, 1),
69
+ ),
70
+ )
71
+ k_scale = cute.make_tensor(
72
+ k_cache.iterator + (self.block_size * self.data_dim),
73
+ layout=cute.make_layout(
74
+ (k_cache.shape[0], self.block_size, 8),
75
+ stride=(k_cache.stride[0], 8, 1),
76
+ ),
77
+ )
78
+
79
+ grid = (out.shape[0], 1024, 1)
80
+ self.kernel(
81
+ out,
82
+ k_data,
83
+ k_scale,
84
+ seq_lens,
85
+ gather_lens,
86
+ block_table,
87
+ offset,
88
+ ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
89
+
90
+ @cute.jit
91
+ def load_g2s(
92
+ self,
93
+ k_data_slice: cute.Tensor,
94
+ k_scale: cute.Tensor,
95
+ block_table: cute.Tensor,
96
+ s_kdata_slice: cute.Tensor,
97
+ s_kscale: cute.Tensor,
98
+ req_id,
99
+ pos,
100
+ lane_id,
101
+ stage_id,
102
+ ):
103
+ # k_data_slice: [num_blocks, block_size, (16, data_dim/16)]
104
+ # s_kdata_slice: [(4, data_dim/16), num_stages]
105
+
106
+ op = cpasync.CopyG2SOp(cute.nvgpu.LoadCacheMode.GLOBAL)
107
+ cp16_atom = cute.make_copy_atom(op, Uint32, num_bits_per_copy=128)
108
+ cp8_atom = cute.make_copy_atom(cpasync.CopyG2SOp(), Uint8, num_bits_per_copy=64)
109
+ page_id = block_table[req_id, pos // self.block_size]
110
+ block_offset = pos % self.block_size
111
+
112
+ # Load the first 512 bytes (32x16B).
113
+ idx = lane_id
114
+ src = k_data_slice[page_id, block_offset, (None, idx)]
115
+ cute.copy(
116
+ cp16_atom,
117
+ cute.recast_tensor(src, Uint32),
118
+ s_kdata_slice[(None, idx), stage_id],
119
+ )
120
+
121
+ # Load the tail 64 bytes.
122
+ idx += 32
123
+ if idx < cutlass.const_expr(self.data_dim // 16):
124
+ src = k_data_slice[page_id, block_offset, (None, idx)]
125
+ cute.copy(
126
+ cp16_atom,
127
+ cute.recast_tensor(src, Uint32),
128
+ s_kdata_slice[(None, idx), stage_id],
129
+ )
130
+ elif idx == cutlass.const_expr(self.data_dim // 16):
131
+ cute.copy(
132
+ cp8_atom,
133
+ k_scale[page_id, block_offset, None],
134
+ s_kscale[None, stage_id],
135
+ )
136
+
137
+ @cute.kernel
138
+ def kernel(
139
+ self,
140
+ out: cute.Tensor,
141
+ k_data: cute.Tensor,
142
+ k_scale: cute.Tensor,
143
+ seq_lens: cute.Tensor,
144
+ gather_lens: cute.Tensor | None,
145
+ block_table: cute.Tensor,
146
+ offset: Int32,
147
+ ):
148
+ req_id, worker_id, _ = cute.arch.block_idx()
149
+ tid, _, _ = cute.arch.thread_idx()
150
+ warp_id = cute.arch.make_warp_uniform(tid // 32)
151
+ lane_id = tid % 32
152
+
153
+ _, num_workers, _ = cute.arch.grid_dim()
154
+
155
+ # Prepare smem.
156
+ smem = cutlass.utils.SmemAllocator()
157
+ s_kdata = smem.allocate_tensor(
158
+ Uint32,
159
+ cute.make_layout((self.data_dim // 4, self.num_warps, self.num_stages)),
160
+ byte_alignment=16,
161
+ )[None, warp_id, None]
162
+ s_kscale = smem.allocate_tensor(
163
+ Uint8,
164
+ cute.make_layout((8, self.num_warps, self.num_stages)),
165
+ byte_alignment=8,
166
+ )[None, warp_id, None]
167
+
168
+ # Prepare for 16B cp.async, also for BF16 smem loads later.
169
+ k_data_slice = cute.logical_divide(k_data, (None, None, 16))
170
+ s_kdata_16B_slice = cute.logical_divide(s_kdata, (4, None))
171
+
172
+ # Load FP8 elems in 8B units, so once dequantized, they are 16B units.
173
+ s_kdata_8B_slice = cute.logical_divide(s_kdata, (2, None))
174
+
175
+ # 16B st.global.
176
+ out_slice = cute.logical_divide(out, (None, None, 8))
177
+
178
+ cp_op = cute.nvgpu.CopyUniversalOp()
179
+ cp8_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64)
180
+ cp16_atom = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=128)
181
+
182
+ seq_len = seq_lens[req_id]
183
+ gather_len = seq_len
184
+ if cutlass.const_expr(gather_lens is not None):
185
+ gather_len = gather_lens[req_id] # type: ignore[index]
186
+ start_pos = seq_len - gather_len
187
+
188
+ # Start prefetch.
189
+ for i in cutlass.range_constexpr(self.num_stages - 1):
190
+ next_pos = (
191
+ start_pos
192
+ + worker_id * self.num_warps
193
+ + warp_id
194
+ + i * num_workers * self.num_warps
195
+ )
196
+ if next_pos < seq_len:
197
+ self.load_g2s(
198
+ k_data_slice,
199
+ k_scale,
200
+ block_table,
201
+ s_kdata_16B_slice,
202
+ s_kscale,
203
+ req_id,
204
+ next_pos,
205
+ lane_id,
206
+ i,
207
+ )
208
+ cute.arch.cp_async_commit_group()
209
+ prefetch_stage = self.num_stages - 1
210
+ compute_stage = 0
211
+
212
+ # Main loop.
213
+ for i in range(
214
+ worker_id * self.num_warps + warp_id,
215
+ gather_len,
216
+ num_workers * self.num_warps,
217
+ ):
218
+ pos = start_pos + i
219
+
220
+ # Prefetch next stage.
221
+ next_pos = pos + num_workers * self.num_warps * (self.num_stages - 1)
222
+ if next_pos < seq_len:
223
+ self.load_g2s(
224
+ k_data_slice,
225
+ k_scale,
226
+ block_table,
227
+ s_kdata_16B_slice,
228
+ s_kscale,
229
+ req_id,
230
+ next_pos,
231
+ lane_id,
232
+ prefetch_stage,
233
+ )
234
+ prefetch_stage = (prefetch_stage + 1) % self.num_stages
235
+ cute.arch.cp_async_commit_group()
236
+
237
+ # Wait for gmem->smem to finish.
238
+ cute.arch.cp_async_wait_group(self.num_stages - 1)
239
+ cute.arch.sync_warp()
240
+
241
+ # There are 512 elems per token. As a warp, data0 holds the first
242
+ # 256 elems and data1 holds the second 256 elems, i.e. each thread
243
+ # holds 8 FP8 elems. This keeps the dequantized 8 BF16 elems as
244
+ # contiguous 16B global stores. On Blackwell, this might not be
245
+ # necessary as we have 32B global stores, but doing it this way
246
+ # does not seem to be slower.
247
+ data0 = cute.make_rmem_tensor((2,), Uint32)
248
+ data1 = cute.make_rmem_tensor((2,), Uint32)
249
+ cute.copy(cp8_atom, s_kdata_8B_slice[(None, lane_id), compute_stage], data0)
250
+ cute.copy(
251
+ cp8_atom,
252
+ s_kdata_8B_slice[(None, lane_id + 32), compute_stage],
253
+ data1,
254
+ )
255
+
256
+ # Convert to bf16x2 via bit manipulation. FP8 scales are per 64
257
+ # elements. An 8-element chunk advances the scale index by
258
+ # chunk_id * 8 // group_size.
259
+ scale0_u32 = Uint32(s_kscale[lane_id * 8 // self.group_size, compute_stage])
260
+ scale0_bf16x2 = (scale0_u32 << Uint32(23)) | (scale0_u32 << Uint32(7))
261
+ scale1_u32 = Uint32(
262
+ s_kscale[(lane_id + 32) * 8 // self.group_size, compute_stage]
263
+ )
264
+ scale1_bf16x2 = (scale1_u32 << Uint32(23)) | (scale1_u32 << Uint32(7))
265
+
266
+ # cvt.rn.scaled::n2::ue8m0.bf16x2.e4m3x2 requires PTX 9.2
267
+ # (CUDA 13.2).
268
+ dequant0 = cute.make_rmem_tensor(4, Uint32)
269
+ dequant1 = cute.make_rmem_tensor(4, Uint32)
270
+ for j in cutlass.range_constexpr(2):
271
+ tmp0 = _fp8x4_to_bf16x4(data0[j])
272
+ tmp1 = _fp8x4_to_bf16x4(data1[j])
273
+
274
+ # BF16 multiply is safe because the scales are exact powers of 2.
275
+ dequant0[j * 2] = _bf16x2_mul(tmp0[0], scale0_bf16x2)
276
+ dequant1[j * 2] = _bf16x2_mul(tmp1[0], scale1_bf16x2)
277
+ dequant0[j * 2 + 1] = _bf16x2_mul(tmp0[1], scale0_bf16x2)
278
+ dequant1[j * 2 + 1] = _bf16x2_mul(tmp1[1], scale1_bf16x2)
279
+
280
+ # Last 64 elems are BF16 tail, corresponds to dequant1 of last
281
+ # 8 threads. We have 448 FP8 + 64 BF16 -> 28x 16B for FP8 +
282
+ # 8x 16B for BF16.
283
+ if lane_id + 32 >= self.fp8_dim // 8:
284
+ idx = self.fp8_dim // 16 + (lane_id + 32) - self.fp8_dim // 8
285
+ cute.copy(
286
+ cp16_atom,
287
+ s_kdata_16B_slice[(None, idx), compute_stage],
288
+ dequant1,
289
+ )
290
+
291
+ # Store two 16B BF16 chunks per lane: first half, then second half.
292
+ dst = out_slice[req_id, offset + i, (None, lane_id)]
293
+ cute.copy(cp16_atom, dequant0, cute.recast_tensor(dst, Uint32))
294
+
295
+ dst = out_slice[req_id, offset + i, (None, lane_id + 32)]
296
+ cute.copy(cp16_atom, dequant1, cute.recast_tensor(dst, Uint32))
297
+
298
+ compute_stage = (compute_stage + 1) % self.num_stages
299
+
300
+ @cache
301
+ @staticmethod
302
+ def compile(
303
+ fp8_dim: int = 448,
304
+ block_size: int = 64,
305
+ has_gather_lens: bool = True,
306
+ ):
307
+ num_reqs = cute.sym_int()
308
+ head_dim = DequantGatherKCacheKernel.head_dim
309
+ head_bytes = fp8_dim + (head_dim - fp8_dim) * 2 + 8
310
+
311
+ out = make_fake_tensor(BFloat16, (num_reqs, cute.sym_int(), head_dim), 16)
312
+ k_cache = cute.runtime.make_fake_tensor(
313
+ Uint8,
314
+ (cute.sym_int(), block_size, head_bytes),
315
+ stride=(cute.sym_int64(divisibility=32), head_bytes, 1),
316
+ assumed_align=32,
317
+ )
318
+ seq_lens = make_fake_tensor(Int32, (num_reqs,))
319
+ gather_lens = make_fake_tensor(Int32, (num_reqs,)) if has_gather_lens else None
320
+ block_table = make_fake_tensor(Int32, (num_reqs, cute.sym_int()))
321
+
322
+ kernel = DequantGatherKCacheKernel(fp8_dim, block_size)
323
+ stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
324
+ return cute.compile(
325
+ kernel,
326
+ out,
327
+ k_cache,
328
+ seq_lens,
329
+ gather_lens,
330
+ block_table,
331
+ Int32(0),
332
+ stream,
333
+ options="--enable-tvm-ffi",
334
+ )
vllm/v1/attention/ops/deepseek_v4_ops/fused_indexer_q_cutedsl.py
+8
-92
@@ -1,7 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-# once we have more CuteDSL kernels in vLLM, we can refactor small helper functions
4
-# to a separate file
3
from functools import cache
4
5
import cutlass
@@ -9,10 +7,16 @@ import cutlass.cute as cute
7
import torch
8
from cuda.bindings.driver import CUstream
9
from cutlass import BFloat16, Float32, Int64, Uint8, Uint32, const_expr
12
-from cutlass._mlir.dialects import llvm
13
-from cutlass.cutlass_dsl import T, dsl_user_op
10
from quack.compile_utils import make_fake_tensor
11
12
+from vllm.v1.attention.ops.deepseek_v4_ops.cutedsl_utils import (
13
+ _bf16x2_abs,
14
+ _bf16x2_max,
15
+ _bf16x2_to_fp32,
16
+ _fp32x2_to_bf16x2,
17
+ _fp32x8_to_fp4x8,
18
+ _recast_val,
19
+)
20
from vllm.vllm_flash_attn.cute import utils as cute_utils
21
22
# MXFP4: 32 elements per block, packed 2 nibbles per byte, ue8m0 block scale.
@@ -61,94 +65,6 @@ def fused_indexer_q_rope_quant_mxfp4_cutedsl(
65
)
66
67
64
-@dsl_user_op
65
-def _recast_val(x, dtype, *, loc=None, ip=None):
66
- return dtype(llvm.bitcast(dtype.mlir_type, x.ir_value(loc=loc, ip=ip)))
67
-
68
-
69
-@dsl_user_op
70
-def _fp32x2_to_bf16x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32:
71
- out = llvm.inline_asm(
72
- T.i32(),
73
- [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
74
- "cvt.rn.bf16x2.f32 $0, $2, $1;",
75
- "=r,f,f",
76
- has_side_effects=False,
77
- is_align_stack=False,
78
- )
79
- return Uint32(out)
80
-
81
-
82
-@dsl_user_op
83
-def _bf16x2_to_fp32(data: Uint32, *, loc=None, ip=None) -> tuple[Float32, Float32]:
84
- out = llvm.inline_asm(
85
- llvm.StructType.get_literal([T.f32(), T.f32()]),
86
- [data.ir_value(loc=loc, ip=ip)],
87
- "shl.b32 $0, $2, 16;\n\tand.b32 $1, $2, 0xFFFF0000;\n",
88
- "=f,=f,r",
89
- has_side_effects=False,
90
- is_align_stack=False,
91
- )
92
- return (
93
- Float32(llvm.extractvalue(T.f32(), out, [0], loc=loc, ip=ip)),
94
- Float32(llvm.extractvalue(T.f32(), out, [1], loc=loc, ip=ip)),
95
- )
96
-
97
-
98
-@dsl_user_op
99
-def _bf16x2_abs(a: Uint32, *, loc=None, ip=None) -> Uint32:
100
- out = llvm.inline_asm(
101
- T.i32(),
102
- [a.ir_value(loc=loc, ip=ip)],
103
- "abs.bf16x2 $0, $1;",
104
- "=r,r",
105
- has_side_effects=False,
106
- is_align_stack=False,
107
- )
108
- return Uint32(out)
109
-
110
-
111
-@dsl_user_op
112
-def _bf16x2_max(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
113
- out = llvm.inline_asm(
114
- T.i32(),
115
- [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
116
- "max.bf16x2 $0, $1, $2;",
117
- "=r,r,r",
118
- has_side_effects=False,
119
- is_align_stack=False,
120
- )
121
- return Uint32(out)
122
-
123
-
124
-@dsl_user_op
125
-def _fp32x8_to_fp4x8(
126
- vals: cute.Tensor,
127
- offset: cutlass.Constexpr[int],
128
- *,
129
- loc=None,
130
- ip=None,
131
-) -> Uint32:
132
- # Pack eight scaled FP32 values into four E2M1x2 bytes, returned as one b32.
133
- assert vals.element_type is Float32
134
- out = llvm.inline_asm(
135
- T.i32(),
136
- [vals[offset + i].ir_value(loc=loc, ip=ip) for i in range(8)],
137
- "{\n\t"
138
- ".reg .b8 x0, x1, x2, x3;\n\t"
139
- "cvt.rn.satfinite.e2m1x2.f32 x0, $2, $1;\n\t"
140
- "cvt.rn.satfinite.e2m1x2.f32 x1, $4, $3;\n\t"
141
- "cvt.rn.satfinite.e2m1x2.f32 x2, $6, $5;\n\t"
142
- "cvt.rn.satfinite.e2m1x2.f32 x3, $8, $7;\n\t"
143
- "mov.b32 $0, {x0, x1, x2, x3};\n\t"
144
- "}\n",
145
- "=r,f,f,f,f,f,f,f,f",
146
- has_side_effects=False,
147
- is_align_stack=False,
148
- )
149
- return Uint32(out)
150
-
151
-
68
class IndexerQMxFp4Kernel:
69
"""Eight-thread subwarps process one ``(token, head)`` row."""
70