Skip to content

Fix SM100 V32 sparse decode FP8 handling #184

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jianglan89 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix SM100 V32 sparse decode FP8 handling #184

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 38 additions & 23 deletions csrc/sm100/decode/head64/kernel.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -675,9 +675,9 @@ KernelTemplate<MODEL_TYPE>
auto process_one_block = [&](int block_idx, auto is_extra_block_t) {
static constexpr bool IS_EXTRA_BLOCK = std::is_same_v<decltype(is_extra_block_t), IsExtraBlock>;
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];
Expand All @@ -691,25 +691,23 @@ KernelTemplate<MODEL_TYPE>
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;
Expand All @@ -718,9 +716,7 @@ KernelTemplate<MODEL_TYPE>
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;
Expand Down Expand Up @@ -783,10 +779,29 @@ KernelTemplate<MODEL_TYPE>
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
Expand Down
172 changes: 172 additions & 0 deletions tests/test_flash_mla_sm100_sparse_decode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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)