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

fix: proper v_major from tensor

Browse Source
This commit is contained in:
biondizzle
2026-05-24 01:55:37 +00:00
parent 384de55dfd
commit 394f08601a
1 changed files with 3 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tests/unit/test_cotiled_diag.py
View File

@@ -39,12 +39,11 @@ def main():
# Derive major modes exactly as FmhaKernel does
from cutlass.utils import LayoutEnum
from cutlass.cute.nvgpu import OperandMajorMode
a_major = LayoutEnum.from_tensor(mQ).mma_major_mode()
b_major = LayoutEnum.from_tensor(mK).mma_major_mode()
# V FMHA layout: use explicit constants
# pv_n_tile=256, s_k=128
v_major = LayoutEnum.COL_MAJOR # layout (256, 128, 1) stride (1, 256, 32768) = col-major
v_major = LayoutEnum.from_tensor(mV).mma_major_mode()
print(f"a_major: {a_major}, b_major: {b_major}, v_major: {v_major}") # layout (256, 128, 1) stride (1, 256, 32768) = col-major
c_layout = LayoutEnum.from_tensor(mC)