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

Single-segment: use normalize=False + per-row normalization from row_sums

Browse Source
This commit is contained in:
biondizzle
2026-05-27 06:48:56 +00:00
parent fe55bf23a0
commit 8f87109f86
1 changed files with 43 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
dsv4/kernels/attention/production.py
View File

@@ -134,7 +134,7 @@ def _attention_single_head_normalized(
sink_bias: torch.Tensor = None,
use_smem_p: bool = False,
) -> torch.Tensor:
"""Run FMHA for a single head with in-kernel normalization (single KV tile)."""
"""Run FMHA for a single head with Python normalization (single KV tile)."""
_, T, hd = q.shape
N = k.shape[1]
apply_swa_mask = swa_len is not None
@@ -142,14 +142,15 @@ def _attention_single_head_normalized(
compiled, kernel = _get_or_compile_kernel(
head_dim=hd, s_k=N, use_smem_p=use_smem_p,
normalize=True, apply_swa_mask=apply_swa_mask,
normalize=False, apply_swa_mask=apply_swa_mask,
is_causal=is_causal, n_comp=n_comp if n_comp > 0 else None,
apply_sink_bias=apply_sink_bias,
)
pv_n_tile = kernel.pv_n_tile
n_pv_tiles = kernel.n_pv_tiles
output = torch.zeros(T, hd, dtype=torch.bfloat16, device='cuda')
output_unnorm = torch.zeros(T, hd, dtype=torch.float32, device='cuda')
lse_val = None
for nt in range(n_pv_tiles):
v_start = nt * pv_n_tile
@@ -157,6 +158,8 @@ def _attention_single_head_normalized(
v_tile = v[0, :, v_start:v_end].contiguous()
v_kernel = v_tile.unsqueeze(-1)
c_tile = torch.zeros(T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
lse_tensor = torch.zeros(T, 1, 1, dtype=torch.float32, device='cuda')
row_sums_tensor = torch.zeros(T, 1, 1, dtype=torch.float32, device='cuda')
q_input = q[0].contiguous().unsqueeze(-1)
k_input = k[0].contiguous().unsqueeze(-1)
@@ -166,13 +169,47 @@ def _attention_single_head_normalized(
mK = ct.from_dlpack(k_input).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_input))
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))
mRS = ct.from_dlpack(row_sums_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums_tensor))
compiled(mQ, mK, mV, mC, stream)
compiled(mQ, mK, mV, mC, stream, lse=mLSE, row_sums=mRS)
torch.cuda.synchronize()
output[:, v_start:v_end] = c_tile[:, :, 0]
output_unnorm[:, v_start:v_end] = c_tile[:, :, 0].float()
if nt == 0:
lse_val = lse_tensor[0, 0, 0].item()
return output.unsqueeze(0) # (1, T, hd)
# Normalize: O_norm = O_unnorm / row_sum
# row_sum is computed from lse: row_sum = exp(lse - row_max * ln2) ... complex
# But we have row_sums from the kernel!
# row_sums_tensor[0,0,0] has the row_sum for row 0 only.
# For per-row normalization, we'd need per-row row_sums.
# The kernel outputs row_sums[sfw_idx, 0, 0] for each of 128 rows.
# But we can also compute from the un-normalized O and the reference.
#
# Simpler: compute row_sum from the un-normalized O and the LSE.
# lse = ln(row_sum) + row_max * ln(2)
# exp(lse) = row_sum * exp(row_max * ln(2)) = row_sum * 2^row_max
# O_unnorm = P @ V where P = softmax * row_sum
# Wait, the kernel's P = exp2(S*scale - row_max) which is NOT softmax.
# softmax = P / row_sum
# O_unnorm = P @ V = softmax * row_sum @ V = O_norm * row_sum
# So: O_norm = O_unnorm / row_sum
#
# We need row_sum per row. The kernel outputs it in row_sums_tensor.
# But row_sums_tensor only has values for 128 rows, and we need all T rows.
# Actually T=128, so row_sums_tensor[0:128, 0, 0] has all rows.
#
# Let me extract per-row row_sums.
# But the kernel only writes row_sums[sfw_idx, 0, 0] for sfw_idx 0..127.
# And row_sums_tensor is (T, 1, 1) = (128, 1, 1).
# So row_sums_tensor[:, 0, 0] should have all 128 rows.
# Re-run to get row_sums (we already have them from the last call)
row_sums_per_row = row_sums_tensor[:, 0, 0].float().unsqueeze(1) # (T, 1)
row_sums_per_row = row_sums_per_row.clamp(min=1e-30)
output_norm = output_unnorm / row_sums_per_row
return output_norm.to(torch.bfloat16).unsqueeze(0) # (1, T, hd)
def _attention_single_head(