[Bugfix][SM120][MLA] Support NoPE sparse MLA (GLM-5.3-Flash) on the FlashInfer SM120 backend (#55277)
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: Kimi Code <noreply@moonshot.cn>
Gabriel Wu committed
Sep 23, 2026 at 17:03 UTC
c843f0cab78b413f2b0a5c2282cb71e3a3937b11
2 files changed
+55
-6
csrc/libtorch_stable/cache_kernels.cu
+11
-6
@@ -549,16 +549,20 @@ __global__ void concat_and_cache_ds_mla_kernel(
549
scalar_t* kv_cache_16bit =
550
reinterpret_cast<scalar_t*>(&kv_cache[dst_idx_start]);
551
552
- // The last warp handles the RoPE part
552
+ // Zero the reserved RoPE tail for NoPE rows.
553
if (threadIdx.x >= 64) {
554
// Each thread handles two elements of RoPE
555
const int8_t pe_idx_start = (threadIdx.x - 64) * 2;
556
- const int64_t src_idx = token_idx * k_pe_stride + pe_idx_start;
557
- // Vectorized load of two 16-bit values, performed as one 32-bit load
558
- const int32_t vals = *reinterpret_cast<const int32_t*>(&k_pe[src_idx]);
556
// RoPE values start after the packed 8-bit NoPE values and the
557
// 32-bit scales
558
const int64_t dst_idx = kv_lora_rank / 2 + 8 + pe_idx_start;
559
+ if (pe_dim == 0) {
560
+ *reinterpret_cast<int32_t*>(&kv_cache_16bit[dst_idx]) = 0;
561
+ return;
562
+ }
563
+ const int64_t src_idx = token_idx * k_pe_stride + pe_idx_start;
564
+ // Vectorized load of two 16-bit values, performed as one 32-bit load
565
+ const int32_t vals = *reinterpret_cast<const int32_t*>(&k_pe[src_idx]);
566
// Vectorized store of two 16-bit values, performed as one 32-bit store
567
*reinterpret_cast<int32_t*>(&kv_cache_16bit[dst_idx]) = vals;
568
return;
@@ -600,7 +604,7 @@ __global__ void concat_and_cache_ds_mla_kernel(
604
}
605
606
// Now all threads in the block scale and write their elements
603
- // NoPE data is packed in the first kv_lora_rank/2 bytes (first 256 bytes)
607
+ // NoPE data occupies the first kv_lora_rank bytes (512 bytes).
608
const int64_t dst_idx_base = dst_idx_start + (threadIdx.x * 8);
609
610
uint8_t result[8];
@@ -936,7 +940,8 @@ void concat_and_cache_mla(
940
if (kv_cache_dtype == "fp8_ds_mla") {
941
STD_TORCH_CHECK(kv_lora_rank == 512,
942
"kv_lora_rank must be 512 for fp8_ds_mla");
939
- STD_TORCH_CHECK(pe_dim == 64, "pe_dim must be 64 for fp8_ds_mla");
943
+ STD_TORCH_CHECK(pe_dim == 64 || pe_dim == 0,
944
+ "pe_dim must be 64 or 0 for fp8_ds_mla");
945
STD_TORCH_CHECK(kv_cache.size(2) == 656 / kv_cache.element_size(),
946
"kv_cache.size(2) must be 656 bytes for fp8_ds_mla");
947
STD_TORCH_CHECK(kv_c.element_size() == 2,
tests/kernels/attention/test_cache.py
+44
@@ -1012,6 +1012,50 @@ def test_concat_and_cache_ds_mla(
1012
torch.testing.assert_close(kv_rope, ref_rope, atol=0.001, rtol=0.1)
1013
1014
1015
+@pytest.mark.parametrize("device", CUDA_DEVICES)
1016
+@pytest.mark.parametrize("block_size", [64, 256])
1017
+@torch.inference_mode()
1018
+def test_concat_and_cache_ds_mla_nope(device: str, block_size: int) -> None:
1019
+ """NoPE matches zero RoPE, clears valid tails, and preserves unused slots."""
1020
+ dtype = torch.bfloat16
1021
+ if current_platform.is_rocm():
1022
+ pytest.skip("concat_and_cache_mla doesn't support fp8_ds_mla on ROCm")
1023
+ num_tokens, num_blocks = 3, 2
1024
+ set_random_seed(0)
1025
+ torch.set_default_device(device)
1026
+ torch.accelerator.set_device_index(device)
1027
+
1028
+ slot_mapping = torch.tensor(
1029
+ [block_size - 1, block_size, -1], dtype=torch.long, device=device
1030
+ )
1031
+ kv_c = torch.randn(num_tokens, 512, dtype=dtype, device=device)
1032
+ k_pe = torch.empty(num_tokens, 0, dtype=dtype, device=device)
1033
+ zero_k_pe = torch.zeros(num_tokens, 64, dtype=dtype, device=device)
1034
+ scale = torch.tensor(1.0, dtype=torch.float32, device=device)
1035
+ kv_cache = torch.full(
1036
+ (num_blocks, block_size, 656), 0xFF, dtype=torch.uint8, device=device
1037
+ )
1038
+ ref_cache = kv_cache.clone()
1039
+
1040
+ opcheck(
1041
+ torch.ops._C_cache_ops.concat_and_cache_mla,
1042
+ (kv_c, k_pe, kv_cache, slot_mapping, "fp8_ds_mla", scale),
1043
+ test_utils=DEFAULT_OPCHECK_TEST_UTILS,
1044
+ )
1045
+ kv_cache.fill_(0xFF)
1046
+ ops.concat_and_cache_mla(kv_c, k_pe, kv_cache, slot_mapping, "fp8_ds_mla", scale)
1047
+ ops.concat_and_cache_mla(
1048
+ kv_c, zero_k_pe, ref_cache, slot_mapping, "fp8_ds_mla", scale
1049
+ )
1050
+
1051
+ torch.testing.assert_close(kv_cache, ref_cache, atol=0, rtol=0)
1052
+ rows = kv_cache.view(num_blocks * block_size, 656)
1053
+ assert (rows[slot_mapping[:2], 528:] == 0).all()
1054
+ untouched = torch.ones(num_blocks * block_size, dtype=torch.bool, device=device)
1055
+ untouched[slot_mapping[:2]] = False
1056
+ assert (rows[untouched] == 0xFF).all()
1057
+
1058
+
1059
# Bytes per token for the nvfp4_ds_mla cache layout (see flashmla_sparse.py).
1060
NVFP4_DS_MLA_ENTRY_SIZE = 352
1061