From 5929e532263f4adbf015e09bbf8a53cb5fb9a0f9 Mon Sep 17 00:00:00 2001 From: luoli Date: Sun, 10 May 2026 15:53:43 +0800 Subject: [PATCH] Fix SM100 V32 sparse decode FP8 handling --- csrc/sm100/decode/head64/kernel.cuh | 61 ++++--- tests/test_flash_mla_sm100_sparse_decode.py | 172 ++++++++++++++++++++ 2 files changed, 210 insertions(+), 23 deletions(-) create mode 100644 tests/test_flash_mla_sm100_sparse_decode.py diff --git a/csrc/sm100/decode/head64/kernel.cuh b/csrc/sm100/decode/head64/kernel.cuh index 7c46921c..ddc16c29 100644 --- a/csrc/sm100/decode/head64/kernel.cuh +++ b/csrc/sm100/decode/head64/kernel.cuh @@ -675,9 +675,9 @@ KernelTemplate auto process_one_block = [&](int block_idx, auto is_extra_block_t) { static constexpr bool IS_EXTRA_BLOCK = std::is_same_v; int cur_block_size = IS_EXTRA_BLOCK ? params.extra_page_block_size : params.page_block_size; - int64_t cur_k_block_stride = IS_EXTRA_BLOCK ? params.stride_extra_kv_block : params.stride_kv_block; + [[maybe_unused]] int64_t cur_k_block_stride = IS_EXTRA_BLOCK ? params.stride_extra_kv_block : params.stride_kv_block; [[maybe_unused]] int cur_k_row_stride = IS_EXTRA_BLOCK ? params.stride_extra_kv_row : params.stride_kv_row; - uint8_t* cur_k_scales_ptr = IS_EXTRA_BLOCK ? extra_k_scales_ptr : k_scales_ptr; + [[maybe_unused]] uint8_t* cur_k_scales_ptr = IS_EXTRA_BLOCK ? extra_k_scales_ptr : k_scales_ptr; int cur_tma_coords_step_per_block = IS_EXTRA_BLOCK ? tma_coords_step_per_extra_block : tma_coords_step_per_block; int abs_pos, my_indices[2]; @@ -691,25 +691,23 @@ KernelTemplate plan.bar_valid_coord_scale_free[rs.index_buf_idx].wait(rs.index_bar_phase^1); int tma_coords[2]; - e8m0 scales[2*NUM_SCALES_EACH_TOKEN]; + [[maybe_unused]] e8m0 scales[2*NUM_SCALES_EACH_TOKEN]; char valid_mask = 0; CUTE_UNROLL for (int i = 0; i < 2; ++i) { - int block_idx, idx_in_block; - block_idx = (unsigned int)my_indices[i] / cur_block_size; - idx_in_block = (unsigned int)my_indices[i] % cur_block_size; - bool is_token_valid = my_indices[i] != -1 && (abs_pos+i < (IS_EXTRA_BLOCK?args.extra_topk_length:args.topk_length)); + int block_idx = 0, idx_in_block = 0; + int max_token_idx = (IS_EXTRA_BLOCK ? params.extra_num_blocks : params.num_blocks) * cur_block_size; + bool is_token_valid = + my_indices[i] >= 0 && + my_indices[i] < max_token_idx && + (abs_pos+i < (IS_EXTRA_BLOCK?args.extra_topk_length:args.topk_length)); + if (is_token_valid) { + block_idx = my_indices[i] / cur_block_size; + idx_in_block = my_indices[i] % cur_block_size; + } valid_mask |= is_token_valid << i; tma_coords[i] = is_token_valid ? block_idx*cur_tma_coords_step_per_block + idx_in_block*tma_coords_step_per_token : -1; // If the token is invalid because it topk position exceeds topk_length, we must manually fill tma_coords with -1 to avoid copying-in NaN. - if constexpr (MODEL_TYPE == ModelType::V32) { - int64_t offset = is_token_valid ? block_idx*cur_k_block_stride + idx_in_block*cur_k_row_stride : 0; - float4 cur_scale_fp32 = __ldg((float4*)(cur_k_scales_ptr + offset)); - e8m0 res[4]; - *(__nv_fp8x2_storage_t*)(res+0) = __nv_cvt_float2_to_e8m0x2(float2{cur_scale_fp32.x, cur_scale_fp32.y}, __NV_NOSAT, cudaRoundZero); - *(__nv_fp8x2_storage_t*)(res+2) = __nv_cvt_float2_to_e8m0x2(float2{cur_scale_fp32.z, cur_scale_fp32.w}, __NV_NOSAT, cudaRoundZero); - if (!is_token_valid) *(uint32_t*)res = (uint32_t)0; - *(uint32_t*)(scales+i*NUM_SCALES_EACH_TOKEN) = *(uint32_t*)(res); - } else { + if constexpr (MODEL_TYPE != ModelType::V32) { int64_t offset = block_idx*cur_k_block_stride + idx_in_block*8; // Each token has 7 scale factors with an extra 1B padding uint64_t scalesx8 = is_token_valid ? __ldg((uint64_t*)(cur_k_scales_ptr + offset)) : 0; *(uint64_t*)(scales+i*NUM_SCALES_EACH_TOKEN) = scalesx8; @@ -718,9 +716,7 @@ KernelTemplate valid_mask <<= lane_idx%4*2; valid_mask |= __shfl_xor_sync(0xFFFFFFFF, valid_mask, 0x1); valid_mask |= __shfl_xor_sync(0xFFFFFFFF, valid_mask, 0x2); - if constexpr (MODEL_TYPE == ModelType::V32) { - *(uint64_t*)(plan.scales[rs.index_buf_idx] + lane_idx*2) = *(uint64_t*)scales; - } else { + if constexpr (MODEL_TYPE != ModelType::V32) { *(__int128_t*)(plan.scales[rs.index_buf_idx] + lane_idx*2) = *(__int128_t*)scales; } *(int2*)(plan.tma_coord[rs.index_buf_idx] + lane_idx*2) = *(int2*)tma_coords; @@ -783,10 +779,29 @@ KernelTemplate for (int local_row_idx = 0; local_row_idx < ROWS_PER_GROUP; ++local_row_idx) { int row_idx = local_row_idx*NUM_GROUPS + group_idx; bf16 scales[4]; - e8m0 scales_e8m0[4]; - *(uint32_t*)scales_e8m0 = *(uint32_t*)plan.scales[rs.index_buf_idx][row_idx]; - *(__nv_bfloat162_raw*)(scales+0) = __nv_cvt_e8m0x2_to_bf162raw(*(unsigned short*)(scales_e8m0+0)); - *(__nv_bfloat162_raw*)(scales+2) = __nv_cvt_e8m0x2_to_bf162raw(*(unsigned short*)(scales_e8m0+2)); + int tma_coord = plan.tma_coord[rs.index_buf_idx][row_idx]; + float4 cur_scale_fp32 = float4{0.0f, 0.0f, 0.0f, 0.0f}; + if (idx_in_group == 0 && tma_coord >= 0) { + uint8_t* cur_k_scales_ptr = + block_idx >= args.num_orig_kv_blocks ? + (uint8_t*)params.extra_kv + D_NOPE : + (uint8_t*)params.kv + D_NOPE; + cur_scale_fp32 = __ldg((float4*)(cur_k_scales_ptr + (int64_t)tma_coord*TMA_K_STRIDE)); + } + uint32_t scale01 = 0, scale23 = 0; + if (idx_in_group == 0) { + scales[0] = (bf16)cur_scale_fp32.x; + scales[1] = (bf16)cur_scale_fp32.y; + scales[2] = (bf16)cur_scale_fp32.z; + scales[3] = (bf16)cur_scale_fp32.w; + scale01 = *(uint32_t*)(scales + 0); + scale23 = *(uint32_t*)(scales + 2); + } + int group_leader_lane = lane_idx - idx_in_group; + scale01 = __shfl_sync(0xFFFFFFFF, scale01, group_leader_lane); + scale23 = __shfl_sync(0xFFFFFFFF, scale23, group_leader_lane); + *(uint32_t*)(scales + 0) = scale01; + *(uint32_t*)(scales + 2) = scale23; uint64_t cur_data_fp8x8 = get_raw_fp8(local_row_idx, 0); CUTE_UNROLL diff --git a/tests/test_flash_mla_sm100_sparse_decode.py b/tests/test_flash_mla_sm100_sparse_decode.py new file mode 100644 index 00000000..35053545 --- /dev/null +++ b/tests/test_flash_mla_sm100_sparse_decode.py @@ -0,0 +1,172 @@ +import os +import sys +import unittest + +import torch + +THIS_DIR = os.path.dirname(os.path.abspath(__file__)) +if THIS_DIR not in sys.path: + sys.path.insert(0, THIS_DIR) + +import flash_mla +import quant + + +class FlashMLASM100SparseDecodeRegressionTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + if not torch.cuda.is_available(): + raise unittest.SkipTest("CUDA is required") + + cls.device = torch.device("cuda:0") + torch.cuda.set_device(cls.device) + torch.set_default_device(cls.device) + torch.set_default_dtype(torch.bfloat16) + torch.set_float32_matmul_precision("high") + cls.cc = torch.cuda.get_device_capability(cls.device) + + def _require_sm100(self): + if self.cc[0] != 10: + self.skipTest("This regression test exercises the SM100 sparse decode path") + + def _make_v32_fp8_kv_cache(self, num_blocks: int, block_size: int) -> torch.Tensor: + device = self.device + torch.manual_seed(20260510) + + k_cache = torch.empty( + (num_blocks, block_size, 1, 656), + dtype=torch.float8_e4m3fn, + device=device, + ) + + nope = (torch.randn((num_blocks, block_size, 1, 512), device=device) * 0.5).clamp_( + -2.0, + 2.0, + ) + k_cache[..., :512] = nope.to(torch.float8_e4m3fn) + + token_ids = torch.arange(num_blocks * block_size, dtype=torch.float32, device=device).view( + num_blocks, + block_size, + 1, + 1, + ) + tile_ids = torch.arange(4, dtype=torch.float32, device=device).view(1, 1, 1, 4) + scales = k_cache[..., 512:528].view(torch.float32) + scales.copy_(0.0137 + (token_ids % 17) * 0.00031 + tile_ids * 0.00103) + + rope = k_cache[..., 528:].view(torch.bfloat16) + rope.copy_((torch.randn_like(rope.float()) * 0.03).to(torch.bfloat16)) + return k_cache + + def _reference_v32_sparse_decode( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + indices: torch.Tensor, + topk_length: torch.Tensor, + sm_scale: float, + ): + batch, seqlen_q, num_heads, head_dim = q.shape + head_dim_v = 512 + topk = indices.shape[-1] + + kv = quant.dequantize_k_cache( + k_cache, + quant.FP8KVCacheLayout.V32_FP8Sparse, + ).view(-1, head_dim) + physical_tokens = kv.shape[0] + + invalid = (indices < 0) | (indices >= physical_tokens) + invalid |= torch.arange(topk, device=q.device).view(1, 1, topk) >= topk_length.view( + batch, + 1, + 1, + ) + + fixed_indices = indices.clamp(0, physical_tokens - 1) + gathered = kv.index_select(0, fixed_indices.reshape(-1)).view( + batch, + seqlen_q, + topk, + head_dim, + ).float() + + logits = torch.einsum("bshd,bstd->bsht", q.float(), gathered) + logits *= sm_scale + logits.masked_fill_(invalid.view(batch, seqlen_q, 1, topk), float("-inf")) + + lse = torch.logsumexp(logits, dim=-1) + probs = torch.exp(logits - lse.unsqueeze(-1)) + out = torch.einsum("bsht,bstd->bshd", probs, gathered[..., :head_dim_v]) + + no_valid_token = lse == float("-inf") + out[no_valid_token.unsqueeze(-1).expand_as(out)] = 0.0 + lse[no_valid_token] = float("inf") + return out.to(torch.bfloat16), lse.transpose(1, 2) + + def test_v32_fp8_scales_and_physical_oob_indices(self): + self._require_sm100() + + batch = 1 + seqlen_q = 2 + num_heads = 64 + head_dim = 576 + head_dim_v = 512 + topk = 64 + num_blocks = 2 + block_size = 64 + physical_tokens = num_blocks * block_size + + torch.manual_seed(129103) + q = (torch.randn((batch, seqlen_q, num_heads, head_dim), device=self.device) * 0.08).to( + torch.bfloat16 + ) + k_cache = self._make_v32_fp8_kv_cache(num_blocks, block_size) + + valid = torch.arange(0, 48, dtype=torch.int32, device=self.device) + physical_oob = torch.arange( + physical_tokens, + physical_tokens + 12, + dtype=torch.int32, + device=self.device, + ) + masked_tail = torch.tensor([7, -1, physical_tokens + 31, 19], dtype=torch.int32, device=self.device) + row0 = torch.cat([valid, physical_oob, masked_tail]) + row1 = torch.cat([valid.flip(0), physical_oob.flip(0), masked_tail.flip(0)]) + indices = torch.stack([row0, row1]).view(batch, seqlen_q, topk) + topk_length = torch.tensor([60], dtype=torch.int32, device=self.device) + + sched_meta, num_splits = flash_mla.get_mla_metadata() + out, lse = flash_mla.flash_mla_with_kvcache( + q, + k_cache, + None, + None, + head_dim_v, + sched_meta, + num_splits, + head_dim ** -0.55, + False, + True, + indices, + None, + None, + None, + topk_length, + None, + ) + ref_out, ref_lse = self._reference_v32_sparse_decode( + q, + k_cache, + indices, + topk_length, + head_dim ** -0.55, + ) + + torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=2.01 / 128) + torch.testing.assert_close(lse, ref_lse, atol=1e-5, rtol=8.01 / 65536) + + +if __name__ == "__main__": + unittest.main(verbosity=2)