{"record":{"id":"ec28c46e1738ccfb","repo":"xai-org/x-algorithm","slug":"the-layout-of-mdk-is-wrong","errorCode":null,"errorMessage":"The layout of mdK is wrong","messagePattern":"The layout of mdK is wrong","errorType":"validation","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/cutedsl/ranker_fa4/flash_bwd_sm100.py","lineNumber":538,"sourceCode":"            self.tiled_mma_dQ,\n        ) = self._get_tiled_mma()\n        self._setup_smem_layout()\n\n        self.cluster_shape_mnk = (*self.cluster_shape_mn, 1)\n        self.cluster_layout_vmnk = cute.tiled_divide(\n            cute.make_layout(self.cluster_shape_mnk),\n            (self.tiled_mma_S.thr_id.shape,),\n        )\n        self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1])\n        self.is_q_do_mcast = self.num_mcast_ctas_b > 1\n\n        if const_expr(not self.dKV_postprocess):\n            self.mdK_layout_enum = LayoutEnum.from_tensor(mdK)\n            self.mdV_layout_enum = LayoutEnum.from_tensor(mdV)\n            dK_major_mode = self.mdK_layout_enum.mma_major_mode()\n            dV_major_mode = self.mdV_layout_enum.mma_major_mode()\n            if const_expr(dK_major_mode != tcgen05.OperandMajorMode.K):\n                raise RuntimeError(\"The layout of mdK is wrong\")\n            if const_expr(dV_major_mode != tcgen05.OperandMajorMode.K):\n                raise RuntimeError(\"The layout of mdV is wrong\")\n\n        if const_expr(self.use_tma_store and not self.dKV_postprocess):\n            tma_copy_op_dKV = cpasync.CopyBulkTensorTileS2GOp()\n            tma_atom_dK, mdK_tma_tensor = cpasync.make_tiled_tma_atom(\n                tma_copy_op_dKV,\n                mdK,\n                cute.select(self.sdK_layout, mode=[0, 1]),\n                self.sdK_epi_tile,\n                1,\n            )\n            tma_atom_dV, mdV_tma_tensor = cpasync.make_tiled_tma_atom(\n                tma_copy_op_dKV,\n                mdV,\n                cute.select(self.sdV_layout, mode=[0, 1]),\n                self.sdV_epi_tile,\n                1,","sourceCodeStart":520,"sourceCodeEnd":556,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/cutedsl/ranker_fa4/flash_bwd_sm100.py#L520-L556","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Allocate mdK with the same layout/stride order as the input K tensor (K-major), e.g. torch.empty_like(k)","Avoid passing stride-permuted views; make the tensor contiguous in the expected order","Check mdK.stride() matches the K-major pattern expected by tcgen05.OperandMajorMode.K"],"exampleFix":"# before\ndK = torch.empty_like(q)                # wrong layout\ndV = torch.empty_like(q)\n\n# after\ndK = torch.empty_like(k)                # K-major, matches mdK layout enum\ndV = torch.empty_like(v)","handlingStrategy":"validation","validationCode":"assert dK.stride() == k.stride() and dK.shape == k.shape, (dK.stride(), k.stride())\nassert dV.stride() == v.stride()","typeGuard":"def k_major_like(k: torch.Tensor, ref: torch.Tensor) -> bool:\n    return k.shape == ref.shape and k.stride() == ref.stride()","tryCatchPattern":null,"preventionTips":["Allocate dK/dV with torch.empty_like(k)/empty_like(v)","Never pass stride-permuted views as gradient outputs"],"tags":["layout-validation","flash-attention","sm100","tensor-strides"],"backgroundTag":"unsupported-dtype","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}