diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 5c5c045a..5f6fdc98 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -108,11 +108,13 @@ class FmhaKernel: def __call__(self, q, k, v, c, stream, lse=None): self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype # For multi-CTA, q has shape (n_h, T, hd, 1). Layout depends on inner (T, hd) dims only. - # Create a view of q without the head dimension for layout computation. + # LayoutEnum.from_tensor needs 2D/3D tensors. For 4D Q, we extract the + # inner dimensions by slicing away the head dim. if const_expr(self.num_ctas > 1): - q_inner = cute.local_tile(q, cute.make_layout(1, stride=1), (0,)) # select head 0's view + # q shape: (n_h, T, hd, 1). Slice head dim to get (T, hd, 1) view. + q_inner = q[0] # select head 0 → (T, hd, 1) self.a_major = LayoutEnum.from_tensor(q_inner).mma_major_mode() - c_inner = cute.local_tile(c, cute.make_layout(1, stride=1), (0,)) + c_inner = c[0] # select head 0 → (T, hd, 1) self.c_layout = LayoutEnum.from_tensor(c_inner) else: self.a_major = LayoutEnum.from_tensor(q).mma_major_mode()