diff --git a/flashinfer/fused_moe/prepare.py b/flashinfer/fused_moe/prepare.py index 8c9aeb6c690..863ba718b57 100644 --- a/flashinfer/fused_moe/prepare.py +++ b/flashinfer/fused_moe/prepare.py @@ -649,6 +649,17 @@ def prepare_trtllm_fp8_block_weights( raise ValueError( "MXFP8 requires hidden_size and intermediate_size divisible by 32." ) + if hidden_size % 128 != 0 or (2 * intermediate_size) % 128 != 0: + # The shuffled-weight layout interleaves scale rows in 128-row + # tiles; the exact (non-padded) reshape below needs both scale + # matrices' row counts (2*intermediate and hidden) to be 128-row + # aligned -- same constraint the legacy pipeline carried + # implicitly (gh #4087). + raise ValueError( + "MXFP8 shuffled-weight layout requires hidden_size and " + "2*intermediate_size divisible by 128." + ) + from ..fp4_quantization import block_scale_interleave from ..quantization.fp8_quantization import mxfp8_quantize from .core import ( _maybe_get_cached_w3_w1_permute_indices, @@ -673,7 +684,21 @@ def prepare_trtllm_fp8_block_weights( is_gated_act_gemm=True, ) w1_q.append(q.view(torch.uint8)[permute.to(device)].view(q.dtype)) - w1_sf.append(sf[permute_sf.to(device)]) + # Row permutation alone is NOT the layout the MxFp8 kernels consume: + # legacy shuffle_matrix_sf_a = row permute THEN 128x4 + # block_scale_interleave (gh #4087 -- dropping the interleave reads + # ~99.8% of scale bytes from the wrong address; errors stay sparse + # only because random E8M0 exponents often coincide). + # block_scale_interleave returns the flattened interleaved buffer; + # reshape restores the logical (rows, cols) view the runner + # validates (the kernel consumes the raw interleaved bytes) -- + # the legacy pipeline's exact idiom. The exact reshape is safe: + # both row counts are multiples of 128, so no padding is added. + w1_sf.append( + block_scale_interleave(sf[permute_sf.to(device)].contiguous()).reshape( + 2 * intermediate_size, hidden_size // 32 + ) + ) q, sf = mxfp8_quantize(w2_bf16[expert], is_sf_swizzled_layout=False) sf = sf.view(torch.uint8).reshape(hidden_size, intermediate_size // 32) @@ -687,7 +712,11 @@ def prepare_trtllm_fp8_block_weights( num_elts_per_sf=32, ) w2_q.append(q.view(torch.uint8)[permute.to(device)].view(q.dtype)) - w2_sf.append(sf[permute_sf.to(device)]) + w2_sf.append( + block_scale_interleave(sf[permute_sf.to(device)].contiguous()).reshape( + hidden_size, intermediate_size // 32 + ) + ) w1_q, w1_sf = torch.stack(w1_q), torch.stack(w1_sf) w2_q, w2_sf = torch.stack(w2_q), torch.stack(w2_sf) diff --git a/tests/moe/test_mxfp8_weight_prepare_layout.py b/tests/moe/test_mxfp8_weight_prepare_layout.py new file mode 100644 index 00000000000..50654fd6b82 --- /dev/null +++ b/tests/moe/test_mxfp8_weight_prepare_layout.py @@ -0,0 +1,88 @@ +"""Layout regression for gh #4087: MxFp8 weight-scale preparation. + +The MxFp8 unified-MoE weight path (``prepare.py``) must produce the SAME +scale-factor layout as the legacy ``shuffle_matrix_sf_a`` pipeline the +kernels were written against: row permutation THEN 128x4 +``block_scale_interleave``. #4026's adapter did only the permutation; the +gap escaped its own unit test because that test hard-coded the incomplete +layout as expected, and escaped the functional test because atol=0.05 +dwarfed the output scale (~0.0025). + +Per review guidance the checks (a) compare directly against the legacy +reference pipeline and (b) use sentinel data where scale bytes are all +distinct — with random weights, misplaced reads often land on an equal +E8M0 exponent and cancel out. +""" + +import pytest +import torch + +from flashinfer.fp4_quantization import block_scale_interleave +from flashinfer.utils import get_compute_capability + + +def _skip_unless_sm100(): + if not torch.cuda.is_available(): + pytest.skip("no CUDA") + if get_compute_capability(torch.device("cuda:0"))[0] != 10: + pytest.skip("MxFp8 trtllm path requires SM100-family") + + +def _sentinel_sf(rows, cols_sf, dev): + # Cycle a prime-period pattern so any misplaced byte changes the result. + return ( + (torch.arange(rows * cols_sf, dtype=torch.int64, device=dev) % 251) + .to(torch.uint8) + .reshape(rows, cols_sf) + ) + + +@pytest.mark.parametrize("rows,cols", [(1024, 768), (256, 512), (2048, 1024)]) +def test_w2_sf_shuffle_matches_legacy(rows, cols): + """Non-gated (w2) scale path: permute+interleave == shuffle_matrix_sf_a.""" + _skip_unless_sm100() + from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache + from flashinfer.quantization.fp4_quantization import shuffle_matrix_sf_a + + dev = torch.device("cuda:0") + sf = _sentinel_sf(rows, cols // 32, dev) + + cache = {} + permute_sf = get_w2_permute_indices_with_cache(cache, sf, 128, num_elts_per_sf=32) + got = block_scale_interleave(sf[permute_sf.to(dev)].contiguous()) + + want = shuffle_matrix_sf_a(sf, 128, num_elts_per_sf=32) + + assert got.shape == want.shape, f"{got.shape} vs {want.shape}" + n_diff = int((got != want).sum()) + assert n_diff == 0, f"{n_diff}/{want.numel()} scale bytes misplaced" + + +@pytest.mark.parametrize("rows,cols", [(2 * 768, 1024), (2 * 256, 512)]) +def test_w31_sf_is_interleaved(rows, cols): + """Gated (w3_w1) scale path: the result must be 128x4-interleaved, i.e. + NOT equal to the permutation-only layout #4026 produced (the exact gh + #4087 regression), and elementwise-preserving (a permutation of bytes).""" + _skip_unless_sm100() + from flashinfer.fused_moe.core import _maybe_get_cached_w3_w1_permute_indices + + dev = torch.device("cuda:0") + sf = _sentinel_sf(rows, cols // 32, dev) + + cache = {} + permute_sf = _maybe_get_cached_w3_w1_permute_indices( + cache, sf, 128, num_elts_per_sf=32, is_gated_act_gemm=True + ) + permuted_only = sf[permute_sf.to(dev)].contiguous() + got = block_scale_interleave(permuted_only) + + # Regression guard: the shipped layout must NOT be the permutation-only + # layout #4026 produced (the exact gh #4087 regression). + assert not torch.equal( + got.flatten()[: permuted_only.numel()], permuted_only.flatten() + ), "interleave was a no-op -- layout regression (gh #4087)" + if got.numel() == permuted_only.numel(): + # Pure relayout (no padding): byte multiset must be preserved. + assert torch.equal( + got.flatten().sort().values, permuted_only.flatten().sort().values + ), "interleave lost/duplicated scale bytes"