diff --git a/flashinfer/fused_moe/__init__.py b/flashinfer/fused_moe/__init__.py index d00c7f65bf6..50c68150e7d 100644 --- a/flashinfer/fused_moe/__init__.py +++ b/flashinfer/fused_moe/__init__.py @@ -44,6 +44,7 @@ TrtllmFp4RoutedRunner, TrtllmFp8BlockRunner, TrtllmFp8PerTensorRunner, + TrtllmMxInt4RoutedRunner, ) # Legacy flat-argument APIs (unchanged, not deprecated) @@ -144,6 +145,7 @@ "TrtllmFp4RoutedRunner", "TrtllmFp8BlockRunner", "TrtllmFp8PerTensorRunner", + "TrtllmMxInt4RoutedRunner", "QuantConfig", "QuantVariant", "RoutingConfig", diff --git a/flashinfer/fused_moe/api.py b/flashinfer/fused_moe/api.py index fc46194bb5d..60811d5afb6 100644 --- a/flashinfer/fused_moe/api.py +++ b/flashinfer/fused_moe/api.py @@ -446,11 +446,34 @@ class TrtllmMxInt4Config: @classmethod def supported(cls, arch: int) -> bool: - # Same trtllm-gen routed batched-GEMM path as the FP4/BF16 backends, so it - # inherits the same manifest coverage. The pinned Rubin manifest includes - # MxInt4 sm107a kernels, although unified MxInt4 still has no runner. + # SM100/SM103 use the forward-compatible sm100f cubins; SM107 selects + # the dedicated Rubin artifact. return arch in _TRTLLM_ROUTED_ARCHS + @staticmethod + def prepare_weights( + w1_bf16, + w2_bf16, + *, + num_local_experts: int, + hidden_size: int, + intermediate_size: int, + device=None, + permute_cache=None, + ): + """Build a ``trtllm_mxint4_routed`` view from canonical BF16 weights.""" + from .prepare import prepare_trtllm_mxint4_weights + + return prepare_trtllm_mxint4_weights( + w1_bf16, + w2_bf16, + num_local_experts=num_local_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + device=device, + permute_cache=permute_cache, + ) + def __repr__(self) -> str: return "TrtllmMxInt4Config()" @@ -754,6 +777,8 @@ class MoEActivationPack: * W4A16 with ``TrtllmFp4Config``: raw ``bfloat16 [M, H]`` values with no activation scale; weights use the MXFP4 preparation contract. * BF16: raw ``bfloat16 [M, H]`` values with no scale tensor. + * MxInt4: raw ``bfloat16 [M, H]`` values with no scale tensor; weights are + packed signed INT4 with BF16 block scales. * DeepSeek FP8: ``float8_e4m3fn [M, H]`` values with transposed ``float32 [H/128, M]`` block scales. * MXFP8: ``float8_e4m3fn [M, H]`` values with token-major @@ -776,8 +801,9 @@ class MoEActivationPack: itself per ``RoutingConfig.method``. ``topk_ids`` / ``topk_weights`` stay ``None`` — the runner allocates internal kernel-filled buffers, and the routing result is not surfaced back through the pack (routing replay is a separate, future capability). TRTLLM FP4, - BF16, block-FP8, and per-tensor-FP8 runners support this mode; ``MoELayer`` dispatches - a logits pack only to capable backends (see each runner's ``supported_routing_modes``). + BF16, block-FP8, per-tensor-FP8, and MxInt4 runners support this mode; + ``MoELayer`` dispatches a logits pack only to capable backends (see each runner's + ``supported_routing_modes``). ``topk_ids`` / ``topk_weights`` follow the routed-MoE naming convention (gh #2425); they keep the field positions of the former ``selected_experts`` / ``final_scales``, so diff --git a/flashinfer/fused_moe/core.py b/flashinfer/fused_moe/core.py index a99c3d7d5cd..ca1f6ea8435 100644 --- a/flashinfer/fused_moe/core.py +++ b/flashinfer/fused_moe/core.py @@ -170,8 +170,14 @@ def _maybe_get_cached_w3_w1_permute_indices( num_elts_per_sf: Union[None, int] = None, is_gated_act_gemm: bool = True, ) -> torch.Tensor: - # Create a unique cache key (weight_type, weight_shape) - cache_key = ("w3_w1", dst_w3_w1_weight.shape) + # Include every parameter that changes the generated permutation. + cache_key = ( + "w3_w1", + dst_w3_w1_weight.shape, + epilogue_tile_m, + num_elts_per_sf, + is_gated_act_gemm, + ) if cache_key not in _cache_permute_indices: # Get permute indices and chain them together if is_gated_act_gemm: @@ -202,8 +208,13 @@ def get_w2_permute_indices_with_cache( epilogue_tile_m: int, num_elts_per_sf: Union[None, int] = None, ) -> torch.Tensor: - # Create a unique cache key (weight_type, weight_shape) - cache_key = ("w2", dst_w2_weight.shape) + # Include every parameter that changes the generated permutation. + cache_key = ( + "w2", + dst_w2_weight.shape, + epilogue_tile_m, + num_elts_per_sf, + ) if cache_key not in _cache_permute_indices: if num_elts_per_sf is None: permute_indices = get_shuffle_matrix_a_row_indices( diff --git a/flashinfer/fused_moe/layer.py b/flashinfer/fused_moe/layer.py index 7770582ed43..df2b70aa9db 100644 --- a/flashinfer/fused_moe/layer.py +++ b/flashinfer/fused_moe/layer.py @@ -38,6 +38,7 @@ TrtllmFp4Config, TrtllmFp8BlockConfig, TrtllmFp8PerTensorConfig, + TrtllmMxInt4Config, ) from .runners import ( B12xNvfp4Runner, @@ -47,6 +48,7 @@ TrtllmFp4RoutedRunner, TrtllmFp8BlockRunner, TrtllmFp8PerTensorRunner, + TrtllmMxInt4RoutedRunner, ) from .utils import map_to_hybrid_bucket @@ -60,6 +62,7 @@ TrtllmBf16RoutedRunner, TrtllmFp8BlockRunner, TrtllmFp8PerTensorRunner, + TrtllmMxInt4RoutedRunner, B12xNvfp4Runner, B12xW4A16Runner, ] @@ -71,6 +74,7 @@ TrtllmBf16Config: TrtllmBf16RoutedRunner, TrtllmFp8BlockConfig: TrtllmFp8BlockRunner, TrtllmFp8PerTensorConfig: TrtllmFp8PerTensorRunner, + TrtllmMxInt4Config: TrtllmMxInt4RoutedRunner, B12xNvfp4Config: B12xNvfp4Runner, B12xW4A16Config: B12xW4A16Runner, } diff --git a/flashinfer/fused_moe/prepare.py b/flashinfer/fused_moe/prepare.py index 5c4a2d0c576..6e97582a6ea 100644 --- a/flashinfer/fused_moe/prepare.py +++ b/flashinfer/fused_moe/prepare.py @@ -44,11 +44,12 @@ ) from ..utils import get_compute_capability -# Module-level permute-index cache. Permute indices depend only on weight -# dims, so the cache is safe to reuse across shapes and calls. +# Module-level permute-index caches. Permute indices depend on weight geometry +# and layout parameters, so matching keys are safe to reuse across calls. _TRTLLM_PERMUTE_CACHE: dict = {} _TRTLLM_FP8_PERMUTE_CACHE: dict = {} _TRTLLM_FP8_PER_TENSOR_PERMUTE_CACHE: dict = {} +_TRTLLM_MXINT4_PERMUTE_CACHE: dict = {} # The E8M0 range clamp and residual-scale factorization are adapted from @@ -971,6 +972,138 @@ def prepare_trtllm_fp8_per_tensor_activations( return quantized.to(torch.float8_e4m3fn), None +def _mxint4_quantize( + weights: torch.Tensor, sf_vec_size: int = 32 +) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize the last dimension to signed packed INT4 with BF16 block scales.""" + blocks = weights.reshape(-1, sf_vec_size) + block_max = blocks.amax(dim=-1, keepdim=True).to(torch.float32) + block_min = blocks.amin(dim=-1, keepdim=True).to(torch.float32) + block_max = block_max * (8.0 / 7.0) + amax = torch.where(block_max > -block_min, block_max, -block_min) + scales = amax / 8.0 + scales = torch.where(scales > 0, scales, torch.ones_like(scales)) + quantized = ( + (blocks * scales.reciprocal()) + .round() + .clamp(-8, 7) + .to(torch.int8) + .reshape(-1, sf_vec_size // 2, 2) + ) + nibbles = (quantized & 0x0F).to(torch.uint8) + packed = nibbles[..., 0] | (nibbles[..., 1] << 4) + return ( + packed.reshape(*weights.shape[:-1], weights.shape[-1] // 2), + scales.to(torch.bfloat16), + ) + + +def prepare_trtllm_mxint4_weights( + w1_bf16: torch.Tensor, + w2_bf16: torch.Tensor, + *, + num_local_experts: int, + hidden_size: int, + intermediate_size: int, + device: Optional[torch.device] = None, + permute_cache: Optional[dict] = None, +) -> Dict[str, torch.Tensor]: + """Build the TRTLLM MxInt4 ``trtllm_mxint4_routed`` weight view. + + Canonical BF16 expert weights are quantized in 32-element K blocks, then + shuffled for fused SwiGLU / transposed-MMA output. Packed INT4 payloads use + BlockMajorK while BF16 scale tensors use TRTLLM's block-scale interleave. + """ + from ..quantization.fp4_quantization import block_scale_interleave + from .core import ( + _maybe_get_cached_w3_w1_permute_indices, + convert_to_block_layout, + get_w2_permute_indices_with_cache, + ) + + if device is None: + device = w1_bf16.device + device = torch.device(device) + if w1_bf16.dtype != torch.bfloat16 or w2_bf16.dtype != torch.bfloat16: + raise ValueError( + "prepare_trtllm_mxint4_weights expects BF16 weights, got " + f"w1={w1_bf16.dtype}, w2={w2_bf16.dtype}." + ) + expected_w1 = (num_local_experts, 2 * intermediate_size, hidden_size) + expected_w2 = (num_local_experts, hidden_size, intermediate_size) + if tuple(w1_bf16.shape) != expected_w1 or tuple(w2_bf16.shape) != expected_w2: + raise ValueError( + f"weight shapes {tuple(w1_bf16.shape)}/{tuple(w2_bf16.shape)} != " + f"expected {expected_w1}/{expected_w2}." + ) + if hidden_size % 256 != 0 or intermediate_size % 256 != 0: + raise ValueError( + "TRTLLM MxInt4 requires hidden_size and intermediate_size divisible by 256." + ) + + w1 = w1_bf16.to(device).contiguous() + w2 = w2_bf16.to(device).contiguous() + w1_q, w1_sf = _mxint4_quantize(w1) + w2_q, w2_sf = _mxint4_quantize(w2) + w1_sf = w1_sf.reshape(num_local_experts, 2 * intermediate_size, hidden_size // 32) + w2_sf = w2_sf.reshape(num_local_experts, hidden_size, intermediate_size // 32) + + if permute_cache is None: + permute_cache = _TRTLLM_MXINT4_PERMUTE_CACHE + epilogue_tile_m = 128 + block_k = 128 + w1_views, w1_scale_views, w2_views, w2_scale_views = [], [], [], [] + for expert in range(num_local_experts): + w1_permute = _maybe_get_cached_w3_w1_permute_indices( + permute_cache, w1_q[expert], epilogue_tile_m + ) + w1_scale_permute = _maybe_get_cached_w3_w1_permute_indices( + permute_cache, + w1_sf[expert], + epilogue_tile_m, + num_elts_per_sf=32, + ) + w2_permute = get_w2_permute_indices_with_cache( + permute_cache, w2_q[expert], epilogue_tile_m + ) + # Keep the established flat-test MxInt4 scale permutation contract; + # preparation parity tests cover this asymmetric GEMM1/GEMM2 setting. + w2_scale_permute = get_w2_permute_indices_with_cache( + permute_cache, + w2_sf[expert], + epilogue_tile_m, + num_elts_per_sf=16, + ) + + w1_views.append( + convert_to_block_layout( + w1_q[expert][w1_permute.to(device)].contiguous(), block_k + ) + ) + w1_scale_views.append( + block_scale_interleave( + w1_sf[expert][w1_scale_permute.to(device)].contiguous() + ) + ) + w2_views.append( + convert_to_block_layout( + w2_q[expert][w2_permute.to(device)].contiguous(), block_k + ) + ) + w2_scale_views.append( + block_scale_interleave( + w2_sf[expert][w2_scale_permute.to(device)].contiguous() + ) + ) + + return { + "gemm1_weights": torch.stack(w1_views), + "gemm1_weights_scale": torch.stack(w1_scale_views).view(torch.bfloat16), + "gemm2_weights": torch.stack(w2_views), + "gemm2_weights_scale": torch.stack(w2_scale_views).view(torch.bfloat16), + } + + def prepare_trtllm_bf16_weights( w1_bf16: torch.Tensor, w2_bf16: torch.Tensor, diff --git a/flashinfer/fused_moe/runners.py b/flashinfer/fused_moe/runners.py index c43f386fb07..bb8939fadd5 100644 --- a/flashinfer/fused_moe/runners.py +++ b/flashinfer/fused_moe/runners.py @@ -1532,6 +1532,300 @@ def __hash__(self): return hash(("trtllm_bf16_routed",)) +# --------------------------------------------------------------------------- +# TRTLLM MxInt4 runner — BF16 activations, packed INT4 BlockMajorK weights +# --------------------------------------------------------------------------- + + +class TrtllmMxInt4RoutedRunner(MoERunner): + """MxInt4 adapter over the canonical TRTLLM MoE runner.""" + + backend_key = "trtllm_mxint4_routed" + supported_routing_modes = ( + RoutingInputMode.PackedPrecomputed, + RoutingInputMode.FromLogits, + ) + supported_quant_variants = (QuantVariant.MxInt4,) + + def check_support(self) -> None: + super().check_support() + if self.config.activation.type is not ActivationType.Swiglu: + raise NotImplementedError( + f"{type(self).__name__} supports only the Swiglu activation." + ) + if not self.config.execution.do_finalize: + raise NotImplementedError( + f"{type(self).__name__} supports only do_finalize=True." + ) + from ..utils import get_compute_capability + from .api import TrtllmMxInt4Config + + major, minor = get_compute_capability(self.device) + arch = major * 10 + minor + if not TrtllmMxInt4Config.supported(arch): + raise NotImplementedError( + f"{type(self).__name__} is enabled only on supported " + f"SM100/SM103/SM107 targets, got sm{arch}." + ) + + def __init__(self, config: MoEConfig, device: torch.device): + from ..tllm_enums import DtypeTrtllmGen, Fp8QuantizationType + from ..utils import device_support_pdl + from .core import get_trtllm_moe_sm100_module + + self.config = config + self.device = device + self._module = get_trtllm_moe_sm100_module() + + routing = config.routing + experts = config.experts + execution = config.execution + self._num_local_experts = experts.local_num_experts or routing.num_experts + self._local_expert_offset = experts.local_expert_offset + self._intermediate_size = experts.intermediate_size + self._activation_type = int(config.activation.type) + self._tune_max_num_tokens = execution.tune_max_num_tokens + + self._dtype_act = DtypeTrtllmGen.Bfloat16 + self._dtype_weights = DtypeTrtllmGen.MxInt4 + self._fp8_quantization_type = Fp8QuantizationType.NoneFp8 + + enable_pdl = execution.enable_pdl + if enable_pdl is None: + enable_pdl = device_support_pdl(device) + self._enable_pdl = enable_pdl + + self._inner: Any = None + self._static_kwargs: dict = {} + self.tuning_config: Any = None + + def _ensure_inner(self, hidden_size: int) -> None: + if self._inner is not None: + return + from ..tllm_enums import WeightLayout + + self._inner = self._module.MoERunner( + top_k=self.config.routing.top_k, + num_local_experts=self._num_local_experts, + dtype_act=self._dtype_act, + dtype_weights=self._dtype_weights, + fp8_quantization_type=self._fp8_quantization_type, + hidden_size=hidden_size, + intermediate_size=self._intermediate_size, + activation_type=self._activation_type, + use_shuffled_weight=True, + weight_layout=int(WeightLayout.BlockMajorK), + use_per_token_scaling=False, + num_experts=self.config.routing.num_experts, + ) + + def get_valid_tactics( # type: ignore[override] + self, inputs: List[torch.Tensor], profile: Any + ) -> List[Any]: + return self._inner.get_valid_tactics(inputs, profile) + + def forward( + self, + inputs: List[torch.Tensor], + tactic: Any = -1, + do_preparation: bool = False, + **kwargs: Any, + ) -> torch.Tensor: + self._inner.forward( + inputs, + tactic=tactic, + do_preparation=do_preparation, + **self._static_kwargs, + ) + return inputs[0] + + def pack_inputs( + self, act: MoEActivationPack, weights: MoEWeightPack + ) -> List[torch.Tensor]: + from .core import MoeRunnerInputs + + view = weights.get_view(self.backend_key) + routing = self.config.routing + hidden_states = act.hidden_states_q + if hidden_states.dtype != torch.bfloat16 or hidden_states.dim() != 2: + raise ValueError(f"{type(self).__name__}: hidden_states_q must be 2D BF16.") + if not hidden_states.is_contiguous(): + raise ValueError( + f"{type(self).__name__}: hidden_states_q must be contiguous." + ) + if act.hidden_states_scale is not None: + raise ValueError( + f"{type(self).__name__}: hidden_states_scale must be None." + ) + num_tokens, hidden_size = hidden_states.shape + if hidden_size % 256 != 0 or self._intermediate_size % 256 != 0: + raise ValueError( + f"{type(self).__name__}: hidden_size and intermediate_size " + "must be divisible by 256." + ) + routing_input_mode = act.routing_input_mode + if routing_input_mode == RoutingInputMode.FromLogits: + _validate_logits_inputs( + act, num_tokens, routing.num_experts, type(self).__name__ + ) + if act.routing_logits.dtype != torch.bfloat16: + raise TypeError( + f"{type(self).__name__}: FromLogits currently requires " + f"bfloat16 routing_logits, got {act.routing_logits.dtype}." + ) + if act.routing_bias is not None: + if act.routing_bias.dtype != torch.bfloat16: + raise TypeError( + f"{type(self).__name__}: routing_bias must be bfloat16, " + f"got {act.routing_bias.dtype}." + ) + routing_logits = act.routing_logits + routing_bias = act.routing_bias + topk_ids = hidden_states.new_empty( + (num_tokens, routing.top_k), dtype=torch.int32 + ) + expert_weights = hidden_states.new_empty( + (num_tokens, routing.top_k), dtype=torch.bfloat16 + ) + elif routing_input_mode == RoutingInputMode.PackedPrecomputed: + _validate_prerouted_inputs( + act, num_tokens, routing.top_k, type(self).__name__ + ) + routing_logits = None + routing_bias = None + topk_ids = _pack_prerouted_topk_ids(act) + expert_weights = act.topk_weights.new_empty( + (num_tokens, routing.top_k), dtype=torch.bfloat16 + ) + else: + raise NotImplementedError( + f"{type(self).__name__} supports only FromLogits and " + "PackedPrecomputed routing." + ) + + required = ( + "gemm1_weights", + "gemm1_weights_scale", + "gemm2_weights", + "gemm2_weights_scale", + ) + missing = [key for key in required if key not in view] + if missing: + raise KeyError(f"{self.backend_key} weight view is missing {missing}.") + for key in required: + tensor = view[key] + if tensor.device != hidden_states.device: + raise ValueError( + f"{type(self).__name__}: {key} is on {tensor.device}, " + f"expected {hidden_states.device}." + ) + if not tensor.is_contiguous(): + raise ValueError(f"{type(self).__name__}: {key} must be contiguous.") + if ( + view["gemm1_weights"].dtype != torch.uint8 + or view["gemm2_weights"].dtype != torch.uint8 + ): + raise TypeError("MxInt4 packed weights must be uint8.") + if ( + view["gemm1_weights_scale"].dtype != torch.bfloat16 + or view["gemm2_weights_scale"].dtype != torch.bfloat16 + ): + raise TypeError("MxInt4 weight scales must be bfloat16.") + expected_shapes = { + "gemm1_weights": ( + self._num_local_experts, + hidden_size // 256, + 2 * self._intermediate_size, + 128, + ), + "gemm1_weights_scale": ( + self._num_local_experts, + 2 * self._intermediate_size * hidden_size // 32, + ), + "gemm2_weights": ( + self._num_local_experts, + self._intermediate_size // 256, + hidden_size, + 128, + ), + "gemm2_weights_scale": ( + self._num_local_experts, + hidden_size * self._intermediate_size // 32, + ), + } + for key, expected in expected_shapes.items(): + if tuple(view[key].shape) != expected: + raise ValueError( + f"{type(self).__name__}: {key} shape " + f"{tuple(view[key].shape)} != expected {expected}." + ) + for key in ("gemm1_alpha", "gemm1_beta", "gemm1_clamp_limit"): + tensor = view.get(key) + if tensor is None: + continue + if tensor.device != hidden_states.device: + raise ValueError( + f"{type(self).__name__}: {key} is on {tensor.device}, " + f"expected {hidden_states.device}." + ) + if tensor.dtype != torch.float32: + raise TypeError( + f"{type(self).__name__}: {key} must be float32, got {tensor.dtype}." + ) + if tuple(tensor.shape) != (self._num_local_experts,): + raise ValueError( + f"{type(self).__name__}: {key} shape {tuple(tensor.shape)} " + f"!= expected ({self._num_local_experts},)." + ) + if not tensor.is_contiguous(): + raise ValueError(f"{type(self).__name__}: {key} must be contiguous.") + + output = hidden_states.new_empty((num_tokens, hidden_size)) + moe_inputs = MoeRunnerInputs( + output=output, + routing_logits=routing_logits, + topk_ids=topk_ids, + expert_weights=expert_weights, + hidden_states=hidden_states, + hidden_states_scale=None, + gemm1_lora_delta=None, + per_token_scale=None, + ) + + self._static_kwargs = dict( + routing_bias=routing_bias, + gemm1_weights=view["gemm1_weights"], + gemm1_weights_scale=view["gemm1_weights_scale"], + gemm1_alpha=view.get("gemm1_alpha"), + gemm1_beta=view.get("gemm1_beta"), + gemm1_clamp_limit=view.get("gemm1_clamp_limit"), + gemm2_weights=view["gemm2_weights"], + gemm2_weights_scale=view["gemm2_weights_scale"], + num_experts=routing.num_experts, + n_group=routing.n_group, + topk_group=routing.topk_group, + local_expert_offset=self._local_expert_offset, + routed_scaling_factor=routing.routed_scaling_factor, + routing_method_type=int(routing.method), + do_finalize=self.config.execution.do_finalize, + enable_pdl=self._enable_pdl, + norm_topk_prob=True, + ) + + self._ensure_inner(hidden_size) + self.tuning_config = self._inner._make_tuning_config( + moe_inputs, + tune_max_num_tokens=self._tune_max_num_tokens, + routing_input_mode=routing_input_mode, + use_cuda_graph=True, + use_cold_l2_cache=True, + ) + return moe_inputs.to_list() + + def __hash__(self): + return hash(("trtllm_mxint4_routed",)) + + # --------------------------------------------------------------------------- # SM12x b12x runners — fixed tactic, existing wrapper delegation # --------------------------------------------------------------------------- diff --git a/tests/moe/test_unified_moe_fuzz.py b/tests/moe/test_unified_moe_fuzz.py index 45c74d7a8ef..46dea48fc9d 100644 --- a/tests/moe/test_unified_moe_fuzz.py +++ b/tests/moe/test_unified_moe_fuzz.py @@ -76,10 +76,10 @@ every mode, so a kernel that routes wrong is caught by check #2. In-kernel routing is single-GPU (non-EP) here; EP + in-kernel routing semantics are a separate validation. -Coverage today: NVFP4, BF16, block-FP8, and per-tensor FP8 on SM100. TRTLLM -FP4 supports packed/unpacked pre-routed and in-kernel routing; BF16, block-FP8, -and per-tensor FP8 support packed pre-routed and in-kernel routing. CuteDSL is -pre-routed-only. +Coverage today: NVFP4, BF16, block/per-tensor FP8, MXFP4/W4A16, and MxInt4. +CuteDSL NVFP4 is pre-routed-only; FromLogits and UnpackedPrecomputed restrict +dispatch to capable TRTLLM runners. MxInt4 covers packed and BF16-FromLogits +routing. OPT-IN: this suite is gated behind FLASHINFER_UMOE_FUZZ (see the pytestmark below) and is SKIPPED unless that env var is set -- waived in CI pending root-cause of a @@ -196,6 +196,7 @@ TrtllmFp4Config, TrtllmFp8BlockConfig, TrtllmFp8PerTensorConfig, + TrtllmMxInt4Config, ) from flashinfer.fused_moe.layer import _BACKEND_RUNNERS from flashinfer.quantization import e2m1_and_ufp8sf_scale_to_float @@ -398,6 +399,16 @@ def _bf16_act_pack_logits(x, routing_logits, routing_bias): ) +def _mxint4_act_pack_logits(x, routing_logits, routing_bias): + return MoEActivationPack( + hidden_states_q=x, + hidden_states_scale=None, + routing_input_mode=RoutingInputMode.FromLogits, + routing_logits=routing_logits, + routing_bias=routing_bias, + ) + + def _bf16_reference( x, w1, w2, selected_experts, final_scales, intermediate_size, expert_offset=0 ): @@ -701,6 +712,43 @@ def _fp8_per_tensor_reference( return out +def _mxint4_quant_dequant(weights): + blocks = weights.float().reshape(-1, 32) + block_max = blocks.amax(dim=-1, keepdim=True) * (8.0 / 7.0) + block_min = blocks.amin(dim=-1, keepdim=True) + scales = torch.maximum(block_max, -block_min) / 8.0 + scales = torch.where(scales > 0, scales, torch.ones_like(scales)) + quantized = (blocks / scales).round().clamp(-8, 7) + stored_scales = scales.to(torch.bfloat16).float() + return (quantized * stored_scales).reshape_as(weights) + + +def _mxint4_reference( + x, + w1, + w2, + selected_experts, + final_scales, + intermediate_size, + expert_offset=0, +): + x32 = x.float() + w1_32 = _mxint4_quant_dequant(w1) + w2_32 = _mxint4_quant_dequant(w2) + final_scales = final_scales.to(torch.bfloat16).float() + out = torch.zeros_like(x32) + for local_e in range(w1.shape[0]): + token, slot = torch.where(selected_experts == local_e + expert_offset) + if token.numel() == 0: + continue + fc1 = x32[token] @ w1_32[local_e].t() + inter = F.silu(fc1[:, intermediate_size:]) * fc1[:, :intermediate_size] + inter = inter.to(torch.bfloat16).float() + expert_out = (inter @ w2_32[local_e].t()).to(torch.bfloat16).float() + out[token] += final_scales[token, slot, None] * expert_out + return out + + _DTYPE = { QuantVariant.NVFP4: DTypeHandler( variant=QuantVariant.NVFP4, @@ -804,7 +852,19 @@ def _fp8_per_tensor_reference( atol_frac=0.05, # provisional; recalibrate over the expanded SM100 sweep rtol=0.3, ), - # MXINT4 adds one entry when its runner is wired upstream. + QuantVariant.MxInt4: DTypeHandler( + variant=QuantVariant.MxInt4, + candidate_configs=(TrtllmMxInt4Config,), + snap=_bf16_snap, + make_act_pack=_bf16_act_pack, + make_act_pack_logits=_mxint4_act_pack_logits, + reference=_mxint4_reference, + poison=_poison_bf16_out, + out_dtype=torch.bfloat16, + # Curated FromLogits observes max|diff| / ||ref||inf ~= 0.0335. + atol_frac=0.04, + rtol=0.3, + ), } # Cfg.variant string <-> handler lookup (labels stay lowercase enum names). @@ -1038,10 +1098,10 @@ def _gen(seed): if unpacked else rng.choice(prerouted_variants) ) - # The legacy TRTLLM MXFP4 modes are validated only with BF16 router logits. + # The legacy TRTLLM MXFP4 and MxInt4 modes are validated only with BF16 logits. logits_dtype = ( "bf16" - if variant in ("mxfp4", "w4a16") + if variant in ("mxfp4", "w4a16", "mxint4") else ("fp32" if rng.random() < 0.25 else "bf16") ) return Cfg( @@ -1308,6 +1368,29 @@ def _gen(seed): routing_input_mode="unpacked", unpacked_weights_dtype="fp32", ), + Cfg( + 64, + 512, + 512, + 16, + 4, + "mxint4", + "imbalanced", + 900_040, + ), # packed MxInt4; seed % 4 == 0 exercises production autotuning + Cfg( + 64, + 512, + 512, + 16, + 4, + "mxint4", + "uniform", + 900_041, + routing_method=RoutingMethodType.Default, + routing_input_mode="fromlogits", + logits_dtype="bf16", + ), ] if _ONLY_SEEDS: # perfect-repro: run only the named seed(s) _curated_by_seed = {c.seed: c for c in _CURATED} @@ -1558,6 +1641,10 @@ def test_unified_moe_fuzz(cfg): dev = torch.device("cuda") if handler.variant is QuantVariant.W4A16 and sm == 103: pytest.skip("TRTLLM MXFP4×BF16 is disabled on SM103") + if handler.variant is QuantVariant.MxInt4 and ( + cfg.hidden % 256 != 0 or cfg.intermediate % 256 != 0 + ): + pytest.skip("TRTLLM MxInt4 requires hidden/intermediate divisible by 256") # Backend *config classes* whose runner is registered in the live MoELayer registry AND valid # on this arch. A newly-wired backend lands here automatically. wired_backends = [ diff --git a/tests/moe/test_unified_moe_mxint4.py b/tests/moe/test_unified_moe_mxint4.py new file mode 100644 index 00000000000..6e8e02c9e91 --- /dev/null +++ b/tests/moe/test_unified_moe_mxint4.py @@ -0,0 +1,526 @@ +"""Unified TRTLLM MxInt4 MoE preparation and runner tests.""" + +import dataclasses +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from flashinfer.autotuner import autotune +from flashinfer.fused_moe import ( + BackendOptions, + ExpertConfig, + MoEActivationPack, + MoEConfig, + MoELayer, + MoEWeightPack, + QuantConfig, + QuantVariant, + RoutingConfig, + RoutingInputMode, + TrtllmMxInt4Config, + TrtllmMxInt4RoutedRunner, +) +from flashinfer.fused_moe.core import ( + _maybe_get_cached_w3_w1_permute_indices, + get_w2_permute_indices_with_cache, +) +from flashinfer.fused_moe.prepare import _mxint4_quantize +from flashinfer.tllm_enums import RoutingMethodType +from flashinfer.utils import get_compute_capability + + +def _is_mxint4_arch() -> bool: + return torch.cuda.is_available() and get_compute_capability( + torch.device("cuda") + ) in ((10, 0), (10, 3), (10, 7)) + + +mxint4_required = pytest.mark.skipif( + not _is_mxint4_arch(), + reason="Unified TRTLLM MxInt4 MoE requires SM100/SM103/SM107", +) + + +def _dequant_mxint4(packed: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + low = (packed & 0x0F).to(torch.int8) + high = ((packed >> 4) & 0x0F).to(torch.int8) + values = torch.stack((low, high), dim=-1).reshape(*packed.shape[:-1], -1) + values = torch.where(values < 8, values, values - 16).float() + return values * scales.float().repeat_interleave(32, dim=-1) + + +def _mxint4_reference( + x, + w1, + w2, + selected_experts, + final_scales, + intermediate_size, + expert_offset=0, +): + out = torch.zeros_like(x.float()) + final_scales = final_scales.to(torch.bfloat16).float() + for local_e in range(w1.shape[0]): + token, slot = torch.where(selected_experts == local_e + expert_offset) + if token.numel() == 0: + continue + fc1 = x[token].float() @ w1[local_e].float().t() + inter = F.silu(fc1[:, intermediate_size:]) * fc1[:, :intermediate_size] + inter = inter.to(torch.bfloat16).float() + expert_out = (inter @ w2[local_e].float().t()).to(torch.bfloat16).float() + out[token] += final_scales[token, slot, None] * expert_out + return out + + +def _make_case( + *, + num_tokens=16, + hidden_size=256, + intermediate_size=256, + num_experts=8, + top_k=2, + local_num_experts=None, + local_expert_offset=0, + routing_input_mode=RoutingInputMode.PackedPrecomputed, + routing_method=None, +): + routing_method = routing_method or RoutingMethodType.Default + device = torch.device("cuda") + local_num_experts = local_num_experts or num_experts + generator = torch.Generator(device=device).manual_seed(42) + x = torch.randn( + num_tokens, + hidden_size, + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + w1 = ( + torch.randn( + local_num_experts, + 2 * intermediate_size, + hidden_size, + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + / 8 + ) + w2 = ( + torch.randn( + local_num_experts, + hidden_size, + intermediate_size, + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + / 8 + ) + logits = torch.randn( + num_tokens, + num_experts, + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + if local_num_experts != num_experts: + logits[:, :local_expert_offset] = -20 + logits[:, local_expert_offset + local_num_experts :] = -20 + routing_bias = None + n_group = topk_group = None + routed_scaling_factor = None + if routing_method is RoutingMethodType.DeepSeekV3: + from tests.moe.trtllm_gen_fused_moe_utils import noaux_tc_ref + + n_group, topk_group, routed_scaling_factor = 4, 2, 1.0 + routing_bias = torch.randn( + num_experts, + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + scores = noaux_tc_ref( + logits.float(), + routing_bias.float(), + n_group=n_group, + topk_group=topk_group, + top_k=top_k, + routed_scaling_factor=routed_scaling_factor, + ) + scales, selected = torch.topk(scores, top_k, dim=-1) + else: + scales, selected = torch.topk( + torch.softmax(logits.float(), dim=-1), top_k, dim=-1 + ) + selected = selected.to(torch.int32) + + view = TrtllmMxInt4Config.prepare_weights( + w1, + w2, + num_local_experts=local_num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + device=device, + ) + q1, s1 = _mxint4_quantize(w1) + q2, s2 = _mxint4_quantize(w2) + w1_dequant = _dequant_mxint4( + q1, s1.reshape(local_num_experts, 2 * intermediate_size, hidden_size // 32) + ) + w2_dequant = _dequant_mxint4( + q2, s2.reshape(local_num_experts, hidden_size, intermediate_size // 32) + ) + reference = _mxint4_reference( + x, + w1_dequant, + w2_dequant, + selected, + scales, + intermediate_size, + expert_offset=local_expert_offset, + ) + + config = MoEConfig( + routing=RoutingConfig( + num_experts=num_experts, + top_k=top_k, + method=routing_method, + n_group=n_group, + topk_group=topk_group, + routed_scaling_factor=routed_scaling_factor, + ), + quant=QuantConfig(variant=QuantVariant.MxInt4), + experts=ExpertConfig( + intermediate_size=intermediate_size, + local_expert_offset=local_expert_offset, + local_num_experts=local_num_experts, + ), + backend=BackendOptions((TrtllmMxInt4Config(),)), + ) + if routing_input_mode is RoutingInputMode.FromLogits: + act = MoEActivationPack( + hidden_states_q=x, + hidden_states_scale=None, + routing_input_mode=routing_input_mode, + routing_logits=logits, + routing_bias=routing_bias, + ) + else: + act = MoEActivationPack( + hidden_states_q=x, + hidden_states_scale=None, + topk_ids=selected, + topk_weights=scales, + ) + weights = MoEWeightPack() + weights.prepare_for("trtllm_mxint4_routed", view) + return act, weights, config, reference, (w1, w2) + + +def _assert_mxint4_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + torch.testing.assert_close(actual.float(), expected.float(), atol=0.05, rtol=0.2) + + +@pytest.mark.parametrize("bad_dtype", [torch.float16, torch.float32]) +def test_mxint4_prepare_rejects_non_bf16(bad_dtype): + w1 = torch.empty(2, 512, 256, dtype=bad_dtype) + w2 = torch.empty(2, 256, 256, dtype=bad_dtype) + with pytest.raises(ValueError, match="BF16"): + TrtllmMxInt4Config.prepare_weights( + w1, + w2, + num_local_experts=2, + hidden_size=256, + intermediate_size=256, + ) + + +def test_mxint4_prepare_rejects_unaligned_geometry(): + w1 = torch.empty(2, 768, 256, dtype=torch.bfloat16) + w2 = torch.empty(2, 256, 384, dtype=torch.bfloat16) + with pytest.raises(ValueError, match="divisible by 256"): + TrtllmMxInt4Config.prepare_weights( + w1, + w2, + num_local_experts=2, + hidden_size=256, + intermediate_size=384, + ) + + +@pytest.mark.parametrize( + ("compute_capability", "supported"), + [((10, 0), True), ((10, 3), True), ((10, 7), True), ((12, 0), False)], +) +def test_mxint4_runner_arch_support(monkeypatch, compute_capability, supported): + import flashinfer.utils as utils + + config = MoEConfig( + routing=RoutingConfig(num_experts=8, top_k=2), + quant=QuantConfig(variant=QuantVariant.MxInt4), + experts=ExpertConfig(intermediate_size=256), + backend=BackendOptions((TrtllmMxInt4Config(),)), + ) + runner = TrtllmMxInt4RoutedRunner.__new__(TrtllmMxInt4RoutedRunner) + runner.config = config + runner.device = torch.device("cuda") + monkeypatch.setattr(utils, "get_compute_capability", lambda _: compute_capability) + if supported: + runner.check_support() + else: + with pytest.raises(NotImplementedError, match="SM100/SM103/SM107"): + runner.check_support() + + +@mxint4_required +def test_mxint4_prepare_matches_flat_test_layout(): + from tests.moe.trtllm_gen_fused_moe_utils import MxInt4BlockScaleMoe + + act, _, _, _, canonical = _make_case(num_experts=2) + w1, w2 = canonical + actual = TrtllmMxInt4Config.prepare_weights( + w1, + w2, + num_local_experts=2, + hidden_size=256, + intermediate_size=256, + device=act.hidden_states_q.device, + permute_cache={}, + ) + implementation = MxInt4BlockScaleMoe() + implementation._cache_permute_indices = {} + quantized = implementation.quantize_weights(w1, w2, act.hidden_states_q) + args = SimpleNamespace( + gemm1_weights=quantized["gemm1_weights"], + gemm2_weights=quantized["gemm2_weights"], + gemm1_scales=quantized["gemm1_scales"], + gemm2_scales=quantized["gemm2_scales"], + ) + expected = implementation.prepare_static_weights_for_kernel( + None, args, w1, w2, 256, 256, 2, None + ) + assert torch.equal(actual["gemm1_weights"], expected["gemm1_weights"]) + assert torch.equal(actual["gemm1_weights_scale"], expected["gemm1_scales"]) + assert torch.equal(actual["gemm2_weights"], expected["gemm2_weights"]) + assert torch.equal(actual["gemm2_weights_scale"], expected["gemm2_scales"]) + + +def test_w2_permute_cache_key_includes_epilogue_tile_m(): + weight = torch.empty(128, 64, dtype=torch.uint8) + cache = {} + tile_64 = get_w2_permute_indices_with_cache(cache, weight, 64) + actual_tile_128 = get_w2_permute_indices_with_cache(cache, weight, 128) + expected_tile_128 = get_w2_permute_indices_with_cache({}, weight, 128) + + assert not torch.equal(tile_64, expected_tile_128) + assert torch.equal(actual_tile_128, expected_tile_128) + + +def test_w3_w1_permute_cache_key_includes_gated_activation_mode(): + weight = torch.empty(256, 64, dtype=torch.uint8) + cache = {} + gated = _maybe_get_cached_w3_w1_permute_indices( + cache, weight, 128, is_gated_act_gemm=True + ) + actual_ungated = _maybe_get_cached_w3_w1_permute_indices( + cache, weight, 128, is_gated_act_gemm=False + ) + expected_ungated = _maybe_get_cached_w3_w1_permute_indices( + {}, weight, 128, is_gated_act_gemm=False + ) + + assert not torch.equal(gated, expected_ungated) + assert torch.equal(actual_ungated, expected_ungated) + + +@mxint4_required +@pytest.mark.parametrize( + "routing_input_mode", + [RoutingInputMode.PackedPrecomputed, RoutingInputMode.FromLogits], + ids=["packed", "from-logits"], +) +def test_mxint4_layer_and_direct_runner_match_reference(routing_input_mode): + act, weights, config, reference, _ = _make_case( + routing_input_mode=routing_input_mode + ) + layer_output = MoELayer(config)(act, weights) + _assert_mxint4_close(layer_output, reference) + + runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + direct_output = runner.forward(runner.pack_inputs(act, weights)) + _assert_mxint4_close(direct_output, reference) + + +@mxint4_required +def test_mxint4_from_logits_deepseek_routing_bias(): + act, weights, config, reference, _ = _make_case( + routing_input_mode=RoutingInputMode.FromLogits, + routing_method=RoutingMethodType.DeepSeekV3, + ) + _assert_mxint4_close(MoELayer(config)(act, weights), reference) + + +@mxint4_required +def test_mxint4_nonzero_expert_offset(): + act, weights, config, reference, _ = _make_case( + num_experts=8, local_num_experts=4, local_expert_offset=4 + ) + _assert_mxint4_close(MoELayer(config)(act, weights), reference) + + +@mxint4_required +def test_mxint4_explicit_autotune_matches_reference(): + act, weights, config, reference, _ = _make_case() + layer = MoELayer(config) + with autotune(True): + output = layer(act, weights) + assert layer.winner_backend == "trtllm_mxint4_routed" + _assert_mxint4_close(output, reference) + + +@mxint4_required +def test_mxint4_from_logits_rejects_fp32_until_validated(): + act, weights, config, _, _ = _make_case( + routing_input_mode=RoutingInputMode.FromLogits + ) + act.routing_logits = act.routing_logits.float() + runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + with pytest.raises(TypeError, match="requires bfloat16 routing_logits"): + runner.pack_inputs(act, weights) + + +@mxint4_required +def test_mxint4_from_logits_rejects_fp32_bias(): + act, weights, config, _, _ = _make_case( + routing_input_mode=RoutingInputMode.FromLogits, + routing_method=RoutingMethodType.DeepSeekV3, + ) + act.routing_bias = act.routing_bias.float() + runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + with pytest.raises(TypeError, match="routing_bias must be bfloat16"): + runner.pack_inputs(act, weights) + + +@mxint4_required +@pytest.mark.parametrize("field", ["hidden_states_q", "routing_logits", "routing_bias"]) +def test_mxint4_runner_rejects_noncontiguous_runtime_inputs(field): + act, weights, config, _, _ = _make_case( + routing_input_mode=RoutingInputMode.FromLogits, + routing_method=RoutingMethodType.DeepSeekV3, + ) + tensor = getattr(act, field) + if tensor.dim() == 2: + tensor = tensor.T.contiguous().T + else: + tensor = tensor.repeat_interleave(2)[::2] + assert not tensor.is_contiguous() + setattr(act, field, tensor) + + runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + with pytest.raises(ValueError, match=rf"{field} must be contiguous"): + runner.pack_inputs(act, weights) + + +@mxint4_required +@pytest.mark.parametrize( + "key", + [ + "gemm1_weights", + "gemm1_weights_scale", + "gemm2_weights", + "gemm2_weights_scale", + ], +) +def test_mxint4_runner_rejects_malformed_prepared_view(key): + act, weights, config, _, _ = _make_case() + view = weights.get_view("trtllm_mxint4_routed") + view[key] = view[key][..., :-1].contiguous() + runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + with pytest.raises(ValueError, match=rf"{key} shape"): + runner.pack_inputs(act, weights) + + +@mxint4_required +@pytest.mark.parametrize("dimension", ["hidden", "intermediate"]) +def test_mxint4_runner_rejects_unaligned_runtime_geometry(dimension): + act, weights, config, _, _ = _make_case() + if dimension == "hidden": + act.hidden_states_q = torch.empty( + act.hidden_states_q.shape[0], + 384, + dtype=torch.bfloat16, + device=act.hidden_states_q.device, + ) + else: + config = dataclasses.replace( + config, + experts=dataclasses.replace(config.experts, intermediate_size=384), + ) + runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + with pytest.raises(ValueError, match="divisible by 256"): + runner.pack_inputs(act, weights) + + +@mxint4_required +@pytest.mark.parametrize( + ("mutation", "error_type", "match"), + [ + ("device", ValueError, "is on"), + ("dtype", TypeError, "must be float32"), + ("shape", ValueError, "shape"), + ("contiguous", ValueError, "must be contiguous"), + ], +) +def test_mxint4_runner_validates_optional_gemm1_params(mutation, error_type, match): + act, weights, config, _, _ = _make_case() + view = weights.get_view("trtllm_mxint4_routed") + num_local_experts = config.experts.local_num_experts or config.routing.num_experts + if mutation == "device": + value = torch.ones(num_local_experts, dtype=torch.float32, device="meta") + elif mutation == "dtype": + value = torch.ones( + num_local_experts, dtype=torch.float16, device=act.hidden_states_q.device + ) + elif mutation == "shape": + value = torch.ones( + num_local_experts - 1, + dtype=torch.float32, + device=act.hidden_states_q.device, + ) + else: + value = torch.ones( + 1, dtype=torch.float32, device=act.hidden_states_q.device + ).expand(num_local_experts) + assert not value.is_contiguous() + view["gemm1_alpha"] = value + + runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + with pytest.raises(error_type, match=match): + runner.pack_inputs(act, weights) + + +@mxint4_required +@pytest.mark.parametrize( + "routing_input_mode", + [RoutingInputMode.PackedPrecomputed, RoutingInputMode.FromLogits], + ids=["packed", "from-logits"], +) +def test_mxint4_cuda_graph_replay(routing_input_mode): + act, weights, config, reference, _ = _make_case( + routing_input_mode=routing_input_mode + ) + runner = TrtllmMxInt4RoutedRunner(config, torch.device("cuda")) + inputs = runner.pack_inputs(act, weights) + eager = runner.forward(inputs).clone() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + runner.forward(inputs) + inputs[0].fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert torch.isfinite(inputs[0]).all() + torch.testing.assert_close(inputs[0], eager, atol=0.05, rtol=0.01) + _assert_mxint4_close(inputs[0], reference)