@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a37e47100c
Fork 0

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, 20:45 UTC

Add CuTe DSL sparse compressor support (#43584)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>

Jie Fang committed May 26, 2026 at 15:11 UTC a37e47100ca1f941671a4418dbdfd7a5afca4811
3 files changed +1419 -39
vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py
+38
@@ -11,6 +11,12 @@ 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 +
20 RoPE is register-based via tl.reshape -> tl.split -> tl.interleave (or the
21 even/odd halves are consumed directly for MXFP4, no interleave needed).
22 FP8 UE8M0 quant uses tl.reshape to tile [N_QUANT_BLOCKS, QUANT_BLOCK] for
@@ -19,11 +25,43 @@ even/odd halves, producing (N_QUANT_BLOCKS, MXFP4_BLOCK/2) packed nibbles
25 and N_QUANT_BLOCKS ue8m0 bytes.
26 """
27
28 +from functools import cache
29 +
30 from vllm.triton_utils import tl, triton
31
32 from .fused_indexer_q import _fp32x2_to_fp4x2
33
34
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,
47 + )
48 +
49 +
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 +
65 # =============================================================================
66 # DeepseekV4 Attention path (head=512, nope=448 FP8 + rope=64 bf16)
67 # =============================================================================
vllm/models/deepseek_v4/common/ops/sparse_attn_compress_cutedsl.py new
+1240
@@ -0,0 +1,1240 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""CuTe DSL sparse-attention compressor for DeepSeek V4.
4 +
5 +The public wrappers provide the C4 fused and C128 split kernels.
6 +"""
7 +
8 +from __future__ import annotations
9 +
10 +from functools import cache
11 +
12 +import cutlass
13 +import cutlass.cute as cute
14 +import torch
15 +from cuda.bindings.driver import CUstream
16 +from cutlass import BFloat16, Float32, Int32, Int64, Uint8, Uint16, Uint32, const_expr
17 +from cutlass._mlir.dialects import llvm
18 +from cutlass.cutlass_dsl import T, dsl_user_op
19 +from quack.compile_utils import make_fake_tensor
20 +
21 +_TORCH_TO_CUTE = {
22 + torch.bfloat16: BFloat16,
23 + torch.float32: Float32,
24 +}
25 +
26 +
27 +@dsl_user_op
28 +def _recast_val(x, dtype, *, loc=None, ip=None):
29 + return dtype(llvm.bitcast(dtype.mlir_type, x.ir_value(loc=loc, ip=ip)))
30 +
31 +
32 +@dsl_user_op
33 +def _fp32x2_to_bf16x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32:
34 + out = llvm.inline_asm(
35 + T.i32(),
36 + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
37 + "cvt.rn.bf16x2.f32 $0, $2, $1;",
38 + "=r,f,f",
39 + has_side_effects=False,
40 + is_align_stack=False,
41 + )
42 + return Uint32(out)
43 +
44 +
45 +@dsl_user_op
46 +def _bf16x2_to_fp32(data: Uint32, *, loc=None, ip=None) -> tuple[Float32, Float32]:
47 + out = llvm.inline_asm(
48 + llvm.StructType.get_literal([T.f32(), T.f32()]),
49 + [data.ir_value(loc=loc, ip=ip)],
50 + "shl.b32 $0, $2, 16;\n\tand.b32 $1, $2, 0xFFFF0000;\n",
51 + "=f,=f,r",
52 + has_side_effects=False,
53 + is_align_stack=False,
54 + )
55 + return (
56 + Float32(llvm.extractvalue(T.f32(), out, [0], loc=loc, ip=ip)),
57 + Float32(llvm.extractvalue(T.f32(), out, [1], loc=loc, ip=ip)),
58 + )
59 +
60 +
61 +@dsl_user_op
62 +def _fp32x2_to_fp8e4m3x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint16:
63 + out = llvm.inline_asm(
64 + T.i16(),
65 + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
66 + "cvt.rn.satfinite.e4m3x2.f32 $0, $2, $1;",
67 + "=h,f,f",
68 + has_side_effects=False,
69 + is_align_stack=False,
70 + )
71 + return Uint16(out)
72 +
73 +
74 +class SparseAttnCompressNormRopeStoreC4Kernel:
75 + min_scale = 1.0e-4
76 + rcp_ln2 = 1.4426950408889634
77 +
78 + def __init__(
79 + self,
80 + head_size: int,
81 + state_width: int,
82 + rope_head_dim: int,
83 + fp8_max: float,
84 + quant_block: int,
85 + token_stride: int,
86 + scale_dim: int,
87 + compress_ratio: int,
88 + overlap: bool,
89 + ):
90 + self.head_dim = head_size
91 + self.state_width = state_width
92 + self.rope_dim = rope_head_dim
93 + self.nope_dim = head_size - rope_head_dim
94 + self.fp8_max = fp8_max
95 + self.quant_block = quant_block
96 + self.token_stride = token_stride
97 + self.scale_dim = scale_dim
98 + self.num_warps = head_size // quant_block
99 + self.nope_blocks = self.nope_dim // quant_block
100 + self.tb_size = head_size // 2
101 + self.compress_ratio = compress_ratio
102 + self.overlap = overlap
103 + self.window = (1 + int(overlap)) * compress_ratio
104 +
105 + @cute.jit
106 + def __call__(
107 + self,
108 + state_cache: cute.Tensor,
109 + token_to_req_indices: cute.Tensor,
110 + positions: cute.Tensor,
111 + slot_mapping: cute.Tensor,
112 + block_table: cute.Tensor,
113 + block_size: Int64,
114 + rms_norm_weight: cute.Tensor,
115 + rms_norm_eps: Float32,
116 + cos_sin_cache: cute.Tensor,
117 + k_cache: cute.Tensor,
118 + kv_slot_mapping: cute.Tensor,
119 + kv_cache_block_size: Int64,
120 + stream: CUstream,
121 + ):
122 + grid = (slot_mapping.shape[0], 1, 1)
123 + self.kernel(
124 + state_cache,
125 + token_to_req_indices,
126 + positions,
127 + slot_mapping,
128 + block_table,
129 + block_size,
130 + rms_norm_weight,
131 + rms_norm_eps,
132 + cos_sin_cache,
133 + k_cache,
134 + kv_slot_mapping,
135 + kv_cache_block_size,
136 + ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
137 +
138 + @cute.kernel
139 + def kernel(
140 + self,
141 + state_cache: cute.Tensor,
142 + token_to_req_indices: cute.Tensor,
143 + positions: cute.Tensor,
144 + slot_mapping: cute.Tensor,
145 + block_table: cute.Tensor,
146 + block_size: Int64,
147 + rms_norm_weight: cute.Tensor,
148 + rms_norm_eps: Float32,
149 + cos_sin_cache: cute.Tensor,
150 + k_cache: cute.Tensor,
151 + kv_slot_mapping: cute.Tensor,
152 + kv_cache_block_size: Int64,
153 + ):
154 + token_idx, _, _ = cute.arch.block_idx()
155 + tid, _, _ = cute.arch.thread_idx()
156 + warp_id = cute.arch.make_warp_uniform(tid // 32)
157 + lane_id = tid % 32
158 + elem0 = tid * 2
159 + elem1 = elem0 + 1
160 +
161 + slot_id = slot_mapping[token_idx]
162 + has_position = token_idx < positions.shape[0]
163 + position = Int64(0)
164 + if has_position:
165 + position = positions[token_idx]
166 + boundary = has_position and (
167 + (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
168 + )
169 + has_req_idx = token_idx < token_to_req_indices.shape[0]
170 + has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
171 + kv_slot_idx = Int64(-1)
172 + if has_kv_slot_idx:
173 + kv_slot_idx = kv_slot_mapping[token_idx]
174 + active = (
175 + slot_id >= Int64(0) and has_req_idx and boundary and kv_slot_idx >= Int64(0)
176 + )
177 +
178 + if active:
179 + req_idx = token_to_req_indices[token_idx]
180 + start = position - Int64(self.window - 1)
181 +
182 + smem = cutlass.utils.SmemAllocator()
183 + s_block_numbers = smem.allocate_tensor(
184 + Int32, cute.make_layout((self.window,)), byte_alignment=4
185 + )
186 + partial_sums = smem.allocate_tensor(
187 + Float32, cute.make_layout((self.num_warps,)), byte_alignment=4
188 + )
189 + rrms_shared = smem.allocate_tensor(
190 + Float32, cute.make_layout((1,)), byte_alignment=4
191 + )
192 +
193 + for row in cutlass.range_constexpr(self.window):
194 + pos = start + Int64(row)
195 + if tid == row:
196 + block_number_i32 = Int32(0)
197 + if pos >= Int64(0):
198 + block_index = pos // block_size
199 + block_number_i32 = block_table[req_idx, block_index]
200 + s_block_numbers[row] = block_number_i32
201 + cute.arch.sync_threads()
202 +
203 + max0 = -Float32.inf
204 + max1 = -Float32.inf
205 + sum0 = Float32(0.0)
206 + sum1 = Float32(0.0)
207 + product0 = Float32(0.0)
208 + product1 = Float32(0.0)
209 +
210 + for row in cutlass.range_constexpr(self.window):
211 + pos = start + Int64(row)
212 + if pos >= Int64(0):
213 + block_index = pos // block_size
214 + block_offset = pos - block_index * block_size
215 + block_number = s_block_numbers[row].to(Int64)
216 + head_offset = Int64((row // self.compress_ratio) * self.head_dim)
217 + row_base = (
218 + block_number * state_cache.stride[0]
219 + + block_offset * state_cache.stride[1]
220 + + head_offset
221 + )
222 +
223 + score0 = state_cache.iterator[
224 + row_base + Int64(self.state_width) + elem0.to(Int64)
225 + ]
226 + kv0 = state_cache.iterator[row_base + elem0.to(Int64)]
227 + new_max0 = cute.arch.fmax(max0, score0)
228 + old_scale0 = cute.math.exp2(
229 + (max0 - new_max0) * Float32(self.rcp_ln2), fastmath=True
230 + )
231 + new_scale0 = cute.math.exp2(
232 + (score0 - new_max0) * Float32(self.rcp_ln2), fastmath=True
233 + )
234 + sum0 = sum0 * old_scale0 + new_scale0
235 + product0 = product0 * old_scale0 + kv0 * new_scale0
236 + max0 = new_max0
237 +
238 + score1 = state_cache.iterator[
239 + row_base + Int64(self.state_width) + elem1.to(Int64)
240 + ]
241 + kv1 = state_cache.iterator[row_base + elem1.to(Int64)]
242 + new_max1 = cute.arch.fmax(max1, score1)
243 + old_scale1 = cute.math.exp2(
244 + (max1 - new_max1) * Float32(self.rcp_ln2), fastmath=True
245 + )
246 + new_scale1 = cute.math.exp2(
247 + (score1 - new_max1) * Float32(self.rcp_ln2), fastmath=True
248 + )
249 + sum1 = sum1 * old_scale1 + new_scale1
250 + product1 = product1 * old_scale1 + kv1 * new_scale1
251 + max1 = new_max1
252 +
253 + x0 = product0 / sum0
254 + x1 = product1 / sum1
255 +
256 + local_sumsq = x0 * x0 + x1 * x1
257 + warp_sum = local_sumsq
258 + for step in cutlass.range_constexpr(5):
259 + offset = const_expr(16 >> step)
260 + warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
261 +
262 + if lane_id == 0:
263 + partial_sums[warp_id] = warp_sum
264 + cute.arch.sync_threads()
265 + if tid == 0:
266 + total = Float32(0.0)
267 + for i in cutlass.range_constexpr(self.num_warps):
268 + total += partial_sums[i]
269 + rrms_shared[0] = cute.math.rsqrt(
270 + total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
271 + )
272 + cute.arch.sync_threads()
273 +
274 + rrms = rrms_shared[0]
275 + x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32)
276 + x1 = x1 * rrms * rms_norm_weight[elem1].to(Float32)
277 +
278 + k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
279 + k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
280 + page = kv_slot_idx // kv_cache_block_size
281 + kv_offset = kv_slot_idx - page * kv_cache_block_size
282 + value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride)
283 + scale_base = (
284 + page * k_cache.stride[0]
285 + + kv_cache_block_size * Int64(self.token_stride)
286 + + kv_offset * Int64(self.scale_dim)
287 + )
288 +
289 + if warp_id == self.nope_blocks:
290 + pair_idx = lane_id
291 + compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
292 + self.compress_ratio
293 + )
294 + cos_v = cos_sin_cache[compressed_pos, pair_idx]
295 + sin_v = cos_sin_cache[
296 + compressed_pos, pair_idx + Int32(self.rope_dim // 2)
297 + ]
298 + real = x0 * cos_v - x1 * sin_v
299 + imag = x0 * sin_v + x1 * cos_v
300 + packed = _fp32x2_to_bf16x2(real, imag)
301 + out_base = value_base + Int64(self.nope_dim) + (lane_id * 4).to(Int64)
302 + k_cache_u32.iterator[out_base // Int64(4)] = packed
303 + else:
304 + q_packed = _fp32x2_to_bf16x2(x0, x1)
305 + q0, q1 = _bf16x2_to_fp32(q_packed)
306 + abs0 = cute.math.absf(q0)
307 + abs1 = cute.math.absf(q1)
308 + local_absmax = cute.arch.fmax(abs0, abs1)
309 + absmax = local_absmax
310 + for step in cutlass.range_constexpr(5):
311 + offset = const_expr(16 >> step)
312 + absmax = cute.arch.fmax(
313 + absmax, cute.arch.shuffle_sync_bfly(absmax, offset)
314 + )
315 + scale_raw = cute.arch.fmax(
316 + Float32(self.min_scale),
317 + absmax / Float32(self.fp8_max),
318 + )
319 + bits = _recast_val(scale_raw, Uint32)
320 + ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
321 + inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32)
322 + y0 = cute.arch.fmin(
323 + cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)),
324 + Float32(self.fp8_max),
325 + )
326 + y1 = cute.arch.fmin(
327 + cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)),
328 + Float32(self.fp8_max),
329 + )
330 + packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
331 + out_base = value_base + (warp_id * self.quant_block + lane_id * 2).to(
332 + Int64
333 + )
334 + k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
335 + if lane_id == 0:
336 + k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to(Uint8)
337 + if warp_id == 0:
338 + k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8(
339 + 0
340 + )
341 +
342 + @cache
343 + @staticmethod
344 + def compile(
345 + head_size: int = 512,
346 + state_width: int = 1024,
347 + rope_head_dim: int = 64,
348 + fp8_max: float = 448.0,
349 + quant_block: int = 64,
350 + token_stride: int = 576,
351 + scale_dim: int = 8,
352 + kv_block_stride: int = 74752,
353 + compress_ratio: int = 4,
354 + overlap: bool = True,
355 + norm_weight_dtype: type[cutlass.Numeric] = Float32,
356 + ):
357 + if compress_ratio != 4 or not overlap:
358 + raise ValueError("CuTe DSL C4 fused sparse-attn requires C4 overlap.")
359 + if head_size != 512:
360 + raise ValueError(
361 + "CuTe DSL C4 fused sparse-attn currently requires head_size=512."
362 + )
363 + if state_width != 2 * head_size:
364 + raise ValueError(
365 + "CuTe DSL C4 fused sparse-attn requires state_width=2*head_size."
366 + )
367 + if quant_block != 64:
368 + raise ValueError(
369 + "CuTe DSL C4 fused sparse-attn currently requires quant_block=64."
370 + )
371 + if rope_head_dim != 64:
372 + raise ValueError(
373 + "CuTe DSL C4 fused sparse-attn currently requires rope_head_dim=64."
374 + )
375 + if token_stride < head_size + rope_head_dim:
376 + raise ValueError("token_stride is too small for the packed FP8/BF16 row.")
377 + expected_scale_dim = (head_size - rope_head_dim) // quant_block + 1
378 + if scale_dim < expected_scale_dim:
379 + raise ValueError("scale_dim is too small for the UE8M0 scale row.")
380 +
381 + num_positions = cute.sym_int()
382 + num_slots = cute.sym_int()
383 + num_req_indices = cute.sym_int()
384 + num_kv_slots = cute.sym_int()
385 + num_state_blocks = cute.sym_int()
386 + num_kv_blocks = cute.sym_int()
387 + state_cache_block_size = cute.sym_int()
388 + block_table_width = cute.sym_int()
389 + max_pos = cute.sym_int()
390 + state_cache_width = state_width * 2
391 +
392 + state_cache = cute.runtime.make_fake_tensor(
393 + Float32,
394 + (num_state_blocks, state_cache_block_size, state_cache_width),
395 + stride=(
396 + cute.sym_int64(divisibility=16),
397 + cute.sym_int64(divisibility=16),
398 + 1,
399 + ),
400 + assumed_align=16,
401 + )
402 + token_to_req_indices = make_fake_tensor(
403 + Int32, (num_req_indices,), divisibility=4
404 + )
405 + positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
406 + slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
407 + block_table = make_fake_tensor(
408 + Int32, (cute.sym_int(), block_table_width), divisibility=1
409 + )
410 + rms_norm_weight = make_fake_tensor(
411 + norm_weight_dtype, (head_size,), divisibility=4
412 + )
413 + cos_sin_cache = cute.runtime.make_fake_tensor(
414 + Float32,
415 + (max_pos, rope_head_dim),
416 + stride=(cute.sym_int64(divisibility=4), 1),
417 + assumed_align=4,
418 + )
419 + k_cache = cute.runtime.make_fake_tensor(
420 + Uint8,
421 + (num_kv_blocks, cute.sym_int(), cute.sym_int()),
422 + stride=(
423 + cute.sym_int64(divisibility=16),
424 + cute.sym_int64(divisibility=8),
425 + 1,
426 + ),
427 + assumed_align=16,
428 + )
429 + kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8)
430 +
431 + kernel = SparseAttnCompressNormRopeStoreC4Kernel(
432 + head_size,
433 + state_width,
434 + rope_head_dim,
435 + fp8_max,
436 + quant_block,
437 + token_stride,
438 + scale_dim,
439 + compress_ratio,
440 + overlap,
441 + )
442 + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
443 + return cute.compile(
444 + kernel,
445 + state_cache,
446 + token_to_req_indices,
447 + positions,
448 + slot_mapping,
449 + block_table,
450 + Int64(0),
451 + rms_norm_weight,
452 + Float32(0.0),
453 + cos_sin_cache,
454 + k_cache,
455 + kv_slot_mapping,
456 + Int64(0),
457 + stream,
458 + options="--enable-tvm-ffi",
459 + )
460 +
461 +
462 +class SparseAttnCompressKernel:
463 + head_tile = 64
464 + rows_per_warp = 8
465 + row_pairs_per_warp = rows_per_warp // 2
466 + elems_per_lane = 4
467 + lanes_per_row = head_tile // elems_per_lane
468 + num_warps = 16
469 + stats_warp_stride = num_warps + 1
470 + tb_size = num_warps * 32
471 + rcp_ln2 = 1.4426950408889634
472 +
473 + def __init__(
474 + self,
475 + head_size: int,
476 + state_width: int,
477 + compress_ratio: int,
478 + overlap: bool,
479 + ):
480 + self.head_dim = head_size
481 + self.num_splits = head_size // self.head_tile
482 + self.state_width = state_width
483 + self.compress_ratio = compress_ratio
484 + self.overlap = overlap
485 + self.window = (1 + int(overlap)) * compress_ratio
486 +
487 + @cute.jit
488 + def __call__(
489 + self,
490 + state_cache: cute.Tensor,
491 + token_to_req_indices: cute.Tensor,
492 + positions: cute.Tensor,
493 + slot_mapping: cute.Tensor,
494 + block_table: cute.Tensor,
495 + block_size: Int64,
496 + compressed_kv: cute.Tensor,
497 + stream: CUstream,
498 + ):
499 + grid = (slot_mapping.shape[0] * self.num_splits, 1, 1)
500 + self.kernel(
501 + state_cache,
502 + token_to_req_indices,
503 + positions,
504 + slot_mapping,
505 + block_table,
506 + block_size,
507 + compressed_kv,
508 + ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
509 +
510 + @cute.kernel
511 + def kernel(
512 + self,
513 + state_cache: cute.Tensor,
514 + token_to_req_indices: cute.Tensor,
515 + positions: cute.Tensor,
516 + slot_mapping: cute.Tensor,
517 + block_table: cute.Tensor,
518 + block_size: Int64,
519 + compressed_kv: cute.Tensor,
520 + ):
521 + block_id, _, _ = cute.arch.block_idx()
522 + tid, _, _ = cute.arch.thread_idx()
523 + warp_id = cute.arch.make_warp_uniform(tid // 32)
524 + lane_id = tid % 32
525 + row_lane = lane_id // self.lanes_per_row
526 + col_group = lane_id % self.lanes_per_row
527 +
528 + token_idx = block_id // self.num_splits
529 + split_idx = block_id - token_idx * self.num_splits
530 + col_base = split_idx * self.head_tile + col_group * self.elems_per_lane
531 +
532 + slot_id = slot_mapping[token_idx]
533 + has_position = token_idx < positions.shape[0]
534 + position = Int64(0)
535 + if has_position:
536 + position = positions[token_idx]
537 + boundary = has_position and (
538 + (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
539 + )
540 + has_req_idx = token_idx < token_to_req_indices.shape[0]
541 + active = slot_id >= Int64(0) and has_req_idx and boundary
542 +
543 + if active:
544 + smem = cutlass.utils.SmemAllocator()
545 + s_max = smem.allocate_tensor(
546 + Float32,
547 + cute.make_layout(
548 + (
549 + self.lanes_per_row,
550 + self.elems_per_lane,
551 + self.stats_warp_stride,
552 + ),
553 + stride=(
554 + self.elems_per_lane * self.stats_warp_stride,
555 + self.stats_warp_stride,
556 + 1,
557 + ),
558 + ),
559 + byte_alignment=4,
560 + )
561 + s_sum = smem.allocate_tensor(
562 + Float32,
563 + cute.make_layout(
564 + (
565 + self.lanes_per_row,
566 + self.elems_per_lane,
567 + self.stats_warp_stride,
568 + ),
569 + stride=(
570 + self.elems_per_lane * self.stats_warp_stride,
571 + self.stats_warp_stride,
572 + 1,
573 + ),
574 + ),
575 + byte_alignment=4,
576 + )
577 + s_product = smem.allocate_tensor(
578 + Float32,
579 + cute.make_layout(
580 + (
581 + self.lanes_per_row,
582 + self.elems_per_lane,
583 + self.stats_warp_stride,
584 + ),
585 + stride=(
586 + self.elems_per_lane * self.stats_warp_stride,
587 + self.stats_warp_stride,
588 + 1,
589 + ),
590 + ),
591 + byte_alignment=4,
592 + )
593 +
594 + row_pair_layout = cute.make_layout(
595 + (self.row_pairs_per_warp, self.elems_per_lane),
596 + stride=(self.elems_per_lane, 1),
597 + )
598 + kv_vals = cute.make_rmem_tensor(row_pair_layout, Float32)
599 + score_vals = cute.make_rmem_tensor(row_pair_layout, Float32)
600 + local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
601 + local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
602 + local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
603 +
604 + for e in cutlass.range_constexpr(self.elems_per_lane):
605 + local_max[e] = -Float32.inf
606 + local_sum[e] = Float32(0.0)
607 + local_product[e] = Float32(0.0)
608 +
609 + req_idx = token_to_req_indices[token_idx]
610 + start = position - Int64(self.window - 1)
611 + cp_f32x4 = cute.make_copy_atom(
612 + cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128
613 + )
614 + row_mask_and_clamp = const_expr(
615 + (cute.arch.WARP_SIZE - self.lanes_per_row) << 8
616 + | (cute.arch.WARP_SIZE - 1)
617 + )
618 +
619 + for i in cutlass.range_constexpr(self.row_pairs_per_warp):
620 + row = warp_id * self.rows_per_warp + i * 2 + row_lane
621 + pos = start + row.to(Int64)
622 + valid = row < self.window and pos >= Int64(0)
623 + head_offset = ((row // self.compress_ratio) * self.head_dim).to(Int64)
624 +
625 + for e in cutlass.range_constexpr(self.elems_per_lane):
626 + kv = Float32(0.0)
627 + score = -Float32.inf
628 + kv_vals[i, e] = kv
629 + score_vals[i, e] = score
630 +
631 + block_index = Int64(0)
632 + block_offset = Int64(0)
633 + block_number_i32 = Int32(0)
634 + if valid:
635 + block_index = pos // block_size
636 + block_offset = pos - block_index * block_size
637 + if col_group == 0:
638 + block_number_i32 = block_table[req_idx, block_index]
639 + block_number_i32 = cute.arch.shuffle_sync(
640 + block_number_i32,
641 + offset=0,
642 + mask_and_clamp=row_mask_and_clamp,
643 + )
644 +
645 + if valid:
646 + block_number = block_number_i32.to(Int64)
647 + row_tensor = state_cache[block_number, block_offset, None]
648 + col_tile = (head_offset + col_base.to(Int64)) // Int64(
649 + self.elems_per_lane
650 + )
651 + kv_src = cute.local_tile(
652 + row_tensor,
653 + tiler=(self.elems_per_lane,),
654 + coord=(col_tile,),
655 + )
656 + score_src = cute.local_tile(
657 + row_tensor,
658 + tiler=(self.elems_per_lane,),
659 + coord=(
660 + col_tile + Int64(self.state_width // self.elems_per_lane),
661 + ),
662 + )
663 + cute.copy(cp_f32x4, kv_src, kv_vals[i, None])
664 + cute.copy(cp_f32x4, score_src, score_vals[i, None])
665 +
666 + for e in cutlass.range_constexpr(self.elems_per_lane):
667 + local_max[e] = cute.arch.fmax(local_max[e], score_vals[i, e])
668 +
669 + for e in cutlass.range_constexpr(self.elems_per_lane):
670 + if local_max[e] > -Float32.inf:
671 + for i in cutlass.range_constexpr(self.row_pairs_per_warp):
672 + exp_score = cute.math.exp2(
673 + (score_vals[i, e] - local_max[e]) * Float32(self.rcp_ln2),
674 + fastmath=True,
675 + )
676 + local_sum[e] += exp_score
677 + local_product[e] += kv_vals[i, e] * exp_score
678 +
679 + for e in cutlass.range_constexpr(self.elems_per_lane):
680 + pair_max = cute.arch.shuffle_sync_bfly(local_max[e], offset=16)
681 + pair_sum = cute.arch.shuffle_sync_bfly(local_sum[e], offset=16)
682 + pair_product = cute.arch.shuffle_sync_bfly(local_product[e], offset=16)
683 + warp_max = cute.arch.fmax(local_max[e], pair_max)
684 + warp_sum = Float32(0.0)
685 + warp_product = Float32(0.0)
686 + if warp_max > -Float32.inf:
687 + local_scale = cute.math.exp2(
688 + (local_max[e] - warp_max) * Float32(self.rcp_ln2),
689 + fastmath=True,
690 + )
691 + pair_scale = cute.math.exp2(
692 + (pair_max - warp_max) * Float32(self.rcp_ln2),
693 + fastmath=True,
694 + )
695 + warp_sum = local_sum[e] * local_scale + pair_sum * pair_scale
696 + warp_product = (
697 + local_product[e] * local_scale + pair_product * pair_scale
698 + )
699 + if lane_id < self.lanes_per_row:
700 + s_max[col_group, e, warp_id] = warp_max
701 + s_sum[col_group, e, warp_id] = warp_sum
702 + s_product[col_group, e, warp_id] = warp_product
703 + cute.arch.sync_threads()
704 +
705 + out_group = tid // self.num_warps
706 + final_lane = tid % self.num_warps
707 + final_groups_per_pass = const_expr(self.tb_size // self.num_warps)
708 + for pass_idx in cutlass.range_constexpr(
709 + self.head_tile // final_groups_per_pass
710 + ):
711 + out_idx = pass_idx * final_groups_per_pass + out_group
712 + out_lane = out_idx // self.elems_per_lane
713 + out_elem = out_idx % self.elems_per_lane
714 +
715 + local_warp_max = s_max[out_lane, out_elem, final_lane]
716 + global_max = local_warp_max
717 + for step in cutlass.range_constexpr(4):
718 + offset = const_expr(8 >> step)
719 + global_max = cute.arch.fmax(
720 + global_max,
721 + cute.arch.shuffle_sync_bfly(
722 + global_max,
723 + offset=offset,
724 + mask_and_clamp=row_mask_and_clamp,
725 + ),
726 + )
727 +
728 + scale = cute.math.exp2(
729 + (local_warp_max - global_max) * Float32(self.rcp_ln2),
730 + fastmath=True,
731 + )
732 + global_sum = s_sum[out_lane, out_elem, final_lane] * scale
733 + global_product = s_product[out_lane, out_elem, final_lane] * scale
734 + for step in cutlass.range_constexpr(4):
735 + offset = const_expr(8 >> step)
736 + global_sum += cute.arch.shuffle_sync_bfly(
737 + global_sum,
738 + offset=offset,
739 + mask_and_clamp=row_mask_and_clamp,
740 + )
741 + global_product += cute.arch.shuffle_sync_bfly(
742 + global_product,
743 + offset=offset,
744 + mask_and_clamp=row_mask_and_clamp,
745 + )
746 +
747 + if final_lane == 0:
748 + compressed_kv.iterator[
749 + token_idx.to(Int64) * compressed_kv.stride[0]
750 + + (split_idx * self.head_tile + out_idx).to(Int64)
751 + ] = global_product / global_sum
752 +
753 + @cache
754 + @staticmethod
755 + def compile(
756 + head_size: int = 512,
757 + state_width: int = 512,
758 + compress_ratio: int = 128,
759 + overlap: bool = False,
760 + ):
761 + if head_size % SparseAttnCompressKernel.head_tile != 0:
762 + raise ValueError("head_size must be divisible by the 64-wide head tile.")
763 + num_positions = cute.sym_int()
764 + num_slots = cute.sym_int()
765 + num_req_indices = cute.sym_int()
766 + num_blocks = cute.sym_int()
767 + state_cache_block_size = cute.sym_int()
768 + block_table_width = cute.sym_int()
769 + state_cache_width = state_width * 2
770 +
771 + state_cache = cute.runtime.make_fake_tensor(
772 + Float32,
773 + (num_blocks, state_cache_block_size, state_cache_width),
774 + stride=(
775 + cute.sym_int64(divisibility=16),
776 + cute.sym_int64(divisibility=16),
777 + 1,
778 + ),
779 + assumed_align=16,
780 + )
781 + token_to_req_indices = make_fake_tensor(
782 + Int32, (num_req_indices,), divisibility=4
783 + )
784 + positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
785 + slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
786 + block_table = make_fake_tensor(
787 + Int32, (cute.sym_int(), block_table_width), divisibility=1
788 + )
789 + compressed_kv = cute.runtime.make_fake_tensor(
790 + Float32,
791 + (num_slots, head_size),
792 + stride=(cute.sym_int64(divisibility=4), 1),
793 + assumed_align=4,
794 + )
795 +
796 + kernel = SparseAttnCompressKernel(
797 + head_size,
798 + state_width,
799 + compress_ratio,
800 + overlap,
801 + )
802 + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
803 + return cute.compile(
804 + kernel,
805 + state_cache,
806 + token_to_req_indices,
807 + positions,
808 + slot_mapping,
809 + block_table,
810 + Int64(0),
811 + compressed_kv,
812 + stream,
813 + options="--enable-tvm-ffi",
814 + )
815 +
816 +
817 +class SparseAttnNormRopeStoreKernel:
818 + min_scale = 1.0e-4
819 +
820 + def __init__(
821 + self,
822 + head_size: int,
823 + rope_head_dim: int,
824 + fp8_max: float,
825 + quant_block: int,
826 + token_stride: int,
827 + scale_dim: int,
828 + compress_ratio: int,
829 + ):
830 + self.head_dim = head_size
831 + self.rope_dim = rope_head_dim
832 + self.nope_dim = head_size - rope_head_dim
833 + self.fp8_max = fp8_max
834 + self.quant_block = quant_block
835 + self.token_stride = token_stride
836 + self.scale_dim = scale_dim
837 + self.num_warps = head_size // quant_block
838 + self.nope_blocks = self.nope_dim // quant_block
839 + self.tb_size = head_size // 2
840 + self.compress_ratio = compress_ratio
841 +
842 + @cute.jit
843 + def __call__(
844 + self,
845 + compressed_kv: cute.Tensor,
846 + positions: cute.Tensor,
847 + slot_mapping: cute.Tensor,
848 + rms_norm_weight: cute.Tensor,
849 + rms_norm_eps: Float32,
850 + cos_sin_cache: cute.Tensor,
851 + k_cache: cute.Tensor,
852 + kv_slot_mapping: cute.Tensor,
853 + kv_cache_block_size: Int64,
854 + stream: CUstream,
855 + ):
856 + grid = (slot_mapping.shape[0], 1, 1)
857 + self.kernel(
858 + compressed_kv,
859 + positions,
860 + slot_mapping,
861 + rms_norm_weight,
862 + rms_norm_eps,
863 + cos_sin_cache,
864 + k_cache,
865 + kv_slot_mapping,
866 + kv_cache_block_size,
867 + ).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
868 +
869 + @cute.kernel
870 + def kernel(
871 + self,
872 + compressed_kv: cute.Tensor,
873 + positions: cute.Tensor,
874 + slot_mapping: cute.Tensor,
875 + rms_norm_weight: cute.Tensor,
876 + rms_norm_eps: Float32,
877 + cos_sin_cache: cute.Tensor,
878 + k_cache: cute.Tensor,
879 + kv_slot_mapping: cute.Tensor,
880 + kv_cache_block_size: Int64,
881 + ):
882 + token_idx, _, _ = cute.arch.block_idx()
883 + tid, _, _ = cute.arch.thread_idx()
884 + warp_id = cute.arch.make_warp_uniform(tid // 32)
885 + lane_id = tid % 32
886 + elem0 = tid * 2
887 +
888 + slot_id = slot_mapping[token_idx]
889 + has_position = token_idx < positions.shape[0]
890 + position = Int64(0)
891 + if has_position:
892 + position = positions[token_idx]
893 + boundary = has_position and (
894 + (position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
895 + )
896 + has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
897 + kv_slot_idx = Int64(-1)
898 + if has_kv_slot_idx:
899 + kv_slot_idx = kv_slot_mapping[token_idx]
900 + active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0)
901 +
902 + if active:
903 + base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64)
904 + x0 = compressed_kv.iterator[base]
905 + x1 = compressed_kv.iterator[base + Int64(1)]
906 +
907 + local_sumsq = x0 * x0 + x1 * x1
908 + warp_sum = local_sumsq
909 + for step in cutlass.range_constexpr(5):
910 + offset = const_expr(16 >> step)
911 + warp_sum += cute.arch.shuffle_sync_bfly(warp_sum, offset)
912 +
913 + smem = cutlass.utils.SmemAllocator()
914 + partial_sums = smem.allocate_tensor(
915 + Float32, cute.make_layout((self.num_warps,)), byte_alignment=4
916 + )
917 + rrms_shared = smem.allocate_tensor(
918 + Float32, cute.make_layout((1,)), byte_alignment=4
919 + )
920 +
921 + if lane_id == 0:
922 + partial_sums[warp_id] = warp_sum
923 + cute.arch.sync_threads()
924 + if tid == 0:
925 + total = Float32(0.0)
926 + for i in cutlass.range_constexpr(self.num_warps):
927 + total += partial_sums[i]
928 + rrms_shared[0] = cute.math.rsqrt(
929 + total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
930 + )
931 + cute.arch.sync_threads()
932 +
933 + rrms = rrms_shared[0]
934 + x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32)
935 + x1 = x1 * rrms * rms_norm_weight[elem0 + 1].to(Float32)
936 +
937 + k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
938 + k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
939 + page = kv_slot_idx // kv_cache_block_size
940 + kv_offset = kv_slot_idx - page * kv_cache_block_size
941 + value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride)
942 + scale_base = (
943 + page * k_cache.stride[0]
944 + + kv_cache_block_size * Int64(self.token_stride)
945 + + kv_offset * Int64(self.scale_dim)
946 + )
947 +
948 + if warp_id == self.nope_blocks:
949 + pair_idx = lane_id
950 + compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
951 + self.compress_ratio
952 + )
953 + cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(Int64)
954 + cos_v = cos_sin_cache.iterator[cs_base]
955 + sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)]
956 + real = x0 * cos_v - x1 * sin_v
957 + imag = x0 * sin_v + x1 * cos_v
958 + packed = _fp32x2_to_bf16x2(real, imag)
959 + out_base = value_base + Int64(self.nope_dim) + (lane_id * 4).to(Int64)
960 + k_cache_u32.iterator[out_base // Int64(4)] = packed
961 + else:
962 + q_packed = _fp32x2_to_bf16x2(x0, x1)
963 + q0, q1 = _bf16x2_to_fp32(q_packed)
964 + abs0 = cute.math.absf(q0)
965 + abs1 = cute.math.absf(q1)
966 + local_absmax = cute.arch.fmax(abs0, abs1)
967 + absmax = local_absmax
968 + for step in cutlass.range_constexpr(5):
969 + offset = const_expr(16 >> step)
970 + absmax = cute.arch.fmax(
971 + absmax, cute.arch.shuffle_sync_bfly(absmax, offset)
972 + )
973 + scale_raw = cute.arch.fmax(
974 + Float32(self.min_scale),
975 + absmax / Float32(self.fp8_max),
976 + )
977 + bits = _recast_val(scale_raw, Uint32)
978 + ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
979 + inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32)
980 + y0 = cute.arch.fmin(
981 + cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)),
982 + Float32(self.fp8_max),
983 + )
984 + y1 = cute.arch.fmin(
985 + cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)),
986 + Float32(self.fp8_max),
987 + )
988 + packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
989 + out_base = value_base + (warp_id * self.quant_block + lane_id * 2).to(
990 + Int64
991 + )
992 + k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
993 + if lane_id == 0:
994 + k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to(Uint8)
995 + if warp_id == 0:
996 + k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8(
997 + 0
998 + )
999 +
1000 + @cache
1001 + @staticmethod
1002 + def compile(
1003 + head_size: int = 512,
1004 + rope_head_dim: int = 64,
1005 + fp8_max: float = 448.0,
1006 + quant_block: int = 64,
1007 + token_stride: int = 576,
1008 + scale_dim: int = 8,
1009 + kv_block_stride: int = 74752,
1010 + compress_ratio: int = 128,
1011 + norm_weight_dtype: type[cutlass.Numeric] = Float32,
1012 + ):
1013 + if quant_block != 64:
1014 + raise ValueError(
1015 + "CuTe DSL sparse-attn store currently requires quant_block=64."
1016 + )
1017 + if rope_head_dim != 64:
1018 + raise ValueError(
1019 + "CuTe DSL sparse-attn store currently requires rope_head_dim=64."
1020 + )
1021 + if head_size % quant_block != 0:
1022 + raise ValueError("head_size must be divisible by quant_block.")
1023 + if token_stride < head_size + rope_head_dim:
1024 + raise ValueError("token_stride is too small for the packed FP8/BF16 row.")
1025 + expected_scale_dim = (head_size - rope_head_dim) // quant_block + 1
1026 + if scale_dim < expected_scale_dim:
1027 + raise ValueError("scale_dim is too small for the UE8M0 scale row.")
1028 + num_positions = cute.sym_int()
1029 + num_slots = cute.sym_int()
1030 + num_kv_slots = cute.sym_int()
1031 + max_pos = cute.sym_int()
1032 + num_blocks = cute.sym_int()
1033 +
1034 + compressed_kv = cute.runtime.make_fake_tensor(
1035 + Float32,
1036 + (num_slots, head_size),
1037 + stride=(cute.sym_int64(divisibility=4), 1),
1038 + assumed_align=4,
1039 + )
1040 + positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
1041 + slot_mapping = make_fake_tensor(Int64, (num_slots,), divisibility=8)
1042 + rms_norm_weight = make_fake_tensor(
1043 + norm_weight_dtype, (head_size,), divisibility=4
1044 + )
1045 + cos_sin_cache = cute.runtime.make_fake_tensor(
1046 + Float32,
1047 + (max_pos, rope_head_dim),
1048 + stride=(cute.sym_int64(divisibility=4), 1),
1049 + assumed_align=4,
1050 + )
1051 + k_cache = cute.runtime.make_fake_tensor(
1052 + Uint8,
1053 + (num_blocks, cute.sym_int(), cute.sym_int()),
1054 + stride=(
1055 + cute.sym_int64(divisibility=16),
1056 + cute.sym_int64(divisibility=8),
1057 + 1,
1058 + ),
1059 + assumed_align=16,
1060 + )
1061 + kv_slot_mapping = make_fake_tensor(Int64, (num_kv_slots,), divisibility=8)
1062 +
1063 + kernel = SparseAttnNormRopeStoreKernel(
1064 + head_size,
1065 + rope_head_dim,
1066 + fp8_max,
1067 + quant_block,
1068 + token_stride,
1069 + scale_dim,
1070 + compress_ratio,
1071 + )
1072 + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
1073 + return cute.compile(
1074 + kernel,
1075 + compressed_kv,
1076 + positions,
1077 + slot_mapping,
1078 + rms_norm_weight,
1079 + Float32(0.0),
1080 + cos_sin_cache,
1081 + k_cache,
1082 + kv_slot_mapping,
1083 + Int64(0),
1084 + stream,
1085 + options="--enable-tvm-ffi",
1086 + )
1087 +
1088 +
1089 +def _compress_kv_sparse_attn_cutedsl(
1090 + state_cache: torch.Tensor,
1091 + token_to_req_indices: torch.Tensor,
1092 + positions: torch.Tensor,
1093 + slot_mapping: torch.Tensor,
1094 + block_table: torch.Tensor,
1095 + block_size: int,
1096 + compressed_kv: torch.Tensor,
1097 + head_size: int = 512,
1098 + state_width: int = 512,
1099 + compress_ratio: int = 128,
1100 + overlap: bool = False,
1101 +) -> None:
1102 + if positions.numel() == 0:
1103 + return
1104 + compiled = SparseAttnCompressKernel.compile(
1105 + head_size=head_size,
1106 + state_width=state_width,
1107 + compress_ratio=compress_ratio,
1108 + overlap=overlap,
1109 + )
1110 + compiled(
1111 + state_cache,
1112 + token_to_req_indices,
1113 + positions,
1114 + slot_mapping,
1115 + block_table,
1116 + block_size,
1117 + compressed_kv,
1118 + )
1119 +
1120 +
1121 +def _norm_rope_insert_sparse_attn_cutedsl(
1122 + compressed_kv: torch.Tensor,
1123 + positions: torch.Tensor,
1124 + slot_mapping: torch.Tensor,
1125 + rms_norm_weight: torch.Tensor,
1126 + rms_norm_eps: float,
1127 + cos_sin_cache: torch.Tensor,
1128 + k_cache: torch.Tensor,
1129 + kv_slot_mapping: torch.Tensor,
1130 + kv_cache_block_size: int,
1131 + kv_block_stride: int,
1132 + head_size: int = 512,
1133 + rope_head_dim: int = 64,
1134 + fp8_max: float = 448.0,
1135 + quant_block: int = 64,
1136 + token_stride: int = 576,
1137 + scale_dim: int = 8,
1138 + compress_ratio: int = 128,
1139 +) -> None:
1140 + if positions.numel() == 0:
1141 + return
1142 + norm_weight_dtype = _TORCH_TO_CUTE.get(rms_norm_weight.dtype)
1143 + if norm_weight_dtype is None:
1144 + raise ValueError(
1145 + "CuTe DSL sparse-attn store supports rms_norm_weight dtype "
1146 + f"bf16/fp32, got {rms_norm_weight.dtype}."
1147 + )
1148 + if k_cache.ndim != 3:
1149 + raise ValueError(
1150 + "CuTe DSL sparse-attn store expects the real DeepSeek V4 "
1151 + f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}."
1152 + )
1153 + compiled = SparseAttnNormRopeStoreKernel.compile(
1154 + head_size=head_size,
1155 + rope_head_dim=rope_head_dim,
1156 + fp8_max=fp8_max,
1157 + quant_block=quant_block,
1158 + token_stride=token_stride,
1159 + scale_dim=scale_dim,
1160 + kv_block_stride=kv_block_stride,
1161 + compress_ratio=compress_ratio,
1162 + norm_weight_dtype=norm_weight_dtype,
1163 + )
1164 + compiled(
1165 + compressed_kv,
1166 + positions,
1167 + slot_mapping,
1168 + rms_norm_weight,
1169 + rms_norm_eps,
1170 + cos_sin_cache,
1171 + k_cache,
1172 + kv_slot_mapping,
1173 + kv_cache_block_size,
1174 + )
1175 +
1176 +
1177 +def _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
1178 + state_cache: torch.Tensor,
1179 + token_to_req_indices: torch.Tensor,
1180 + positions: torch.Tensor,
1181 + slot_mapping: torch.Tensor,
1182 + block_table: torch.Tensor,
1183 + block_size: int,
1184 + rms_norm_weight: torch.Tensor,
1185 + rms_norm_eps: float,
1186 + cos_sin_cache: torch.Tensor,
1187 + k_cache: torch.Tensor,
1188 + kv_slot_mapping: torch.Tensor,
1189 + kv_cache_block_size: int,
1190 + kv_block_stride: int,
1191 + head_size: int = 512,
1192 + state_width: int = 1024,
1193 + rope_head_dim: int = 64,
1194 + fp8_max: float = 448.0,
1195 + quant_block: int = 64,
1196 + token_stride: int = 576,
1197 + scale_dim: int = 8,
1198 + compress_ratio: int = 4,
1199 + overlap: bool = True,
1200 +) -> None:
1201 + if positions.numel() == 0:
1202 + return
1203 + norm_weight_dtype = _TORCH_TO_CUTE.get(rms_norm_weight.dtype)
1204 + if norm_weight_dtype is None:
1205 + raise ValueError(
1206 + "CuTe DSL sparse-attn fused store supports rms_norm_weight dtype "
1207 + f"bf16/fp32, got {rms_norm_weight.dtype}."
1208 + )
1209 + if k_cache.ndim != 3:
1210 + raise ValueError(
1211 + "CuTe DSL sparse-attn fused store expects the real DeepSeek V4 "
1212 + f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}."
1213 + )
1214 + compiled = SparseAttnCompressNormRopeStoreC4Kernel.compile(
1215 + head_size=head_size,
1216 + state_width=state_width,
1217 + rope_head_dim=rope_head_dim,
1218 + fp8_max=fp8_max,
1219 + quant_block=quant_block,
1220 + token_stride=token_stride,
1221 + scale_dim=scale_dim,
1222 + kv_block_stride=kv_block_stride,
1223 + compress_ratio=compress_ratio,
1224 + overlap=overlap,
1225 + norm_weight_dtype=norm_weight_dtype,
1226 + )
1227 + compiled(
1228 + state_cache,
1229 + token_to_req_indices,
1230 + positions,
1231 + slot_mapping,
1232 + block_table,
1233 + block_size,
1234 + rms_norm_weight,
1235 + rms_norm_eps,
1236 + cos_sin_cache,
1237 + k_cache,
1238 + kv_slot_mapping,
1239 + kv_cache_block_size,
1240 + )
vllm/models/deepseek_v4/compressor.py
+141 -39
@@ -13,9 +13,11 @@ 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,
18 - _fused_kv_compress_norm_rope_insert_sparse_attn,
19 + _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl,
20 + _norm_rope_insert_sparse_attn_cutedsl,
21 )
22 from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE
23 from vllm.platforms import current_platform
@@ -171,6 +173,33 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
173
174
175 class DeepseekCompressor(nn.Module):
176 + _compressed_kv_buffers: ClassVar[dict[tuple[str, int, int], torch.Tensor]] = {}
177 +
178 + @classmethod
179 + def _get_compressed_kv_buffer(
180 + cls,
181 + device: str,
182 + max_num_tokens: int,
183 + head_dim: int,
184 + ) -> torch.Tensor:
185 + if device == "cuda" and torch.accelerator.is_available():
186 + device_key = f"cuda:{torch.accelerator.current_device_index()}"
187 + alloc_device = torch.device(device_key)
188 + else:
189 + device_key = str(device)
190 + alloc_device = torch.device(device)
191 +
192 + key = (device_key, max_num_tokens, head_dim)
193 + buffer = cls._compressed_kv_buffers.get(key)
194 + if buffer is None:
195 + buffer = torch.empty(
196 + (max_num_tokens, head_dim),
197 + dtype=torch.float32,
198 + device=alloc_device,
199 + )
200 + cls._compressed_kv_buffers[key] = buffer
201 + return buffer
202 +
203 def __init__(
204 self,
205 vllm_config: VllmConfig,
@@ -240,12 +269,24 @@ class DeepseekCompressor(nn.Module):
269 assert not use_fp4_cache, (
270 "MXFP4 cache is only supported for indexer (head=128)"
271 )
243 - self._fused_kernel = _fused_kv_compress_norm_rope_insert_sparse_attn
272 + self._use_cutedsl_sparse_compressor = True
273 + self._use_cutedsl_fused_sparse_compressor = self.compress_ratio == 4
274 + self._compress_kernel = _compress_kv_sparse_attn_cutedsl
275 + self._norm_rope_store_kernel = _norm_rope_insert_sparse_attn_cutedsl
276 + self._fused_sparse_kernel = (
277 + _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl
278 + )
279 + self._compressed_kv_buffer = self._get_compressed_kv_buffer(
280 + self.device,
281 + vllm_config.scheduler_config.max_num_batched_tokens,
282 + self.head_dim,
283 + )
284 self._quant_block = 64
285 self._token_stride = self.nope_head_dim + self.rope_head_dim * 2
286 self._scale_dim = self.nope_head_dim // 64 + 1 # 7 real + 1 pad
287 self._num_warps = 4
288 elif self.head_dim == 128:
289 + self._use_cutedsl_sparse_compressor = False
290 if use_fp4_cache:
291 self._fused_kernel = (
292 _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
@@ -339,43 +380,104 @@ class DeepseekCompressor(nn.Module):
380 k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix])
381 kv_cache = self._static_forward_context[self.k_cache_prefix].kv_cache
382
342 - self._fused_kernel[(num_actual,)](
343 - # state cache
344 - state_cache,
345 - state_cache.stride(0),
346 - state_cache.stride(1),
347 - # metadata
348 - token_to_req_indices,
349 - positions,
350 - slot_mapping,
351 - block_table,
352 - block_table.stride(0),
353 - block_size,
354 - # RMSNorm
355 - self.norm.weight,
356 - self.rms_norm_eps,
357 - # RoPE
358 - cos_sin_cache,
359 - cos_sin_cache.stride(0),
360 - # KV cache
361 - kv_cache,
362 - k_cache_metadata.slot_mapping,
363 - kv_cache.shape[1], # paged KV cache block size (tokens per block)
364 - # constexprs
365 - HEAD_SIZE=self.head_dim,
366 - TRITON_BLOCK_SIZE=triton.next_power_of_2(self.head_dim),
367 - STATE_WIDTH=state_width,
368 - COMPRESS_RATIO=self.compress_ratio,
369 - OVERLAP=self.overlap,
370 - ROPE_HEAD_DIM=self.rope_head_dim,
371 - FP8_MAX=448.0,
372 - QUANT_BLOCK=self._quant_block,
373 - TOKEN_STRIDE=self._token_stride,
374 - SCALE_DIM=self._scale_dim,
375 - KV_BLOCK_STRIDE=kv_cache.stride(0),
376 - num_warps=self._num_warps,
377 - **pdl_kwargs,
378 - )
383 + if self._use_cutedsl_sparse_compressor:
384 + if self._use_cutedsl_fused_sparse_compressor:
385 + self._fused_sparse_kernel(
386 + state_cache,
387 + token_to_req_indices,
388 + positions,
389 + slot_mapping,
390 + block_table,
391 + block_size,
392 + self.norm.weight,
393 + self.rms_norm_eps,
394 + cos_sin_cache,
395 + kv_cache,
396 + k_cache_metadata.slot_mapping,
397 + kv_cache.shape[1], # paged KV cache block size
398 + kv_cache.stride(0),
399 + head_size=self.head_dim,
400 + state_width=state_width,
401 + rope_head_dim=self.rope_head_dim,
402 + fp8_max=448.0,
403 + quant_block=self._quant_block,
404 + token_stride=self._token_stride,
405 + scale_dim=self._scale_dim,
406 + compress_ratio=self.compress_ratio,
407 + overlap=self.overlap,
408 + )
409 + else:
410 + compressed_kv = self._compressed_kv_buffer[:num_actual]
411 + self._compress_kernel(
412 + state_cache,
413 + token_to_req_indices,
414 + positions,
415 + slot_mapping,
416 + block_table,
417 + block_size,
418 + compressed_kv,
419 + head_size=self.head_dim,
420 + state_width=state_width,
421 + compress_ratio=self.compress_ratio,
422 + overlap=self.overlap,
423 + )
424 + self._norm_rope_store_kernel(
425 + compressed_kv,
426 + positions,
427 + slot_mapping,
428 + self.norm.weight,
429 + self.rms_norm_eps,
430 + cos_sin_cache,
431 + kv_cache,
432 + k_cache_metadata.slot_mapping,
433 + kv_cache.shape[1], # paged KV cache block size
434 + kv_cache.stride(0),
435 + head_size=self.head_dim,
436 + rope_head_dim=self.rope_head_dim,
437 + fp8_max=448.0,
438 + quant_block=self._quant_block,
439 + token_stride=self._token_stride,
440 + scale_dim=self._scale_dim,
441 + compress_ratio=self.compress_ratio,
442 + )
443 + else:
444 + self._fused_kernel[(num_actual,)](
445 + # state cache
446 + state_cache,
447 + state_cache.stride(0),
448 + state_cache.stride(1),
449 + # metadata
450 + token_to_req_indices,
451 + positions,
452 + slot_mapping,
453 + block_table,
454 + block_table.stride(0),
455 + block_size,
456 + # RMSNorm
457 + self.norm.weight,
458 + self.rms_norm_eps,
459 + # RoPE
460 + cos_sin_cache,
461 + cos_sin_cache.stride(0),
462 + # KV cache
463 + kv_cache,
464 + k_cache_metadata.slot_mapping,
465 + kv_cache.shape[1], # paged KV cache block size (tokens per block)
466 + # constexprs
467 + HEAD_SIZE=self.head_dim,
468 + TRITON_BLOCK_SIZE=triton.next_power_of_2(self.head_dim),
469 + STATE_WIDTH=state_width,
470 + COMPRESS_RATIO=self.compress_ratio,
471 + OVERLAP=self.overlap,
472 + ROPE_HEAD_DIM=self.rope_head_dim,
473 + FP8_MAX=448.0,
474 + QUANT_BLOCK=self._quant_block,
475 + TOKEN_STRIDE=self._token_stride,
476 + SCALE_DIM=self._scale_dim,
477 + KV_BLOCK_STRIDE=kv_cache.stride(0),
478 + num_warps=self._num_warps,
479 + **pdl_kwargs,
480 + )
481
482
483 @triton.jit