--- 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::D_QK == 576); + if (topk == 2176) { + if constexpr (MT == ModelType::GLM_NSA) { + if (num_heads != 64) return false; + launch_prefill_mg( + 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