File size: 1,411 Bytes
963e33f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
--- flashinfer/mla/_sparse_mla_sm120.py
+++ flashinfer/mla/_sparse_mla_sm120.py
@@ -112,6 +112,7 @@
         (64, 512),
         (64, 1024),
         (64, 2048),
+        (64, 2176),
         (128, 128),
         (128, 512),
         (128, 1024),
--- flashinfer/data/csrc/sparse_mla_sm120_decode_dsv3_2.cu
+++ flashinfer/data/csrc/sparse_mla_sm120_decode_dsv3_2.cu
@@ -162,6 +164,7 @@
   DSV3_2_DISPATCH(64, 512)
   DSV3_2_DISPATCH(64, 1024)
   DSV3_2_DISPATCH(64, 2048)
+  DSV3_2_DISPATCH(64, 2176)
   DSV3_2_DISPATCH(128, 128)
   DSV3_2_DISPATCH(128, 512)
   DSV3_2_DISPATCH(128, 1024)
--- flashinfer/data/csrc/sparse_mla_sm120_prefill.cu
+++ flashinfer/data/csrc/sparse_mla_sm120_prefill.cu
@@ -217,6 +217,16 @@
                          float* out_lse, float sm_scale, int num_tokens, size_t stride_kv_block,
                          const int* topk_length_ptr, cudaStream_t stream) {
   static_assert(KVCacheTraits<MT>::D_QK == 576);
+  if (topk == 2176) {
+    if constexpr (MT == ModelType::GLM_NSA) {
+      if (num_heads != 64) return false;
+      launch_prefill_mg<MT, ComputeMode::FP8, 64, 2176, 64>(
+          Q, KV, indices, attn_sink, output, out_lse, sm_scale, num_tokens,
+          stride_kv_block, topk_length_ptr, stream);
+      return true;
+    }
+    return false;
+  }
   if (topk != 2048) return false;
 
   // PBS=64 matches the V32 decode (`decode_dsv3_2_kernel.cuh`). NH=8 covers