diff --git a/kernels/mixed_moe_gemm_2stage.py b/kernels/mixed_moe_gemm_2stage.py index 712291931..3af440ab5 100644 --- a/kernels/mixed_moe_gemm_2stage.py +++ b/kernels/mixed_moe_gemm_2stage.py @@ -138,8 +138,8 @@ def compile_mixed_moe_gemm1( allocator_ping = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem1") _state = {} - if a_dtype not in ("fp8", "fp16", "int8", "fp4"): - raise ValueError(f"a_dtype must be one of ('fp8','fp16','int8','fp4'), got {a_dtype!r}") + if a_dtype not in ("fp8", "fp16", "int8", "fp4", "fp6"): + raise ValueError(f"a_dtype must be one of ('fp8','fp16','int8','fp4','fp6'), got {a_dtype!r}") if b_dtype not in ("fp8", "fp16", "int8", "int4", "fp4"): raise ValueError(f"b_dtype must be one of ('fp8','fp16','int8','int4','fp4'), got {b_dtype!r}") @@ -147,6 +147,7 @@ def compile_mixed_moe_gemm1( is_f16_b = b_dtype == "fp16" is_f8_a = a_dtype == "fp8" is_f4_a = a_dtype == "fp4" + is_f6_a = a_dtype == "fp6" is_f4_b = b_dtype == "fp4" sort_block_m = max(32, tile_m) @@ -161,8 +162,11 @@ def compile_mixed_moe_gemm1( a_elem_bytes = 2 if is_f16_a else 1 b_elem_bytes = 1 tile_k_bytes = int(tile_k) * int(a_elem_bytes) + # fp4: 2 elements per byte (packed); fp6: FP8-padded layout = 1 byte/element (like fp8). a_elem_vec_pack = 2 if is_f4_a else 1 - cbsz = 0 if is_f8_a else 4 + # fp6: 32 B per K=32 chunk (24 B packed codes + 8 B zero pad); fp4/fp8: 16 B. + a_per_lane_kpack_bytes = 32 if is_f6_a else 16 + cbsz = 0 if is_f8_a else (2 if is_f6_a else 4) blgp = 4 if (tile_k_bytes % 64) != 0: @@ -409,7 +413,7 @@ def x_lds_elem(): _pp_b_loads = [p["b_loads"] for p in _pipe_phases] _pp_has_scale = [p["has_scale"] for p in _pipe_phases] - fp4_ratio = 2 if a_dtype == "fp4" else 1 + fp4_ratio = 2 if is_f4_a else 1 # fp6 uses FP8-padded layout (1 byte/elem), same count as fp8 gui_ratio = 1 if gate_up_interleave else 2 _vmcnt_before_barrier = tile_m // 32 // fp4_ratio + tile_n // 32 * gui_ratio @@ -736,7 +740,7 @@ def load_x_tile(base_k): lane_div_16 = layout_get(coord_l16, 0) lane_mod_16 = layout_get(coord_l16, 1) row_a_lds = lane_mod_16 - col_offset_base = lane_div_16 * arith.constant(16, index=True) + col_offset_base = lane_div_16 * arith.constant(a_per_lane_kpack_bytes, index=True) num_acc_n = n_per_wave // 16 c_n_per_wave = arith.constant(n_per_wave, index=True) @@ -1034,6 +1038,11 @@ def prefetch_full_a_from_lds(lds_buffer, ku_limit=k_unroll): if const_expr(is_f8_a): a2, a3 = lds_load_packs_k64(curr_row, col_base + 64, lds_buffer) a_regs.append((a0, a1, a2, a3)) + elif const_expr(is_f6_a): + # fp6: 32B FP8-padded slot; 3rd 16B chunk carries codes, + # 4th 8B is zero pad (ignored by cbsz=2 MFMA). + a2, _ = lds_load_packs_k64(curr_row, col_base + 16, lds_buffer) + a_regs.append((a0, a1, a2)) else: a_regs.append((a0, a1)) return a_regs @@ -1156,6 +1165,10 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): if const_expr(is_f8_a): a0, a1, a2, a3 = a_tile_regs[_a_reg_idx] a128 = pack_i64x4_to_i32x8(a0, a1, a2, a3) + elif const_expr(is_f6_a): + # fp6: 3x16B loads; 4th slot is zero pad (cbsz=2 ignores it) + a0, a1, a2 = a_tile_regs[_a_reg_idx] + a128 = pack_i64x4_to_i32x8(a0, a1, a2, c0_i64) else: a0, a1 = a_tile_regs[_a_reg_idx] a128 = pack_i64x4_to_i32x8(a0, a1, c0_i64, c0_i64) @@ -1200,6 +1213,10 @@ def load_a_subtile(k_idx, mi_idx, lds_buffer): if const_expr(is_f8_a): a2, a3 = lds_load_packs_k64(curr_row, col_base + 64, lds_buffer) return (a0, a1, a2, a3) + elif const_expr(is_f6_a): + # fp6: 3x16B loads; 4th slot is zero pad (cbsz=2 ignores it) + a2, _ = lds_load_packs_k64(curr_row, col_base + 16, lds_buffer) + return (a0, a1, a2) else: return (a0, a1) @@ -1250,6 +1267,9 @@ def _pack(x0, x1, x2, x3): if const_expr(is_f8_a): a128 = _pack(a_reg[0], a_reg[1], a_reg[2], a_reg[3]) + elif const_expr(is_f6_a): + # fp6: 3x16B loads; 4th slot is zero pad (cbsz=2 ignores it) + a128 = _pack(a_reg[0], a_reg[1], a_reg[2], c0_i64) else: a128 = _pack(a_reg[0], a_reg[1], c0_i64, c0_i64) @@ -2479,8 +2499,8 @@ def compile_mixed_moe_gemm2( allocator_ping = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem1") _state = {} - if a_dtype not in ("fp8", "fp16", "int8", "fp4"): - raise ValueError(f"a_dtype must be one of ('fp8','fp16','int8','fp4'), got {a_dtype!r}") + if a_dtype not in ("fp8", "fp16", "int8", "fp4", "fp6"): + raise ValueError(f"a_dtype must be one of ('fp8','fp16','int8','fp4','fp6'), got {a_dtype!r}") if b_dtype not in ("fp8", "fp16", "int8", "int4", "fp4"): raise ValueError(f"b_dtype must be one of ('fp8','fp16','int8','int4','fp4'), got {b_dtype!r}") @@ -2489,6 +2509,8 @@ def compile_mixed_moe_gemm2( is_f8_a = a_dtype == "fp8" is_f4_a = a_dtype == "fp4" + is_f6_a = a_dtype == "fp6" + is_f4_or_f6_a = is_f4_a or is_f6_a is_f4_b = b_dtype == "fp4" _scale_pack_m = 2 # physical mn_pack in preshuffle microscale layout @@ -2505,8 +2527,11 @@ def compile_mixed_moe_gemm2( b_elem_bytes = 1 tile_k_bytes = int(tile_k) * int(a_elem_bytes) + # fp4: 2 elements per byte (packed); fp6: FP8-padded layout = 1 byte/element (like fp8). a_elem_vec_pack = 2 if is_f4_a else 1 - cbsz = 0 if is_f8_a else 4 + # fp6: 32 B per K=32 chunk (24 B packed codes + 8 B zero pad); fp4/fp8: 16 B. + a_per_lane_kpack_bytes = 32 if is_f6_a else 16 + cbsz = 0 if is_f8_a else (2 if is_f6_a else 4) blgp = 4 # ---- Static B preshuffle strides (compile-time) ---- @@ -2548,7 +2573,7 @@ def compile_mixed_moe_gemm2( mfma_i32_k32 = getattr(rocdl, "mfma_i32_16x16x32i8", None) or getattr(rocdl, "mfma_i32_16x16x32_i8", None) if mfma_i32_k32 is None: raise AttributeError( - "INT8 K32 MFMA op not found: expected `rocdl.mfma_i32_16x16x32i8` " "(or `rocdl.mfma_i32_16x16x32_i8`)." + "INT8 K32 MFMA op not found: expected `rocdl.mfma_i32_16x16x32i8` (or `rocdl.mfma_i32_16x16x32_i8`)." ) def _x_elem_type(): @@ -2819,7 +2844,7 @@ def check_c_k_valid_gate(base_k): sx_rsrc = 1 sw_rsrc = 1 if const_expr(not is_f16_a): - if const_expr(is_f4_a or is_f8_a): + if const_expr(is_f4_or_f6_a or is_f8_a): # A2 microscale: e8m0 in sorted layout [sorted_size, K/32]. # Caller must pre-scatter a2_scale via moe_mxfp4_sort. kblk = _div_pow2(k_in, 32) @@ -3083,7 +3108,7 @@ def load_x_tile(base_k): row_a_lds = lane_mod_16 - col_offset_base = lane_div_16 * arith.constant(16, index=True) + col_offset_base = lane_div_16 * arith.constant(a_per_lane_kpack_bytes, index=True) # Dynamic N tiling within block. num_waves = 4 @@ -3448,6 +3473,10 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): col_base1 = col_base + 64 a2, a3 = lds_load_packs_k64(curr_row_a_lds, col_base1, lds_buffer) a128 = pack_i64x4_to_i32x8(a0, a1, a2, a3) + elif const_expr(is_f6_a): + # fp6: 3x16B loads; 4th slot is zero pad (cbsz=2 ignores it) + a2, _ = lds_load_packs_k64(curr_row_a_lds, col_base0 + 16, lds_buffer) + a128 = pack_i64x4_to_i32x8(a0, a1, a2, c0_i64) else: a128 = pack_i64x4_to_i32x8(a0, a1, c0_i64, c0_i64) diff --git a/tests/kernels/test_moe_gemm.py b/tests/kernels/test_moe_gemm.py index e956f83b5..34bd3f9aa 100644 --- a/tests/kernels/test_moe_gemm.py +++ b/tests/kernels/test_moe_gemm.py @@ -416,7 +416,7 @@ def run_moe_stage1( # f"{' (even)' if even_dispatch else ' (random)'}" # ) - if in_dtype not in ("fp8", "fp16", "bf16", "int8", "int8smooth", "int4", "int4_bf16", "fp4", "a8w4"): + if in_dtype not in ("fp8", "fp16", "bf16", "int8", "int8smooth", "int4", "int4_bf16", "fp4", "a8w4", "a6w4"): raise ValueError( f"in_dtype must be one of ('fp8','fp16','bf16','int8','int8smooth','int4','int4_bf16','fp4','a8w4'), got {in_dtype!r}" ) @@ -799,7 +799,7 @@ def launch(o, x, w, sx, sw, st, eids, sw_sorted): print( f"FlyDSL MoE stage1[{in_dtype}]: " f"{us:.1f} us, " - f"{tflops:.2f} TFLOPS(logical, M={tokens*topk}), " + f"{tflops:.2f} TFLOPS(logical, M={tokens * topk}), " f"{tbps:.3f} TB/s (doweight_stage1={doweight_stage1})" ) # Compare + benchmark vs aiter stage1 (optional; enabled by default when aiter is runnable). @@ -912,6 +912,7 @@ def run_moe_stage2( routing_in: Optional[RoutingBuffers] = None, a2_fp8_in: Optional[torch.Tensor] = None, a2_scale_in: Optional[torch.Tensor] = None, + a2_ref_in: Optional[torch.Tensor] = None, return_outputs: bool = False, skip_ref: bool = False, init_scale: float = 0.2, @@ -1042,7 +1043,7 @@ def run_moe_stage2( # f"{' (even)' if even_dispatch else ' (random)'}" # ) - if in_dtype not in ("fp8", "fp16", "bf16", "int8", "int8smooth", "int4", "int4_bf16", "fp4", "a8w4"): + if in_dtype not in ("fp8", "fp16", "bf16", "int8", "int8smooth", "int4", "int4_bf16", "fp4", "a8w4", "a6w4"): raise ValueError( f"in_dtype must be one of ('fp8','fp16','bf16','int8','int8smooth','int4','int4_bf16','fp4','a8w4'), got {in_dtype!r}" ) @@ -1051,8 +1052,9 @@ def run_moe_stage2( is_int8smooth = in_dtype == "int8smooth" is_fp4 = in_dtype == "fp4" is_a8w4 = in_dtype == "a8w4" # MX-FP8 activation + MX-FP4 weight + is_a6w4 = in_dtype == "a6w4" # MXFP6-E2M3 activation + MXFP4 weight # Share the FP4 stage2 path (W2 shuffle / scale sort / mixed kernel). - is_fp4_path = is_fp4 or is_a8w4 + is_fp4_path = is_fp4 or is_a8w4 or is_a6w4 use_packed_int4 = is_int4 or is_int4_bf16 # Quantize inputs / weights. @@ -1090,13 +1092,13 @@ def run_moe_stage2( w1_q, scale_w1 = pertoken_quant(w1_fp32, quant_dtype=torch.int8, dtypeMax=7) w2_q, scale_w2 = pertoken_quant(w2_fp32, quant_dtype=torch.int8, dtypeMax=7) scale_x = None - elif in_dtype in ("fp4", "a8w4"): + elif in_dtype in ("fp4", "a8w4", "a6w4"): from tests.kernels.utils import fp4_utils if fp4_utils is None: pytest.skip("fp4_utils not available (triton not installed)") if "gfx95" not in ARCH: - pytest.skip(f"FP4 MFMA requires gfx950+, got {ARCH}") + pytest.skip(f"FP4/FP6 MFMA requires gfx950+, got {ARCH}") # FP4 / A8W4 share the MXFP4 W2 path; quantize W2 only here. A2 comes # from `a2_fp8_in` (FP4 packed bytes for 'fp4', MX-FP8 e4m3fn for 'a8w4'). w2_flat_fp32 = w2_fp32.view(experts * model_dim, inter_dim) @@ -1263,7 +1265,7 @@ def run_moe_stage2( if is_fp4_path: fp4_accumulate = not bool(use_reduce) - a_dtype_kernel = "fp8" if is_a8w4 else "fp4" + a_dtype_kernel = "fp8" if is_a8w4 else ("fp6" if is_a6w4 else "fp4") exe = compile_mixed_moe_gemm2( model_dim=model_dim, inter_dim=inter_dim, @@ -1485,8 +1487,12 @@ def launch(o, x, w, sx, sw, st, eids, sw_sorted): torch.cuda.synchronize() if not bool(skip_ref): + # For a6w4: use a2_ref_in (low-6-bit E2M3 codes) instead of a2_q + # (FP8-padded kernel layout), and tell the reference to decode as mxfp6. + _a2_ref = a2_ref_in if a2_ref_in is not None else a2_q + _a2_kind = "mxfp6" if a2_ref_in is not None else None ref2 = torch_moe_gemm2( - a2_q, + _a2_ref, w2_q, a2_scale, scale_w2, @@ -1496,6 +1502,7 @@ def launch(o, x, w, sx, sw, st, eids, sw_sorted): doweight_stage2=doweight_stage2, group_size=group_size, scale_w2_groups=scale_w2_groups, + a2_kind=_a2_kind, ) assert verify_output(out.to(torch.float32), ref2, rtol=0.5, atol=0.5) @@ -1525,6 +1532,7 @@ def launch(o, x, w, sx, sw, st, eids, sw_sorted): "int4_bf16": (16, 4, "none", "per_row_f32"), "fp4": (4, 4, "per_block_e8m0_32", "per_block_e8m0_32"), "a8w4": (8, 4, "per_block_e8m0_32", "per_block_e8m0_32"), + "a6w4": (6, 4, "per_block_e8m0_32", "per_block_e8m0_32"), } a_bits, w_bits, a_scale_mode, w_scale_mode = _BYTES_SPEC[in_dtype] if use_groupwise_scale: @@ -1559,7 +1567,7 @@ def launch(o, x, w, sx, sw, st, eids, sw_sorted): tbps = float("nan") if us <= 0 else bytes_moved / 1e12 / (us / 1e6) print( f"FlyDSL MoE stage2 [{kernel_name}] {in_dtype} {'reduce' if use_reduce else 'atomic'} | " - f"{model_dim}x{inter_dim}, E={experts}, K={topk}, M_eff={tokens*topk} | " + f"{model_dim}x{inter_dim}, E={experts}, K={topk}, M_eff={tokens * topk} | " f"{us:.1f} us, {tflops:.2f} TFLOPS, {tbps:.3f} TB/s" ) # Optional compare vs aiter stage2. @@ -1622,7 +1630,7 @@ def launch_ck(o, a2_, w1_, w2_, sorted_ids_, sorted_eids_, num_valid_, w2_scale_ tflops_ck = flops / (us_ck / 1e6) / 1e12 print( f"[aiter] stage2: {us_ck:.1f} us, " - f"{tflops_ck:.2f} TFLOPS(logical, M={tokens*topk}), FlyDSL vs aiter speedups: {tflops / tflops_ck:.2f}x" + f"{tflops_ck:.2f} TFLOPS(logical, M={tokens * topk}), FlyDSL vs aiter speedups: {tflops / tflops_ck:.2f}x" ) # Correctness run (best-effort; do not fail perf comparison if aiter diverges). @@ -1873,7 +1881,11 @@ def test_moe_gemm_2stage( # a8w4 -> MX-FP8 e4m3fn (1 B/elem) # run_moe_stage2 sorts the raw E8M0 scale [tokens*topk, inter_dim//32] internally. out1_fp32 = out1_fp16.to(torch.float32).view(tokens * topk, inter_dim) - quantize_a2 = _per_1x32_mxfp8_quant if in_dtype == "a8w4" else _per_1x32_fp4_quant + quantize_a2 = ( + _per_1x32_mxfp8_quant + if in_dtype == "a8w4" + else _per_1x32_mxfp6_quant if in_dtype == "a6w4" else _per_1x32_fp4_quant + ) a2_q, a2_scale = quantize_a2(out1_fp32) elif w_fp4_kernel: a2_q = out1_fp16.to(torch.float32) @@ -1985,6 +1997,21 @@ def _per_1x32_mxfp8_quant(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: return x_q, scale_bytes +def _per_1x32_mxfp6_quant( + x: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize a tensor to MXFP6-E2M3 with per-1x32 E8M0 block scaling. + + Returns: + a_pad [..., K] uint8 FP8-padded kernel input (cbsz=2 layout). + scale_e8m0 [..., K//32] uint8 E8M0 block scale. + a_unpacked [..., K] uint8 low-6-bit E2M3 codes used by the reference. + """ + from tests.kernels.utils.fp4_utils import per_1x32_f6_quant + + return per_1x32_f6_quant(x.bfloat16()) + + # Test Helpers for MoE GEMM2 Mode Comparison def _make_reduce_mode_compile_fn( use_flydsl_reduce: bool = True, use_valid_mask: bool = False, scale_dtype: str = "f32" @@ -2269,6 +2296,83 @@ def test_moe_gemm_w4a16_groupwise_scale(scale_dtype): pytest.mark.skipif("gfx95" not in ARCH, reason="FP4 shape requires gfx950+"), ], ), + # ── Realworld MoE model shapes (stage2 down-projection) ────────────── + # tokens, model_dim, inter_dim, experts, topk, tile_m, tile_n, tile_k + # Mixtral-8x7B: hidden=4096, inter=14336, E=8, topk=2 + pytest.param( + 128, + 4096, + 256, + 8, + 2, + 64, + 256, + 256, + id="Mixtral8x7B-T128", + marks=pytest.mark.skipif("gfx95" not in ARCH, reason="FP4 requires gfx950+"), + ), + pytest.param( + 512, + 4096, + 256, + 8, + 2, + 64, + 256, + 256, + id="Mixtral8x7B-T512", + marks=pytest.mark.skipif("gfx95" not in ARCH, reason="FP4 requires gfx950+"), + ), + # Qwen3-30B-A3B: hidden=2048, inter=768, E=64, topk=8 + pytest.param( + 128, + 2048, + 256, + 64, + 8, + 64, + 256, + 256, + id="Qwen3-30B-A3B-T128", + marks=pytest.mark.skipif("gfx95" not in ARCH, reason="FP4 requires gfx950+"), + ), + pytest.param( + 512, + 2048, + 256, + 64, + 8, + 64, + 256, + 256, + id="Qwen3-30B-A3B-T512", + marks=pytest.mark.skipif("gfx95" not in ARCH, reason="FP4 requires gfx950+"), + ), + # DeepSeek-V3: hidden=7168, inter=2048, E=256, topk=8 + pytest.param( + 128, + 7168, + 256, + 256, + 8, + 64, + 256, + 128, + id="DeepSeekV3-T128", + marks=pytest.mark.skipif("gfx95" not in ARCH, reason="FP4 requires gfx950+"), + ), + pytest.param( + 512, + 7168, + 256, + 256, + 8, + 64, + 256, + 128, + id="DeepSeekV3-T512", + marks=pytest.mark.skipif("gfx95" not in ARCH, reason="FP4 requires gfx950+"), + ), ], ) @pytest.mark.parametrize( @@ -2277,6 +2381,7 @@ def test_moe_gemm_w4a16_groupwise_scale(scale_dtype): "fp8", pytest.param("fp4", marks=pytest.mark.skipif("gfx95" not in ARCH, reason="FP4 requires gfx950+")), pytest.param("a8w4", marks=pytest.mark.skipif("gfx95" not in ARCH, reason="A8W4 requires gfx950+")), + pytest.param("a6w4", marks=pytest.mark.skipif("gfx95" not in ARCH, reason="A6W4 requires gfx950+")), ], ) def test_moe_stage2_standalone( @@ -2302,7 +2407,7 @@ def test_moe_stage2_standalone( 3. Reduce mode (FlyDSL): GEMM2 + FlyDSL reduce kernel For FP4 / A8W4: atomic mode + torch reduce mode (MXFP4 weight path). """ - is_fp4_path = in_dtype in ("fp4", "a8w4") + is_fp4_path = in_dtype in ("fp4", "a8w4", "a6w4") if is_fp4_path: from tests.kernels.utils import fp4_utils @@ -2341,12 +2446,26 @@ def test_moe_stage2_standalone( _A2_QUANTIZERS = { "fp4": _per_1x32_fp4_quant, "a8w4": _per_1x32_mxfp8_quant, + "a6w4": _per_1x32_mxfp6_quant, } device = torch.device("cuda") torch.manual_seed(seed) a2_fp32 = torch.randn((tokens * topk, inter_dim), device=device, dtype=torch.float32) * 0.2 - a2_q, a2_scale = _A2_QUANTIZERS[in_dtype](a2_fp32) - fp4_args = dict(common_args, a2_fp8_in=a2_q, a2_scale_in=a2_scale) + quant_result = _A2_QUANTIZERS[in_dtype](a2_fp32) + if in_dtype == "a6w4": + # _per_1x32_mxfp6_quant returns (a_pad, scale, a_unpacked). + # Pass a_pad to the kernel, a_unpacked + a2_kind="mxfp6" to the + # reference so torch_moe_gemm2 can dequantize correctly. + a2_q, a2_scale, a2_unpacked = quant_result + fp4_args = dict( + common_args, + a2_fp8_in=a2_q, + a2_scale_in=a2_scale, + a2_ref_in=a2_unpacked, + ) + else: + a2_q, a2_scale = quant_result + fp4_args = dict(common_args, a2_fp8_in=a2_q, a2_scale_in=a2_scale) run_moe_stage2(**fp4_args, kernel_name=f"moe_gemm2_atomic_{in_dtype}") run_moe_stage2(**fp4_args, use_reduce=True, kernel_name=f"moe_gemm2_reduce_torch_{in_dtype}") return diff --git a/tests/kernels/test_ref.py b/tests/kernels/test_ref.py index 718140ba3..083f16357 100644 --- a/tests/kernels/test_ref.py +++ b/tests/kernels/test_ref.py @@ -36,7 +36,7 @@ def _detect_scale_kind(x: torch.Tensor, scale: torch.Tensor | None) -> str: def _logical_k(x: torch.Tensor, kind: str) -> int: """Return the logical K dim (unpacked) given the detected scale kind.""" k = int(x.shape[-1]) - return k * 2 if kind == "mxfp4" else k + return k * 2 if kind == "mxfp4" else k # mxfp6/mxfp8/scalar: 1 byte per element def _dequant_mxfp4_per_1x32(x_fp4: torch.Tensor, scale_e8m0: torch.Tensor) -> torch.Tensor: @@ -73,6 +73,27 @@ def _dequant_mxfp8_per_1x32(x_fp8: torch.Tensor, scale_e8m0: torch.Tensor) -> to return x_f32.view(*x_fp8.shape[:-1], logical_k) +def _dequant_mxfp6_per_1x32(x_fp6: torch.Tensor, scale_e8m0: torch.Tensor) -> torch.Tensor: + """Dequantize MXFP6-E2M3 codes with per-1x32 E8M0 block scales to fp32. + + x_fp6: uint8 tensor where the low 6 bits of each byte are the E2M3 code + (the ``a_unpacked`` layout returned by ``per_1x32_f6_quant``). + scale_e8m0: uint8 tensor, one byte per 32-element K block. + """ + from tests.kernels.utils.fp4_utils import e8m0_to_f32, fp6_e2m3_to_f32 + + shape = x_fp6.shape + logical_k = int(shape[-1]) + if logical_k % 32 != 0: + raise ValueError(f"MXFP6 logical K must be divisible by 32, got {logical_k}") + + x_flat = x_fp6.reshape(-1, logical_k) + s_flat = scale_e8m0.view(torch.uint8).reshape(-1, logical_k // 32) + decoded = fp6_e2m3_to_f32(x_flat) # [N, K] float32 + scales_f32 = e8m0_to_f32(s_flat).repeat_interleave(32, dim=1) # [N, K] float32 + return (decoded * scales_f32).view(shape) + + def _dequant(x: torch.Tensor, scale: torch.Tensor | None, kind: str) -> torch.Tensor: """Unified fp32 dequantization driven by the detected ``kind``. @@ -84,6 +105,8 @@ def _dequant(x: torch.Tensor, scale: torch.Tensor | None, kind: str) -> torch.Te return _dequant_mxfp4_per_1x32(x, scale) if kind == "mxfp8": return _dequant_mxfp8_per_1x32(x, scale) + if kind == "mxfp6": + return _dequant_mxfp6_per_1x32(x, scale) x_f32 = x.to(torch.float32) return x_f32 if scale is None else x_f32 * scale @@ -172,6 +195,7 @@ def torch_moe_gemm2( doweight_stage2: bool, group_size: int = -1, scale_w2_groups: torch.Tensor | None = None, + a2_kind: str | None = None, ) -> torch.Tensor: """Return [tokens, model_dim] fp32. @@ -185,17 +209,22 @@ def torch_moe_gemm2( group_size: -1 for per-row scale (uses scale_w2), >0 for group-wise scale. scale_w2_groups: Group-wise scale tensor of shape [E, inter_dim//group_size, model_dim] (Opt 0 layout). Required when group_size > 0; ignored otherwise. + a2_kind: Override auto-detection of the A2 format. Use ``"mxfp6"`` when + passing MXFP6-E2M3 unpacked codes (low-6-bit uint8 per element). + When ``None`` the kind is inferred from the tensor dtype/shape. """ assert a2_q.is_cuda and w2_q.is_cuda tokens, topk = topk_ids.shape # Independent per-1x32 block-scale detection for a2 and w2; see # ``_detect_scale_kind`` for the classification rules. - a_kind = _detect_scale_kind(a2_q, scale_a2) + # a2_kind may be passed explicitly when auto-detection is ambiguous (e.g. + # mxfp6 and mxfp8 share the same byte-per-element ratio). + a_kind = a2_kind or _detect_scale_kind(a2_q, scale_a2) w_kind = _detect_scale_kind(w2_q, scale_w2) inter_dim = _logical_k(a2_q, a_kind) - if a_kind in ("mxfp4", "mxfp8"): + if a_kind in ("mxfp4", "mxfp8", "mxfp6"): if a2_q.dim() == 2: a2_q = a2_q.view(tokens, topk, -1) scale_a2 = scale_a2.view(tokens, topk, -1)