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

D5c: add row_sum output for proper external normalization

Browse Source 
The kernel's O_unnorm is max-shifted (divided by 2^row_max), so
O_norm != O_unnorm * exp(-LSE). Instead, O_norm = O_unnorm / row_sum.
Added mRowSums output tensor to enable correct normalization.
This commit is contained in:
biondizzle
2026-05-26 15:07:22 +00:00
parent 31e6426049
commit 016edbcc97
2 changed files with 25 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -109,7 +109,7 @@ class FmhaKernel:
cute.size_in_bytes(self.q_dtype, v_s)) * cta
@cute.jit
def __call__(self, q, k, v, c, stream, lse=None, swa_len=None, sink_bias=None):
def __call__(self, q, k, v, c, stream, lse=None, swa_len=None, sink_bias=None, row_sums=None):
self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype
self.a_major = LayoutEnum.from_tensor(q).mma_major_mode()
self.b_major = LayoutEnum.from_tensor(k).mma_major_mode()
@@ -152,10 +152,13 @@ class FmhaKernel:
# else: sink_bias is already a CuTe tensor (caller must pass via ct.from_dlpack)
# Grid: (M_tiles, 1, batch) where M = n_h * T packed into M dimension
# For single-head (n_h=1): grid=(1,1,1) — backward compatible
self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.p_smem_s,self.c_smem_s,self.epi_tile,lse,swa_len,sink_bias).launch(grid=(1,1,self.batch_size),block=[self.threads_per_cta,1,1],stream=stream)
if const_expr(row_sums is None):
row_sums = cute.make_tensor(lse.iterator, cute.make_layout((1,), stride=(0,)))
self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.p_smem_s,self.c_smem_s,self.epi_tile,lse,swa_len,sink_bias,row_sums).launch(grid=(1,1,self.batch_size),block=[self.threads_per_cta,1,1],stream=stream)
@cute.kernel
def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, p_smem_s, c_smem_s, epi_tile, mLSE, swa_len, mSinkBias):
def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, p_smem_s, c_smem_s, epi_tile, mLSE, swa_len, mSinkBias, mRowSums):
warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
tidx,_,_ = cute.arch.thread_idx()
if warp_idx == self.tma_warp_id:
@@ -583,6 +586,8 @@ class FmhaKernel:
_ln2 = Float32(0.6931471805599453) # ln(2)
lse_val = cute.math.log(row_sum, fastmath=True) + _row_max_safe * _ln2
mLSE[sfw_idx, Int32(0), Int32(0)] = lse_val
# Also output row_sum for external normalization (D5c)
mRowSums[sfw_idx, Int32(0), Int32(0)] = row_sum
tmem.relinquish_alloc_permit()
tmem.free(tmem_ptr)

31
tests/unit/test_d5c_fused.py
View File

@@ -162,6 +162,7 @@ def test_d5c_combined():
# Allocate output
c_out = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
lse_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
row_sum_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
# Prepare CuTe tensors
def to_cute(t):
@@ -172,31 +173,32 @@ def test_d5c_combined():
mV = to_cute(v_combined.unsqueeze(-1).contiguous())
mC = to_cute(c_out)
mLSE = to_cute(lse_out)
mRowSums = to_cute(row_sum_out)
# Compile
print('Compiling D5c kernel (combined KV + sink bias)...', flush=True)
mSinkBias = to_cute(attn_sink)
compiled = cute.compile(
kernel, mQ, mK, mV, mC, stream, mLSE,
swa_len=swa_len, sink_bias=mSinkBias,
swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums,
)
# Run
print('Running D5c kernel...', flush=True)
compiled(
mQ, mK, mV, mC, stream, mLSE,
swa_len=swa_len, sink_bias=mSinkBias,
swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums,
)
torch.cuda.synchronize()
# Check results
# Kernel outputs UN-NORMALIZED O (normalize=False). Normalize using per-row LSE.
# O_norm[i] = O_unnorm[i] * exp(-lse[i])
# Kernel outputs UN-NORMALIZED O (normalize=False). Normalize using per-row row_sum.
# O_norm[i] = O_unnorm[i] / row_sum[i]
o_kernel_unnorm = c_out[:, :, 0].float() # (m, hd)
lse_kernel = lse_out[:, 0, 0].float() # (m,)
row_sums = row_sum_out[:, 0, 0].float() # (m,)
# Normalize each row
o_kernel = o_kernel_unnorm * (-lse_kernel.unsqueeze(1)).exp() # (m, hd)
# Normalize each row by its row_sum
o_kernel = o_kernel_unnorm / row_sums.unsqueeze(1).clamp(min=1e-30)
cos = torch.nn.functional.cosine_similarity(
o_kernel.flatten().unsqueeze(0),
@@ -210,7 +212,8 @@ def test_d5c_combined():
if cos < 0.99:
print(f' kernel[0,:4]={o_kernel[0,:4].tolist()}')
print(f' ref[0,:4]={ref_combined[0,:4].tolist()}')
print(f' LSE range: {lse_kernel.min().item():.4f} to {lse_kernel.max().item():.4f}')
print(f' row_sum range: {row_sums.min().item():.4f} to {row_sums.max().item():.4f}')
print(f' LSE range: {lse_out[:,0,0].min().item():.4f} to {lse_out[:,0,0].max().item():.4f}')
def test_d5c_with_causal():
@@ -258,6 +261,7 @@ def test_d5c_with_causal():
c_out = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda')
lse_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
row_sum_out = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
def to_cute(t):
return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t))
@@ -267,22 +271,21 @@ def test_d5c_with_causal():
mV = to_cute(v_combined.unsqueeze(-1).contiguous())
mC = to_cute(c_out)
mLSE = to_cute(lse_out)
print('Compiling D5c kernel (causal + sink bias)...', flush=True)
mRowSums = to_cute(row_sum_out)
mSinkBias = to_cute(attn_sink)
compiled = cute.compile(
kernel, mQ, mK, mV, mC, stream, mLSE,
swa_len=swa_len, sink_bias=mSinkBias,
swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums,
)
compiled(
mQ, mK, mV, mC, stream, mLSE,
swa_len=swa_len, sink_bias=mSinkBias,
swa_len=swa_len, sink_bias=mSinkBias, row_sums=mRowSums,
)
torch.cuda.synchronize()
o_kernel_unnorm = c_out[:, :, 0].float()
lse_kernel = lse_out[:, 0, 0].float()
o_kernel = o_kernel_unnorm * (-lse_kernel.unsqueeze(1)).exp()
row_sums = row_sum_out[:, 0, 0].float()
o_kernel = o_kernel_unnorm / row_sums.unsqueeze(1).clamp(min=1e-30)
cos = torch.nn.functional.cosine_similarity(
o_kernel.flatten().unsqueeze(0),