NVFP4 mega MoE: sf_id=0 fix for scale_vec::4X + UINT8 TMA + SF pipeline + interleaving

Root cause of ILLEGAL_INSTRUCTION: make_runtime_instr_desc_with_sf_id(instr_desc, k, k)
passed sf_id=1 for k=1 (second UMMA atom), but mxf4nvf4 with scale_vec::4X requires
sf_id=0 always — the hardware implicitly reads 4 SF positions per atom from a single
TMEM region. Non-zero sf_id causes the hardware to access invalid TMEM offsets.

Also includes:
- UINT8 TMA for packed FP4 (avoids 16U4 driver bugs)
- NVFP4 SF pipeline: 2 K-columns per BLOCK_K for group_size=16
- MN-major SF interleaving for gate/up L1 weights
- Fix contiguous copy for SF byte view
- Preserve MN-major layout in SF interleave
- Force contiguous on SF tensors before C++ call
- Unpack weight tuples before printing
- Single transpose back to MN-major (don't double-transpose)
This commit is contained in:
2026-05-12 20:26:13 +00:00
parent 26a8ab75a1
commit 74bf612771
5 changed files with 138 additions and 10 deletions

View File

@@ -882,6 +882,8 @@ sm100_fp8_nvfp4_mega_moe_impl(void* y,
// NVFP4: group=16 → 2 SF K-columns per BLOCK_K (128/16/4=2)
// Each UTCCP call moves 128 int32s → 4 TMEM cols
// We need 2 UTCCP calls per SF: one per K-column
// NOTE: No SMEM warp transpose needed — transform_sf_token_idx
// pre-arranges the data in the correct UTCCP layout via global memory
using cute_utccp_t = cute::SM100_UTCCP_4x32dp128bit_2cta;
#pragma unroll
@@ -906,8 +908,14 @@ sm100_fp8_nvfp4_mega_moe_impl(void* y,
// Issue UMMA
#pragma unroll
for (uint32_t k = 0; k < BLOCK_K / UMMA_K; ++ k) {
// NVFP4 scale_vec::4X: sf_id must always be 0.
// The hardware implicitly reads 4 SF positions per UMMA atom
// from the single TMEM region [scale_A_tmem]/[scale_B_tmem].
// Unlike scale_vec::1X (MXFP4) where each atom needs a unique sf_id
// to index sub-columns, scale_vec::4X ignores sf_id or requires 0.
// Passing sf_id=k (k=1 for second UMMA atom) was the ILLEGAL_INSTRUCTION bug.
const auto runtime_instr_desc =
mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, k, k);
mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, 0, 0);
a_desc.lo = mma::sm100::advance_umma_desc_lo<
cute::UMMA::Major::K, LOAD_BLOCK_M, kSwizzleAMode, uint8_t>(a_desc_base_lo, 0, k * (UMMA_K / 2));
b_desc.lo = mma::sm100::advance_umma_desc_lo<