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

Match tensor shapes from working test_d1_kv_merge

Browse Source
This commit is contained in:
biondizzle
2026-05-27 06:56:04 +00:00
parent 3a25c7feff
commit 6ee61717c0
1 changed files with 9 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -162,17 +162,20 @@ def _attention_single_head(
for nt in range(n_pv_tiles):
v_start = nt * pv_n_tile
v_end = v_start + pv_n_tile
v_tile = v_seg[0, :, v_start:v_end].contiguous()
v_kernel = v_tile.unsqueeze(-1)
v_tile = v_seg[0, :, v_start:v_end].contiguous() # (T, pv_n_tile) 2D
v_kernel = v_tile.unsqueeze(-1) # (T, pv_n_tile, 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')
q_input = q[0].contiguous().unsqueeze(-1)
k_input = k_seg[0].contiguous().unsqueeze(-1)
# Match test_d1_kv_merge shapes exactly
q_input = q[0].contiguous() # (T, hd) 2D
k_input = k_seg[0].contiguous() # (s_k, hd) 2D
q_3d = q_input.unsqueeze(-1) # (T, hd, 1)
k_3d = k_input.unsqueeze(-1) # (s_k, hd, 1)
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
mQ = ct.from_dlpack(q_input).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_input))
mK = ct.from_dlpack(k_input).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_input))
mQ = ct.from_dlpack(q_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_3d))
mK = ct.from_dlpack(k_3d).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_3d))
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))