xai-org/x-algorithm · error · RuntimeError

The layout of mdV is wrong

Error message

The layout of mdV is wrong

What it means

In the SM100 flash-attention backward kernel, mdV (the dV accumulator tensor) is expected to be laid out in K-major (column-major) mode as required by tcgen05 MMA operands. The kernel derives the layout via LayoutEnum.from_tensor(mdV) and checks its mma_major_mode(); if it is not OperandMajorMode.K it aborts because the MMA instruction descriptor would be invalid.

Source

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

        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,
            )
        else:

View on GitHub (pinned to 24c60942c5)

Solutions

1. Check how mdV is allocated/passed on the host side and ensure its contiguous dimension is the K (head-dim/inner) axis so the layout enum resolves to K-major
2. Transpose or re-layout mdV (cute.select or .T equivalent) before invoking the kernel, mirroring how mdK is prepared
3. Verify mdK and mdV are constructed identically; the mdK check above passing while mdV fails points at asymmetry in their creation
4. If tensors come from PyTorch, ensure the corresponding dV tensor is made contiguous in the same dim as dK

Example fix

// before
mdV = cute.make_tensor(dV_ptr, mD_layout)  # M-major
// after
mdV = cute.make_tensor(dV_ptr, cute.select(mD_layout, mode=[1, 0, 2, 3]))  # K-major

from phoenix.xrex.cutedsl.ranker_fa4.flash_bwd_sm100 import LayoutEnum
import cutlass.tcgen05 as tcgen05
assert LayoutEnum.from_tensor(mdV).mma_major_mode() == tcgen05.OperandMajorMode.K, 'mdV must be K-major'

When it happens

Trigger: Calling the flash backward kernel __call__ with an mdV tensor whose memory layout is M-major/row-major (e.g. created transposed or with a swapped stride order) instead of K-major, so mma_major_mode() returns MN mode.

Common situations: Passing dV output tensors allocated with wrong stride order, reusing tensors from the forward pass without transposing, or changes in how the host-side wrapper constructs mdV after a refactor of the kernel signature.

Related errors

• block-sparse arrays cover {bs_num_blocks} m-tiles but the ke
• Unsupported tcgen05 MMA op kind: {type(op).__name__}
• Block sparse tensors{context} must have shapes (B, H, M) and
• Only Float16 or BFloat16 is supported
• The layout of mdK is wrong

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