From 8441b88136cdac0b14e8d87e1f8c497d368b8c3d Mon Sep 17 00:00:00 2001 From: bigSheep123 <1554869970@qq.com> Date: Fri, 14 Aug 2026 13:52:24 +0800 Subject: [PATCH 1/2] [BugFix][Hopper][Blackwell] Read the WGMMA/UMMA K-panel stride from the layout A shared operand sliced along MN silently computed wrong results whenever block_K spanned more than one swizzle atom: T.gemm(a_s, b_s[0:64, :], c0, transpose_B=True) # wrong past the first K atom When block_K exceeds one swizzle atom the shared layout splits K into panels stored panel-major -- a (128, 128) bf16 tile decodes as (128,(64,2)):(64,(1,8192)) -- so stepping to the next K panel moves by the *buffer's* MN extent. The atom offset formulas instead reconstructed that step as `mn_extent * swizzle_atom_elems` using the extent of the operand, which is smaller for a slice. The slice origin was already layout-derived and correct, so only K atoms past the first were affected: the two halves of a split buffer read each other's rows, in bounds and without any diagnostic. block_K == 64 was correct only by accident -- at 16-bit dtypes that is exactly one 128B atom, so `ki // k_atom_size` is always zero and the bad multiplier was multiplied away. MN-major operands never reach the term at all. Take the panel stride from the decoded layout instead of reconstructing it: cute.restrict already returns the sublayout carrying it, and slicing does not change a stride. Expose it as `k_panel_elems` on WGMMADescriptorParams / TCGEN05DescriptorParams, None when the layout has a single K panel or the operand is MN-major -- exactly the cases where the term is multiplied by zero, so those keep the previous expression and behaviour. The same defect was duplicated across all seven atom emitters: dense Hopper (ss/rs), sparse Hopper (ss/rs) and Blackwell (ss/ts/block-scaled). Verified on H200 (sm_90a): - Whole-buffer path is unchanged: across 120 K-major layouts (fp16/bf16/fp8_e4m3/ fp32 x 5 MN extents x 6 K extents), k_panel_elems is either None or exactly equal to the old reconstruction, for both descriptor builders. The fix is a strict generalization of the path that already worked. - New test_tilelang_kernel_gemm_sliced_operand.py cross-checks a sliced operand against one-buffer-per-slice over operand x majorness x block_K; it fails on the four K-major block_K>64 combinations before this change and passes after. - A layout-level assertion covers the Blackwell descriptor, which has no numeric coverage here for lack of the hardware. Closes #2964 --- ...st_tilelang_cuda_tcgen05_operand_layout.py | 30 +++ ...est_tilelang_kernel_gemm_sliced_operand.py | 178 ++++++++++++++++++ .../macro/tcgen05_macro_generator.py | 49 +++-- .../intrinsics/macro/wgmma_macro_generator.py | 39 +++- .../macro/wgmma_sp_macro_generator.py | 14 +- 5 files changed, 287 insertions(+), 23 deletions(-) create mode 100644 testing/python/kernel/test_tilelang_kernel_gemm_sliced_operand.py diff --git a/testing/python/cuda/test_tilelang_cuda_tcgen05_operand_layout.py b/testing/python/cuda/test_tilelang_cuda_tcgen05_operand_layout.py index 362055f3d5..c20d63eb84 100644 --- a/testing/python/cuda/test_tilelang_cuda_tcgen05_operand_layout.py +++ b/testing/python/cuda/test_tilelang_cuda_tcgen05_operand_layout.py @@ -367,3 +367,33 @@ def plain_2d(m, k): reference.is_k_major, ) assert staged.slice_byte_offset - reference.slice_byte_offset == 2 * m_extent * k_extent * 2 + + +def test_tcgen05_descriptor_k_panel_stride_follows_the_buffer(): + """An MN slice keeps the buffer's K-panel spacing, not the slice's. + + ``block_K`` past one swizzle atom stores the operand as (MN, (atom, panels)), + so stepping to the next K panel moves by ``buffer_MN * atom`` elements. That + step must be read off the layout: reconstructing it from the operand extent + (``mn_extent * atom``) halves the stride for a half-height slice and makes + every K atom past the first read the wrong panel. + """ + + # Two M atoms tall so the slice origin stays on the 128-row atom grid that + # validate_tcgen05_operand_regions requires, and two K panels wide. + m_extent, k_extent = 256, 128 + atom_elems = 64 # 128B swizzle atom over bfloat16 + + def plain_2d(m, k): + return m * atom_elems + k % atom_elems + (k // atom_elems) * m_extent * atom_elems + + layout = _swizzle_128b((m_extent, k_extent), plain_2d) + buffer = tvm.tirx.decl_buffer((m_extent, k_extent), "bfloat16", scope="shared") + + whole = compute_umma_descriptor(layout, buffer, False, region=_ranges((0, 0), (m_extent, k_extent))) + sliced = compute_umma_descriptor(layout, buffer, False, region=_ranges((128, 0), (128, k_extent))) + + assert whole.k_panel_elems == m_extent * atom_elems + # Same spacing for the slice, even though its MN extent is half as large. + assert sliced.k_panel_elems == m_extent * atom_elems + assert sliced.slice_byte_offset == 128 * atom_elems * 2 diff --git a/testing/python/kernel/test_tilelang_kernel_gemm_sliced_operand.py b/testing/python/kernel/test_tilelang_kernel_gemm_sliced_operand.py new file mode 100644 index 0000000000..35964b6d01 --- /dev/null +++ b/testing/python/kernel/test_tilelang_kernel_gemm_sliced_operand.py @@ -0,0 +1,178 @@ +"""A shared operand that is a slice of a wider buffer must match the same GEMM +written with one buffer per slice. + +Regression coverage for the WGMMA K-panel stride: a K-major operand whose +``block_K`` spans more than one swizzle atom is stored as (MN, (atom, panels)), +and stepping from one K panel to the next moves by the *buffer's* MN extent. The +offset formulas used to reconstruct that step from the operand's own extent, +which is smaller for a slice, so ``T.gemm(A_s, B_s[0:64, :], ...)`` silently read +the wrong panel for every ``ki`` past the first atom. ``block_K == 64`` (bf16 = +one 128B atom, single panel) never exercised the term and stayed correct. +""" + +import pytest +import torch + +import tilelang +import tilelang.language as T +import tilelang.testing + +ACCUM = T.float32 + + +def _matmul_sliced(M, N, K, block_M, block_N, block_K, dtype, split, operand, transposed, sliced): + """Split one GEMM into two half-width ones over a shared operand. + + ``operand`` selects which operand is split ("A" splits M, "B" splits N), + ``transposed`` gives that operand's ``transpose_`` flag (which decides whether + the split axis is the buffer's leading or trailing one), and ``sliced`` picks + the shape under test (one buffer consumed as two slices) versus the reference + (one buffer per half). + """ + trans_a = operand == "A" and transposed + trans_b = operand == "B" and transposed + a_shape = (K, M) if trans_a else (M, K) + b_shape = (N, K) if trans_b else (K, N) + split_full = block_N if operand == "B" else block_M + rest = split_full - split + + def a_tile(extent): + return (block_K, extent) if trans_a else (extent, block_K) + + def b_tile(extent): + return (extent, block_K) if trans_b else (block_K, extent) + + @T.prim_func + def main( + A: T.Tensor(a_shape, dtype), + B: T.Tensor(b_shape, dtype), + C: T.Tensor((M, N), dtype), + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): + C0 = T.alloc_fragment((split, block_N) if operand == "A" else (block_M, split), ACCUM) + C1 = T.alloc_fragment((rest, block_N) if operand == "A" else (block_M, rest), ACCUM) + T.clear(C0) + T.clear(C1) + + if operand == "B": + A_s = T.alloc_shared(a_tile(block_M), dtype) + if sliced: + B_s = T.alloc_shared(b_tile(block_N), dtype) + else: + B_s0 = T.alloc_shared(b_tile(split), dtype) + B_s1 = T.alloc_shared(b_tile(rest), dtype) + else: + B_s = T.alloc_shared(b_tile(block_N), dtype) + if sliced: + A_s = T.alloc_shared(a_tile(block_M), dtype) + else: + A_s0 = T.alloc_shared(a_tile(split), dtype) + A_s1 = T.alloc_shared(a_tile(rest), dtype) + + for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=0): + if operand == "B": + if trans_a: + T.copy(A[ko * block_K, by * block_M], A_s) + else: + T.copy(A[by * block_M, ko * block_K], A_s) + if sliced: + if trans_b: + T.copy(B[bx * block_N, ko * block_K], B_s) + T.gemm(A_s, B_s[0:split, :], C0, trans_a, trans_b) + T.gemm(A_s, B_s[split:block_N, :], C1, trans_a, trans_b) + else: + T.copy(B[ko * block_K, bx * block_N], B_s) + T.gemm(A_s, B_s[:, 0:split], C0, trans_a, trans_b) + T.gemm(A_s, B_s[:, split:block_N], C1, trans_a, trans_b) + else: + if trans_b: + T.copy(B[bx * block_N, ko * block_K], B_s0) + T.copy(B[bx * block_N + split, ko * block_K], B_s1) + else: + T.copy(B[ko * block_K, bx * block_N], B_s0) + T.copy(B[ko * block_K, bx * block_N + split], B_s1) + T.gemm(A_s, B_s0, C0, trans_a, trans_b) + T.gemm(A_s, B_s1, C1, trans_a, trans_b) + else: + if trans_b: + T.copy(B[bx * block_N, ko * block_K], B_s) + else: + T.copy(B[ko * block_K, bx * block_N], B_s) + if sliced: + if trans_a: + T.copy(A[ko * block_K, by * block_M], A_s) + T.gemm(A_s[:, 0:split], B_s, C0, trans_a, trans_b) + T.gemm(A_s[:, split:block_M], B_s, C1, trans_a, trans_b) + else: + T.copy(A[by * block_M, ko * block_K], A_s) + T.gemm(A_s[0:split, :], B_s, C0, trans_a, trans_b) + T.gemm(A_s[split:block_M, :], B_s, C1, trans_a, trans_b) + else: + if trans_a: + T.copy(A[ko * block_K, by * block_M], A_s0) + T.copy(A[ko * block_K, by * block_M + split], A_s1) + else: + T.copy(A[by * block_M, ko * block_K], A_s0) + T.copy(A[by * block_M + split, ko * block_K], A_s1) + T.gemm(A_s0, B_s, C0, trans_a, trans_b) + T.gemm(A_s1, B_s, C1, trans_a, trans_b) + + if operand == "B": + T.copy(C0, C[by * block_M, bx * block_N]) + T.copy(C1, C[by * block_M, bx * block_N + split]) + else: + T.copy(C0, C[by * block_M, bx * block_N]) + T.copy(C1, C[by * block_M + split, bx * block_N]) + + return main + + +def run_sliced_operand(M, N, K, block_M, block_N, block_K, dtype, split, operand, transposed): + torch_dtype = getattr(torch, dtype) + trans_a = operand == "A" and transposed + trans_b = operand == "B" and transposed + + torch.manual_seed(0) + a = torch.randn((K, M) if trans_a else (M, K), device="cuda", dtype=torch_dtype) + b = torch.randn((N, K) if trans_b else (K, N), device="cuda", dtype=torch_dtype) + ref = (a.T if trans_a else a).float() @ (b.T if trans_b else b).float() + + outputs = {} + for sliced in (False, True): + kernel = tilelang.compile( + _matmul_sliced(M, N, K, block_M, block_N, block_K, dtype, split, operand, transposed, sliced), + out_idx=[2], + target="cuda", + pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True}, + ) + outputs[sliced] = kernel(a, b).float() + + # The two shapes issue the same instructions over the same values, so they + # agree bit-for-bit; the fp32 reference only bounds the accumulation error. + torch.testing.assert_close(outputs[True], outputs[False], rtol=0, atol=0) + torch.testing.assert_close(outputs[True], ref, rtol=1e-2, atol=1) + + +@tilelang.testing.requires_cuda +@tilelang.testing.requires_cuda_compute_version_ge(9, 0) +@pytest.mark.parametrize("operand", ["A", "B"]) +@pytest.mark.parametrize("transposed", [False, True]) +@pytest.mark.parametrize("block_K", [64, 128, 256]) +def test_gemm_sliced_shared_operand(operand, transposed, block_K): + """block_K 128/256 put more than one swizzle atom on the K axis.""" + run_sliced_operand( + M=256, + N=256, + K=512, + block_M=128, + block_N=128, + block_K=block_K, + dtype="bfloat16", + split=64, + operand=operand, + transposed=transposed, + ) + + +if __name__ == "__main__": + tilelang.testing.main() diff --git a/tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py b/tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py index 5a2f521b4b..e541fb8e3c 100644 --- a/tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py +++ b/tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py @@ -58,6 +58,14 @@ class TCGEN05DescriptorParams: buffer; passed to ``T.increase_descriptor_offset`` after building the descriptor from the buffer base. ``0`` for a whole-buffer / base-origin operand. Computed by :func:`compute_umma_descriptor` from the CuTe layout.""" + k_panel_elems: int | None = None + """Element stride between consecutive K swizzle-atom panels (K-major only), + read off the decoded layout. ``None`` when the layout has a single K panel, + or for an MN-major operand, in which case the ``ki // k_atom_size`` term the + offset formulas multiply it into is always zero / unused. Callers must not + reconstruct this as ``mn_extent * swizzle_atom_elems``: for an operand that + is a slice of a wider buffer (``B[64:128, :]``) the panels stay spaced by the + *buffer's* MN extent, not the operand's.""" @dataclass(frozen=True) @@ -187,6 +195,20 @@ def compute_umma_descriptor(tl_layout, buffer, transposed: bool, micro_size_k: i f"only row-major operand layouts are supported." ) + # Element stride between consecutive K swizzle-atom panels, taken from the + # decoded layout: a K-major operand wider than one swizzle atom decodes as + # (MN, (atom, panels)), and the panel sub-mode's stride is what the offset + # formulas need. Reconstructing it as `mn_dim * swizzle_atom_elems` only holds + # when the operand covers its whole buffer -- for a slice, mn_dim shrinks while + # the panel spacing does not. + k_panel_elems = None + if k_major: + k_shape = tile[k_idx].shape + if isinstance(k_shape, tuple) and len(k_shape) == 2: + panel_stride = tile[k_idx].stride[-1] + if isinstance(panel_stride, int): + k_panel_elems = panel_stride + # SwizzleAtomMNSize per UMMA LayoutType (NONE->1, B32->2, B64->4, B128->8) = 1 << b_bits. W = 1 << byte_swizzle.b_bits swizzled = not swizzle_mode.is_none() @@ -236,6 +258,7 @@ def compute_umma_descriptor(tl_layout, buffer, transposed: bool, micro_size_k: i elem_bits=elem_bits, is_k_major=k_major, slice_byte_offset=slice_byte_offset, + k_panel_elems=k_panel_elems, ) @@ -951,6 +974,11 @@ def tcgen05_ss_atom( bk_atom_size = b_params.k_atom_size a_swizzle_atom_elems = a_params.swizzle_atom_elems b_swizzle_atom_elems = b_params.swizzle_atom_elems + # K-panel stride from the decoded layout; m_dim / n_dim_per_cta are operand + # extents and only match the panel spacing for a whole-buffer operand + # (see TCGEN05DescriptorParams.k_panel_elems). + a_k_panel_elems = a_params.k_panel_elems if a_params.k_panel_elems is not None else m_dim * a_swizzle_atom_elems + b_k_panel_elems = b_params.k_panel_elems if b_params.k_panel_elems is not None else n_dim_per_cta * b_swizzle_atom_elems mask_zero = T.cast(0, T.int32) # Pre-compute offsets @@ -958,16 +986,14 @@ def tcgen05_ss_atom( A_elem_offset = ( (ki % ak_atom_size) * micro_size_k + inst_m_idx * atom_m_per_cta * a_swizzle_atom_elems - + (ki // ak_atom_size) * m_dim * a_swizzle_atom_elems + + (ki // ak_atom_size) * a_k_panel_elems ) else: A_elem_offset = inst_m_idx * atom_m_per_cta * k_dim + ki * a_swizzle_atom_elems * micro_size_k if b_params.is_k_major: B_elem_offset = ( - (ki // bk_atom_size) * n_dim_per_cta * b_swizzle_atom_elems - + (ki % bk_atom_size) * micro_size_k - + inst_n_idx * atom_n * b_swizzle_atom_elems + (ki // bk_atom_size) * b_k_panel_elems + (ki % bk_atom_size) * micro_size_k + inst_n_idx * atom_n * b_swizzle_atom_elems ) else: B_elem_offset = ki * b_swizzle_atom_elems * micro_size_k + inst_n_idx * atom_n * ( @@ -1052,15 +1078,15 @@ def tcgen05_ts_atom( b_elem_bits = b_params.elem_bits bk_atom_size = b_params.k_atom_size b_swizzle_atom_elems = b_params.swizzle_atom_elems + # See TCGEN05DescriptorParams.k_panel_elems. + b_k_panel_elems = b_params.k_panel_elems if b_params.k_panel_elems is not None else n_dim_per_cta * b_swizzle_atom_elems mask_zero = T.cast(0, T.int32) A_tmem_offset = a_params.atom_offset(inst_m_idx * atom_m_per_cta, ki * meta.atom_k) if b_params.is_k_major: B_elem_offset = ( - (ki // bk_atom_size) * n_dim_per_cta * b_swizzle_atom_elems - + (ki % bk_atom_size) * micro_size_k - + inst_n_idx * atom_n * b_swizzle_atom_elems + (ki // bk_atom_size) * b_k_panel_elems + (ki % bk_atom_size) * micro_size_k + inst_n_idx * atom_n * b_swizzle_atom_elems ) else: B_elem_offset = ki * b_swizzle_atom_elems * micro_size_k + inst_n_idx * atom_n * ( @@ -1144,21 +1170,22 @@ def tcgen05_blockscaled_atom( bk_atom_size = b_params.k_atom_size a_swizzle_atom_elems = a_params.swizzle_atom_elems b_swizzle_atom_elems = b_params.swizzle_atom_elems + # See TCGEN05DescriptorParams.k_panel_elems. + a_k_panel_elems = a_params.k_panel_elems if a_params.k_panel_elems is not None else m_dim * a_swizzle_atom_elems + b_k_panel_elems = b_params.k_panel_elems if b_params.k_panel_elems is not None else n_dim_per_cta * b_swizzle_atom_elems if a_params.is_k_major: A_elem_offset = ( (ki % ak_atom_size) * micro_size_k + inst_m_idx * atom_m_per_cta * a_swizzle_atom_elems - + (ki // ak_atom_size) * m_dim * a_swizzle_atom_elems + + (ki // ak_atom_size) * a_k_panel_elems ) else: A_elem_offset = inst_m_idx * atom_m_per_cta * k_dim + ki * a_swizzle_atom_elems * micro_size_k if b_params.is_k_major: B_elem_offset = ( - (ki // bk_atom_size) * n_dim_per_cta * b_swizzle_atom_elems - + (ki % bk_atom_size) * micro_size_k - + inst_n_idx * atom_n * b_swizzle_atom_elems + (ki // bk_atom_size) * b_k_panel_elems + (ki % bk_atom_size) * micro_size_k + inst_n_idx * atom_n * b_swizzle_atom_elems ) else: B_elem_offset = ki * b_swizzle_atom_elems * micro_size_k + inst_n_idx * atom_n * ( diff --git a/tilelang/cuda/intrinsics/macro/wgmma_macro_generator.py b/tilelang/cuda/intrinsics/macro/wgmma_macro_generator.py index 5a9a450acd..db649f9a86 100644 --- a/tilelang/cuda/intrinsics/macro/wgmma_macro_generator.py +++ b/tilelang/cuda/intrinsics/macro/wgmma_macro_generator.py @@ -80,6 +80,14 @@ class WGMMADescriptorParams: buffer; passed to ``T.increase_descriptor_offset`` after building the descriptor from the buffer base. ``0`` for a whole-buffer / base-origin operand. Computed by :func:`compute_gmma_descriptor` from the CuTe layout.""" + k_panel_elems: int | None = None + """Element stride between consecutive K swizzle-atom panels (K-major only), + read off the decoded layout. ``None`` when the layout has a single K panel, + or for an MN-major operand, in which case the ``ki // k_atom_size`` term the + offset formulas multiply it into is always zero / unused. Callers must not + reconstruct this as ``mn_extent * swizzle_atom_elems``: for an operand that + is a slice of a wider buffer (``B[64:128, :]``) the panels stay spaced by the + *buffer's* MN extent, not the operand's.""" def compute_gmma_descriptor(tl_layout, buffer, transposed: bool, micro_size_k: int = 16, region=None) -> WGMMADescriptorParams: @@ -149,6 +157,21 @@ def compute_gmma_descriptor(tl_layout, buffer, transposed: bool, micro_size_k: i f"only row-major operand layouts are supported." ) + # Element stride between consecutive K swizzle-atom panels, taken from the + # decoded layout in element space. A K-major operand wider than one swizzle + # atom decodes as (MN, (atom, panels)), and the panel sub-mode's stride is + # what the atom offset formulas need to step by. Reconstructing it as + # `mn_dim * swizzle_atom_elems` is only valid when the operand covers the + # whole buffer -- for a slice (`B[64:128, :]`) mn_dim is the slice extent + # while the panels stay spaced by the buffer's MN extent. + k_panel_elems = None + if k_major: + k_shape = tile[k_idx].shape + if isinstance(k_shape, tuple) and len(k_shape) == 2: + panel_stride = tile[k_idx].stride[-1] + if isinstance(panel_stride, int): + k_panel_elems = panel_stride + # W per CuTe LayoutType (INTERLEAVE->1, B32->2, B64->4, B128->8) = 1 << b_bits. W = 1 << byte_swizzle.b_bits swizzled = not swizzle_mode.is_none() @@ -203,6 +226,7 @@ def compute_gmma_descriptor(tl_layout, buffer, transposed: bool, micro_size_k: i elems_in_bytes=elems_in_bytes, is_k_major=k_major, slice_byte_offset=slice_byte_offset, + k_panel_elems=k_panel_elems, ) @@ -643,6 +667,7 @@ def wgmma_rs_atom( elems_in_bytes = b_params.elems_in_bytes bk_atom_size = b_params.k_atom_size b_swizzle_atom_elems = b_params.swizzle_atom_elems + b_k_panel_elems = b_params.k_panel_elems if b_params.k_panel_elems is not None else n_dim * b_swizzle_atom_elems thread_binding = self.get_thread_binding() @@ -656,9 +681,7 @@ def _rs_atom(A_buf, desc_b, C_buf): scale_out = T.Select(ki != 0, 1, T.Select(clear_accum, 0, 1)) B_offset = ( - (ki // bk_atom_size) * n_dim * b_swizzle_atom_elems - + warp_j * wgmma_inst_n * b_swizzle_atom_elems - + (ki % bk_atom_size) * micro_size_k + (ki // bk_atom_size) * b_k_panel_elems + warp_j * wgmma_inst_n * b_swizzle_atom_elems + (ki % bk_atom_size) * micro_size_k if b_params.is_k_major else ( ki * b_swizzle_atom_elems * micro_size_k + warp_j * wgmma_inst_n * (k_dim if n_dim // b_swizzle_atom_elems > 1 else 1) @@ -744,6 +767,8 @@ def wgmma_ss_atom( bk_atom_size = b_params.k_atom_size a_swizzle_atom_elems = a_params.swizzle_atom_elems b_swizzle_atom_elems = b_params.swizzle_atom_elems + a_k_panel_elems = a_params.k_panel_elems if a_params.k_panel_elems is not None else m_dim * a_swizzle_atom_elems + b_k_panel_elems = b_params.k_panel_elems if b_params.k_panel_elems is not None else n_dim * b_swizzle_atom_elems thread_binding = self.get_thread_binding() @@ -757,16 +782,12 @@ def _ss_atom(desc_a, desc_b, C_buf): warp_j = warp_n * num_inst_n + inst_n_idx A_offset = ( - (ki % ak_atom_size) * micro_size_k - + warp_i * 64 * a_swizzle_atom_elems - + (ki // ak_atom_size) * m_dim * a_swizzle_atom_elems + (ki % ak_atom_size) * micro_size_k + warp_i * 64 * a_swizzle_atom_elems + (ki // ak_atom_size) * a_k_panel_elems if a_is_k_major else warp_i * 64 * k_dim + ki * a_swizzle_atom_elems * micro_size_k ) B_offset = ( - (ki // bk_atom_size) * n_dim * b_swizzle_atom_elems - + (ki % bk_atom_size) * micro_size_k - + warp_j * wgmma_inst_n * b_swizzle_atom_elems + (ki // bk_atom_size) * b_k_panel_elems + (ki % bk_atom_size) * micro_size_k + warp_j * wgmma_inst_n * b_swizzle_atom_elems if b_is_k_major else ( ki * b_swizzle_atom_elems * micro_size_k + warp_j * wgmma_inst_n * (k_dim if n_dim // b_swizzle_atom_elems > 1 else 1) diff --git a/tilelang/cuda/intrinsics/macro/wgmma_sp_macro_generator.py b/tilelang/cuda/intrinsics/macro/wgmma_sp_macro_generator.py index 06af8d7797..6cf888c28f 100644 --- a/tilelang/cuda/intrinsics/macro/wgmma_sp_macro_generator.py +++ b/tilelang/cuda/intrinsics/macro/wgmma_sp_macro_generator.py @@ -184,6 +184,11 @@ def wgmma_ss( # where max specially handles the case when n_dim is 8. ak_atom_size = a_params.k_atom_size bk_atom_size = b_params.k_atom_size + # K-panel stride comes from the decoded layout; m_dim/n_dim are the operand + # extents and only coincide with the panel spacing when the operand covers + # its whole buffer (see WGMMADescriptorParams.k_panel_elems). + a_k_panel_elems = a_params.k_panel_elems if a_params.k_panel_elems is not None else m_dim * a_swizzle_atom_elems + b_k_panel_elems = b_params.k_panel_elems if b_params.k_panel_elems is not None else n_dim * b_swizzle_atom_elems wgmma_inst_m, wgmma_inst_n = self.wgmma_inst_m, self.wgmma_inst_n num_inst_m = 4 * self.warp_row_tiles // wgmma_inst_m num_inst_n = self.warp_col_tiles // wgmma_inst_n @@ -253,13 +258,13 @@ def _warp_mma(A_base_ptr, B_base_ptr, C_buf): A_offset = ( (ki % ak_atom_size) * (micro_size_k // self.SPARSE_FACTOR) + warp_i * 64 * a_swizzle_atom_elems - + (ki // ak_atom_size) * m_dim * a_swizzle_atom_elems + + (ki // ak_atom_size) * a_k_panel_elems if a_is_k_major else warp_i * 64 * (k_dim // self.SPARSE_FACTOR) + ki * a_swizzle_atom_elems * (micro_size_k // self.SPARSE_FACTOR) ) B_offset = ( - (ki // bk_atom_size) * n_dim * b_swizzle_atom_elems + (ki // bk_atom_size) * b_k_panel_elems + (ki % bk_atom_size) * micro_size_k + warp_j * wgmma_inst_n * b_swizzle_atom_elems if b_is_k_major @@ -339,6 +344,9 @@ def wgmma_rs( b_swizzle_atom_elems = b_params.swizzle_atom_elems b_slice_byte_offset = b_params.slice_byte_offset bk_atom_size = b_params.k_atom_size + # See WGMMADescriptorParams.k_panel_elems: n_dim is the operand extent, not + # the panel spacing, once the operand is a slice of a wider buffer. + b_k_panel_elems = b_params.k_panel_elems if b_params.k_panel_elems is not None else n_dim * b_swizzle_atom_elems wgmma_inst_m, wgmma_inst_n = self.wgmma_inst_m, self.wgmma_inst_n num_inst_m = 4 * self.warp_row_tiles // wgmma_inst_m num_inst_n = self.warp_col_tiles // wgmma_inst_n @@ -395,7 +403,7 @@ def _warp_mma(A_buf, B_base_ptr, C_buf): warp_j = warp_n * num_inst_n + j A_offset = ki * warp_rows * local_size_a + i * local_size_a B_offset = ( - (ki // bk_atom_size) * n_dim * b_swizzle_atom_elems + (ki // bk_atom_size) * b_k_panel_elems + warp_j * wgmma_inst_n * b_swizzle_atom_elems + (ki % bk_atom_size) * micro_size_k if b_is_k_major From 5a84c00cb0e2245de275e5220f5cff37b4712f3d Mon Sep 17 00:00:00 2001 From: bigSheep123 <1554869970@qq.com> Date: Fri, 14 Aug 2026 14:13:04 +0800 Subject: [PATCH 2/2] [BugFix][Hopper][Blackwell] Refuse a multi-panel K mode with no constant stride Review follow-up on the K-panel stride decode. The decode returned None whenever the K mode was not a canonical 2-tuple, and every consumer then fell back to `mn_extent * swizzle_atom_elems` -- the exact reconstruction this series removes. The fallback is only sound for a single-panel layout, where `ki // k_atom_size` is always zero, and nothing checked that. So a swizzled multi-panel layout that failed to decode would silently re-enter the wrong-code path: the same shape as the original defect, where a formula's precondition was assumed but never enforced. Move the decode into a shared `decode_k_panel_elems` used by both descriptor builders, and make it assert instead of returning None when a swizzled operand spans more than one panel without exposing a constant panel stride. Panel count is rounded up, so a K extent that is not a whole number of atoms cannot slip through as single-panel. Also tighten the accept condition: reading `stride[-1]` as the panel step is only valid once the leading sub-mode is confirmed to be the contiguous atom. On a panels-first `(2, atom):(panel_stride, 1)` ordering the old condition would have returned 1 -- a silently wrong step rather than a refusal. Consumers now read the step through `k_panel_stride(mn_extent)` on the params dataclass, so the contract and its single-panel default live in one place instead of being restated at eleven call sites. New test_tilelang_cuda_wgmma_operand_layout.py covers the decode directly (hardware-independent): canonical multi-panel strides, the single-panel and MN-major None cases, and each refusal -- panels-first, flat multi-atom, a three-sub-mode K, and a partial-atom extent. Re-verified on H200: the whole-buffer invariant is unchanged by the tightening (120 K-major layouts x 2 builders, 60 single-panel, 0 disagreements), every multi-panel K slice is still accepted with the buffer-derived stride, and the full suite set passes (174 passed, 40 skipped). --- ...test_tilelang_cuda_wgmma_operand_layout.py | 71 +++++++++++++++ .../macro/tcgen05_macro_generator.py | 54 +++++------ .../intrinsics/macro/wgmma_macro_generator.py | 90 ++++++++++++++----- .../macro/wgmma_sp_macro_generator.py | 11 +-- 4 files changed, 164 insertions(+), 62 deletions(-) create mode 100644 testing/python/cuda/test_tilelang_cuda_wgmma_operand_layout.py diff --git a/testing/python/cuda/test_tilelang_cuda_wgmma_operand_layout.py b/testing/python/cuda/test_tilelang_cuda_wgmma_operand_layout.py new file mode 100644 index 0000000000..6ecd559dce --- /dev/null +++ b/testing/python/cuda/test_tilelang_cuda_wgmma_operand_layout.py @@ -0,0 +1,71 @@ +"""Unit tests for the WGMMA/UMMA K-panel stride decode. + +``decode_k_panel_elems`` is the single source of truth for how far an atom offset +steps to reach the next K swizzle-atom panel. Reconstructing that step from the +operand's MN extent is wrong for an operand that is a slice of a wider buffer, so +the decode must either return the layout's own stride or refuse — never fall back +silently, which is the wrong-code path the decode exists to remove. +""" + +import pytest + +import tilelang.testing # noqa: F401 (registers the test main) +from tilelang.layout import SwizzleMode, cute +from tilelang.cuda.intrinsics.macro.wgmma_macro_generator import decode_k_panel_elems + +ATOM = 64 # elements per 128B swizzle atom at 16-bit dtypes +B128 = SwizzleMode.SWIZZLE_128B +NONE = SwizzleMode.NONE + + +def _k_mode(shape, stride): + return cute.make_layout(shape, stride) + + +def test_canonical_multi_panel_returns_the_layout_stride(): + """(atom, panels):(1, panel_stride) -- the step is read off the layout.""" + assert decode_k_panel_elems(_k_mode([ATOM, 2], [1, 8192]), True, B128, ATOM, "WGMMA") == 8192 + assert decode_k_panel_elems(_k_mode([ATOM, 4], [1, 8192]), True, B128, ATOM, "WGMMA") == 8192 + # The stride is the buffer's, so a taller buffer gives a larger step for the + # same operand shape -- exactly what a sliced operand needs. + assert decode_k_panel_elems(_k_mode([ATOM, 2], [1, 16384]), True, B128, ATOM, "WGMMA") == 16384 + + +def test_single_panel_and_mn_major_need_no_step(): + """None is only returned where ``ki // k_atom_size`` is always zero.""" + # Flat K mode: one atom wide, so there is no second panel to step to. + assert decode_k_panel_elems(_k_mode(ATOM, 1), True, B128, ATOM, "WGMMA") is None + # MN-major operands never reach the panel term. + assert decode_k_panel_elems(_k_mode([ATOM, 2], [1, 8192]), False, B128, ATOM, "WGMMA") is None + # Unswizzled layouts have no panels at all. + assert decode_k_panel_elems(_k_mode(128, 1), True, NONE, 128, "WGMMA") is None + + +def test_panels_first_ordering_is_rejected_not_misread(): + """A (panels, atom) K mode must not have its trailing atom stride read as the step. + + Taking ``stride[-1]`` without checking that the leading sub-mode is the + contiguous atom would return 1 here — a silently wrong panel step. + """ + with pytest.raises(AssertionError, match=r"spans 2 swizzle-atom panels"): + decode_k_panel_elems(_k_mode([2, ATOM], [8192, 1]), True, B128, ATOM, "WGMMA") + + +def test_multi_panel_without_a_constant_stride_asserts(): + """Refusing beats falling back to the extent-based reconstruction.""" + # Flat but two atoms wide: no panel sub-mode to read a stride from. + with pytest.raises(AssertionError, match=r"spans 2 swizzle-atom panels"): + decode_k_panel_elems(_k_mode(2 * ATOM, 1), True, B128, ATOM, "UMMA") + # Three sub-modes: panel spacing is not expressible as one scalar step. + with pytest.raises(AssertionError, match=r"spans 4 swizzle-atom panels"): + decode_k_panel_elems(_k_mode([ATOM, 2, 2], [1, 8192, 16384]), True, B128, ATOM, "WGMMA") + + +def test_partial_atom_extent_still_counts_as_two_panels(): + """A K extent that is not a whole number of atoms must not read as single-panel.""" + with pytest.raises(AssertionError, match=r"spans 2 swizzle-atom panels"): + decode_k_panel_elems(_k_mode(ATOM + 16, 1), True, B128, ATOM, "WGMMA") + + +if __name__ == "__main__": + tilelang.testing.main() diff --git a/tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py b/tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py index e541fb8e3c..1aeebc1f7e 100644 --- a/tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py +++ b/tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py @@ -2,6 +2,7 @@ from dataclasses import dataclass import tilelang.language as T from .mma_macro_generator import TensorCoreIntrinEmitter as MMAIntrinEmitter +from .wgmma_macro_generator import decode_k_panel_elems from ..layout.mma_sm100_layout import ( TCGEN05Meta, make_tmem_frg_a, @@ -60,12 +61,22 @@ class TCGEN05DescriptorParams: operand. Computed by :func:`compute_umma_descriptor` from the CuTe layout.""" k_panel_elems: int | None = None """Element stride between consecutive K swizzle-atom panels (K-major only), - read off the decoded layout. ``None`` when the layout has a single K panel, - or for an MN-major operand, in which case the ``ki // k_atom_size`` term the - offset formulas multiply it into is always zero / unused. Callers must not - reconstruct this as ``mn_extent * swizzle_atom_elems``: for an operand that - is a slice of a wider buffer (``B[64:128, :]``) the panels stay spaced by the - *buffer's* MN extent, not the operand's.""" + read off the decoded layout by :func:`decode_k_panel_elems`. ``None`` only + when no panel step is needed -- see :meth:`k_panel_stride`, which is how + callers should read it.""" + + def k_panel_stride(self, mn_extent: int) -> int: + """K-panel step for the atom offset formulas. + + ``mn_extent`` is the operand's own MN extent, used only for the + single-panel / MN-major case where the ``ki // k_atom_size`` factor is + always zero and any value works. Never reconstruct the multi-panel step + from it: for a slice of a wider buffer the panels stay spaced by the + *buffer's* MN extent, not the operand's. + """ + if self.k_panel_elems is not None: + return self.k_panel_elems + return mn_extent * self.swizzle_atom_elems @dataclass(frozen=True) @@ -195,20 +206,6 @@ def compute_umma_descriptor(tl_layout, buffer, transposed: bool, micro_size_k: i f"only row-major operand layouts are supported." ) - # Element stride between consecutive K swizzle-atom panels, taken from the - # decoded layout: a K-major operand wider than one swizzle atom decodes as - # (MN, (atom, panels)), and the panel sub-mode's stride is what the offset - # formulas need. Reconstructing it as `mn_dim * swizzle_atom_elems` only holds - # when the operand covers its whole buffer -- for a slice, mn_dim shrinks while - # the panel spacing does not. - k_panel_elems = None - if k_major: - k_shape = tile[k_idx].shape - if isinstance(k_shape, tuple) and len(k_shape) == 2: - panel_stride = tile[k_idx].stride[-1] - if isinstance(panel_stride, int): - k_panel_elems = panel_stride - # SwizzleAtomMNSize per UMMA LayoutType (NONE->1, B32->2, B64->4, B128->8) = 1 << b_bits. W = 1 << byte_swizzle.b_bits swizzled = not swizzle_mode.is_none() @@ -249,6 +246,8 @@ def compute_umma_descriptor(tl_layout, buffer, transposed: bool, micro_size_k: i # Elements per swizzle atom along the non-K (MN) dimension; the unswizzled # case spans the whole MN tile (bit-based so sub-byte dtypes stay packed). swizzle_atom_elems = mn_dim if swizzle_mode.is_none() else _bytes_to_elements(swizzle_mode.swizzle_byte_size(), elem_bits) + + k_panel_elems = decode_k_panel_elems(tile[k_idx], k_major, swizzle_mode, swizzle_atom_elems, "UMMA") return TCGEN05DescriptorParams( swizzle_mode=swizzle_mode, leading_byte_offset=int(lbo), @@ -974,11 +973,8 @@ def tcgen05_ss_atom( bk_atom_size = b_params.k_atom_size a_swizzle_atom_elems = a_params.swizzle_atom_elems b_swizzle_atom_elems = b_params.swizzle_atom_elems - # K-panel stride from the decoded layout; m_dim / n_dim_per_cta are operand - # extents and only match the panel spacing for a whole-buffer operand - # (see TCGEN05DescriptorParams.k_panel_elems). - a_k_panel_elems = a_params.k_panel_elems if a_params.k_panel_elems is not None else m_dim * a_swizzle_atom_elems - b_k_panel_elems = b_params.k_panel_elems if b_params.k_panel_elems is not None else n_dim_per_cta * b_swizzle_atom_elems + a_k_panel_elems = a_params.k_panel_stride(m_dim) + b_k_panel_elems = b_params.k_panel_stride(n_dim_per_cta) mask_zero = T.cast(0, T.int32) # Pre-compute offsets @@ -1078,8 +1074,7 @@ def tcgen05_ts_atom( b_elem_bits = b_params.elem_bits bk_atom_size = b_params.k_atom_size b_swizzle_atom_elems = b_params.swizzle_atom_elems - # See TCGEN05DescriptorParams.k_panel_elems. - b_k_panel_elems = b_params.k_panel_elems if b_params.k_panel_elems is not None else n_dim_per_cta * b_swizzle_atom_elems + b_k_panel_elems = b_params.k_panel_stride(n_dim_per_cta) mask_zero = T.cast(0, T.int32) A_tmem_offset = a_params.atom_offset(inst_m_idx * atom_m_per_cta, ki * meta.atom_k) @@ -1170,9 +1165,8 @@ def tcgen05_blockscaled_atom( bk_atom_size = b_params.k_atom_size a_swizzle_atom_elems = a_params.swizzle_atom_elems b_swizzle_atom_elems = b_params.swizzle_atom_elems - # See TCGEN05DescriptorParams.k_panel_elems. - a_k_panel_elems = a_params.k_panel_elems if a_params.k_panel_elems is not None else m_dim * a_swizzle_atom_elems - b_k_panel_elems = b_params.k_panel_elems if b_params.k_panel_elems is not None else n_dim_per_cta * b_swizzle_atom_elems + a_k_panel_elems = a_params.k_panel_stride(m_dim) + b_k_panel_elems = b_params.k_panel_stride(n_dim_per_cta) if a_params.is_k_major: A_elem_offset = ( diff --git a/tilelang/cuda/intrinsics/macro/wgmma_macro_generator.py b/tilelang/cuda/intrinsics/macro/wgmma_macro_generator.py index db649f9a86..160c5f7d0f 100644 --- a/tilelang/cuda/intrinsics/macro/wgmma_macro_generator.py +++ b/tilelang/cuda/intrinsics/macro/wgmma_macro_generator.py @@ -30,6 +30,51 @@ def _min_leaf_stride(stride) -> int: return int(stride) +def decode_k_panel_elems(k_mode, k_major: bool, swizzle_mode, swizzle_atom_elems: int, kind: str) -> int | None: + """Element stride between consecutive K swizzle-atom panels, off the layout. + + Shared by the WGMMA and UMMA descriptor builders. A K-major operand wider + than one swizzle atom decodes with its K mode split into ``(atom, panels)``, + and that panel sub-mode's stride is the step the atom offset formulas need. + It cannot be reconstructed as ``mn_extent * swizzle_atom_elems``, which only + holds when the operand covers its whole buffer -- for a slice the MN extent + shrinks while the panel spacing does not. + + ``k_mode`` is the operand tile's K mode in element space (post-``restrict``). + Returns ``None`` only when no panel step is needed: an MN-major operand, or a + layout with a single K panel, where the ``ki // k_atom_size`` factor the step + is multiplied by is always zero. + + A swizzled multi-panel layout whose K mode is not the canonical + ``(atom, panels):(1, panel_stride)`` asserts rather than falling back to the + extent-based reconstruction: that fallback is precisely the defect this + decode exists to remove, so re-entering it silently would reintroduce a + wrong-code path. + """ + if not k_major: + return None + shape, stride = k_mode.shape, k_mode.stride + if isinstance(shape, tuple) and len(shape) == 2: + # Only read the trailing stride as the panel step once the leading + # sub-mode is confirmed to be the contiguous atom -- on a panels-first + # ordering the trailing stride would be the atom's, i.e. silently wrong. + panel_stride = stride[-1] + if stride[0] == 1 and shape[0] == swizzle_atom_elems and isinstance(panel_stride, int): + return panel_stride + if swizzle_mode.is_none(): + # Panels are a swizzle concept; an unswizzled operand has none. + return None + # Round up: a K extent that is not a whole number of atoms still spans two + # panels, and must not slip through as "single panel". + n_panels = -(-int(cute.size(k_mode)) // swizzle_atom_elems) + assert n_panels <= 1, ( + f"{kind} K-major operand spans {n_panels} swizzle-atom panels but its decoded K mode " + f"{shape}:{stride} is not the canonical (atom, panels):(1, panel_stride); the atom " + f"offset cannot be formed without a constant panel stride." + ) + return None + + def select_wgmma_inst_n(warp_col_tiles: int) -> int: """Widest legal WGMMA ``N`` that tiles ``warp_col_tiles`` exactly. @@ -82,12 +127,22 @@ class WGMMADescriptorParams: operand. Computed by :func:`compute_gmma_descriptor` from the CuTe layout.""" k_panel_elems: int | None = None """Element stride between consecutive K swizzle-atom panels (K-major only), - read off the decoded layout. ``None`` when the layout has a single K panel, - or for an MN-major operand, in which case the ``ki // k_atom_size`` term the - offset formulas multiply it into is always zero / unused. Callers must not - reconstruct this as ``mn_extent * swizzle_atom_elems``: for an operand that - is a slice of a wider buffer (``B[64:128, :]``) the panels stay spaced by the - *buffer's* MN extent, not the operand's.""" + read off the decoded layout by :func:`decode_k_panel_elems`. ``None`` only + when no panel step is needed -- see :meth:`k_panel_stride`, which is how + callers should read it.""" + + def k_panel_stride(self, mn_extent: int) -> int: + """K-panel step for the atom offset formulas. + + ``mn_extent`` is the operand's own MN extent, used only for the + single-panel / MN-major case where the ``ki // k_atom_size`` factor is + always zero and any value works. Never reconstruct the multi-panel step + from it: for a slice of a wider buffer the panels stay spaced by the + *buffer's* MN extent, not the operand's. + """ + if self.k_panel_elems is not None: + return self.k_panel_elems + return mn_extent * self.swizzle_atom_elems def compute_gmma_descriptor(tl_layout, buffer, transposed: bool, micro_size_k: int = 16, region=None) -> WGMMADescriptorParams: @@ -157,21 +212,6 @@ def compute_gmma_descriptor(tl_layout, buffer, transposed: bool, micro_size_k: i f"only row-major operand layouts are supported." ) - # Element stride between consecutive K swizzle-atom panels, taken from the - # decoded layout in element space. A K-major operand wider than one swizzle - # atom decodes as (MN, (atom, panels)), and the panel sub-mode's stride is - # what the atom offset formulas need to step by. Reconstructing it as - # `mn_dim * swizzle_atom_elems` is only valid when the operand covers the - # whole buffer -- for a slice (`B[64:128, :]`) mn_dim is the slice extent - # while the panels stay spaced by the buffer's MN extent. - k_panel_elems = None - if k_major: - k_shape = tile[k_idx].shape - if isinstance(k_shape, tuple) and len(k_shape) == 2: - panel_stride = tile[k_idx].stride[-1] - if isinstance(panel_stride, int): - k_panel_elems = panel_stride - # W per CuTe LayoutType (INTERLEAVE->1, B32->2, B64->4, B128->8) = 1 << b_bits. W = 1 << byte_swizzle.b_bits swizzled = not swizzle_mode.is_none() @@ -217,6 +257,8 @@ def compute_gmma_descriptor(tl_layout, buffer, transposed: bool, micro_size_k: i # Elements per swizzle atom along the non-K (MN) dimension; the unswizzled # case spans the whole MN tile. swizzle_atom_elems = mn_dim if swizzle_mode.is_none() else swizzle_mode.swizzle_byte_size() // elems_in_bytes + + k_panel_elems = decode_k_panel_elems(tile[k_idx], k_major, swizzle_mode, swizzle_atom_elems, "WGMMA") return WGMMADescriptorParams( swizzle_mode=swizzle_mode, leading_byte_offset=int(lbo), @@ -667,7 +709,7 @@ def wgmma_rs_atom( elems_in_bytes = b_params.elems_in_bytes bk_atom_size = b_params.k_atom_size b_swizzle_atom_elems = b_params.swizzle_atom_elems - b_k_panel_elems = b_params.k_panel_elems if b_params.k_panel_elems is not None else n_dim * b_swizzle_atom_elems + b_k_panel_elems = b_params.k_panel_stride(n_dim) thread_binding = self.get_thread_binding() @@ -767,8 +809,8 @@ def wgmma_ss_atom( bk_atom_size = b_params.k_atom_size a_swizzle_atom_elems = a_params.swizzle_atom_elems b_swizzle_atom_elems = b_params.swizzle_atom_elems - a_k_panel_elems = a_params.k_panel_elems if a_params.k_panel_elems is not None else m_dim * a_swizzle_atom_elems - b_k_panel_elems = b_params.k_panel_elems if b_params.k_panel_elems is not None else n_dim * b_swizzle_atom_elems + a_k_panel_elems = a_params.k_panel_stride(m_dim) + b_k_panel_elems = b_params.k_panel_stride(n_dim) thread_binding = self.get_thread_binding() diff --git a/tilelang/cuda/intrinsics/macro/wgmma_sp_macro_generator.py b/tilelang/cuda/intrinsics/macro/wgmma_sp_macro_generator.py index 6cf888c28f..219b1b7dad 100644 --- a/tilelang/cuda/intrinsics/macro/wgmma_sp_macro_generator.py +++ b/tilelang/cuda/intrinsics/macro/wgmma_sp_macro_generator.py @@ -184,11 +184,8 @@ def wgmma_ss( # where max specially handles the case when n_dim is 8. ak_atom_size = a_params.k_atom_size bk_atom_size = b_params.k_atom_size - # K-panel stride comes from the decoded layout; m_dim/n_dim are the operand - # extents and only coincide with the panel spacing when the operand covers - # its whole buffer (see WGMMADescriptorParams.k_panel_elems). - a_k_panel_elems = a_params.k_panel_elems if a_params.k_panel_elems is not None else m_dim * a_swizzle_atom_elems - b_k_panel_elems = b_params.k_panel_elems if b_params.k_panel_elems is not None else n_dim * b_swizzle_atom_elems + a_k_panel_elems = a_params.k_panel_stride(m_dim) + b_k_panel_elems = b_params.k_panel_stride(n_dim) wgmma_inst_m, wgmma_inst_n = self.wgmma_inst_m, self.wgmma_inst_n num_inst_m = 4 * self.warp_row_tiles // wgmma_inst_m num_inst_n = self.warp_col_tiles // wgmma_inst_n @@ -344,9 +341,7 @@ def wgmma_rs( b_swizzle_atom_elems = b_params.swizzle_atom_elems b_slice_byte_offset = b_params.slice_byte_offset bk_atom_size = b_params.k_atom_size - # See WGMMADescriptorParams.k_panel_elems: n_dim is the operand extent, not - # the panel spacing, once the operand is a slice of a wider buffer. - b_k_panel_elems = b_params.k_panel_elems if b_params.k_panel_elems is not None else n_dim * b_swizzle_atom_elems + b_k_panel_elems = b_params.k_panel_stride(n_dim) wgmma_inst_m, wgmma_inst_n = self.wgmma_inst_m, self.wgmma_inst_n num_inst_m = 4 * self.warp_row_tiles // wgmma_inst_m num_inst_n = self.warp_col_tiles // wgmma_inst_n