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 1, 23:45 UTC
optimize the compressor 128 split cutedsl kernel (#44230)
Signed-off-by: Jie Fang <jief@nvidia.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Jie Fang committed
Jun 4, 2026 at 11:22 UTC
f0cd590d62ea8190a4f00b49e981f5f9972a8d56
1 file changed
+233
-228
vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py
+233
-228
@@ -508,30 +508,28 @@ class SparseAttnCompressNormRopeStoreC4Kernel:
508
)
509
510
511
-class SparseAttnCompressKernel:
511
+class SparseAttnCompressC128Block8Kernel:
512
head_tile = 64
513
rows_per_warp = 16
514
- row_pairs_per_warp = rows_per_warp // 2
515
- elems_per_lane = 4
514
+ elems_per_lane = 2
515
lanes_per_row = head_tile // elems_per_lane
516
num_warps = 8
518
- stats_warp_stride = num_warps + 1
517
+ stats_lane_stride = lanes_per_row + 1
518
+ final_reduce_steps = 3
519
+ final_reduce_initial_offset = 4
520
tb_size = num_warps * 32
521
+ compress_ratio = 128
522
+ state_block_size = 8
523
rcp_ln2 = 1.4426950408889634
524
525
def __init__(
526
self,
527
head_size: int,
528
state_width: int,
526
- compress_ratio: int,
527
- overlap: bool,
529
):
530
self.head_dim = head_size
531
self.num_splits = head_size // self.head_tile
532
self.state_width = state_width
532
- self.compress_ratio = compress_ratio
533
- self.overlap = overlap
534
- self.window = (1 + int(overlap)) * compress_ratio
533
534
@cute.jit
535
def __call__(
@@ -541,7 +539,6 @@ class SparseAttnCompressKernel:
539
positions: cute.Tensor,
540
slot_mapping: cute.Tensor,
541
block_table: cute.Tensor,
544
- block_size: Int64,
542
compressed_kv: cute.Tensor,
543
stream: CUstream,
544
):
@@ -552,7 +549,6 @@ class SparseAttnCompressKernel:
549
positions,
550
slot_mapping,
551
block_table,
555
- block_size,
552
compressed_kv,
553
).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
554
@@ -564,29 +560,36 @@ class SparseAttnCompressKernel:
560
positions: cute.Tensor,
561
slot_mapping: cute.Tensor,
562
block_table: cute.Tensor,
567
- block_size: Int64,
563
compressed_kv: cute.Tensor,
564
):
565
block_id, _, _ = cute.arch.block_idx()
566
tid, _, _ = cute.arch.thread_idx()
567
warp_id = cute.arch.make_warp_uniform(tid // 32)
568
lane_id = tid % 32
574
- row_lane = lane_id // self.lanes_per_row
569
col_group = lane_id % self.lanes_per_row
570
571
token_idx = block_id // self.num_splits
572
split_idx = block_id - token_idx * self.num_splits
573
col_base = split_idx * self.head_tile + col_group * self.elems_per_lane
574
581
- slot_id = slot_mapping[token_idx]
582
- has_position = token_idx < positions.shape[0]
575
position = Int64(0)
584
- if has_position:
576
+ req_idx = Int32(0)
577
+ slot_id = Int64(-1)
578
+ has_position = token_idx < positions.shape[0]
579
+ has_req_idx = token_idx < token_to_req_indices.shape[0]
580
+ if lane_id == 0:
581
+ slot_id = slot_mapping[token_idx]
582
+ if lane_id == 0 and has_position:
583
position = positions[token_idx]
584
+ if lane_id == 0 and has_req_idx:
585
+ req_idx = token_to_req_indices[token_idx]
586
+ slot_id = cute.arch.shuffle_sync(slot_id, offset=0)
587
+ position = cute.arch.shuffle_sync(position, offset=0)
588
+ req_idx = cute.arch.shuffle_sync(req_idx, offset=0)
589
boundary = has_position and (
590
(position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
591
)
589
- has_req_idx = token_idx < token_to_req_indices.shape[0]
592
+ start = position - Int64(self.compress_ratio - 1)
593
active = slot_id >= Int64(0) and has_req_idx and boundary
594
595
if active:
@@ -595,13 +598,13 @@ class SparseAttnCompressKernel:
598
Float32,
599
cute.make_layout(
600
(
601
+ self.num_warps,
602
self.lanes_per_row,
603
self.elems_per_lane,
600
- self.stats_warp_stride,
604
),
605
stride=(
603
- self.elems_per_lane * self.stats_warp_stride,
604
- self.stats_warp_stride,
606
+ self.stats_lane_stride * self.elems_per_lane,
607
+ self.elems_per_lane,
608
1,
609
),
610
),
@@ -611,13 +614,13 @@ class SparseAttnCompressKernel:
614
Float32,
615
cute.make_layout(
616
(
617
+ self.num_warps,
618
self.lanes_per_row,
619
self.elems_per_lane,
616
- self.stats_warp_stride,
620
),
621
stride=(
619
- self.elems_per_lane * self.stats_warp_stride,
620
- self.stats_warp_stride,
622
+ self.stats_lane_stride * self.elems_per_lane,
623
+ self.elems_per_lane,
624
1,
625
),
626
),
@@ -627,25 +630,25 @@ class SparseAttnCompressKernel:
630
Float32,
631
cute.make_layout(
632
(
633
+ self.num_warps,
634
self.lanes_per_row,
635
self.elems_per_lane,
632
- self.stats_warp_stride,
636
),
637
stride=(
635
- self.elems_per_lane * self.stats_warp_stride,
636
- self.stats_warp_stride,
638
+ self.stats_lane_stride * self.elems_per_lane,
639
+ self.elems_per_lane,
640
1,
641
),
642
),
643
byte_alignment=4,
644
)
645
643
- row_pair_layout = cute.make_layout(
644
- (self.row_pairs_per_warp, self.elems_per_lane),
646
+ row_layout = cute.make_layout(
647
+ (self.rows_per_warp, self.elems_per_lane),
648
stride=(self.elems_per_lane, 1),
649
)
647
- kv_vals = cute.make_rmem_tensor(row_pair_layout, Float32)
648
- score_vals = cute.make_rmem_tensor(row_pair_layout, Float32)
650
+ kv_vals = cute.make_rmem_tensor(row_layout, Float32)
651
+ score_vals = cute.make_rmem_tensor(row_layout, Float32)
652
local_max = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
653
local_sum = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
654
local_product = cute.make_rmem_tensor((self.elems_per_lane,), Float32)
@@ -655,100 +658,61 @@ class SparseAttnCompressKernel:
658
local_sum[e] = Float32(0.0)
659
local_product[e] = Float32(0.0)
660
658
- req_idx = token_to_req_indices[token_idx]
659
- start = position - Int64(self.window - 1)
660
- cp_f32x4 = cute.make_copy_atom(
661
- cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128
661
+ first_block_index = start // Int64(self.state_block_size)
662
+ warp_block_index = first_block_index + (warp_id * 2).to(Int64)
663
+ block0_i32 = Int32(0)
664
+ block1_i32 = Int32(0)
665
+ if lane_id == 0:
666
+ block0_i32 = block_table[req_idx, warp_block_index]
667
+ block1_i32 = block_table[req_idx, warp_block_index + Int64(1)]
668
+ block0_i32 = cute.arch.shuffle_sync(block0_i32, offset=0)
669
+ block1_i32 = cute.arch.shuffle_sync(block1_i32, offset=0)
670
+
671
+ cp_f32x2 = cute.make_copy_atom(
672
+ cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=64
673
)
663
- row_mask_and_clamp = const_expr(
664
- (cute.arch.WARP_SIZE - self.lanes_per_row) << 8
665
- | (cute.arch.WARP_SIZE - 1)
674
+ final_mask_and_clamp = const_expr(
675
+ (cute.arch.WARP_SIZE - self.num_warps) << 8 | (cute.arch.WARP_SIZE - 1)
676
)
667
-
668
- for i in cutlass.range_constexpr(self.row_pairs_per_warp):
669
- row = warp_id * self.rows_per_warp + i * 2 + row_lane
670
- pos = start + row.to(Int64)
671
- valid = row < self.window and pos >= Int64(0)
672
- head_offset = ((row // self.compress_ratio) * self.head_dim).to(Int64)
673
-
674
- for e in cutlass.range_constexpr(self.elems_per_lane):
675
- kv = Float32(0.0)
676
- score = -Float32.inf
677
- kv_vals[i, e] = kv
678
- score_vals[i, e] = score
679
-
680
- block_index = Int64(0)
681
- block_offset = Int64(0)
682
- block_number_i32 = Int32(0)
683
- if valid:
684
- block_index = pos // block_size
685
- block_offset = pos - block_index * block_size
686
- if col_group == 0:
687
- block_number_i32 = block_table[req_idx, block_index]
688
- block_number_i32 = cute.arch.shuffle_sync(
689
- block_number_i32,
690
- offset=0,
691
- mask_and_clamp=row_mask_and_clamp,
677
+ col_tile = col_base.to(Int64) // Int64(self.elems_per_lane)
678
+ score_col_tile = col_tile + Int64(self.state_width // self.elems_per_lane)
679
+
680
+ for i in cutlass.range_constexpr(self.rows_per_warp):
681
+ block_number_i32 = block0_i32
682
+ block_offset = Int64(i)
683
+ if const_expr(i >= self.state_block_size):
684
+ block_number_i32 = block1_i32
685
+ block_offset = Int64(i - self.state_block_size)
686
+ row_tensor = state_cache[block_number_i32.to(Int64), block_offset, None]
687
+ kv_src = cute.local_tile(
688
+ row_tensor,
689
+ tiler=(self.elems_per_lane,),
690
+ coord=(col_tile,),
691
)
693
-
694
- if valid:
695
- block_number = block_number_i32.to(Int64)
696
- row_tensor = state_cache[block_number, block_offset, None]
697
- col_tile = (head_offset + col_base.to(Int64)) // Int64(
698
- self.elems_per_lane
699
- )
700
- kv_src = cute.local_tile(
701
- row_tensor,
702
- tiler=(self.elems_per_lane,),
703
- coord=(col_tile,),
704
- )
705
- score_src = cute.local_tile(
706
- row_tensor,
707
- tiler=(self.elems_per_lane,),
708
- coord=(
709
- col_tile + Int64(self.state_width // self.elems_per_lane),
710
- ),
711
- )
712
- cute.copy(cp_f32x4, kv_src, kv_vals[i, None])
713
- cute.copy(cp_f32x4, score_src, score_vals[i, None])
692
+ score_src = cute.local_tile(
693
+ row_tensor,
694
+ tiler=(self.elems_per_lane,),
695
+ coord=(score_col_tile,),
696
+ )
697
+ cute.copy(cp_f32x2, kv_src, kv_vals[i, None])
698
+ cute.copy(cp_f32x2, score_src, score_vals[i, None])
699
700
for e in cutlass.range_constexpr(self.elems_per_lane):
701
local_max[e] = cute.arch.fmax(local_max[e], score_vals[i, e])
702
703
for e in cutlass.range_constexpr(self.elems_per_lane):
719
- if local_max[e] > -Float32.inf:
720
- for i in cutlass.range_constexpr(self.row_pairs_per_warp):
721
- exp_score = cute.math.exp2(
722
- (score_vals[i, e] - local_max[e]) * Float32(self.rcp_ln2),
723
- fastmath=True,
724
- )
725
- local_sum[e] += exp_score
726
- local_product[e] += kv_vals[i, e] * exp_score
727
-
728
- for e in cutlass.range_constexpr(self.elems_per_lane):
729
- pair_max = cute.arch.shuffle_sync_bfly(local_max[e], offset=16)
730
- pair_sum = cute.arch.shuffle_sync_bfly(local_sum[e], offset=16)
731
- pair_product = cute.arch.shuffle_sync_bfly(local_product[e], offset=16)
732
- warp_max = cute.arch.fmax(local_max[e], pair_max)
733
- warp_sum = Float32(0.0)
734
- warp_product = Float32(0.0)
735
- if warp_max > -Float32.inf:
736
- local_scale = cute.math.exp2(
737
- (local_max[e] - warp_max) * Float32(self.rcp_ln2),
738
- fastmath=True,
739
- )
740
- pair_scale = cute.math.exp2(
741
- (pair_max - warp_max) * Float32(self.rcp_ln2),
704
+ for i in cutlass.range_constexpr(self.rows_per_warp):
705
+ exp_score = cute.math.exp2(
706
+ (score_vals[i, e] - local_max[e]) * Float32(self.rcp_ln2),
707
fastmath=True,
708
)
744
- warp_sum = local_sum[e] * local_scale + pair_sum * pair_scale
745
- warp_product = (
746
- local_product[e] * local_scale + pair_product * pair_scale
747
- )
748
- if lane_id < self.lanes_per_row:
749
- s_max[col_group, e, warp_id] = warp_max
750
- s_sum[col_group, e, warp_id] = warp_sum
751
- s_product[col_group, e, warp_id] = warp_product
709
+ local_sum[e] += exp_score
710
+ local_product[e] += kv_vals[i, e] * exp_score
711
+
712
+ for e in cutlass.range_constexpr(self.elems_per_lane):
713
+ s_max[warp_id, col_group, e] = local_max[e]
714
+ s_sum[warp_id, col_group, e] = local_sum[e]
715
+ s_product[warp_id, col_group, e] = local_product[e]
716
cute.arch.sync_threads()
717
718
out_group = tid // self.num_warps
@@ -761,16 +725,16 @@ class SparseAttnCompressKernel:
725
out_lane = out_idx // self.elems_per_lane
726
out_elem = out_idx % self.elems_per_lane
727
764
- local_warp_max = s_max[out_lane, out_elem, final_lane]
728
+ local_warp_max = s_max[final_lane, out_lane, out_elem]
729
global_max = local_warp_max
766
- for step in cutlass.range_constexpr(3):
767
- offset = const_expr(4 >> step)
730
+ for step in cutlass.range_constexpr(self.final_reduce_steps):
731
+ offset = const_expr(self.final_reduce_initial_offset >> step)
732
global_max = cute.arch.fmax(
733
global_max,
734
cute.arch.shuffle_sync_bfly(
735
global_max,
736
offset=offset,
773
- mask_and_clamp=row_mask_and_clamp,
737
+ mask_and_clamp=final_mask_and_clamp,
738
),
739
)
740
@@ -778,19 +742,19 @@ class SparseAttnCompressKernel:
742
(local_warp_max - global_max) * Float32(self.rcp_ln2),
743
fastmath=True,
744
)
781
- global_sum = s_sum[out_lane, out_elem, final_lane] * scale
782
- global_product = s_product[out_lane, out_elem, final_lane] * scale
783
- for step in cutlass.range_constexpr(3):
784
- offset = const_expr(4 >> step)
745
+ global_sum = s_sum[final_lane, out_lane, out_elem] * scale
746
+ global_product = s_product[final_lane, out_lane, out_elem] * scale
747
+ for step in cutlass.range_constexpr(self.final_reduce_steps):
748
+ offset = const_expr(self.final_reduce_initial_offset >> step)
749
global_sum += cute.arch.shuffle_sync_bfly(
750
global_sum,
751
offset=offset,
788
- mask_and_clamp=row_mask_and_clamp,
752
+ mask_and_clamp=final_mask_and_clamp,
753
)
754
global_product += cute.arch.shuffle_sync_bfly(
755
global_product,
756
offset=offset,
793
- mask_and_clamp=row_mask_and_clamp,
757
+ mask_and_clamp=final_mask_and_clamp,
758
)
759
760
if final_lane == 0:
@@ -804,10 +768,8 @@ class SparseAttnCompressKernel:
768
def compile(
769
head_size: int = 512,
770
state_width: int = 512,
807
- compress_ratio: int = 128,
808
- overlap: bool = False,
771
):
810
- if head_size % SparseAttnCompressKernel.head_tile != 0:
772
+ if head_size % SparseAttnCompressC128Block8Kernel.head_tile != 0:
773
raise ValueError("head_size must be divisible by the 64-wide head tile.")
774
num_positions = cute.sym_int()
775
num_slots = cute.sym_int()
@@ -838,15 +800,13 @@ class SparseAttnCompressKernel:
800
compressed_kv = cute.runtime.make_fake_tensor(
801
Float32,
802
(num_slots, head_size),
841
- stride=(cute.sym_int64(divisibility=4), 1),
803
+ stride=(head_size, 1),
804
assumed_align=4,
805
)
806
845
- kernel = SparseAttnCompressKernel(
807
+ kernel = SparseAttnCompressC128Block8Kernel(
808
head_size,
809
state_width,
848
- compress_ratio,
849
- overlap,
810
)
811
stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
812
return cute.compile(
@@ -856,7 +816,6 @@ class SparseAttnCompressKernel:
816
positions,
817
slot_mapping,
818
block_table,
859
- Int64(0),
819
compressed_kv,
820
stream,
821
options="--enable-tvm-ffi",
@@ -875,6 +834,7 @@ class SparseAttnNormRopeStoreKernel:
834
token_stride: int,
835
scale_dim: int,
836
compress_ratio: int,
837
+ static_kv_cache_block_size: int,
838
):
839
self.head_dim = head_size
840
self.rope_dim = rope_head_dim
@@ -887,6 +847,7 @@ class SparseAttnNormRopeStoreKernel:
847
self.nope_blocks = self.nope_dim // quant_block
848
self.tb_size = head_size // 2
849
self.compress_ratio = compress_ratio
850
+ self.static_kv_cache_block_size = static_kv_cache_block_size
851
852
@cute.jit
853
def __call__(
@@ -899,7 +860,6 @@ class SparseAttnNormRopeStoreKernel:
860
cos_sin_cache: cute.Tensor,
861
k_cache: cute.Tensor,
862
kv_slot_mapping: cute.Tensor,
902
- kv_cache_block_size: Int64,
863
stream: CUstream,
864
):
865
grid = (slot_mapping.shape[0], 1, 1)
@@ -912,7 +872,6 @@ class SparseAttnNormRopeStoreKernel:
872
cos_sin_cache,
873
k_cache,
874
kv_slot_mapping,
915
- kv_cache_block_size,
875
).launch(grid=grid, block=(self.tb_size, 1, 1), stream=stream)
876
877
@cute.kernel
@@ -926,7 +885,6 @@ class SparseAttnNormRopeStoreKernel:
885
cos_sin_cache: cute.Tensor,
886
k_cache: cute.Tensor,
887
kv_slot_mapping: cute.Tensor,
929
- kv_cache_block_size: Int64,
888
):
889
token_idx, _, _ = cute.arch.block_idx()
890
tid, _, _ = cute.arch.thread_idx()
@@ -934,21 +892,41 @@ class SparseAttnNormRopeStoreKernel:
892
lane_id = tid % 32
893
elem0 = tid * 2
894
937
- slot_id = slot_mapping[token_idx]
938
- has_position = token_idx < positions.shape[0]
895
position = Int64(0)
940
- if has_position:
896
+ kv_slot_idx = Int64(-1)
897
+ has_position = token_idx < positions.shape[0]
898
+ slot_id = Int64(-1)
899
+ if lane_id == 0:
900
+ slot_id = slot_mapping[token_idx]
901
+ if lane_id == 0 and has_position:
902
position = positions[token_idx]
903
+ has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
904
+ if lane_id == 0 and has_kv_slot_idx:
905
+ kv_slot_idx = kv_slot_mapping[token_idx]
906
+ slot_id = cute.arch.shuffle_sync(slot_id, offset=0)
907
+ position = cute.arch.shuffle_sync(position, offset=0)
908
+ kv_slot_idx = cute.arch.shuffle_sync(kv_slot_idx, offset=0)
909
boundary = has_position and (
910
(position + Int64(1)) % Int64(self.compress_ratio) == Int64(0)
911
)
945
- has_kv_slot_idx = token_idx < kv_slot_mapping.shape[0]
946
- kv_slot_idx = Int64(-1)
947
- if has_kv_slot_idx:
948
- kv_slot_idx = kv_slot_mapping[token_idx]
912
active = slot_id >= Int64(0) and boundary and kv_slot_idx >= Int64(0)
913
914
if active:
915
+ k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
916
+ k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
917
+ static_block_size = Int64(self.static_kv_cache_block_size)
918
+ page = kv_slot_idx // static_block_size
919
+ kv_offset = kv_slot_idx - page * static_block_size
920
+ scale_row_offset = static_block_size * Int64(self.token_stride)
921
+ value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride)
922
+ scale_base = (
923
+ page * k_cache.stride[0]
924
+ + scale_row_offset
925
+ + kv_offset * Int64(self.scale_dim)
926
+ )
927
+ weight0 = rms_norm_weight[elem0].to(Float32)
928
+ weight1 = rms_norm_weight[elem0 + 1].to(Float32)
929
+
930
base = token_idx.to(Int64) * compressed_kv.stride[0] + elem0.to(Int64)
931
x0 = compressed_kv.iterator[base]
932
x1 = compressed_kv.iterator[base + Int64(1)]
@@ -963,42 +941,32 @@ class SparseAttnNormRopeStoreKernel:
941
partial_sums = smem.allocate_tensor(
942
Float32, cute.make_layout((self.num_warps,)), byte_alignment=4
943
)
966
- rrms_shared = smem.allocate_tensor(
967
- Float32, cute.make_layout((1,)), byte_alignment=4
968
- )
944
945
if lane_id == 0:
946
partial_sums[warp_id] = warp_sum
947
cute.arch.sync_threads()
973
- if tid == 0:
974
- total = Float32(0.0)
975
- for i in cutlass.range_constexpr(self.num_warps):
976
- total += partial_sums[i]
977
- rrms_shared[0] = cute.math.rsqrt(
978
- total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
979
- )
980
- cute.arch.sync_threads()
948
982
- rrms = rrms_shared[0]
983
- x0 = x0 * rrms * rms_norm_weight[elem0].to(Float32)
984
- x1 = x1 * rrms * rms_norm_weight[elem0 + 1].to(Float32)
949
+ total = partial_sums[lane_id % self.num_warps]
950
+ sum_mask_and_clamp = const_expr(
951
+ (cute.arch.WARP_SIZE - self.num_warps) << 8 | (cute.arch.WARP_SIZE - 1)
952
+ )
953
+ for step in cutlass.range_constexpr(3):
954
+ offset = const_expr(4 >> step)
955
+ total += cute.arch.shuffle_sync_bfly(
956
+ total,
957
+ offset,
958
+ mask_and_clamp=sum_mask_and_clamp,
959
+ )
960
986
- k_cache_u16 = cute.recast_tensor(k_cache, Uint16)
987
- k_cache_u32 = cute.recast_tensor(k_cache, Uint32)
988
- page = kv_slot_idx // kv_cache_block_size
989
- kv_offset = kv_slot_idx - page * kv_cache_block_size
990
- value_base = page * k_cache.stride[0] + kv_offset * Int64(self.token_stride)
991
- scale_base = (
992
- page * k_cache.stride[0]
993
- + kv_cache_block_size * Int64(self.token_stride)
994
- + kv_offset * Int64(self.scale_dim)
961
+ rrms = cute.math.rsqrt(
962
+ total / Float32(self.head_dim) + rms_norm_eps, fastmath=True
963
)
964
+ x0 = x0 * rrms * weight0
965
+ x1 = x1 * rrms * weight1
966
967
if warp_id == self.nope_blocks:
968
pair_idx = lane_id
999
- compressed_pos = (position // Int64(self.compress_ratio)) * Int64(
1000
- self.compress_ratio
1001
- )
969
+ compressed_pos = position - Int64(self.compress_ratio - 1)
970
cs_base = compressed_pos * cos_sin_cache.stride[0] + pair_idx.to(Int64)
971
cos_v = cos_sin_cache.iterator[cs_base]
972
sin_v = cos_sin_cache.iterator[cs_base + Int64(self.rope_dim // 2)]
@@ -1058,6 +1026,7 @@ class SparseAttnNormRopeStoreKernel:
1026
kv_block_stride: int = 74752,
1027
compress_ratio: int = 128,
1028
norm_weight_dtype: type[cutlass.Numeric] = Float32,
1029
+ static_kv_cache_block_size: int = 0,
1030
):
1031
if quant_block != 64:
1032
raise ValueError(
@@ -1074,6 +1043,11 @@ class SparseAttnNormRopeStoreKernel:
1043
expected_scale_dim = (head_size - rope_head_dim) // quant_block + 1
1044
if scale_dim < expected_scale_dim:
1045
raise ValueError("scale_dim is too small for the UE8M0 scale row.")
1046
+ if static_kv_cache_block_size <= 0:
1047
+ raise ValueError(
1048
+ "CuTe DSL sparse-attn store requires a positive static "
1049
+ "kv_cache_block_size."
1050
+ )
1051
num_positions = cute.sym_int()
1052
num_slots = cute.sym_int()
1053
num_kv_slots = cute.sym_int()
@@ -1083,7 +1057,7 @@ class SparseAttnNormRopeStoreKernel:
1057
compressed_kv = cute.runtime.make_fake_tensor(
1058
Float32,
1059
(num_slots, head_size),
1086
- stride=(cute.sym_int64(divisibility=4), 1),
1060
+ stride=(head_size, 1),
1061
assumed_align=4,
1062
)
1063
positions = make_fake_tensor(Int64, (num_positions,), divisibility=8)
@@ -1094,14 +1068,14 @@ class SparseAttnNormRopeStoreKernel:
1068
cos_sin_cache = cute.runtime.make_fake_tensor(
1069
Float32,
1070
(max_pos, rope_head_dim),
1097
- stride=(cute.sym_int64(divisibility=4), 1),
1071
+ stride=(rope_head_dim, 1),
1072
assumed_align=4,
1073
)
1074
k_cache = cute.runtime.make_fake_tensor(
1075
Uint8,
1076
(num_blocks, cute.sym_int(), cute.sym_int()),
1077
stride=(
1104
- cute.sym_int64(divisibility=16),
1078
+ kv_block_stride,
1079
cute.sym_int64(divisibility=8),
1080
1,
1081
),
@@ -1117,6 +1091,7 @@ class SparseAttnNormRopeStoreKernel:
1091
token_stride,
1092
scale_dim,
1093
compress_ratio,
1094
+ static_kv_cache_block_size,
1095
)
1096
stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
1097
return cute.compile(
@@ -1129,48 +1104,66 @@ class SparseAttnNormRopeStoreKernel:
1104
cos_sin_cache,
1105
k_cache,
1106
kv_slot_mapping,
1132
- Int64(0),
1107
stream,
1108
options="--enable-tvm-ffi",
1109
)
1110
1111
1138
-def compress_kv_sparse_attn_cutedsl(
1139
- state_cache: torch.Tensor,
1140
- token_to_req_indices: torch.Tensor,
1141
- positions: torch.Tensor,
1142
- slot_mapping: torch.Tensor,
1143
- block_table: torch.Tensor,
1112
+def compile_split_sparse_attn_cutedsl(
1113
+ head_size: int,
1114
+ state_width: int,
1115
block_size: int,
1145
- compressed_kv: torch.Tensor,
1146
- head_size: int = 512,
1147
- state_width: int = 512,
1148
- compress_ratio: int = 128,
1149
- overlap: bool = False,
1150
-) -> None:
1151
- if positions.numel() == 0:
1152
- return
1153
- compiled = SparseAttnCompressKernel.compile(
1116
+ rope_head_dim: int,
1117
+ fp8_max: float,
1118
+ quant_block: int,
1119
+ token_stride: int,
1120
+ scale_dim: int,
1121
+ kv_cache_block_size: int,
1122
+ kv_block_stride: int,
1123
+ compress_ratio: int,
1124
+ overlap: bool,
1125
+ rms_norm_weight_dtype: torch.dtype,
1126
+):
1127
+ if not (
1128
+ head_size == 512
1129
+ and state_width == head_size
1130
+ and compress_ratio == 128
1131
+ and not overlap
1132
+ and block_size == 8
1133
+ ):
1134
+ raise ValueError(
1135
+ "CuTe DSL split sparse-attn wrapper only supports the real "
1136
+ "DeepSeek V4 C128 layout: head_size=512, state_width=512, "
1137
+ "compress_ratio=128, overlap=False, block_size=8."
1138
+ )
1139
+ compress = SparseAttnCompressC128Block8Kernel.compile(
1140
head_size=head_size,
1141
state_width=state_width,
1156
- compress_ratio=compress_ratio,
1157
- overlap=overlap,
1142
)
1159
- compiled(
1160
- state_cache,
1161
- token_to_req_indices,
1162
- positions,
1163
- slot_mapping,
1164
- block_table,
1165
- block_size,
1166
- compressed_kv,
1143
+ norm_weight_dtype = _TORCH_TO_CUTE[rms_norm_weight_dtype]
1144
+ store = SparseAttnNormRopeStoreKernel.compile(
1145
+ head_size,
1146
+ rope_head_dim,
1147
+ fp8_max,
1148
+ quant_block,
1149
+ token_stride,
1150
+ scale_dim,
1151
+ kv_block_stride,
1152
+ compress_ratio,
1153
+ norm_weight_dtype,
1154
+ kv_cache_block_size,
1155
)
1156
+ return compress, store
1157
1158
1170
-def norm_rope_insert_sparse_attn_cutedsl(
1171
- compressed_kv: torch.Tensor,
1159
+def split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
1160
+ state_cache: torch.Tensor,
1161
+ token_to_req_indices: torch.Tensor,
1162
positions: torch.Tensor,
1163
slot_mapping: torch.Tensor,
1164
+ block_table: torch.Tensor,
1165
+ block_size: int,
1166
+ compressed_kv: torch.Tensor,
1167
rms_norm_weight: torch.Tensor,
1168
rms_norm_eps: float,
1169
cos_sin_cache: torch.Tensor,
@@ -1179,38 +1172,58 @@ def norm_rope_insert_sparse_attn_cutedsl(
1172
kv_cache_block_size: int,
1173
kv_block_stride: int,
1174
head_size: int = 512,
1175
+ state_width: int = 512,
1176
rope_head_dim: int = 64,
1177
fp8_max: float = 448.0,
1178
quant_block: int = 64,
1179
token_stride: int = 576,
1180
scale_dim: int = 8,
1181
compress_ratio: int = 128,
1182
+ overlap: bool = False,
1183
) -> None:
1184
+ if k_cache.ndim != 3:
1185
+ raise ValueError(
1186
+ "CuTe DSL sparse-attn store expects the real DeepSeek V4 "
1187
+ f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}."
1188
+ )
1189
+ if kv_cache_block_size != k_cache.shape[1]:
1190
+ raise ValueError(
1191
+ "CuTe DSL split sparse-attn wrapper expected kv_cache_block_size "
1192
+ f"to match k_cache.shape[1], got {kv_cache_block_size} and "
1193
+ f"{k_cache.shape[1]}."
1194
+ )
1195
if positions.numel() == 0:
1196
return
1191
- norm_weight_dtype = _TORCH_TO_CUTE.get(rms_norm_weight.dtype)
1192
- if norm_weight_dtype is None:
1197
+ if rms_norm_weight.dtype not in _TORCH_TO_CUTE:
1198
raise ValueError(
1199
"CuTe DSL sparse-attn store supports rms_norm_weight dtype "
1200
f"bf16/fp32, got {rms_norm_weight.dtype}."
1201
)
1197
- if k_cache.ndim != 3:
1198
- raise ValueError(
1199
- "CuTe DSL sparse-attn store expects the real DeepSeek V4 "
1200
- f"3D k_cache layout [num_blocks, block_size, 584], got ndim={k_cache.ndim}."
1201
- )
1202
- compiled = SparseAttnNormRopeStoreKernel.compile(
1203
- head_size=head_size,
1204
- rope_head_dim=rope_head_dim,
1205
- fp8_max=fp8_max,
1206
- quant_block=quant_block,
1207
- token_stride=token_stride,
1208
- scale_dim=scale_dim,
1209
- kv_block_stride=kv_block_stride,
1210
- compress_ratio=compress_ratio,
1211
- norm_weight_dtype=norm_weight_dtype,
1202
+ compress, store = compile_split_sparse_attn_cutedsl(
1203
+ head_size,
1204
+ state_width,
1205
+ block_size,
1206
+ rope_head_dim,
1207
+ fp8_max,
1208
+ quant_block,
1209
+ token_stride,
1210
+ scale_dim,
1211
+ kv_cache_block_size,
1212
+ kv_block_stride,
1213
+ compress_ratio,
1214
+ overlap,
1215
+ rms_norm_weight.dtype,
1216
)
1213
- compiled(
1217
+ compress(
1218
+ state_cache,
1219
+ token_to_req_indices,
1220
+ positions,
1221
+ slot_mapping,
1222
+ block_table,
1223
+ compressed_kv,
1224
+ )
1225
+
1226
+ store(
1227
compressed_kv,
1228
positions,
1229
slot_mapping,
@@ -1219,7 +1232,6 @@ def norm_rope_insert_sparse_attn_cutedsl(
1232
cos_sin_cache,
1233
k_cache,
1234
kv_slot_mapping,
1222
- kv_cache_block_size,
1235
)
1236
1237
@@ -1346,7 +1358,7 @@ def compress_norm_rope_store_cutedsl(
1358
dtype=torch.float32,
1359
device=state_cache.device,
1360
)
1349
- compress_kv_sparse_attn_cutedsl(
1361
+ split_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
1362
state_cache,
1363
token_to_req_indices,
1364
positions,
@@ -1354,15 +1366,6 @@ def compress_norm_rope_store_cutedsl(
1366
block_table,
1367
block_size,
1368
compressed_kv,
1357
- head_size=head_dim,
1358
- state_width=state_width,
1359
- compress_ratio=compress_ratio,
1360
- overlap=overlap,
1361
- )
1362
- norm_rope_insert_sparse_attn_cutedsl(
1363
- compressed_kv,
1364
- positions,
1365
- slot_mapping,
1369
rms_norm_weight,
1370
rms_norm_eps,
1371
cos_sin_cache,
@@ -1371,10 +1374,12 @@ def compress_norm_rope_store_cutedsl(
1374
kv_cache.shape[1], # paged KV cache block size
1375
kv_cache.stride(0),
1376
head_size=head_dim,
1377
+ state_width=state_width,
1378
rope_head_dim=rope_head_dim,
1379
fp8_max=448.0,
1380
quant_block=quant_block,
1381
token_stride=token_stride,
1382
scale_dim=scale_dim,
1383
compress_ratio=compress_ratio,
1384
+ overlap=overlap,
1385
)