glm-5.3-flash-tiny-random-nvfp4 / vllm_patch /flashinfer-glm53.patch
cyijun2k's picture
Add architecture-faithful GLM-5.3 Flash NVFP4 tiny-random fixture
dc386a5 verified
Raw
History Blame Contribute Delete
1.41 kB
--- 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