biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: v_major derivation in cotiled diag

Browse Source
This commit is contained in:
biondizzle
2026-05-24 01:53:40 +00:00
parent 2318bc3d77
commit 60fa7fb8f4
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tests/unit/test_cotiled_diag.py
View File

@@ -47,12 +47,12 @@ def main():
a_major = LayoutEnum.from_tensor(mQ).mma_major_mode()
b_major = LayoutEnum.from_tensor(mK).mma_major_mode()
# v_fmha layout
v_fmha_layout = cute.make_layout(
(pv_n_tile, s_k, 1),
stride=(1, pv_n_tile, pv_n_tile * s_k),
)
v_major = LayoutEnum.ROW_MAJOR # based on the v_fmha layout
# v_major: must match FmhaKernel's derivation
# In FmhaKernel, v_fmha is created with layout (pv_n_tile, s_k, 1) stride (1, pv_n_tile, ...)
# We create a temporary tensor to get the LayoutEnum
v_fmha_tensor = torch.randn(pv_n_tile, s_k, 1, dtype=torch.bfloat16, device='cuda')
mV_fmha = ct.from_dlpack(v_fmha_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_fmha_tensor))
v_major = LayoutEnum.from_tensor(mV_fmha).mma_major_mode()
qk_mma = utils.sm100.make_trivial_tiled_mma(
q_dtype, q_dtype, a_major, b_major, Float32,