xai-org/x-algorithm · error · RuntimeError

The layout of mdK is wrong

Error message

The layout of mdK is wrong

What it means

In flash_bwd_sm100.__call__, when dKV postprocess is disabled the K-gradient output tensor mdK must have an MMA K-major layout (as derived via LayoutEnum.from_tensor). A row-major/other layout raises RuntimeError because TMA/MMA stores assume K-major.

Source

Thrown at phoenix/xrex/cutedsl/ranker_fa4/flash_bwd_sm100.py:538

            self.tiled_mma_dQ,
        ) = self._get_tiled_mma()
        self._setup_smem_layout()

        self.cluster_shape_mnk = (*self.cluster_shape_mn, 1)
        self.cluster_layout_vmnk = cute.tiled_divide(
            cute.make_layout(self.cluster_shape_mnk),
            (self.tiled_mma_S.thr_id.shape,),
        )
        self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1])
        self.is_q_do_mcast = self.num_mcast_ctas_b > 1

        if const_expr(not self.dKV_postprocess):
            self.mdK_layout_enum = LayoutEnum.from_tensor(mdK)
            self.mdV_layout_enum = LayoutEnum.from_tensor(mdV)
            dK_major_mode = self.mdK_layout_enum.mma_major_mode()
            dV_major_mode = self.mdV_layout_enum.mma_major_mode()
            if const_expr(dK_major_mode != tcgen05.OperandMajorMode.K):
                raise RuntimeError("The layout of mdK is wrong")
            if const_expr(dV_major_mode != tcgen05.OperandMajorMode.K):
                raise RuntimeError("The layout of mdV is wrong")

        if const_expr(self.use_tma_store and not self.dKV_postprocess):
            tma_copy_op_dKV = cpasync.CopyBulkTensorTileS2GOp()
            tma_atom_dK, mdK_tma_tensor = cpasync.make_tiled_tma_atom(
                tma_copy_op_dKV,
                mdK,
                cute.select(self.sdK_layout, mode=[0, 1]),
                self.sdK_epi_tile,
                1,
            )
            tma_atom_dV, mdV_tma_tensor = cpasync.make_tiled_tma_atom(
                tma_copy_op_dKV,
                mdV,
                cute.select(self.sdV_layout, mode=[0, 1]),
                self.sdV_epi_tile,
                1,

View on GitHub (pinned to 24c60942c5)

Solutions

1. Allocate mdK with the same layout/stride order as the input K tensor (K-major), e.g. torch.empty_like(k)
2. Avoid passing stride-permuted views; make the tensor contiguous in the expected order
3. Check mdK.stride() matches the K-major pattern expected by tcgen05.OperandMajorMode.K

Example fix

# before
dK = torch.empty_like(q)                # wrong layout
dV = torch.empty_like(q)

# after
dK = torch.empty_like(k)                # K-major, matches mdK layout enum
dV = torch.empty_like(v)

assert dK.stride() == k.stride() and dK.shape == k.shape, (dK.stride(), k.stride())
assert dV.stride() == v.stride()

def k_major_like(k: torch.Tensor, ref: torch.Tensor) -> bool:
    return k.shape == ref.shape and k.stride() == ref.stride()

When it happens

Trigger: Calling the SM100 backward kernel with a transposed or non-K-major mdK tensor — e.g. allocating dK with shape/layout mirroring a row-major Q tensor instead of the K-major layout the kernel writes, or passing a view with permuted strides.

Common situations: Allocating dK/dV outputs with torch.empty_like(q) instead of matching K's layout; permuting strides for a fused epilogue; a container/refactor changing tensor stride order.

Related errors

• The layout of mdV is wrong
• Block sparsity + paged KV not supported on SM100
• block-sparse arrays cover {bs_num_blocks} m-tiles but the ke
• Block sparse tensors{context} must have shapes (B, H, M) and
• Only Float16 or BFloat16 is supported

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/ec28c46e1738ccfb. Report an issue: GitHub.