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:
@@ -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<
|
||||
|
||||
Reference in New Issue
Block a user