biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

D1: fix per-row LSE output + add KV merge test v2 with per-row LSE

Browse Source
This commit is contained in:
biondizzle
2026-05-24 22:21:51 +00:00
parent 18f3274c0b
commit 674c5b9c18
2 changed files with 128 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
dsv4/kernels/attention/fmha.py
View File

@@ -465,7 +465,7 @@ class FmhaKernel:
)
cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO_i)
for k in cutlass.range(cute.size(tTMrO_i), vectorize=True):
tTMrO_i[k] = tTMrO_i[k] * Float32(1.0) # DEBUG: NO-OP round-trip test
tTMrO_i[k] = tTMrO_i[k] * acc_scale
cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i)
cute.arch.fence_view_async_tmem_store()
@@ -505,14 +505,15 @@ class FmhaKernel:
# Compute LSE: lse = ln(row_sum) + row_max * ln(2)
# Only when emitting un-normalized output (D5a path).
# When normalize=True, LSE is not needed (in-kernel normalization).
# Each thread writes its row's LSE. With 128 softmax threads and 128 rows,
# each thread (sfw_idx) owns exactly one row.
if const_expr(not self.normalize):
_row_max_safe = row_max
if row_max == -cutlass.Float32.inf:
_row_max_safe = Float32(0.0)
if sfw_idx == 0:
_ln2 = Float32(0.6931471805599453) # ln(2)
lse_val = cute.math.log(row_sum, fastmath=True) + _row_max_safe * _ln2
mLSE[0] = lse_val
_ln2 = Float32(0.6931471805599453) # ln(2)
lse_val = cute.math.log(row_sum, fastmath=True) + _row_max_safe * _ln2
mLSE[sfw_idx] = lse_val
tmem.relinquish_alloc_permit()
tmem.free(tmem_ptr)

122
tests/unit/test_d1_kv_merge_v2.py Normal file
View File

@@ -0,0 +1,122 @@
"""
D1: Multi-KV-tile merge using per-row LSE and log-sum-exp.
Strategy: Run s_k=128 kernel per KV segment, get per-row O and LSE.
Merge using the D5 formula:
O = (exp(lse_0) * O_0 + exp(lse_1) * O_1) / (exp(lse_0) + exp(lse_1))
This avoids the broken TMEM round-trip O rescale.
"""
import torch, math
import cutlass.cute as cute
import cutlass.torch as ct
import cuda.bindings.driver as cuda
from dsv4.kernels.attention.fmha import FmhaKernel
def test_multi_kv_merge(hd=64, s_k=256):
m = 128
n_kv_segments = s_k // 128
torch.manual_seed(42)
q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda')
v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
# FP32 reference (full attention)
qf = q[:, :, 0].float()
kf = k[:, :, 0].float()
scale = 1.0 / math.sqrt(hd)
attn_scores = qf @ kf.T * scale
attn_max = attn_scores.max(dim=-1, keepdim=True)[0]
attn_exp = torch.exp(attn_scores - attn_max)
attn_sum = attn_exp.sum(dim=-1, keepdim=True)
ref_norm = (attn_exp / attn_sum) @ v.float()
# Run s_k=128 kernel per KV segment
kernel = FmhaKernel(head_dim=hd, s_k=128, use_smem_p=False, normalize=False)
pv_n_tile = kernel.pv_n_tile
n_pv_tiles = kernel.n_pv_tiles
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
# Compile once
k_seg0 = k[:128]
v_tile0 = v[:128, 0:pv_n_tile].contiguous()
v_kernel0 = v_tile0.unsqueeze(-1)
c_tile0 = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
# Per-row LSE: shape (m,)
lse_tensor = torch.zeros(m, dtype=torch.float32, device='cuda')
mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))
mK = ct.from_dlpack(k_seg0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_seg0))
mV = ct.from_dlpack(v_kernel0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel0))
mC = ct.from_dlpack(c_tile0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile0))
mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))
print(f' Compiling (hd={hd}, s_k=128, {n_kv_segments} segments)...', flush=True)
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE)
# Accumulate across KV segments using log-sum-exp merge
o_accum = None # Will be (m, hd) FP32
lse_accum = None # Will be (m,) FP32
for seg in range(n_kv_segments):
k_start = seg * 128
k_end = k_start + 128
k_seg = k[k_start:k_end]
v_seg = v[k_start:k_end]
seg_o = torch.zeros(m, hd, dtype=torch.float32, device='cuda')
for nt in range(n_pv_tiles):
v_start = nt * pv_n_tile
v_end = v_start + pv_n_tile
v_tile = v_seg[:, v_start:v_end].contiguous()
v_kernel = v_tile.unsqueeze(-1)
c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
lse_tensor.zero_()
mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))
mK = ct.from_dlpack(k_seg).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_seg))
mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel))
mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile))
mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))
compiled(mQ, mK, mV, mC, stream, mLSE)
torch.cuda.synchronize()
seg_o[:, v_start:v_end] = c_tile[:, :, 0].float()
seg_lse = lse_tensor.clone() # (m,) per-row LSE
# Log-sum-exp merge with accumulator
if o_accum is None:
o_accum = seg_o
lse_accum = seg_lse
else:
e_old = torch.exp(lse_accum) # (m,)
e_new = torch.exp(seg_lse) # (m,)
e_sum = e_old + e_new # (m,)
o_accum = (e_old.unsqueeze(-1) * o_accum + e_new.unsqueeze(-1) * seg_o) / e_sum.unsqueeze(-1)
lse_accum = torch.log(e_sum)
cos = torch.nn.functional.cosine_similarity(
o_accum.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0)
).item()
print(f' hd={hd}, s_k={s_k} ({n_kv_segments} segments): cos_norm {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}')
return cos
def test():
print("=== D1: Multi-KV Merge via Per-Row Log-Sum-Exp ===\n")
test_multi_kv_merge(64, 256)
test_multi_kv_merge(64, 384)
test_multi_kv_merge(64, 512)
test_multi_kv_merge(64, 1024)
if __name__ == '__main__':
test()