From 6ee61717c075d640ab6f3f6ac34a0732da9c15ee Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 27 May 2026 06:56:04 +0000 Subject: [PATCH] Match tensor shapes from working test_d1_kv_merge --- dsv4/kernels/attention/production.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/dsv4/kernels/attention/production.py b/dsv4/kernels/attention/production.py index fb527c40..58a63573 100644 --- a/dsv4/kernels/attention/production.py +++ b/dsv4/kernels/attention/production.py @@ -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))