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, 07:15 UTC
[Kernel][DSv4] Optimize sparse FP8 compressor kernels (#44161)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Yongye Zhu committed
Jun 1, 2026 at 12:18 UTC
035733515f25764cfa828b269cd762d38e4959b9
1 file changed
+139
-91
vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py
+139
-91
@@ -96,9 +96,16 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
96
self.quant_block = quant_block
97
self.token_stride = token_stride
98
self.scale_dim = scale_dim
99
- self.num_warps = head_size // quant_block
99
+ self.elems_per_lane = 8
100
+ self.copy_elems = 4
101
+ self.copy_chunks = self.elems_per_lane // self.copy_elems
102
+ self.lanes_per_group = quant_block // self.elems_per_lane
103
+ self.groups_per_warp = 32 // self.lanes_per_group
104
+ self.scale_reduce_steps = self.lanes_per_group.bit_length() - 1
105
+ self.scale_reduce_offset = self.lanes_per_group // 2
106
+ self.num_warps = (head_size // quant_block) // self.groups_per_warp
107
self.nope_blocks = self.nope_dim // quant_block
101
- self.tb_size = head_size // 2
108
+ self.tb_size = self.num_warps * 32
109
self.compress_ratio = compress_ratio
110
self.overlap = overlap
111
self.window = (1 + int(overlap)) * compress_ratio
@@ -156,8 +163,9 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
163
tid, _, _ = cute.arch.thread_idx()
164
warp_id = cute.arch.make_warp_uniform(tid // 32)
165
lane_id = tid % 32
159
- elem0 = tid * 2
160
- elem1 = elem0 + 1
166
+ group_lane = lane_id % self.lanes_per_group
167
+ group_idx = warp_id * self.groups_per_warp + lane_id // self.lanes_per_group
168
+ elem_base = group_idx * self.quant_block + group_lane * self.elems_per_lane
169
170
slot_id = slot_mapping[token_idx]
171
has_position = token_idx < positions.shape[0]
@@ -201,12 +209,24 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
209
s_block_numbers[row] = block_number_i32
210
cute.arch.sync_threads()
211
204
- max0 = -Float32.inf
205
- max1 = -Float32.inf
206
- sum0 = Float32(0.0)
207
- sum1 = Float32(0.0)
208
- product0 = Float32(0.0)
209
- product1 = Float32(0.0)
212
+ local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
213
+ local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
214
+ local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
215
+
216
+ for e in cutlass.range_constexpr(self.elems_per_lane):
217
+ local_max[e] = -Float32.inf
218
+ local_sum[e] = Float32(0.0)
219
+ local_product[e] = Float32(0.0)
220
+
221
+ cp_f32x4 = cute.make_copy_atom(
222
+ cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128
223
+ )
224
+ copy_layout = cute.make_layout(
225
+ (self.copy_chunks, self.copy_elems),
226
+ stride=(self.copy_elems, 1),
227
+ )
228
+ kv_vals = cute.make_rmem_tensor(copy_layout, Float32)
229
+ score_vals = cute.make_rmem_tensor(copy_layout, Float32)
230
231
for row in cutlass.range_constexpr(self.window):
232
pos = start + Int64(row)
@@ -215,46 +235,51 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
235
block_offset = pos - block_index * block_size
236
block_number = s_block_numbers[row].to(Int64)
237
head_offset = Int64((row // self.compress_ratio) * self.head_dim)
218
- row_base = (
219
- block_number * state_cache.stride[0]
220
- + block_offset * state_cache.stride[1]
221
- + head_offset
222
- )
223
-
224
- score0 = state_cache.iterator[
225
- row_base + Int64(self.state_width) + elem0.to(Int64)
226
- ]
227
- kv0 = state_cache.iterator[row_base + elem0.to(Int64)]
228
- new_max0 = cute.arch.fmax(max0, score0)
229
- old_scale0 = cute.math.exp2(
230
- (max0 - new_max0) * Float32(self.rcp_ln2), fastmath=True
231
- )
232
- new_scale0 = cute.math.exp2(
233
- (score0 - new_max0) * Float32(self.rcp_ln2), fastmath=True
234
- )
235
- sum0 = sum0 * old_scale0 + new_scale0
236
- product0 = product0 * old_scale0 + kv0 * new_scale0
237
- max0 = new_max0
238
-
239
- score1 = state_cache.iterator[
240
- row_base + Int64(self.state_width) + elem1.to(Int64)
241
- ]
242
- kv1 = state_cache.iterator[row_base + elem1.to(Int64)]
243
- new_max1 = cute.arch.fmax(max1, score1)
244
- old_scale1 = cute.math.exp2(
245
- (max1 - new_max1) * Float32(self.rcp_ln2), fastmath=True
246
- )
247
- new_scale1 = cute.math.exp2(
248
- (score1 - new_max1) * Float32(self.rcp_ln2), fastmath=True
249
- )
250
- sum1 = sum1 * old_scale1 + new_scale1
251
- product1 = product1 * old_scale1 + kv1 * new_scale1
252
- max1 = new_max1
238
+ row_tensor = state_cache[block_number, block_offset, None]
239
+ for chunk in cutlass.range_constexpr(self.copy_chunks):
240
+ copy_elem = const_expr(chunk * self.copy_elems)
241
+ col_tile = (
242
+ head_offset + (elem_base + Int32(copy_elem)).to(Int64)
243
+ ) // Int64(self.copy_elems)
244
+ kv_src = cute.local_tile(
245
+ row_tensor,
246
+ tiler=(self.copy_elems,),
247
+ coord=(col_tile,),
248
+ )
249
+ score_src = cute.local_tile(
250
+ row_tensor,
251
+ tiler=(self.copy_elems,),
252
+ coord=(
253
+ col_tile + Int64(self.state_width // self.copy_elems),
254
+ ),
255
+ )
256
+ cute.copy(cp_f32x4, kv_src, kv_vals[chunk, None])
257
+ cute.copy(cp_f32x4, score_src, score_vals[chunk, None])
258
+
259
+ for e in cutlass.range_constexpr(self.elems_per_lane):
260
+ chunk = const_expr(e // self.copy_elems)
261
+ copy_elem = const_expr(e % self.copy_elems)
262
+ score = score_vals[chunk, copy_elem]
263
+ kv = kv_vals[chunk, copy_elem]
264
+ new_max = cute.arch.fmax(local_max[e], score)
265
+ old_scale = cute.math.exp2(
266
+ (local_max[e] - new_max) * Float32(self.rcp_ln2),
267
+ fastmath=True,
268
+ )
269
+ new_scale = cute.math.exp2(
270
+ (score - new_max) * Float32(self.rcp_ln2),
271
+ fastmath=True,
272
+ )
273
+ local_sum[e] = local_sum[e] * old_scale + new_scale
274
+ local_product[e] = local_product[e] * old_scale + kv * new_scale
275
+ local_max[e] = new_max
276
254
- x0 = product0 / sum0
255
- x1 = product1 / sum1
277
+ x = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
278
+ local_sumsq = Float32(0.0)
279
+ for e in cutlass.range_constexpr(self.elems_per_lane):
280
+ x[e] = local_product[e] / local_sum[e]
281
+ local_sumsq += x[e] * x[e]
282
257
- local_sumsq = x0 * x0 + x1 * x1
283
warp_sum = local_sumsq
284
for step in cutlass.range_constexpr(5):
285
offset = const_expr(16 >> step)
@@ -273,8 +298,9 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
298
cute.arch.sync_threads()
299
300
rrms = rrms_shared[0]
276
- x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32)
277
- x1 = x1 * rrms * rms_norm_weight[elem1].to(Float32)
301
+ for e in cutlass.range_constexpr(self.elems_per_lane):
302
+ elem = elem_base + e
303
+ x[e] = x[e] * rrms * rms_norm_weight[elem].to(Float32)
304
305
k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
306
k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
@@ -287,31 +313,53 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
313
+ kv_offset * Int64(self.scale_dim)
314
)
315
290
- if warp_id == self.nope_blocks:
291
- pair_idx = lane_id
316
+ if group_idx == self.nope_blocks:
317
compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
318
self.compress_ratio
319
)
295
- cos_v = cos_sin_cache[compressed_pos, pair_idx]
296
- sin_v = cos_sin_cache[
297
- compressed_pos, pair_idx + Int32(self.rope_dim // 2)
298
- ]
299
- real = x0 * cos_v - x1 * sin_v
300
- imag = x0 * sin_v + x1 * cos_v
301
- packed = _fp32x2_to_bf16x2(real, imag)
302
- out_base = value_base + Int64(self.nope_dim) + (lane_id * 4).to(Int64)
303
- k_cache_u32.iterator[out_base // Int64(4)] = packed
320
+ for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
321
+ elem = const_expr(pair * 2)
322
+ pair_idx = (elem_base - self.nope_dim) // 2 + Int32(pair)
323
+ cos_v = cos_sin_cache[compressed_pos, pair_idx]
324
+ sin_v = cos_sin_cache[
325
+ compressed_pos, pair_idx + Int32(self.rope_dim // 2)
326
+ ]
327
+ real = x[elem] * cos_v - x[elem + 1] * sin_v
328
+ imag = x[elem] * sin_v + x[elem + 1] * cos_v
329
+ packed = _fp32x2_to_bf16x2(real, imag)
330
+ out_base = (
331
+ value_base
332
+ + Int64(self.nope_dim)
333
+ + ((elem_base - self.nope_dim + Int32(elem)) * 2).to(Int64)
334
+ )
335
+ k_cache_u32.iterator[out_base // Int64(4)] = packed
336
else:
305
- q_packed = _fp32x2_to_bf16x2(x0, x1)
306
- q0, q1 = _bf16x2_to_fp32(q_packed)
307
- abs0 = cute.math.absf(q0)
308
- abs1 = cute.math.absf(q1)
309
- local_absmax = cute.arch.fmax(abs0, abs1)
337
+ q = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
338
+ local_absmax = Float32(0.0)
339
+ for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
340
+ elem = const_expr(pair * 2)
341
+ q_packed = _fp32x2_to_bf16x2(x[elem], x[elem + 1])
342
+ q0, q1 = _bf16x2_to_fp32(q_packed)
343
+ q[elem] = q0
344
+ q[elem + 1] = q1
345
+ local_absmax = cute.arch.fmax(
346
+ local_absmax,
347
+ cute.arch.fmax(cute.math.absf(q0), cute.math.absf(q1)),
348
+ )
349
absmax = local_absmax
311
- for step in cutlass.range_constexpr(5):
312
- offset = const_expr(16 >> step)
350
+ group_mask_and_clamp = const_expr(
351
+ (cute.arch.WARP_SIZE - self.lanes_per_group) << 8
352
+ | (cute.arch.WARP_SIZE - 1)
353
+ )
354
+ for step in cutlass.range_constexpr(self.scale_reduce_steps):
355
+ offset = const_expr(self.scale_reduce_offset >> step)
356
absmax = cute.arch.fmax(
314
- absmax, cute.arch.shuffle_sync_bfly(absmax, offset)
357
+ absmax,
358
+ cute.arch.shuffle_sync_bfly(
359
+ absmax,
360
+ offset=offset,
361
+ mask_and_clamp=group_mask_and_clamp,
362
+ ),
363
)
364
scale_raw = cute.arch.fmax(
365
Float32(self.min_scale),
@@ -320,22 +368,22 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
368
bits = _recast_val(scale_raw, Uint32)
369
ue8m0 = ((bits + Uint32(0x7FFFFF)) >> Uint32(23)) & Uint32(0xFF)
370
inv_scale = _recast_val((Uint32(254) - ue8m0) << Uint32(23), Float32)
323
- y0 = cute.arch.fmin(
324
- cute.arch.fmax(q0 * inv_scale, Float32(-self.fp8_max)),
325
- Float32(self.fp8_max),
326
- )
327
- y1 = cute.arch.fmin(
328
- cute.arch.fmax(q1 * inv_scale, Float32(-self.fp8_max)),
329
- Float32(self.fp8_max),
330
- )
331
- packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
332
- out_base = value_base + (warp_id * self.quant_block + lane_id * 2).to(
333
- Int64
334
- )
335
- k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
336
- if lane_id == 0:
337
- k_cache.iterator[scale_base + warp_id.to(Int64)] = ue8m0.to(Uint8)
338
- if warp_id == 0:
371
+ for pair in cutlass.range_constexpr(self.elems_per_lane // 2):
372
+ elem = const_expr(pair * 2)
373
+ y0 = cute.arch.fmin(
374
+ cute.arch.fmax(q[elem] * inv_scale, Float32(-self.fp8_max)),
375
+ Float32(self.fp8_max),
376
+ )
377
+ y1 = cute.arch.fmin(
378
+ cute.arch.fmax(q[elem + 1] * inv_scale, Float32(-self.fp8_max)),
379
+ Float32(self.fp8_max),
380
+ )
381
+ packed_fp8 = _fp32x2_to_fp8e4m3x2(y0, y1)
382
+ out_base = value_base + (elem_base + Int32(elem)).to(Int64)
383
+ k_cache_u16.iterator[out_base // Int64(2)] = packed_fp8
384
+ if group_lane == 0:
385
+ k_cache.iterator[scale_base + group_idx.to(Int64)] = ue8m0.to(Uint8)
386
+ if group_idx == 0:
387
k_cache.iterator[scale_base + Int64(self.nope_blocks)] = Uint8(
388
0
389
)
@@ -462,11 +510,11 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
510
511
class SparseAttnCompressKernel:
512
head_tile = 64
465
- rows_per_warp = 8
513
+ rows_per_warp = 16
514
row_pairs_per_warp = rows_per_warp // 2
515
elems_per_lane = 4
516
lanes_per_row = head_tile // elems_per_lane
469
- num_warps = 16
517
+ num_warps = 8
518
stats_warp_stride = num_warps + 1
519
tb_size = num_warps * 32
520
rcp_ln2 = 1.4426950408889634
@@ -715,8 +763,8 @@ class SparseAttnCompressKernel:
763
764
local_warp_max = s_max[out_lane, out_elem, final_lane]
765
global_max = local_warp_max
718
- for step in cutlass.range_constexpr(4):
719
- offset = const_expr(8 >> step)
766
+ for step in cutlass.range_constexpr(3):
767
+ offset = const_expr(4 >> step)
768
global_max = cute.arch.fmax(
769
global_max,
770
cute.arch.shuffle_sync_bfly(
@@ -732,8 +780,8 @@ class SparseAttnCompressKernel:
780
)
781
global_sum = s_sum[out_lane, out_elem, final_lane] * scale
782
global_product = s_product[out_lane, out_elem, final_lane] * scale
735
- for step in cutlass.range_constexpr(4):
736
- offset = const_expr(8 >> step)
783
+ for step in cutlass.range_constexpr(3):
784
+ offset = const_expr(4 >> step)
785
global_sum += cute.arch.shuffle_sync_bfly(
786
global_sum,
787
offset=offset,