diff --git a/src/models/model_type.h b/src/models/model_type.h index 19d9892568..90fe98f7a9 100644 --- a/src/models/model_type.h +++ b/src/models/model_type.h @@ -17,7 +17,7 @@ namespace Generators { struct ModelType { inline static bool IsLLM(const std::string& model_type) { // Large-language model (LLM) - static constexpr std::array LLM = {"chatglm", "decoder", "ernie4_5", "gemma", "gemma2", "gemma3_text", "gemma3_vl_text", "gemma4_text", "gpt2", "gptoss", "granite", "granitemoehybrid", "hunyuandensev1", "internlm2", "lfm2", "lfm2_vl_text", "llama", "mistral", "mistral3_text", "nemotron", "olmo", "phi", "phimoe", "phi3", "phi3small", "qwen2", "qwen2_5_vl_text", "qwen3", "qwen3_vl_text", "qwen3_5_moe_text", "qwen3_5_text", "smollm3"}; + static constexpr std::array LLM = {"chatglm", "decoder", "ernie4_5", "gemma", "gemma2", "gemma3_text", "gemma3_vl_text", "gemma4_text", "gpt2", "gptoss", "granite", "granitemoehybrid", "hunyuandensev1", "internlm2", "lfm2", "lfm2_moe", "lfm2_vl_text", "llama", "mistral", "mistral3_text", "nemotron", "olmo", "phi", "phimoe", "phi3", "phi3small", "qwen2", "qwen2_5_vl_text", "qwen3", "qwen3_vl_text", "qwen3_5_moe_text", "qwen3_5_text", "smollm3"}; return std::find(LLM.begin(), LLM.end(), model_type) != LLM.end(); } @@ -75,9 +75,10 @@ struct ModelType { inline static bool IsLFM2(const std::string& model_type) { // Liquid Foundation Model 2: hybrid attention/conv architecture with conv state cache. - // "lfm2_vl_text" is an LFM2-VL checkpoint exported as a plain text decoder; it has the same - // conv/attention layer mix and needs the same conv state cache. - return model_type == "lfm2" || model_type == "lfm2_vl_text"; + // The MoE variant (LFM2-8B-A1B, LFM2.5-8B-A1B, LFM2-24B-A2B) has the same runtime shape, and + // "lfm2_vl_text" is an LFM2-VL checkpoint exported as a plain text decoder; both have the same + // conv/attention layer mix and need the same conv state cache. + return model_type == "lfm2" || model_type == "lfm2_moe" || model_type == "lfm2_vl_text"; } }; diff --git a/src/python/py/models/README.md b/src/python/py/models/README.md index 60d030efce..11c8d377c5 100644 --- a/src/python/py/models/README.md +++ b/src/python/py/models/README.md @@ -75,6 +75,7 @@ The tool currently supports the following model architectures. - HunYuan Dense V1 - InternLM2 - LFM2 (text and the decoder of LFM2-VL / LFM2.5-VL) +- LFM2 MoE - Llama - Mistral - Nemotron @@ -599,7 +600,8 @@ python builder.py -m model_name -o path_to_output_folder -p int4 -e execution_pr ##### QMoE Block Size This scenario is for when you want to set the block size for QMoE expert weights. -Set `qmoe_block_size` to `0` or a negative value for per-channel quantization. CUDA block-wise QMoE supports only `32`, `64`, or `128`; the default is `32` except for TRT-RTX, which defaults to `128`. +Set `qmoe_block_size` to `0` or a negative value for per-channel quantization. Block-wise QMoE on CPU, CUDA, and WebGPU supports only `32`, `64`, or `128`; TRT-RTX also accepts `16` and `256`. The default is `32` except for TRT-RTX, which defaults to `128`. +WebGPU block-wise QMoE requires both `hidden_size` and `moe_intermediate_size` to be divisible by `qmoe_block_size`. Raw block-wise INT4 QMoE requires both dimensions to be even. ```bash # From wheel: @@ -609,6 +611,12 @@ python -m onnxruntime_genai.models.builder -m model_name -o path_to_output_folde python builder.py -m model_name -o path_to_output_folder -p int4 -e execution_provider --extra_options qmoe_block_size=128 ``` +On CPU and WebGPU the block-wise expert weights use ONNX Runtime's MatMulNBits encoding with signed block scales: each block's max-magnitude element maps exactly to the lowest quantized value, so no extreme is clipped. This encoding has been validated end to end against the Hugging Face reference on CPU, and the WebGPU QMoE kernel consumes the MatMulNBits layout directly. TRT-RTX keeps the builder's original symmetric encoding with positive block scales until the signed-scale grid has been measured on that execution provider. + +The CPU `QMoE` kernel dequantizes the experts to fp32 on every call unless the environment variable `ORT_USE_MLAS_Q4_GEMM_MOE=1` is set at runtime, which enables its AVX-512 MLAS Q4 fast path. That fast path re-quantizes to the same grid as this encoding, so enabling it only changes results at fp32 rounding level. + +ONNX Runtime builds that include [microsoft/onnxruntime#32644](https://github.com/microsoft/onnxruntime/pull/32644) run block-wise experts on the MLAS QNBit GEMM (`MatMulNBits`) kernels directly, without the environment variable. For LFM2.5-8B-A1B int4 on a 30-core x86 server, CPU decode went from about 0.8 s per token with the default dequantize-then-GEMM path to 35 ms per token with `ORT_USE_MLAS_Q4_GEMM_MOE=1` and 15 ms per token with #32644. + ##### QMoE Weights Prepacked This scenario is for when you want to control the CUDA QMoE expert weight layout. The default value is `-1`, which lets the builder choose the layout automatically. Use `0` to export raw weights and let CUDA prepack them at runtime, or `1` to export CUTLASS-prepacked weights. diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py index 821c21585b..622d60daeb 100644 --- a/src/python/py/models/builder.py +++ b/src/python/py/models/builder.py @@ -29,6 +29,7 @@ HunyuanDenseV1Model, InternLM2Model, LFM2Model, + LFM2MoEModel, LlamaModel, Mistral3TextModel, MistralModel, @@ -562,6 +563,9 @@ def create_model( onnx_model = InternLM2Model(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options) elif config.architectures[0] == "Lfm2ForCausalLM": onnx_model = LFM2Model(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options) + elif config.architectures[0] == "Lfm2MoeForCausalLM": + onnx_model = LFM2MoEModel(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options) + onnx_model.model_type = "lfm2_moe" elif config.architectures[0] == "Lfm2VlForConditionalGeneration": onnx_model = LFM2Model(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options) # With the embedding layer excluded the decoder is one stage of the LFM2-VL vision pipeline; @@ -728,7 +732,9 @@ def get_args(): Default value is 32. qmoe_block_size = <=0/16/32/64/128/256: Specify the block size for QMoE expert weights quantization. Set <= 0 for per-channel quantization. Default is 128 for TRT-RTX, 32 for others. - CUDA block-wise QMoE supports 32/64/128 only. + CPU, CUDA, and WebGPU block-wise QMoE support 32/64/128 only; TRT-RTX also accepts 16/256. + WebGPU requires hidden_size and moe_intermediate_size to be divisible by qmoe_block_size. + Raw block-wise INT4 QMoE requires both dimensions to be even. Supported EPs: CPU, CUDA, WebGPU, TRT-RTX. qmoe_weights_prepacked = -1/0/1: Specify the CUDA QMoE expert weight layout. -1 lets the builder choose automatically, 0 exports raw weights for runtime prepacking, and 1 exports CUTLASS-prepacked weights. diff --git a/src/python/py/models/builders/__init__.py b/src/python/py/models/builders/__init__.py index 9544a51862..0925693898 100644 --- a/src/python/py/models/builders/__init__.py +++ b/src/python/py/models/builders/__init__.py @@ -14,7 +14,7 @@ from .granite import GraniteModel, GraniteMoEHybridModel from .hunyuan import HunyuanDenseV1Model from .internlm import InternLM2Model -from .lfm2 import LFM2Model +from .lfm2 import LFM2Model, LFM2MoEModel from .llama import LlamaModel from .mistral import Mistral3TextModel, MistralModel from .mtp import MTPModel @@ -58,6 +58,7 @@ "HunyuanDenseV1Model", "InternLM2Model", "LFM2Model", + "LFM2MoEModel", "LlamaModel", "MTPModel", "Mistral3TextModel", diff --git a/src/python/py/models/builders/base.py b/src/python/py/models/builders/base.py index 01432113ca..9763ca12c6 100644 --- a/src/python/py/models/builders/base.py +++ b/src/python/py/models/builders/base.py @@ -404,32 +404,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): self.make_quant_config_init() # MoE-specific variables - num_experts = ( - config.num_local_experts - if hasattr(config, "num_local_experts") - else config.num_experts - if hasattr(config, "num_experts") - else 0 - ) - top_k_experts = config.num_experts_per_tok if hasattr(config, "num_experts_per_tok") else 0 - swiglu_limit = config.swiglu_limit if hasattr(config, "swiglu_limit") else None - self.moe_attrs = { - "op_type": "MoE", # MoE op to use - "num_experts": num_experts, # Number of experts in MoE layer - "top_k": top_k_experts, # Number of experts to select in MoE layer - "activation_alpha": 1.0, # Alpha parameter used in activation function - "activation_beta": 0.0, # Beta parameter used in activation function - "activation_type": self.activation, # Activation function for MoE layer - "expert_weight_bits": -1, # Number of bits used in quantized MoE weights (only INT4 or INT8 are supported). - "normalize_routing_weights": False, # Normalize routing weights in MoE layer - "swiglu_fusion": 0, # Fusion level for SwiGLU activation function - "swiglu_limit": swiglu_limit, # Value used to clamp results into a certain range in SwiGLU activation function - "use_sparse_mixer": False, # Use SparseMixer in MoE layer (used in Phi-3.5 MoE) - "weights_prepacked": 0, # CUDA QMoE layout: -1=auto/omit, 0=raw, 1=CUTLASS-prepacked - "quant_type": "int", # QMoE quantization type: "int" (INT4/INT8), "fp4" (MXFP4), or "nvfp4" (NVFP4). - "global_scale_names": {}, # Per-layer QMoE global-scale initializer names, when required. - "zero_point_names": {}, # Per-layer QMoE zero-point initializer names, when required. - } + self.make_moe_attrs_init(config) self.make_moe_init() # LM head-specific variables @@ -1018,6 +993,38 @@ def make_quant_config_init(self): elif not isinstance(self.quant_config, QuantConfig): raise TypeError("_quant_config must be a QuantConfig instance") + def make_moe_attrs_init(self, config): + num_experts = ( + config.num_local_experts + if hasattr(config, "num_local_experts") + else config.num_experts + if hasattr(config, "num_experts") + else 0 + ) + top_k_experts = config.num_experts_per_tok if hasattr(config, "num_experts_per_tok") else 0 + swiglu_limit = config.swiglu_limit if hasattr(config, "swiglu_limit") else None + self.moe_attrs = { + "op_type": "MoE", # MoE op to use + "num_experts": num_experts, # Number of experts in MoE layer + "top_k": top_k_experts, # Number of experts to select in MoE layer + "activation_alpha": 1.0, # Alpha parameter used in activation function + "activation_beta": 0.0, # Beta parameter used in activation function + "activation_type": self.activation, # Activation function for MoE layer + "expert_weight_bits": -1, # Number of bits used in quantized MoE weights (only INT4 or INT8 are supported). + "normalize_routing_weights": False, # Normalize routing weights in MoE layer + "swiglu_fusion": 0, # Fusion level for SwiGLU activation function + "swiglu_limit": swiglu_limit, # Value used to clamp results into a certain range in SwiGLU activation function + "use_sparse_mixer": False, # Use SparseMixer in MoE layer (used in Phi-3.5 MoE) + "router_sentinel": None, # Router score given to unselected experts when the model selects them in-graph (e.g. LFM2-MoE) + "num_dense_layers": 0, # Leading decoder layers that keep a dense MLP instead of a MoE layer (e.g. LFM2-MoE) + "use_expert_bias": False, # Select experts on scores + a load-balancing bias but mix with the unbiased scores (e.g. LFM2-MoE) + "routed_scaling_factor": 1.0, # Multiplier applied to the routed experts' output (e.g. LFM2-MoE) + "weights_prepacked": 0, # CUDA QMoE layout: -1=auto/omit, 0=raw, 1=CUTLASS-prepacked + "quant_type": "int", # QMoE quantization type: "int" (INT4/INT8), "fp4" (MXFP4), or "nvfp4" (NVFP4). + "global_scale_names": {}, # Per-layer QMoE global-scale initializer names, when required. + "zero_point_names": {}, # Per-layer QMoE zero-point initializer names, when required. + } + def make_moe_init(self): # MoE quantization scheme comes from `quant_config.moe.type` ("int4"/"int8"/"mxfp4"/"nvfp4"), which maps to # (expert_weight_bits, QMoE quant_type): @@ -1034,11 +1041,11 @@ def make_moe_init(self): # MXFP4 and NVFP4 both resolve to the "mx" kind; the QMoE op tells them apart by dtype name # ("mxfp4" -> op "fp4", "nvfp4" -> op "nvfp4"). Integer dtypes use the plain "int" QMoE path. - self.moe_attrs["moe_op_type"] = "QMoE" if moe_descriptor.is_quantized else "MoE" + self.moe_attrs["op_type"] = "QMoE" if moe_descriptor.is_quantized else "MoE" if moe_descriptor.kind == "mx": - self.moe_attrs["qmoe_quant_type"] = "nvfp4" if moe_descriptor.name == "nvfp4" else "fp4" + self.moe_attrs["quant_type"] = "nvfp4" if moe_descriptor.name == "nvfp4" else "fp4" else: - self.moe_attrs["qmoe_quant_type"] = "int" + self.moe_attrs["quant_type"] = "int" # weights_prepacked is a CUDA-only QMoE layout contract. Non-CUDA EPs omit the attribute and use # their normal blockwise QMoE encoding, so CUDA-prepacked exports are not intended to be shared @@ -1046,6 +1053,11 @@ def make_moe_init(self): # raw [E, N, K/pack] weights and let the CUDA runtime PrePack hook transform them). self.moe_attrs["weights_prepacked"] = self.quant_config.moe.weights_prepacked + if self.moe_attrs["swiglu_limit"] is None and self.ep == "trt-rtx": + # TRT-RTX EP builds currently require QMoE swiglu_limit to be present on every MoE model; + # use +inf to preserve the "no clamp" behavior when the model omits it. + self.moe_attrs["swiglu_limit"] = float("inf") + def make_lm_head_init(self, config): pass @@ -2082,6 +2094,23 @@ def make_gather(self, name, inputs, dtype, shape, axis): self.make_node("Gather", inputs=inputs, outputs=[output], name=name, axis=axis) self.make_value(output, dtype, shape=shape) + def make_gather_elements(self, name, inputs, dtype, shape, axis): + output = f"{name}/output_0" + self.make_node("GatherElements", inputs=inputs, outputs=[output], name=name, axis=axis) + self.make_value(output, dtype, shape=shape) + + def make_scatter_elements(self, name, inputs, dtype, shape, axis): + output = f"{name}/output_0" + self.make_node("ScatterElements", inputs=inputs, outputs=[output], name=name, axis=axis) + self.make_value(output, dtype, shape=shape) + + def make_topk(self, name, inputs, dtype, shape, axis=-1, largest=True): + # Emits both TopK outputs: `{name}/output_0` holds the values and `{name}/output_1` the indices. + outputs = [f"{name}/output_0", f"{name}/output_1"] + self.make_node("TopK", inputs=inputs, outputs=outputs, name=name, axis=axis, largest=largest) + self.make_value(outputs[0], dtype, shape=shape) + self.make_value(outputs[1], ir.DataType.INT64, shape=shape) + def make_reshape(self, name, inputs, dtype, shape): output = f"{name}/output_0" self.make_node("Reshape", inputs=inputs, outputs=[output], name=name) @@ -2258,6 +2287,26 @@ def make_sigmoid(self, name, root_input, dtype, shape): self.make_node("Sigmoid", inputs=[root_input], outputs=[output], name=name) self.make_value(output, dtype, shape=shape) + def make_log(self, name, root_input, dtype, shape): + output = f"{name}/output_0" + self.make_node("Log", inputs=[root_input], outputs=[output], name=name) + self.make_value(output, dtype, shape=shape) + + def make_exp(self, name, root_input, dtype, shape): + output = f"{name}/output_0" + self.make_node("Exp", inputs=[root_input], outputs=[output], name=name) + self.make_value(output, dtype, shape=shape) + + def make_neg(self, name, root_input, dtype, shape): + output = f"{name}/output_0" + self.make_node("Neg", inputs=[root_input], outputs=[output], name=name) + self.make_value(output, dtype, shape=shape) + + def make_reciprocal(self, name, root_input, dtype, shape): + output = f"{name}/output_0" + self.make_node("Reciprocal", inputs=[root_input], outputs=[output], name=name) + self.make_value(output, dtype, shape=shape) + def make_cos(self, name, root_input, dtype, shape): output = f"{name}/output_0" self.make_node("Cos", inputs=[root_input], outputs=[output], name=name) @@ -5097,13 +5146,56 @@ def make_moe(self, layer_id, moe, root_input): def make_moe_preprocessing(self, layer_id, moe, root_input): raise NotImplementedError("MoE weight preprocessing must be implemented by the model class.") + def make_moe_expert_names(self, layer_id): + """Initializer names of the packed expert tensors that the MoE/QMoE op consumes. + + Returns a dict with the keys ``gate_up_weight``, ``gate_up_scales``, ``gate_up_bias``, + ``down_weight``, ``down_scales`` and ``down_bias``. The weight name ends in ``qweight`` for + QMoE and ``weight`` for MoE. + """ + weight_type = "qweight" if self.moe_attrs["op_type"] == "QMoE" else "weight" + names = {} + for key, proj in (("gate_up", "gate_up_proj"), ("down", "down_proj")): + prefix = f"model.layers.{layer_id}.moe.experts.{proj}" + names[f"{key}_weight"] = f"{prefix}.{weight_type}" + names[f"{key}_scales"] = f"{prefix}.scales" + names[f"{key}_bias"] = f"{prefix}.bias" + return names + + def make_interleaved_swiglu_moe_preprocessing(self, layer_id, moe): + """Emit the expert initializers for a bias-free SwiGLU MoE layer with fused gate/up experts. + + Hugging Face stores ``[gate | up]`` concatenated along the output axis of ``gate_up_proj`` + ``[E, 2I, H]``; the fused SwiGLU (``swiglu_fusion=1``) wants the rows interleaved + ``[g0, u0, g1, u1, ...]``. Pre-quantized checkpoints (``gate_up_proj`` is ``None``) already + carry QMoE-packed tensors and are passed through. The MoE/QMoE op takes the expert biases as + separate inputs, so zero biases are emitted for the models that have none. + """ + gate_up_weight = None + down_weight = None + if getattr(moe.experts, "gate_up_proj", None) is not None: + raw_gate_up = moe.experts.gate_up_proj + half = raw_gate_up.shape[1] // 2 + gate_up_weight = torch.stack([raw_gate_up[:, :half, :], raw_gate_up[:, half:, :]], dim=2).reshape_as( + raw_gate_up + ) + down_weight = moe.experts.down_proj + self.make_moe_expert_initializers(layer_id, moe.experts, gate_up_weight, down_weight) + + names = self.make_moe_expert_names(layer_id) + num_experts = self.moe_attrs["num_experts"] + self.make_initializer( + torch.zeros(num_experts, 2 * self.moe_intermediate_size), names["gate_up_bias"], to=self.io_dtype + ) + self.make_initializer(torch.zeros(num_experts, self.hidden_size), names["down_bias"], to=self.io_dtype) + def make_moe_expert_initializers(self, layer_id, experts, gate_up_weight=None, down_weight=None): op_type = self.moe_attrs["op_type"] - weight_type = f"{'q' if op_type == 'QMoE' else ''}weight" - gate_up_name = f"model.layers.{layer_id}.moe.experts.gate_up_proj.{weight_type}" - gate_up_scales_name = f"model.layers.{layer_id}.moe.experts.gate_up_proj.scales" - down_name = f"model.layers.{layer_id}.moe.experts.down_proj.{weight_type}" - down_scales_name = f"model.layers.{layer_id}.moe.experts.down_proj.scales" + names = self.make_moe_expert_names(layer_id) + gate_up_name = names["gate_up_weight"] + gate_up_scales_name = names["gate_up_scales"] + down_name = names["down_weight"] + down_scales_name = names["down_scales"] native_quant_type = getattr(experts, "quant_type", None) if native_quant_type is not None: @@ -5164,7 +5256,18 @@ def make_moe_expert_initializers(self, layer_id, experts, gate_up_weight=None, d def make_moe_router(self, layer_id, moe, root_input): raise NotImplementedError("MoE router construction must be implemented by the model class.") - def make_moe_subgraph(self, layer_id, moe, root_input): + def make_moe_router_shape(self, last_dim=None): + """Shape of a per-token router tensor: one row per token, `num_experts` (or `last_dim`) columns. + + The row dim follows `make_hidden_state_shape`: paged attention flattens tokens to `num_tokens`, + so the router tensors must declare the same symbolic dim as the MoE op's input. + """ + rows = "num_tokens" if self.use_paged_attention else "batch_size * sequence_length" + return [rows, self.moe_attrs["num_experts"] if last_dim is None else last_dim] + + def make_moe_subgraph(self, layer_id, moe, root_input, router_probs=None, output_scale=None): + # `router_probs` and `output_scale` are whatever the model's `make_moe_router` returned; models + # whose router output is addressed by name (e.g. `.../router/Reshape/output_0`) leave them as None. raise NotImplementedError("MoE subgraph construction must be implemented by the model class.") def make_moe_op(self, name, **kwargs): @@ -5361,27 +5464,44 @@ def make_qmoe_weights(self, weights): except Exception as e: raise RuntimeError(f"Per-channel QMoE quantization failed: {e}") from e - if self.ep == "cuda" and self.quant_attrs["qmoe_block_size"] > 0: + supported_blockwise_eps = ["cpu", "cuda", "webgpu", "trt-rtx"] + block_size = self.quant_attrs["qmoe_block_size"] + if self.ep not in supported_blockwise_eps: + raise RuntimeError(f"Please use a supported EP ({', '.join(supported_blockwise_eps)}) " + "for QMoE expert weights quantization. " + f"Got qmoe_block_size={block_size} and ep={self.ep}.") + + if self.ep == "trt-rtx": + # TRT-RTX keeps the builder's original symmetric blockwise encoding (positive block scales, + # so the block extreme can be clipped to 7/8 of its magnitude). The signed-scale MatMulNBits + # grid used below has not been measured against the TRT-RTX kernel yet, so the EP stays on + # the encoding it shipped with until that parity run exists. + try: + qweight, scales = self._symmetric_blockwise_quantize(weights, block_size) + self.moe_attrs["block_size"] = block_size + return qweight, scales.to(torch.float16) + except Exception as e: + raise RuntimeError(f"Block-wise quantization failed with block_size={block_size}: {e}") from e + + if block_size not in (32, 64, 128): + raise ValueError(f"Block-wise QMoE only supports block_size 32, 64, or 128, got {block_size}.") + + if self.ep == "cuda": # CUDA QMoE consumes CUTLASS-prepacked expert weights (the kernel's fpA_intB mixed GEMM # layout). For weights_prepacked=-1 (auto) or 1, produce them offline so the QMoE op reads # them directly: quantize with ONNX Runtime's blockwise quantizer, keep the signed scales, # then run pack_weights_for_cuda_mixed_gemm. This is the encoding validated by the - # com.microsoft QMoE CUDA parity tests. The builder's own _symmetric_blockwise_quantize uses - # a different scale/packing convention the kernel cannot consume. + # com.microsoft QMoE CUDA parity tests. # # weights_prepacked=0 ships raw [N, K/pack] weights with ONNX Runtime's MatMulNBits-compatible # blockwise quantizer. This is the exact encoding the CUDA QMoE PrePack hook expects: raw # bytes + blockwise scales, which it lays out into the CUTLASS fpA_intB format at load time. - block_size = self.quant_attrs["qmoe_block_size"] quantize_method = ( self._matmulnbits_blockwise_quantize if weights_prepacked == 0 else self._cutlass_prepacked_blockwise_quantize ) descriptor = "MatMulNBits-compatible" if weights_prepacked == 0 else "CUTLASS-prepacked" - - if block_size not in (32, 64, 128): - raise ValueError(f"CUDA QMoE only supports block_size 32, 64, or 128, got {block_size}.") try: qweight, scales = quantize_method(weights) self.moe_attrs["block_size"] = block_size @@ -5389,22 +5509,16 @@ def make_qmoe_weights(self, weights): except Exception as e: raise RuntimeError(f"{descriptor} QMoE quantization failed with block_size={block_size}: {e}") from e - # Use block-wise quantization for supported EPs when qmoe_block_size > 0. - supported_blockwise_eps = ["cpu", "cuda", "webgpu", "trt-rtx"] - use_blockwise_quant = self.ep in supported_blockwise_eps and self.quant_attrs["qmoe_block_size"] > 0 - - if use_blockwise_quant: - block_size = self.quant_attrs["qmoe_block_size"] - try: - qweight, scales = self._symmetric_blockwise_quantize(weights, block_size) - self.moe_attrs["block_size"] = block_size - return qweight, scales.to(torch.float16) - except Exception as e: - raise RuntimeError(f"Block-wise quantization failed with block_size={block_size}: {e}") from e - - raise RuntimeError(f"Please use a supported EP ({', '.join(supported_blockwise_eps)}) " - "for QMoE expert weights quantization. " - f"Got qmoe_block_size={self.quant_attrs['qmoe_block_size']} and ep={self.ep}.") + # CPU and WebGPU QMoE ship raw [N, K/pack] weights on ONNX Runtime's MatMulNBits grid (signed block + # scales, so no extreme is clipped). The CPU kernel's MLAS Q4 fast path (ORT_USE_MLAS_Q4_GEMM_MOE=1) + # re-quantizes to this same grid, which makes it lossless, and the WebGPU kernel consumes the + # MatMulNBits layout directly. + try: + qweight, scales = self._matmulnbits_blockwise_quantize(weights) + self.moe_attrs["block_size"] = block_size + return qweight, scales.to(torch.float16) + except Exception as e: + raise RuntimeError(f"Block-wise quantization failed with block_size={block_size}: {e}") from e # TODO: replace all five CudaQuantizer methods with calls to native ORT APIs def _symmetric_per_channel_quantize(self, weights): @@ -5470,22 +5584,36 @@ def _matmulnbits_blockwise_quantize(self, weights): ``weights`` is a single expert's weight of logical shape ``[N, K]`` (quantized along the last/``K`` axis). Returns ``(qweight, scales)`` where - ``qweight`` is ``[N, K/pack]`` uint8 (2 INT4 elements per byte; INT8 is - one element per byte) and ``scales`` is ``[N, K/block_size]`` float scales - (SIGNED by default on this blockwise path — the MLAS ``default`` - convention). Layout matches ``quantize_matmul_{4,8}bits``. + ``qweight`` is ``[N, K/pack]`` uint8 (2 INT4 elements per byte; INT8 + is one element per byte) and ``scales`` is ``[N, ceil(K/block_size)]`` + float scales (SIGNED by default on this blockwise path — the MLAS + ``default`` convention). Layout matches ``quantize_matmul_{4,8}bits``. """ bits = int(self.moe_attrs["expert_weight_bits"]) block_size = self.quant_attrs["qmoe_block_size"] - return CudaQuantizer.matmulnbits_blockwise_quantize( + pack = 8 // bits + k = weights.shape[-1] + if k % pack != 0: + raise ValueError(f"INT{bits} QMoE requires expert input dimension K ({k}) to be divisible by {pack}.") + if self.ep == "webgpu" and k % block_size != 0: + # WebGPU indexes raw weights with a whole-block stride, while QMoE requires unpadded storage. + raise ValueError( + f"WebGPU QMoE requires expert input dimension K ({k}) to be divisible by " + f"qmoe_block_size ({block_size}); partial blocks are unsupported." + ) + qweight, scales = CudaQuantizer.matmulnbits_blockwise_quantize( weights, bits, block_size, unsigned_full_range=True, signed_scale=True, ) + # QMoE validates raw storage as [E, N, K/pack]. Drop the quantizer's whole-block padding; + # the scales retain ceil(K/block_size) columns. WebGPU partial blocks are rejected above. + return qweight[:, : k // pack], scales def _symmetric_blockwise_quantize(self, weights, block_size): + """Original symmetric blockwise encoding with positive scales; kept for the TRT-RTX EP.""" bits = self.moe_attrs["expert_weight_bits"] return CudaQuantizer.symmetric_blockwise_quantize( weights, diff --git a/src/python/py/models/builders/lfm2.py b/src/python/py/models/builders/lfm2.py index 524163fc1d..c7e4dce6e0 100644 --- a/src/python/py/models/builders/lfm2.py +++ b/src/python/py/models/builders/lfm2.py @@ -17,6 +17,11 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): # LFM2-specific attributes self.layernorm_attrs["epsilon"] = config.norm_eps + self.make_intermediate_size_init(config) + + self.conv_L_cache = config.conv_L_cache + + def make_intermediate_size_init(self, config): # Calculate the dynamic intermediate_size for the MLP. intermediate_size = config.intermediate_size if config.block_auto_adjust_ff_dim: @@ -28,8 +33,6 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): ) self.intermediate_size = intermediate_size - self.conv_L_cache = config.conv_L_cache - def make_attention_init(self, config): self.attention_attrs["q_norm"] = True self.attention_attrs["k_norm"] = True @@ -160,19 +163,350 @@ def make_layer(self, layer_id, layer): location="ffn", ) - # Alias MLP attribute names for compatibility with the base class - layer.mlp = layer.feed_forward - layer.mlp.gate_proj = layer.mlp.w1 - layer.mlp.up_proj = layer.mlp.w3 - layer.mlp.down_proj = layer.mlp.w2 - self.make_mlp(layer_id, layer.mlp, root_input=self.layernorm_attrs["output_0"]) + self.make_feed_forward(layer_id, layer, root_input=self.layernorm_attrs["output_0"]) self.layernorm_attrs["first_layernorm"] = False if layer_id == self.num_layers - 1: # Norm after last decoder layer of model (last layer --> norm) self.layernorm_attrs["last_layernorm"] = True + def get_feed_forward_module(self, layer): + # Hugging Face names the module `feed_forward`; the quantized-checkpoint IR exposes it as `mlp`. + feed_forward = getattr(layer, "feed_forward", None) + return layer.mlp if feed_forward is None else feed_forward + + def make_feed_forward(self, layer_id, layer, root_input): + mlp = self.get_feed_forward_module(layer) + if hasattr(mlp, "w1"): + # Alias Hugging Face's MLP attribute names for compatibility with the base class + mlp.gate_proj = mlp.w1 + mlp.up_proj = mlp.w3 + mlp.down_proj = mlp.w2 + self.make_mlp(layer_id, mlp, root_input=root_input) + def update_genai_config(self, genai_config): decoder = genai_config["model"]["decoder"] decoder["layer_types"] = self.layer_types decoder["conv_cache_size"] = self.conv_L_cache - 1 + + +class LFM2MoEModel(LFM2Model): + """LFM2-MoE builder (LFM2-8B-A1B, LFM2.5-8B-A1B, LFM2-24B-A2B). + + Same hybrid conv/attention stack as LFM2. The first ``num_dense_layers`` layers keep the dense + SwiGLU MLP; every later layer routes each token to ``num_experts_per_tok`` of ``num_experts`` + experts. The router (``Lfm2MoeSparseMoeBlock.route_tokens_to_experts`` in transformers) is + sigmoid based with an auxiliary-loss-free load-balancing bias: experts are *selected* by the + top-k of ``sigmoid(logits) + expert_bias`` but *mixed* with + ``sigmoid_i / (sum_selected(sigmoid_j) + 1e-6)``. + + The fused MoE/QMoE op only knows softmax-over-top-k routing, so the selection is done in the + graph and the op is fed ``log(sigmoid(logits))`` at the selected experts and a large negative + sentinel everywhere else. Its softmax over the surviving top-k entries then equals + ``sigmoid_i / sum_selected(sigmoid_j)``. The remaining per-token factor + ``sum_selected / (sum_selected + 1e-6)`` is applied to the op output, together with + ``routed_scaling_factor``, so tokens whose selected scores are all tiny (or flushed to zero) + keep the HF output magnitude. + """ + + def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): + super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options) + + if not config.norm_topk_prob: + raise NotImplementedError( + "LFM2-MoE export requires norm_topk_prob=True: the fused MoE op always normalizes the " + "routing weights over the selected experts." + ) + + self.moe_intermediate_size = config.moe_intermediate_size + self.moe_attrs["num_dense_layers"] = config.num_dense_layers + self.moe_attrs["use_expert_bias"] = config.use_expert_bias + self.moe_attrs["routed_scaling_factor"] = config.routed_scaling_factor + + # Router score for the experts that were not selected. Any finite value far below every + # plausible log-sigmoid works: it must fit in fp16 and exp(sentinel - max) must underflow to 0. + self.moe_attrs["router_sentinel"] = -10000.0 + self.moe_attrs["activation_type"] = "swiglu" + self.moe_attrs["swiglu_fusion"] = 1 + self.moe_attrs["normalize_routing_weights"] = True + + def make_intermediate_size_init(self, config): + # Lfm2MoeConfig has no block_auto_adjust_ff_dim: the dense layers use intermediate_size as is. + self.intermediate_size = config.intermediate_size + + def make_feed_forward(self, layer_id, layer, root_input): + if layer_id < self.moe_attrs["num_dense_layers"]: + super().make_feed_forward(layer_id, layer, root_input) + else: + self.make_moe(layer_id, self.get_feed_forward_module(layer), root_input) + + def make_moe(self, layer_id, moe, root_input): + self.make_moe_preprocessing(layer_id, moe, root_input) + router_probs, output_scale = self.make_moe_router(layer_id, moe, root_input) + self.make_moe_subgraph(layer_id, moe, root_input, router_probs, output_scale) + + def make_moe_preprocessing(self, layer_id, moe, root_input): + # Keep the router in floating point so int4 rounding cannot flip an expert choice. + moe.gate.exclude_from_quantization = True + self.make_interleaved_swiglu_moe_preprocessing(layer_id, moe) + + def make_moe_router(self, layer_id, moe, root_input): + """Emit the in-graph expert selection. + + Returns ``(router_probs, output_scale)``: the masked router scores fed to the MoE/QMoE op and + the per-token factor that the op output is multiplied by. + + root_input --> MatMul --> Reshape --> Cast(fp32) --> sigmoid --+--> Add(expert_bias) --> TopK + | | + +--> GatherElements <------+ + | | | | + | | ReduceSum --> ... --> output_scale + | | + | Clip --> Log + | | + +--> Shape --> ConstantOfShape(sentinel) + | + Cast(io_dtype) <-- ScatterElements <---+ + """ + scores_name = self.make_moe_router_scores(layer_id, moe, root_input) + indices_name = self.make_moe_router_selection(layer_id, moe, scores_name) + selected_name = self.make_moe_router_selected_scores(layer_id, scores_name, indices_name) + output_scale = self.make_moe_router_output_scale(layer_id, root_input, selected_name) + router_probs = self.make_moe_router_mask(layer_id, scores_name, indices_name, selected_name) + return router_probs, output_scale + + def make_moe_router_scores(self, layer_id, moe, root_input): + """Emit the per-expert router scores (`sigmoid(logits)`) in fp32 and return their name.""" + basename = f"/model/layers.{layer_id}/moe/router" + logits_shape = self.make_moe_router_shape() + + matmul_name = self.make_matmul(moe.gate, f"{basename}/MatMul", root_input) + reshape_name = f"{basename}/Reshape" + self.make_reshape( + reshape_name, + [f"{matmul_name}/output_0", f"/model/constants/INT64/{[-1, self.moe_attrs['num_experts']]}"], + dtype=self.io_dtype, + shape=logits_shape, + ) + logits_name = f"{reshape_name}/output_0" + if self.io_dtype != ir.DataType.FLOAT: + # The scores drive a discrete choice and are logged below, so keep them in fp32. + cast_name = f"{basename}/Cast" + self.make_cast(cast_name, logits_name, ir.DataType.FLOAT, shape=logits_shape) + logits_name = f"{cast_name}/output_0" + + # sigmoid(x) = 1 / (1 + exp(-x)), spelled out. ONNX Runtime's CPU Sigmoid kernel carries an absolute + # error of about 6e-8 (it is evaluated as 1 - sigmoid(-x)), so it returns 5.96e-8 for x = -16 + # instead of 1.125e-7 and 0 below x = -17. The scores are logged and summed below, where that + # is a large relative error; this form matches torch.sigmoid to fp32 rounding. + sigmoid_basename = f"{basename}/sigmoid" + self.make_neg(f"{sigmoid_basename}/Neg", logits_name, ir.DataType.FLOAT, shape=logits_shape) + self.make_exp( + f"{sigmoid_basename}/Exp", f"{sigmoid_basename}/Neg/output_0", ir.DataType.FLOAT, shape=logits_shape + ) + self.make_add( + f"{sigmoid_basename}/Add", + [f"{sigmoid_basename}/Exp/output_0", f"/model/constants/FLOAT/{1.0}"], + dtype=ir.DataType.FLOAT, + shape=logits_shape, + ) + reciprocal_name = f"{sigmoid_basename}/Reciprocal" + self.make_reciprocal(reciprocal_name, f"{sigmoid_basename}/Add/output_0", ir.DataType.FLOAT, shape=logits_shape) + return f"{reciprocal_name}/output_0" + + def make_moe_router_selection(self, layer_id, moe, scores_name): + """Emit the top-k over the bias-corrected scores and return the selected expert indices.""" + basename = f"/model/layers.{layer_id}/moe/router" + top_k = self.moe_attrs["top_k"] + logits_shape = self.make_moe_router_shape() + + selection_name = scores_name + if self.moe_attrs["use_expert_bias"]: + expert_bias_name = f"model.layers.{layer_id}.moe.expert_bias" + self.make_initializer(moe.expert_bias, expert_bias_name, to=ir.DataType.FLOAT) + add_name = f"{basename}/Add" + self.make_add(add_name, [scores_name, expert_bias_name], dtype=ir.DataType.FLOAT, shape=logits_shape) + selection_name = f"{add_name}/output_0" + + topk_name = f"{basename}/TopK" + self.make_topk( + topk_name, + [selection_name, f"/model/constants/INT64/[{top_k}]"], + dtype=ir.DataType.FLOAT, + shape=self.make_moe_router_shape(last_dim=top_k), + ) + return f"{topk_name}/output_1" + + def make_moe_router_selected_scores(self, layer_id, scores_name, indices_name): + """Emit the unbiased scores of the selected experts (`[rows, top_k]`) and return their name.""" + gather_name = f"/model/layers.{layer_id}/moe/router/GatherElements" + self.make_gather_elements( + gather_name, + [scores_name, indices_name], + dtype=ir.DataType.FLOAT, + shape=self.make_moe_router_shape(last_dim=self.moe_attrs["top_k"]), + axis=1, + ) + return f"{gather_name}/output_0" + + def make_moe_router_output_scale(self, layer_id, root_input, selected_name): + """Emit the per-token factor applied to the MoE op output and return its name. + + HF mixes the selected experts with `sigmoid_i / (sum_selected + 1e-6)` while the op normalizes + to `sigmoid_i / sum_selected`, so the op output is scaled by `sum_selected / (sum_selected + 1e-6)`. + The sum is taken before the clamp in `make_moe_router_mask`, so scores that flushed to zero give + a zero factor, matching HF exactly. `routed_scaling_factor` is folded in here as well. + """ + basename = f"/model/layers.{layer_id}/moe/router/scale" + rows_shape = self.make_moe_router_shape(last_dim=1) + + sum_name = f"{basename}/ReduceSum" + self.make_reduce_sum( + sum_name, + [selected_name, "/model/constants/INT64/[-1]"], + dtype=ir.DataType.FLOAT, + shape=rows_shape, + keepdims=True, + ) + add_name = f"{basename}/Add" + self.make_add( + add_name, + [f"{sum_name}/output_0", f"/model/constants/FLOAT/{1e-6}"], + dtype=ir.DataType.FLOAT, + shape=rows_shape, + ) + div_name = f"{basename}/Div" + self.make_div( + div_name, [f"{sum_name}/output_0", f"{add_name}/output_0"], dtype=ir.DataType.FLOAT, shape=rows_shape + ) + scale_name = f"{div_name}/output_0" + routed_scaling_factor = self.moe_attrs["routed_scaling_factor"] + if routed_scaling_factor != 1.0: + mul_name = f"{basename}/Mul" + self.make_mul( + mul_name, + [scale_name, f"/model/constants/FLOAT/{routed_scaling_factor}"], + dtype=ir.DataType.FLOAT, + shape=rows_shape, + ) + scale_name = f"{mul_name}/output_0" + + # Reshape [rows, 1] to the MoE output's leading dims + [1] so it broadcasts over hidden_size. + output_shape = self.make_hidden_state_shape(last_dim=1) + shape_name = f"{basename}/Shape" + self.make_shape(shape_name, root_input, shape=[len(output_shape)]) + slice_name = f"{basename}/Slice" + self.make_slice( + slice_name, + [f"{shape_name}/output_0", "/model/constants/INT64/[0]", "/model/constants/INT64/[-1]"], + dtype=ir.DataType.INT64, + shape=[len(output_shape) - 1], + ) + concat_name = f"{basename}/Concat" + self.make_concat( + concat_name, + [f"{slice_name}/output_0", "/model/constants/INT64/[1]"], + dtype=ir.DataType.INT64, + shape=[len(output_shape)], + axis=0, + ) + reshape_name = f"{basename}/Reshape" + self.make_reshape( + reshape_name, [scale_name, f"{concat_name}/output_0"], dtype=ir.DataType.FLOAT, shape=output_shape + ) + scale_name = f"{reshape_name}/output_0" + + if self.io_dtype != ir.DataType.FLOAT: + cast_name = f"{basename}/Cast" + self.make_cast(cast_name, scale_name, self.io_dtype, shape=output_shape) + scale_name = f"{cast_name}/output_0" + return scale_name + + def make_moe_router_mask(self, layer_id, scores_name, indices_name, selected_name): + """Emit `log(selected scores)` at the selected experts and a sentinel elsewhere, in `io_dtype`. + + The MoE/QMoE op takes the softmax of the top-k of this tensor, which equals + `sigmoid_i / sum_selected(sigmoid_j)`; `make_moe_router_output_scale` supplies the rest of HF's + `sigmoid_i / (sum_selected + 1e-6)` mixing weight. + """ + basename = f"/model/layers.{layer_id}/moe/router" + logits_shape = self.make_moe_router_shape() + selected_shape = self.make_moe_router_shape(last_dim=self.moe_attrs["top_k"]) + + # A sigmoid that flushed to 0 would give log = -inf, below the sentinel, and let the op's own + # top-k pick a different expert than the graph selected. Clamp so log stays far above it. + clip_name = f"{basename}/Clip" + self.make_clip( + clip_name, + [selected_name, f"/model/constants/FLOAT/{1e-30}", ""], + dtype=ir.DataType.FLOAT, + shape=selected_shape, + ) + log_name = f"{basename}/Log" + self.make_log(log_name, f"{clip_name}/output_0", ir.DataType.FLOAT, shape=selected_shape) + + shape_name = f"{basename}/Shape" + self.make_shape(shape_name, scores_name, shape=[2]) + sentinel_name = f"{basename}/ConstantOfShape" + self.make_constant_of_shape( + sentinel_name, + f"{shape_name}/output_0", + value=ir.tensor([self.moe_attrs["router_sentinel"]], dtype=ir.DataType.FLOAT), + dtype=ir.DataType.FLOAT, + shape=logits_shape, + ) + scatter_name = f"{basename}/ScatterElements" + self.make_scatter_elements( + scatter_name, + [f"{sentinel_name}/output_0", indices_name, f"{log_name}/output_0"], + dtype=ir.DataType.FLOAT, + shape=logits_shape, + axis=1, + ) + router_probs_name = f"{scatter_name}/output_0" + + if self.io_dtype != ir.DataType.FLOAT: + cast_name = f"{basename}/Cast_1" + self.make_cast(cast_name, router_probs_name, self.io_dtype, shape=logits_shape) + router_probs_name = f"{cast_name}/output_0" + return router_probs_name + + def make_moe_subgraph(self, layer_id, moe, root_input, router_probs=None, output_scale=None): + # `make_moe` always passes the pair returned by `make_moe_router`; the defaults only keep the + # base-class signature. + basename = f"/model/layers.{layer_id}/moe" + op_type = self.moe_attrs["op_type"] + names = self.make_moe_expert_names(layer_id) + gate_up_proj_zero_points, down_proj_zero_points = self.moe_attrs.get("zero_point_names", {}).get( + layer_id, ("", "") + ) + use_zero_points = bool(gate_up_proj_zero_points and down_proj_zero_points) and self.ep != "trt-rtx" + gate_up_proj_global_scales, down_proj_global_scales = self.moe_attrs.get("global_scale_names", {}).get( + layer_id, ("", "") + ) + + moe_name = f"{basename}/{op_type}" + self.make_moe_op( + moe_name, + root_input=root_input, + router_probs=router_probs, + weight1=names["gate_up_weight"], + scales1=names["gate_up_scales"] if op_type == "QMoE" else "", + bias1=names["gate_up_bias"], + weight2=names["down_weight"], + scales2=names["down_scales"] if op_type == "QMoE" else "", + bias2=names["down_bias"], + zero_points1=gate_up_proj_zero_points if use_zero_points else "", + zero_points2=down_proj_zero_points if use_zero_points else "", + global_scales1=gate_up_proj_global_scales, + global_scales2=down_proj_global_scales, + ) + + # Apply the routing-mass correction (and routed_scaling_factor) from the router. + mul_name = f"{basename}/Mul" + self.make_mul( + mul_name, [f"{moe_name}/output_0", output_scale], dtype=self.io_dtype, shape=self.make_hidden_state_shape() + ) + + # Assign the MoE output as the residual input of the next SkipLayerNorm + self.layernorm_attrs["skip_input"] = f"{mul_name}/output_0" diff --git a/src/python/py/models/builders/qwen.py b/src/python/py/models/builders/qwen.py index 1cce97ef7f..ff7519e19c 100644 --- a/src/python/py/models/builders/qwen.py +++ b/src/python/py/models/builders/qwen.py @@ -788,10 +788,6 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): self.moe_attrs["activation_type"] = "swiglu" self.moe_attrs["swiglu_fusion"] = 1 self.moe_attrs["normalize_routing_weights"] = True - if self.moe_attrs.get("swiglu_limit") is None and self.ep == "trt-rtx": - # TRT-RTX EP builds currently require QMoE swiglu_limit to be present; - # use +inf to preserve the "no clamp" behavior when the model omits it. - self.moe_attrs["swiglu_limit"] = float("inf") self.moe_intermediate_size = getattr(config, "moe_intermediate_size", 512) self.shared_expert_intermediate_size = getattr( @@ -802,24 +798,7 @@ def get_moe_module(self, layer_id, layer): return layer.mlp def make_moe_preprocessing(self, layer_id, moe, root_input): - gate_up_proj_bias = f"model.layers.{layer_id}.moe.experts.gate_up_proj.bias" - down_proj_bias = f"model.layers.{layer_id}.moe.experts.down_proj.bias" - - gate_up_weight = None - down_weight = None - if getattr(moe.experts, "gate_up_proj", None) is not None: - # Repack HF concatenated [gate|up] to ORT interleaved [g0,u0,g1,u1,...]. - raw_gate_up = moe.experts.gate_up_proj - half = raw_gate_up.shape[1] // 2 - gate_up_weight = torch.stack([raw_gate_up[:, :half, :], raw_gate_up[:, half:, :]], dim=2).reshape_as( - raw_gate_up - ) - down_weight = moe.experts.down_proj - self.make_moe_expert_initializers(layer_id, moe.experts, gate_up_weight, down_weight) - - num_e = self.moe_attrs["num_experts"] - self.make_initializer(torch.zeros(num_e, 2 * self.moe_intermediate_size), gate_up_proj_bias, to=self.io_dtype) - self.make_initializer(torch.zeros(num_e, self.hidden_size), down_proj_bias, to=self.io_dtype) + self.make_interleaved_swiglu_moe_preprocessing(layer_id, moe) def make_moe_router(self, layer_id, moe, root_input): basename = f"/model/layers.{layer_id}/moe" @@ -833,19 +812,13 @@ def make_moe_router(self, layer_id, moe, root_input): f"/model/constants/INT64/{[-1, self.moe_attrs['num_experts']]}", ], dtype=self.io_dtype, - shape=["batch_size * sequence_length", self.moe_attrs["num_experts"]], + shape=self.make_moe_router_shape(), ) def make_moe_subgraph(self, layer_id, moe, root_input): basename = f"/model/layers.{layer_id}/moe" op_type = self.moe_attrs["op_type"] - moe_weight_type = f"{'q' if op_type == 'QMoE' else ''}weight" - gate_up_proj_weight = f"model.layers.{layer_id}.moe.experts.gate_up_proj.{moe_weight_type}" - gate_up_proj_scales = f"model.layers.{layer_id}.moe.experts.gate_up_proj.scales" - gate_up_proj_bias = f"model.layers.{layer_id}.moe.experts.gate_up_proj.bias" - down_proj_weight = f"model.layers.{layer_id}.moe.experts.down_proj.{moe_weight_type}" - down_proj_scales = f"model.layers.{layer_id}.moe.experts.down_proj.scales" - down_proj_bias = f"model.layers.{layer_id}.moe.experts.down_proj.bias" + names = self.make_moe_expert_names(layer_id) gate_up_proj_global_scales, down_proj_global_scales = self.moe_attrs.get("global_scale_names", {}).get( layer_id, ("", "") ) @@ -855,12 +828,12 @@ def make_moe_subgraph(self, layer_id, moe, root_input): moe_name, root_input=root_input, router_probs=f"{basename}/router/Reshape/output_0", - weight1=gate_up_proj_weight, - scales1=gate_up_proj_scales if op_type == "QMoE" else "", - bias1=gate_up_proj_bias, - weight2=down_proj_weight, - scales2=down_proj_scales if op_type == "QMoE" else "", - bias2=down_proj_bias, + weight1=names["gate_up_weight"], + scales1=names["gate_up_scales"] if op_type == "QMoE" else "", + bias1=names["gate_up_bias"], + weight2=names["down_weight"], + scales2=names["down_scales"] if op_type == "QMoE" else "", + bias2=names["down_bias"], global_scales1=gate_up_proj_global_scales, global_scales2=down_proj_global_scales, ) diff --git a/src/python/py/models/loaders/quark.py b/src/python/py/models/loaders/quark.py index 5a73a448e6..cccde8d31e 100644 --- a/src/python/py/models/loaders/quark.py +++ b/src/python/py/models/loaders/quark.py @@ -247,7 +247,7 @@ def combine_and_repack_gate_up(self, experts): def repack_qweight(self, weights, bits) -> torch.Tensor: """ Repacks unpacked uint8 weights (representing 4-bit values) into a packed uint8 tensor. - This mirrors the packing logic from builder.py's _symmetric_blockwise_quantize. + This mirrors the packing logic of CudaQuantizer.symmetric_blockwise_quantize. """ if bits != 4: raise NotImplementedError("This repacking function is specifically for 4-bit weights.") @@ -341,7 +341,7 @@ def unpack_qweight_quark(self, module, reorder=True, dtype="uint4"): """ Unpack `qweight` to standard format and reorder for OGA. This is based on the packing logic from Quark's Pack_4_bits with reorder=True - and the unpacking logic from ORT GenAI's _symmetric_blockwise_quantize. + and the unpacking logic of ORT GenAI's CudaQuantizer.symmetric_blockwise_quantize. """ to_unpack = module.qweight diff --git a/test/cpp/lfm2_cache_layer_types_test.cpp b/test/cpp/lfm2_cache_layer_types_test.cpp index 6f7a91c79a..fb146c6cb4 100644 --- a/test/cpp/lfm2_cache_layer_types_test.cpp +++ b/test/cpp/lfm2_cache_layer_types_test.cpp @@ -98,6 +98,39 @@ void ExpectLayerTypesValidationMessage(const std::string& message, size_t actual EXPECT_NE(message.find("num_hidden_layers (" + std::to_string(expected_layers) + ")"), std::string::npos) << message; } +// LFM2 variants that the exporter labels with their own model type still have to run through the +// LFM2 runtime: the same conv state cache, and no rewind. +void ExpectLfm2Runtime(const std::string& model_type, const std::string& temp_suffix) { + const auto src_dir = GetLfm2ModelPath(); + const auto model_dir = MakeTempDir(temp_suffix); + fs_std::copy(src_dir, model_dir, fs_std::copy_options::recursive | fs_std::copy_options::overwrite_existing); + std::string config = ReadFile(model_dir / "genai_config.json"); + ReplaceFirst(config, "\"type\": \"lfm2\"", "\"type\": \"" + model_type + "\""); + WriteFile(model_dir / "genai_config.json", config); + + auto model = OgaModel::Create(model_dir.string().c_str()); + auto params = OgaGeneratorParams::Create(*model); + params->SetSearchOption("max_length", 8); + auto generator = OgaGenerator::Create(*model, *params); + + const std::vector input_ids{0, 0, 195, 731}; + generator->AppendTokens(input_ids.data(), input_ids.size()); + generator->GenerateNextToken(); + const size_t sequence_length = generator->GetSequenceCount(0); + EXPECT_GT(sequence_length, input_ids.size()); + + // Conv state cannot be rewound, exactly like the dense LFM2 models. The generator's up-front guard + // has to reject it; a throw from deeper in the state would come after the sequence was rewound. + std::string message; + try { + generator->RewindTo(0); + } catch (const std::runtime_error& e) { + message = e.what(); + } + EXPECT_NE(message.find("RewindTo is currently not supported for " + model_type), std::string::npos) << message; + EXPECT_EQ(generator->GetSequenceCount(0), sequence_length); +} + } // namespace TEST(LFM2CacheLayerTypesValidationTest, UndersizedLayerTypesArray) { @@ -138,37 +171,15 @@ TEST(LFM2CacheLayerTypesValidationTest, ValidMatchingLayers) { }); } +// The MoE variant (LFM2-8B-A1B, LFM2.5-8B-A1B, LFM2-24B-A2B) is exported as "lfm2_moe". +TEST(LFM2CacheLayerTypesValidationTest, Lfm2MoeTypeUsesLfm2Runtime) { + SkipIfModelUnavailable(); + ExpectLfm2Runtime("lfm2_moe", "moe_type"); +} + // An LFM2-VL checkpoint exported without exclude_embeds is labeled "lfm2_vl_text": a plain LFM2 -// decoder that must run through the LFM2 runtime (conv state cache, no rewind). +// decoder rather than a stage of the vision pipeline. TEST(LFM2CacheLayerTypesValidationTest, Lfm2VlTextTypeUsesLfm2Runtime) { SkipIfModelUnavailable(); - - const auto src_dir = GetLfm2ModelPath(); - const auto model_dir = MakeTempDir("vl_text_type"); - fs_std::copy(src_dir, model_dir, fs_std::copy_options::recursive | fs_std::copy_options::overwrite_existing); - std::string config = ReadFile(model_dir / "genai_config.json"); - ReplaceFirst(config, "\"type\": \"lfm2\"", "\"type\": \"lfm2_vl_text\""); - WriteFile(model_dir / "genai_config.json", config); - - auto model = OgaModel::Create(model_dir.string().c_str()); - auto params = OgaGeneratorParams::Create(*model); - params->SetSearchOption("max_length", 8); - auto generator = OgaGenerator::Create(*model, *params); - - const std::vector input_ids{0, 0, 195, 731}; - generator->AppendTokens(input_ids.data(), input_ids.size()); - generator->GenerateNextToken(); - const size_t sequence_length = generator->GetSequenceCount(0); - EXPECT_GT(sequence_length, input_ids.size()); - - // Conv state cannot be rewound, exactly like the dense LFM2 models. The generator's up-front guard - // has to reject it; a throw from deeper in the state would come after the sequence was rewound. - std::string message; - try { - generator->RewindTo(0); - } catch (const std::runtime_error& e) { - message = e.what(); - } - EXPECT_NE(message.find("RewindTo is currently not supported for lfm2_vl_text"), std::string::npos) << message; - EXPECT_EQ(generator->GetSequenceCount(0), sequence_length); + ExpectLfm2Runtime("lfm2_vl_text", "vl_text_type"); } diff --git a/test/python/builder/_builder_test_utils.py b/test/python/builder/_builder_test_utils.py new file mode 100644 index 0000000000..81472f4229 --- /dev/null +++ b/test/python/builder/_builder_test_utils.py @@ -0,0 +1,103 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License + +"""Shared helpers for the model-builder unit tests. + +The builder modules import each other as ``models.builders.`` and pull in ``onnx_ir``, +``onnxruntime.quantization`` and ``transformers``. ``load_builder_module`` loads one of them from +source without installing the ``onnxruntime_genai`` package, stubbing whichever of those +dependencies is missing so the pure-Python parts of the builder can still be exercised. +""" + +from __future__ import annotations + +import importlib.util +import sys +import types +from pathlib import Path + +BUILDERS_DIR = Path(__file__).parents[3] / "src" / "python" / "py" / "models" / "builders" + + +def _module_available(module_name): + try: + return importlib.util.find_spec(module_name) is not None + except (ModuleNotFoundError, ValueError): + return False + + +def stub_missing_builder_dependencies(): + if not _module_available("onnx_ir"): + onnx_ir = types.ModuleType("onnx_ir") + onnx_ir.DataType = types.SimpleNamespace(INT4=object(), FLOAT=object(), FLOAT16=object(), BFLOAT16=object()) + tensor_adapters = types.ModuleType("onnx_ir.tensor_adapters") + tensor_adapters.TorchTensor = object + tensor_adapters.to_torch_dtype = lambda dtype: dtype + sys.modules["onnx_ir"] = onnx_ir + sys.modules["onnx_ir.tensor_adapters"] = tensor_adapters + + if not _module_available("onnxruntime.quantization.matmul_nbits_quantizer"): + # Prefer the real onnxruntime package when it is installed; only fabricate a + # top-level stub when the package truly isn't available. This avoids shadowing a + # real onnxruntime wheel (which would break other tests in the session) and only + # supplies the specific submodule the builder needs. + if _module_available("onnxruntime"): + import onnxruntime # noqa: PLC0415 + else: + onnxruntime = sys.modules.setdefault("onnxruntime", types.ModuleType("onnxruntime")) + quantization = getattr(onnxruntime, "quantization", None) + if quantization is None: + quantization = types.ModuleType("onnxruntime.quantization") + matmul_nbits_quantizer = types.ModuleType("onnxruntime.quantization.matmul_nbits_quantizer") + for class_name in ( + "KQuantWeightOnlyQuantConfig", + "MatMulNBitsQuantizer", + "QuantFormat", + "RTNWeightOnlyQuantConfig", + ): + setattr(matmul_nbits_quantizer, class_name, type(class_name, (), {})) + onnxruntime.quantization = quantization + quantization.matmul_nbits_quantizer = matmul_nbits_quantizer + sys.modules["onnxruntime.quantization"] = quantization + sys.modules["onnxruntime.quantization.matmul_nbits_quantizer"] = matmul_nbits_quantizer + + if not _module_available("tqdm"): + tqdm_module = types.ModuleType("tqdm") + tqdm_module.tqdm = lambda iterable=None, *args, **kwargs: iterable + sys.modules["tqdm"] = tqdm_module + + if not _module_available("transformers"): + transformers = types.ModuleType("transformers") + for class_name in ( + "AutoConfig", + "AutoModelForCausalLM", + "AutoModelForSpeechSeq2Seq", + "AutoTokenizer", + "GenerationConfig", + ): + setattr(transformers, class_name, type(class_name, (), {})) + sys.modules["transformers"] = transformers + + +stub_missing_builder_dependencies() + + +def _register_builders_package(): + if str(BUILDERS_DIR.parent) not in sys.path: + sys.path.insert(0, str(BUILDERS_DIR.parent)) + sys.modules.setdefault("models", types.ModuleType("models")) + builders_package = sys.modules.setdefault("models.builders", types.ModuleType("models.builders")) + builders_package.__path__ = [str(BUILDERS_DIR)] + + +def load_builder_module(module_name): + """Load ``src/python/py/models/builders/.py`` as ``models.builders.``.""" + _register_builders_package() + qualified_name = f"models.builders.{module_name}" + if qualified_name in sys.modules: + return sys.modules[qualified_name] + spec = importlib.util.spec_from_file_location(qualified_name, BUILDERS_DIR / f"{module_name}.py") + module = importlib.util.module_from_spec(spec) + sys.modules[qualified_name] = module + spec.loader.exec_module(module) + return module diff --git a/test/python/builder/test_lfm2_moe.py b/test/python/builder/test_lfm2_moe.py new file mode 100644 index 0000000000..a740256e16 --- /dev/null +++ b/test/python/builder/test_lfm2_moe.py @@ -0,0 +1,570 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License + +"""Unit tests for the LFM2-MoE model builder. + +The LFM2-MoE router selects experts by top-k of ``sigmoid(logits) + expert_bias`` but mixes them +with ``sigmoid_i / (sum_selected + 1e-6)``. The fused MoE/QMoE op only knows softmax-over-top-k +routing, so the builder performs the selection in the graph, feeds the op ``log(sigmoid)`` at the +selected experts and a sentinel elsewhere, and scales the op output by +``sum_selected / (sum_selected + 1e-6)``. These tests pin that graph and execute it. +""" + +from __future__ import annotations + +import types + +import numpy as np +import onnxruntime as ort +import pytest +import torch +from _builder_test_utils import load_builder_module + +base_module = load_builder_module("base") +lfm2_module = load_builder_module("lfm2") +ir = base_module.ir +Model = base_module.Model +LFM2Model = lfm2_module.LFM2Model +LFM2MoEModel = lfm2_module.LFM2MoEModel + + +def _stub_base_init(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): + """Stand-in for Model.__init__ that runs only the real MoE attribute hooks (unquantized experts).""" + self.ep = ep + self.activation = "silu" + self.quant_config = types.SimpleNamespace(moe=types.SimpleNamespace(type="none", weights_prepacked=-1)) + self.make_moe_attrs_init(config) + self.make_moe_init() + + +def _moe_attrs(num_experts, top_k, use_expert_bias, routed_scaling_factor): + """`moe_attrs` as the real constructors build them: base defaults plus the LFM2-MoE overrides.""" + config = types.SimpleNamespace( + num_experts=num_experts, + num_experts_per_tok=top_k, + norm_topk_prob=True, + num_dense_layers=0, + moe_intermediate_size=32, + use_expert_bias=use_expert_bias, + routed_scaling_factor=routed_scaling_factor, + ) + with pytest.MonkeyPatch.context() as monkeypatch: + monkeypatch.setattr(LFM2Model, "__init__", _stub_base_init) + return LFM2MoEModel(config, ir.DataType.FLOAT, ir.DataType.FLOAT, "cpu", None, {}).moe_attrs + + +def _recording_model(io_dtype, *, num_experts=8, top_k=2, use_expert_bias=True, routed_scaling_factor=1.0): + """LFM2MoEModel whose graph emitters record instead of building an onnx_ir graph.""" + model = LFM2MoEModel.__new__(LFM2MoEModel) + model.io_dtype = io_dtype + model.ep = "cpu" + model.hidden_size = 64 + model.moe_intermediate_size = 32 + model.use_paged_attention = False + model.moe_attrs = _moe_attrs(num_experts, top_k, use_expert_bias, routed_scaling_factor) + model.layernorm_attrs = {} + model.quant_attrs = {"nodes_to_exclude": []} + model.nodes = [] + model.initializers = {} + model.make_node = lambda op_type, inputs, outputs, name, domain="", **attrs: model.nodes.append( + types.SimpleNamespace(op_type=op_type, inputs=list(inputs), outputs=list(outputs), name=name, attrs=attrs) + ) + model.make_value = lambda name, dtype=None, shape=None: None + model.make_initializer = lambda tensor, name, to=None, raw=False: model.initializers.__setitem__(name, (tensor, to)) + model.make_matmul = lambda matmul, basename, root_input, **kwargs: ( + model.make_node("MatMul", [root_input, f"{basename}.weight"], [f"{basename}/output_0"], name=basename), + basename, + )[1] + return model + + +def _moe_module(num_experts=8, hidden=64, inter=32): + gate = types.SimpleNamespace(weight=torch.randn(num_experts, hidden)) + experts = types.SimpleNamespace( + gate_up_proj=torch.randn(num_experts, 2 * inter, hidden), + down_proj=torch.randn(num_experts, hidden, inter), + ) + return types.SimpleNamespace(gate=gate, experts=experts, expert_bias=torch.randn(num_experts)) + + +def test_lfm2_moe_routes_dense_and_moe_layers(monkeypatch): + model = LFM2MoEModel.__new__(LFM2MoEModel) + model.moe_attrs = {"num_dense_layers": 2} + calls = [] + monkeypatch.setattr(model, "make_mlp", lambda layer_id, mlp, root_input: calls.append(("mlp", layer_id, mlp))) + monkeypatch.setattr(model, "make_moe", lambda layer_id, moe, root_input: calls.append(("moe", layer_id, moe))) + + dense_ffn = types.SimpleNamespace(w1="w1", w2="w2", w3="w3") + moe_ffn = object() + model.make_feed_forward(1, types.SimpleNamespace(feed_forward=dense_ffn), "x") + model.make_feed_forward(2, types.SimpleNamespace(feed_forward=moe_ffn), "x") + + assert calls == [("mlp", 1, dense_ffn), ("moe", 2, moe_ffn)] + # The dense path aliases the HF w1/w3/w2 names onto the base-class gate/up/down names. + assert (dense_ffn.gate_proj, dense_ffn.up_proj, dense_ffn.down_proj) == ("w1", "w3", "w2") + + +def test_lfm2_feed_forward_accepts_both_loader_layouts(monkeypatch): + """Hugging Face layers expose `feed_forward`; the quantized-checkpoint IR exposes `mlp`.""" + model = LFM2Model.__new__(LFM2Model) + seen = [] + monkeypatch.setattr(model, "make_mlp", lambda layer_id, mlp, root_input: seen.append(mlp)) + + hf_ffn = types.SimpleNamespace(w1="w1", w2="w2", w3="w3") + model.make_feed_forward(0, types.SimpleNamespace(feed_forward=hf_ffn), "x") + assert (hf_ffn.gate_proj, hf_ffn.up_proj, hf_ffn.down_proj) == ("w1", "w3", "w2") + + # The quantized IR already carries gate/up/down names and has no w1/w3/w2 to alias. + quantized_ffn = types.SimpleNamespace(gate_proj="g", up_proj="u", down_proj="d") + model.make_feed_forward(1, types.SimpleNamespace(mlp=quantized_ffn), "x") + + assert seen == [hf_ffn, quantized_ffn] + + +def test_lfm2_moe_feed_forward_accepts_quantized_layer_layout(monkeypatch): + model = LFM2MoEModel.__new__(LFM2MoEModel) + model.moe_attrs = {"num_dense_layers": 0} + moe = object() + seen = [] + monkeypatch.setattr(model, "make_moe", lambda layer_id, module, root_input: seen.append(module)) + model.make_feed_forward(3, types.SimpleNamespace(mlp=moe), "x") + assert seen == [moe] + + +def test_lfm2_dense_model_still_owns_the_mlp(monkeypatch): + model = LFM2Model.__new__(LFM2Model) + calls = [] + monkeypatch.setattr(model, "make_mlp", lambda layer_id, mlp, root_input: calls.append((layer_id, mlp))) + ffn = types.SimpleNamespace(w1="w1", w2="w2", w3="w3") + model.make_feed_forward(0, types.SimpleNamespace(feed_forward=ffn), "x") + assert calls == [(0, ffn)] + + +@pytest.mark.parametrize("io_dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16]) +@pytest.mark.parametrize("routed_scaling_factor", [1.0, 2.5]) +def test_lfm2_moe_router_selects_with_bias_and_mixes_without(io_dtype, routed_scaling_factor): + model = _recording_model(io_dtype, routed_scaling_factor=routed_scaling_factor) + moe = _moe_module() + router_probs, output_scale = model.make_moe_router(3, moe, "hidden") + + ops = [n.op_type for n in model.nodes] + casts = ["Cast"] if io_dtype != ir.DataType.FLOAT else [] + scaling = ["Mul"] if routed_scaling_factor != 1.0 else [] + assert ops == [ + "MatMul", + "Reshape", + *casts, + # sigmoid spelled out: ORT's Sigmoid kernel loses all relative precision below ~1e-6 + "Neg", + "Exp", + "Add", + "Reciprocal", + "Add", + "TopK", + "GatherElements", + # output scale: sum_selected / (sum_selected + 1e-6) [* routed_scaling_factor], reshaped to 3D + "ReduceSum", + "Add", + "Div", + *scaling, + "Shape", + "Slice", + "Concat", + "Reshape", + *casts, + # masked router scores + "Clip", + "Log", + "Shape", + "ConstantOfShape", + "ScatterElements", + *casts, + ] + r = "/model/layers.3/moe/router" + by_name = {n.name[len(r) + 1 :]: n for n in model.nodes} + + # Scores: 1 / (1 + exp(-logits)) in fp32. + logits = f"{r}/Cast/output_0" if casts else f"{r}/Reshape/output_0" + assert by_name["sigmoid/Neg"].inputs == [logits] + assert by_name["sigmoid/Exp"].inputs == [f"{r}/sigmoid/Neg/output_0"] + assert by_name["sigmoid/Add"].inputs == [f"{r}/sigmoid/Exp/output_0", "/model/constants/FLOAT/1.0"] + assert by_name["sigmoid/Reciprocal"].inputs == [f"{r}/sigmoid/Add/output_0"] + scores = f"{r}/sigmoid/Reciprocal/output_0" + + # Selection: TopK over sigmoid + expert_bias (fp32). + assert by_name["Add"].inputs == [scores, "model.layers.3.moe.expert_bias"] + assert model.initializers["model.layers.3.moe.expert_bias"][1] == ir.DataType.FLOAT + assert by_name["TopK"].inputs == [f"{r}/Add/output_0", "/model/constants/INT64/[2]"] + assert by_name["TopK"].attrs == {"axis": -1, "largest": True} + indices = f"{r}/TopK/output_1" + + # The unbiased selected scores feed both the mass factor and the mask. + assert by_name["GatherElements"].inputs == [scores, indices] + assert by_name["GatherElements"].attrs == {"axis": 1} + selected = f"{r}/GatherElements/output_0" + + # Mass factor: the sum is taken before the clamp so flushed-to-zero scores give a zero factor. + assert by_name["scale/ReduceSum"].inputs == [selected, "/model/constants/INT64/[-1]"] + assert by_name["scale/ReduceSum"].attrs == {"keepdims": True} + assert by_name["scale/Add"].inputs == [f"{r}/scale/ReduceSum/output_0", "/model/constants/FLOAT/1e-06"] + assert by_name["scale/Div"].inputs == [f"{r}/scale/ReduceSum/output_0", f"{r}/scale/Add/output_0"] + scale = f"{r}/scale/Div/output_0" + if scaling: + assert by_name["scale/Mul"].inputs == [scale, f"/model/constants/FLOAT/{routed_scaling_factor}"] + scale = f"{r}/scale/Mul/output_0" + assert by_name["scale/Shape"].inputs == ["hidden"] + assert by_name["scale/Slice"].inputs == [ + f"{r}/scale/Shape/output_0", + "/model/constants/INT64/[0]", + "/model/constants/INT64/[-1]", + ] + assert by_name["scale/Concat"].inputs == [f"{r}/scale/Slice/output_0", "/model/constants/INT64/[1]"] + assert by_name["scale/Reshape"].inputs == [scale, f"{r}/scale/Concat/output_0"] + expected_scale = f"{r}/scale/Cast/output_0" if casts else f"{r}/scale/Reshape/output_0" + assert output_scale == expected_scale + + # Mask: clamp, log, scatter over a sentinel row. The clamp keeps a flushed-to-zero sigmoid from + # producing -inf below the sentinel, which would let the op's own top-k pick another expert. + assert by_name["Clip"].inputs == [selected, "/model/constants/FLOAT/1e-30", ""] + assert by_name["Log"].inputs == [f"{r}/Clip/output_0"] + sentinel = by_name["ConstantOfShape"].attrs["value"] + assert sentinel.dtype == ir.DataType.FLOAT + router_sentinel = model.moe_attrs["router_sentinel"] + assert sentinel.numpy().tolist() == [router_sentinel] + assert router_sentinel < -1000 and torch.finfo(torch.float16).min < router_sentinel + assert torch.log(torch.tensor(1e-30)).item() > router_sentinel + assert by_name["ScatterElements"].inputs == [f"{r}/ConstantOfShape/output_0", indices, f"{r}/Log/output_0"] + assert by_name["ScatterElements"].attrs == {"axis": 1} + + expected_probs = f"{r}/Cast_1/output_0" if casts else f"{r}/ScatterElements/output_0" + assert router_probs == expected_probs + if casts: + assert by_name["Cast"].attrs == {"to": ir.DataType.FLOAT} + assert by_name["scale/Cast"].attrs == {"to": io_dtype} + assert by_name["Cast_1"].attrs == {"to": io_dtype} + + +def test_lfm2_moe_threads_router_outputs_to_subgraph(monkeypatch): + model = LFM2MoEModel.__new__(LFM2MoEModel) + calls = [] + monkeypatch.setattr(model, "make_moe_preprocessing", lambda *args: calls.append(("pre", args))) + monkeypatch.setattr(model, "make_moe_router", lambda *args: (calls.append(("router", args)), ("probs", "scale"))[1]) + monkeypatch.setattr(model, "make_moe_subgraph", lambda *args: calls.append(("subgraph", args))) + moe = object() + model.make_moe(5, moe, "hidden") + assert calls == [ + ("pre", (5, moe, "hidden")), + ("router", (5, moe, "hidden")), + ("subgraph", (5, moe, "hidden", "probs", "scale")), + ] + + +def test_lfm2_moe_router_without_expert_bias_selects_on_sigmoid(): + model = _recording_model(ir.DataType.FLOAT, use_expert_bias=False) + model.make_moe_router(3, _moe_module(), "hidden") + + ops = [n.op_type for n in model.nodes] + assert ops.count("Add") == 2 # the sigmoid's 1 + exp(-x) and the mass factor's epsilon, no bias add + topk = next(n for n in model.nodes if n.op_type == "TopK") + assert topk.inputs[0] == "/model/layers.3/moe/router/sigmoid/Reciprocal/output_0" + assert "model.layers.3.moe.expert_bias" not in model.initializers + + +def test_lfm2_moe_preprocessing_interleaves_gate_and_up(monkeypatch): + model = _recording_model(ir.DataType.FLOAT16, num_experts=2) + moe = _moe_module(num_experts=2, hidden=4, inter=3) + captured = {} + monkeypatch.setattr( + model, + "make_moe_expert_initializers", + lambda layer_id, experts, gate_up_weight=None, down_weight=None: captured.update( + layer_id=layer_id, experts=experts, gate_up=gate_up_weight, down=down_weight + ), + ) + model.make_moe_preprocessing(2, moe, "hidden") + + raw = moe.experts.gate_up_proj + gate, up = raw[:, :3, :], raw[:, 3:, :] + assert captured["experts"] is moe.experts + assert captured["gate_up"].shape == raw.shape + torch.testing.assert_close(captured["gate_up"][:, 0::2, :], gate) + torch.testing.assert_close(captured["gate_up"][:, 1::2, :], up) + assert captured["down"] is moe.experts.down_proj + + gate_up_bias, _ = model.initializers["model.layers.2.moe.experts.gate_up_proj.bias"] + down_bias, _ = model.initializers["model.layers.2.moe.experts.down_proj.bias"] + assert gate_up_bias.shape == (2, 2 * model.moe_intermediate_size) and not gate_up_bias.any() + assert down_bias.shape == (2, model.hidden_size) and not down_bias.any() + + # The router MatMul is excluded from int4 quantization so rounding cannot flip an expert choice. + assert moe.gate.exclude_from_quantization is True + + +@pytest.mark.parametrize("op_type", ["MoE", "QMoE"]) +def test_make_moe_expert_names_follows_op_type(op_type): + model = Model.__new__(Model) + model.moe_attrs = {"op_type": op_type} + weight = "qweight" if op_type == "QMoE" else "weight" + assert model.make_moe_expert_names(7) == { + "gate_up_weight": f"model.layers.7.moe.experts.gate_up_proj.{weight}", + "gate_up_scales": "model.layers.7.moe.experts.gate_up_proj.scales", + "gate_up_bias": "model.layers.7.moe.experts.gate_up_proj.bias", + "down_weight": f"model.layers.7.moe.experts.down_proj.{weight}", + "down_scales": "model.layers.7.moe.experts.down_proj.scales", + "down_bias": "model.layers.7.moe.experts.down_proj.bias", + } + + +@pytest.mark.parametrize("use_paged_attention,rows", [(False, "batch_size * sequence_length"), (True, "num_tokens")]) +def test_make_moe_router_shape_follows_the_token_layout(use_paged_attention, rows): + model = Model.__new__(Model) + model.moe_attrs = {"num_experts": 8} + model.use_paged_attention = use_paged_attention + assert model.make_moe_router_shape() == [rows, 8] + assert model.make_moe_router_shape(last_dim=2) == [rows, 2] + + +@pytest.mark.parametrize("op_type", ["MoE", "QMoE"]) +def test_lfm2_moe_subgraph_feeds_masked_router_probs_and_scales_the_output(monkeypatch, op_type): + model = _recording_model(ir.DataType.FLOAT16) + model.moe_attrs["op_type"] = op_type + captured = {} + monkeypatch.setattr(model, "make_moe_op", lambda name, **kwargs: captured.update(name=name, **kwargs)) + + model.make_moe_subgraph(4, _moe_module(), "hidden", "masked_probs", "output_scale") + + weight = "qweight" if op_type == "QMoE" else "weight" + scales = ".scales" if op_type == "QMoE" else "" + assert captured["name"] == f"/model/layers.4/moe/{op_type}" + assert captured["root_input"] == "hidden" + assert captured["router_probs"] == "masked_probs" + assert captured["weight1"] == f"model.layers.4.moe.experts.gate_up_proj.{weight}" + assert captured["weight2"] == f"model.layers.4.moe.experts.down_proj.{weight}" + assert captured["scales1"] == (f"model.layers.4.moe.experts.gate_up_proj{scales}" if scales else "") + assert captured["scales2"] == (f"model.layers.4.moe.experts.down_proj{scales}" if scales else "") + assert captured["bias1"] == "model.layers.4.moe.experts.gate_up_proj.bias" + assert captured["bias2"] == "model.layers.4.moe.experts.down_proj.bias" + + (mul,) = model.nodes + assert mul.op_type == "Mul" + assert mul.inputs == [f"/model/layers.4/moe/{op_type}/output_0", "output_scale"] + assert model.layernorm_attrs["skip_input"] == "/model/layers.4/moe/Mul/output_0" + + +def _executable_model(num_experts, top_k, hidden, inter, routed_scaling_factor=1.0): + """LFM2MoEModel with the real emitters, writing into an onnx_ir graph that ORT can execute.""" + model = LFM2MoEModel.__new__(LFM2MoEModel) + model.io_dtype = ir.DataType.FLOAT + model.onnx_dtype = ir.DataType.FLOAT + model.ep = "cpu" + model.hidden_size = hidden + model.moe_intermediate_size = inter + model.use_paged_attention = False + model.moe_attrs = _moe_attrs(num_experts, top_k, True, routed_scaling_factor) + model.layernorm_attrs = {} + model.quant_attrs = {"nodes_to_exclude": []} + model.values = {} + model.node_names = set() + graph = ir.Graph(inputs=(), outputs=(), nodes=(), opset_imports={"": 21, "com.microsoft": 1}, name="lfm2_moe") + model.model = ir.Model(graph, ir_version=10) + graph.inputs.append(model.make_value("hidden", ir.DataType.FLOAT, ["batch_size", "sequence_length", hidden])) + return model, graph + + +def _hf_reference(hidden_states, moe, top_k, routed_scaling_factor, eps=1e-6): + """Mirror of transformers' Lfm2MoeSparseMoeBlock (route_tokens_to_experts + experts). + + `eps=0` gives the plain sum-to-one normalization that the fused MoE op applies on its own. + """ + x = hidden_states.reshape(-1, hidden_states.shape[-1]) + logits = x @ moe.gate.weight.T + scores = torch.sigmoid(logits) + _, selected = torch.topk(scores + moe.expert_bias, k=top_k, dim=-1) + weights = torch.gather(scores, 1, selected) + weights = weights / (weights.sum(dim=-1, keepdim=True) + eps) + out = torch.zeros_like(x) + inter = moe.experts.down_proj.shape[-1] + for token in range(x.shape[0]): + for slot in range(top_k): + expert = selected[token, slot] + gate_up = moe.experts.gate_up_proj[expert] @ x[token] + act = torch.nn.functional.silu(gate_up[:inter]) * gate_up[inter:] + out[token] += weights[token, slot] * (moe.experts.down_proj[expert] @ act) + return (out * routed_scaling_factor).reshape(hidden_states.shape) + + +@pytest.mark.parametrize("routed_scaling_factor", [1.0, 2.5]) +def test_lfm2_moe_executed_graph_matches_hf_routing(tmp_path, routed_scaling_factor): + """Run the emitted MoE layer in ORT and compare with HF's routing math, including the cases where + the fused op's forced sum-to-one normalization differs from HF's `sum + 1e-6`.""" + num_experts, top_k, hidden, inter = 8, 4, 32, 64 + generator = torch.Generator().manual_seed(7) + model, graph = _executable_model(num_experts, top_k, hidden, inter, routed_scaling_factor) + + # Token t reads router logits from column t of the gate weight (one-hot hidden states below). The + # expert bias selects experts 0-3 whatever their sigmoid, so the routing mass can be made arbitrarily small. + at_eps = float(np.log(1e-6 / top_k)) # four selected sigmoids that sum to HF's epsilon + gate_weight = torch.zeros(num_experts, hidden) + gate_weight[:, 0] = torch.tensor([1.5, -0.5, 0.3, 2.0, -1.0, 0.8, -2.0, 0.1]) # ordinary scores + gate_weight[:, 1] = torch.tensor([-16.0, -16.0, -16.0, -16.0, -10.0, -11.0, -12.0, -13.0]) # tiny selected + gate_weight[:, 2] = torch.tensor([-110.0, -111.0, -112.0, -113.0, -10.0, -11.0, -12.0, -13.0]) # flushed + gate_weight[:, 3] = torch.tensor([-5.0, -110.0, -111.0, -112.0, -10.0, -11.0, -12.0, -13.0]) # partly flushed + gate_weight[:, 4] = torch.tensor([at_eps] * top_k + [-10.0, -11.0, -12.0, -13.0]) # sum == eps + expert_bias = torch.tensor([2.0, 2.0, 2.0, 2.0, 0.0, 0.0, 0.0, 0.0]) + moe = types.SimpleNamespace( + gate=types.SimpleNamespace(weight=gate_weight), + experts=types.SimpleNamespace( + gate_up_proj=torch.randn(num_experts, 2 * inter, hidden, generator=generator) / hidden**0.5, + down_proj=torch.randn(num_experts, hidden, inter, generator=generator) / inter**0.5, + ), + expert_bias=expert_bias, + ) + model.make_moe(0, moe, "hidden") + graph.outputs.append(model.values[model.layernorm_attrs["skip_input"]]) + model_path = tmp_path / "lfm2_moe.onnx" + ir.save(model.model, model_path) + session = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"]) + + num_tokens = 5 + hidden_states = torch.zeros(1, num_tokens, hidden) + for token in range(num_tokens): + hidden_states[0, token, token] = 1.0 + (actual,) = session.run(None, {"hidden": hidden_states.numpy()}) + expected = _hf_reference(hidden_states, moe, top_k, routed_scaling_factor).numpy() + np.testing.assert_allclose(actual, expected, rtol=1e-4, atol=1e-6) + + # The allclose above is what catches a regression: tokens 1, 2 and 4 trip it if the graph drops the output + # scale (the fused op alone normalizes like HF with eps=0), and token 3 trips it if the clamp floor leaks + # routing mass to flushed experts. The checks below keep the fixture in the regime where that happens. + without_eps = _hf_reference(hidden_states, moe, top_k, routed_scaling_factor, eps=0.0).numpy() + # Token 1: the selected sigmoids sum to ~4.5e-7, so HF keeps only ~31% of the routing mass. + assert np.isclose(np.linalg.norm(without_eps[0, 1]) / np.linalg.norm(expected[0, 1]), 3.2215, rtol=1e-3) + # Token 2: the selected sigmoids flush to 0 in fp32, so HF returns exactly zero (eps=0 would be 0/0). + assert np.all(expected[0, 2] == 0) + assert np.all(actual[0, 2] == 0) + # Token 3: one selected sigmoid survives and three flush to 0; the clamped three must get no weight. + # Token 4: the selected sigmoids sum to exactly HF's epsilon, so HF keeps half the routing mass. + assert np.isclose(np.linalg.norm(without_eps[0, 4]) / np.linalg.norm(expected[0, 4]), 2.0, rtol=1e-3) + + +@pytest.mark.parametrize("bits", [4, 8]) +@pytest.mark.parametrize("hidden,inter", [(40, 64), (64, 40)]) +def test_lfm2_qmoe_cpu_executes_partial_blocks(tmp_path, bits, hidden, inter): + model, graph = _executable_model(2, 1, hidden, inter) + model.moe_attrs.update(op_type="QMoE", expert_weight_bits=bits) + model.quant_attrs["qmoe_block_size"] = 32 + moe = _moe_module(2, hidden, inter) + # Constant rows with exactly representable scales isolate the storage layout from quantization error. + for weight in (moe.experts.gate_up_proj, moe.experts.down_proj): + rows = torch.arange(weight.shape[0] * weight.shape[1]).reshape(*weight.shape[:2], 1) + weight.copy_(((rows % 7) + 1) / 128.0) + model.make_moe(0, moe, "hidden") + graph.outputs.append(model.values[model.layernorm_attrs["skip_input"]]) + model_path = tmp_path / "lfm2_qmoe.onnx" + ir.save(model.model, model_path) + session = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"]) + hidden_states = torch.linspace(-0.1, 0.2, 2 * hidden).reshape(1, 2, hidden) + (actual,) = session.run(None, {"hidden": hidden_states.numpy()}) + expected = _hf_reference(hidden_states, moe, 1, 1.0).numpy() + np.testing.assert_allclose(actual, expected, rtol=1e-4, atol=1e-6) + + +@pytest.mark.parametrize("hidden,inter", [(33, 64), (64, 33)]) +def test_lfm2_qmoe_export_rejects_odd_int4_dimensions(hidden, inter): + model, _ = _executable_model(2, 1, hidden, inter) + model.moe_attrs.update(op_type="QMoE", expert_weight_bits=4) + model.quant_attrs["qmoe_block_size"] = 32 + with pytest.raises(RuntimeError, match=r"K \(33\) to be divisible by 2"): + model.make_moe(0, _moe_module(2, hidden, inter), "hidden") + + +def test_lfm2_moe_rejects_unnormalized_topk(monkeypatch): + monkeypatch.setattr(LFM2Model, "__init__", _stub_base_init) + config = types.SimpleNamespace(norm_topk_prob=False) + with pytest.raises(NotImplementedError, match="norm_topk_prob"): + LFM2MoEModel(config, ir.DataType.FLOAT, ir.DataType.FLOAT, "cpu", None, {}) + + +def test_lfm2_dense_intermediate_size_honors_block_auto_adjust(): + # Dense LFM2 keeps transformers' FFN width adjustment and does not default the flag away. + model = LFM2Model.__new__(LFM2Model) + config = types.SimpleNamespace( + intermediate_size=12288, + block_auto_adjust_ff_dim=True, + block_ffn_dim_multiplier=1.0, + block_multiple_of=256, + ) + model.make_intermediate_size_init(config) + assert model.intermediate_size == 8192 + + with pytest.raises(AttributeError, match="block_auto_adjust_ff_dim"): + model.make_intermediate_size_init(types.SimpleNamespace(intermediate_size=12288)) + + +def test_lfm2_moe_intermediate_size_uses_config_value(): + # Lfm2MoeConfig has no block_auto_adjust_ff_dim; the dense layers use intermediate_size as is. + model = LFM2MoEModel.__new__(LFM2MoEModel) + model.make_intermediate_size_init(types.SimpleNamespace(intermediate_size=7168)) + assert model.intermediate_size == 7168 + + +def test_lfm2_moe_init_configures_fused_swiglu(monkeypatch): + monkeypatch.setattr(LFM2Model, "__init__", _stub_base_init) + config = types.SimpleNamespace( + norm_topk_prob=True, + num_dense_layers=2, + moe_intermediate_size=1792, + use_expert_bias=True, + routed_scaling_factor=1.0, + ) + model = LFM2MoEModel(config, ir.DataType.FLOAT, ir.DataType.FLOAT, "cpu", None, {}) + assert model.moe_attrs["activation_type"] == "swiglu" + assert model.moe_attrs["swiglu_fusion"] == 1 + assert model.moe_attrs["normalize_routing_weights"] is True + assert model.moe_attrs["router_sentinel"] == -10000.0 + assert model.moe_attrs["num_dense_layers"] == 2 + assert model.moe_attrs["use_expert_bias"] is True + assert model.moe_attrs["routed_scaling_factor"] == 1.0 + assert model.moe_intermediate_size == 1792 + + +@pytest.mark.parametrize( + "ep,swiglu_limit,expected", + [ + ("cpu", None, None), + ("trt-rtx", None, float("inf")), + ("trt-rtx", 7.0, 7.0), + ], +) +def test_make_moe_init_fills_swiglu_limit_for_trt_rtx(ep, swiglu_limit, expected): + # TRT-RTX requires swiglu_limit on QMoE for every MoE model, so the base init supplies +inf. + model = Model.__new__(Model) + model.ep = ep + model.moe_attrs = {"swiglu_limit": swiglu_limit} + model.quant_config = types.SimpleNamespace(moe=types.SimpleNamespace(type="int4", weights_prepacked=-1)) + model.make_moe_init() + assert model.moe_attrs["swiglu_limit"] == expected + + +@pytest.mark.parametrize( + "moe_type,expected", + [ + ("int4", ("QMoE", "int", 4)), + ("int8", ("QMoE", "int", 8)), + ("mxfp4", ("QMoE", "fp4", 4)), + ("nvfp4", ("QMoE", "nvfp4", 4)), + ("none", ("MoE", "int", 0)), + ], +) +def test_make_moe_init_selects_qmoe_for_quantized_experts(moe_type, expected): + # Regression: these keys are what make_moe_op / make_moe_expert_initializers read. + model = Model.__new__(Model) + model.ep = "cpu" + model.moe_attrs = {"swiglu_limit": None} + model.quant_config = types.SimpleNamespace(moe=types.SimpleNamespace(type=moe_type, weights_prepacked=-1)) + model.make_moe_init() + assert ( + model.moe_attrs["op_type"], + model.moe_attrs["quant_type"], + model.moe_attrs["expert_weight_bits"], + ) == expected diff --git a/test/python/builder/test_qmoe_weights.py b/test/python/builder/test_qmoe_weights.py index 3a47e4ecc1..7e09a01405 100644 --- a/test/python/builder/test_qmoe_weights.py +++ b/test/python/builder/test_qmoe_weights.py @@ -14,6 +14,10 @@ 4. restore per-channel quantization when ``qmoe_block_size <= 0``, and 5. keep the SIGNED blockwise scales (taking ``abs()`` reintroduces the garbage-output bug). + +The CPU and WebGPU paths share the raw MatMulNBits encoding and must keep the +``[N, K/pack]`` storage contract the QMoE op validates, for INT4 and INT8. +WebGPU additionally requires whole blocks; INT4 requires even K. """ from __future__ import annotations @@ -21,97 +25,18 @@ import importlib.util import sys import types -from pathlib import Path import pytest import torch +from _builder_test_utils import BUILDERS_DIR, load_builder_module - -def _module_available(module_name): - try: - return importlib.util.find_spec(module_name) is not None - except (ModuleNotFoundError, ValueError): - return False - - -def _stub_missing_builder_dependencies(): - if not _module_available("onnx_ir"): - onnx_ir = types.ModuleType("onnx_ir") - onnx_ir.DataType = types.SimpleNamespace(INT4=object(), FLOAT=object(), FLOAT16=object(), BFLOAT16=object()) - tensor_adapters = types.ModuleType("onnx_ir.tensor_adapters") - tensor_adapters.TorchTensor = object - tensor_adapters.to_torch_dtype = lambda dtype: dtype - sys.modules["onnx_ir"] = onnx_ir - sys.modules["onnx_ir.tensor_adapters"] = tensor_adapters - - if not _module_available("onnxruntime.quantization.matmul_nbits_quantizer"): - # Prefer the real onnxruntime package when it is installed; only fabricate a - # top-level stub when the package truly isn't available. This avoids shadowing a - # real onnxruntime wheel (which would break other tests in the session) and only - # supplies the specific submodule the builder needs. - if _module_available("onnxruntime"): - import onnxruntime # noqa: PLC0415 - else: - onnxruntime = sys.modules.setdefault("onnxruntime", types.ModuleType("onnxruntime")) - quantization = getattr(onnxruntime, "quantization", None) - if quantization is None: - quantization = types.ModuleType("onnxruntime.quantization") - matmul_nbits_quantizer = types.ModuleType("onnxruntime.quantization.matmul_nbits_quantizer") - for class_name in ( - "KQuantWeightOnlyQuantConfig", - "MatMulNBitsQuantizer", - "QuantFormat", - "RTNWeightOnlyQuantConfig", - ): - setattr(matmul_nbits_quantizer, class_name, type(class_name, (), {})) - onnxruntime.quantization = quantization - quantization.matmul_nbits_quantizer = matmul_nbits_quantizer - sys.modules["onnxruntime.quantization"] = quantization - sys.modules["onnxruntime.quantization.matmul_nbits_quantizer"] = matmul_nbits_quantizer - - if not _module_available("tqdm"): - tqdm_module = types.ModuleType("tqdm") - tqdm_module.tqdm = lambda iterable=None, *args, **kwargs: iterable - sys.modules["tqdm"] = tqdm_module - - if not _module_available("transformers"): - transformers = types.ModuleType("transformers") - for class_name in ( - "AutoConfig", - "AutoModelForCausalLM", - "AutoModelForSpeechSeq2Seq", - "AutoTokenizer", - "GenerationConfig", - ): - setattr(transformers, class_name, type(class_name, (), {})) - sys.modules["transformers"] = transformers - - -_stub_missing_builder_dependencies() - -BUILDERS_DIR = Path(__file__).parents[3] / "src" / "python" / "py" / "models" / "builders" -sys.path.insert(0, str(BUILDERS_DIR.parent)) - - -def _load_builder_module(module_name): - spec = importlib.util.spec_from_file_location(f"models.builders.{module_name}", BUILDERS_DIR / f"{module_name}.py") - module = importlib.util.module_from_spec(spec) - sys.modules[f"models.builders.{module_name}"] = module - spec.loader.exec_module(module) - return module - - -sys.modules.setdefault("models", types.ModuleType("models")) -builders_package = sys.modules.setdefault("models.builders", types.ModuleType("models.builders")) -builders_package.__path__ = [str(BUILDERS_DIR)] - -base_module = _load_builder_module("base") +base_module = load_builder_module("base") Model = base_module.Model cuda_quantizer_module = sys.modules["quantization.cuda_quantizer"] qmoe_symmetric_per_channel_quantize = cuda_quantizer_module.CudaQuantizer.qmoe_symmetric_per_channel_quantize -gptoss_module = _load_builder_module("gptoss") +gptoss_module = load_builder_module("gptoss") GPTOSSModel = gptoss_module.GPTOSSModel -phi_module = _load_builder_module("phi") +phi_module = load_builder_module("phi") Phi3MoELongRoPEModel = phi_module.Phi3MoELongRoPEModel @@ -207,6 +132,7 @@ def _load_builder_cli_module(monkeypatch): "HunyuanDenseV1Model", "InternLM2Model", "LFM2Model", + "LFM2MoEModel", "LlamaModel", "Mistral3TextModel", "MistralModel", @@ -473,14 +399,110 @@ def test_cuda_raw_path_for_zero(): assert model.calls == [("matmulnbits", 128)] -def test_non_cuda_does_not_use_cuda_only_paths(): - """The CUDA-only encodings must not be used on other EPs, even when - weights_prepacked is set.""" - model = _FakeMoEModel("cpu", 128, 0) +@pytest.mark.parametrize("ep", ["cpu", "webgpu"]) +@pytest.mark.parametrize("weights_prepacked", [-1, 0, 1]) +def test_cpu_and_webgpu_use_signed_scale_blockwise_quantizer(ep, weights_prepacked): + """CPU and WebGPU ship raw MatMulNBits-convention blockwise weights (signed + block scales) regardless of the CUDA-only weights_prepacked knob, and never + the CUTLASS-prepacked encoding.""" + model = _FakeMoEModel(ep, 128, weights_prepacked) + model.make_qmoe_weights(_W) + assert model.calls == [("matmulnbits", 128)] + assert model.moe_attrs["block_size"] == 128 + + +@pytest.mark.parametrize("weights_prepacked", [-1, 0, 1]) +def test_trt_rtx_keeps_the_original_symmetric_blockwise_encoding(weights_prepacked): + """The signed-scale grid has not been measured on the TRT-RTX kernel, so that + EP stays on the encoding it shipped with.""" + model = _FakeMoEModel("trt-rtx", 128, weights_prepacked) model.make_qmoe_weights(_W) - assert ("cutlass", 128) not in model.calls - assert ("matmulnbits", 128) not in model.calls assert model.calls == [("symmetric", 128)] + assert model.moe_attrs["block_size"] == 128 + + +@pytest.mark.parametrize("ep", ["cpu", "webgpu", "cuda"]) +def test_matmulnbits_blockwise_paths_validate_block_size(ep): + """Every EP on the MatMulNBits grid shares the CUDA block-size constraint.""" + model = _FakeMoEModel(ep, 16, -1) + with pytest.raises(ValueError, match="block_size 32, 64, or 128"): + model.make_qmoe_weights(_W) + assert model.calls == [] + + +@pytest.mark.parametrize("ep,weights_prepacked", [("cpu", -1), ("cuda", 0)]) +@pytest.mark.parametrize("bits,k,expected_columns", [(4, 40, 20), (8, 40, 40), (8, 33, 33)]) +def test_raw_blockwise_storage_drops_the_block_padding(ep, weights_prepacked, bits, k, expected_columns): + """The MatMulNBits quantizer pads K up to whole blocks; the QMoE op validates raw storage as + [E, N, K/pack], so the padding must not reach the initializer.""" + model = _RealMoEModel(ep, 32, weights_prepacked, bits=bits) + model._matmulnbits_blockwise_quantize = Model._matmulnbits_blockwise_quantize.__get__(model) + torch.manual_seed(0) + weights = torch.randn(3, k) + qweight, scales = model.make_qmoe_weights(weights) + assert tuple(qweight.shape) == (3, expected_columns) + assert tuple(scales.shape) == (3, 2) + padded, _ = cuda_quantizer_module.CudaQuantizer.matmulnbits_blockwise_quantize( + weights, bits, 32, unsigned_full_range=True, signed_scale=True + ) + assert tuple(padded.shape) == (3, 2 * (32 // (8 // bits))) + assert torch.equal(qweight, padded[:, :expected_columns]) + + +@pytest.mark.parametrize("ep,weights_prepacked", [("cpu", -1), ("webgpu", -1), ("cuda", 0)]) +def test_raw_blockwise_int4_rejects_odd_input_dimension(ep, weights_prepacked): + model = _RealMoEModel(ep, 32, weights_prepacked) + model._matmulnbits_blockwise_quantize = Model._matmulnbits_blockwise_quantize.__get__(model) + with pytest.raises(RuntimeError, match=r"INT4 QMoE requires expert input dimension K \(33\) to be divisible by 2"): + model.make_qmoe_weights(torch.zeros(4, 33)) + + +@pytest.mark.parametrize("bits", [4, 8]) +@pytest.mark.parametrize("block_size", [32, 64, 128]) +def test_webgpu_blockwise_rejects_partial_blocks(bits, block_size): + model = _RealMoEModel("webgpu", block_size, -1, bits=bits) + model._matmulnbits_blockwise_quantize = Model._matmulnbits_blockwise_quantize.__get__(model) + with pytest.raises(RuntimeError, match=rf"WebGPU QMoE.*K \(40\).*qmoe_block_size \({block_size}\)"): + model.make_qmoe_weights(torch.zeros(4, 40)) + + +@pytest.mark.parametrize("ep,weights_prepacked", [("cpu", -1), ("webgpu", -1)]) +def test_non_cuda_blockwise_int8_uses_offset_128_storage(ep, weights_prepacked): + """INT8 QMoE bytes are q + 128 (the CPU kernel dequantizes with default_zp_8bit = 128), on the + same signed-scale grid as INT4: the block's max-magnitude element maps exactly to qmin.""" + model = _RealMoEModel(ep, 32, weights_prepacked, bits=8) + model._matmulnbits_blockwise_quantize = Model._matmulnbits_blockwise_quantize.__get__(model) + weights = torch.zeros(1, 32) + weights[0, 0] = 8.0 # positive extreme of block 0 + weights[0, 1] = -4.0 + weights[0, 2] = 0.0625 + qweight, scales = model.make_qmoe_weights(weights) + assert qweight.dtype == torch.uint8 and tuple(qweight.shape) == (1, 32) + assert scales.shape == (1, 1) and scales[0, 0] == -0.0625 # signed: 8.0 / qmin(-128) + assert qweight[0, 0].item() == 0 # 8.0 -> qmin (-128) + 128, exact + assert qweight[0, 1].item() == 64 + 128 # -4.0 / -0.0625 = 64 + assert qweight[0, 2].item() == 127 # 0.0625 / -0.0625 = -1 + assert qweight[0, 3].item() == 128 # zero -> the zero point + dequantized = (qweight[0].to(torch.float32) - 128.0) * scales[0, 0].to(torch.float32) + assert torch.equal(dequantized, weights[0]) + + +@pytest.mark.parametrize("ep", ["cpu", "webgpu"]) +def test_non_cuda_blockwise_scales_are_signed_and_do_not_clip_the_extreme(ep): + """The signed-scale grid maps each block's max-magnitude element exactly to + qmin, so a positive extreme is no longer clipped to 7/8 of its value. This is + also the grid the CPU QMoE MLAS Q4 fast path re-quantizes to, keeping that path + lossless.""" + model = _RealMoEModel(ep, 32, -1, bits=4) + model._matmulnbits_blockwise_quantize = Model._matmulnbits_blockwise_quantize.__get__(model) + weights = torch.zeros(1, 32) + weights[0, 0] = 8.0 # positive extreme of block 0 + weights[0, 1] = -4.0 + qweight, scales = model.make_qmoe_weights(weights) + assert scales.shape == (1, 1) and scales[0, 0] == -1.0 # signed: 8.0 / qmin(-8) + q = qweight[0, 0].item() + assert (q & 0xF) == 0 # 8.0 -> qmin (-8) + 8 == 0, exact + assert (q >> 4) == 12 # -4.0 / -1.0 = 4 -> 4 + 8 @pytest.mark.parametrize( diff --git a/test/python/builder/test_qwen_vl.py b/test/python/builder/test_qwen_vl.py index fa74f94e98..0ffb3530d8 100644 --- a/test/python/builder/test_qwen_vl.py +++ b/test/python/builder/test_qwen_vl.py @@ -861,9 +861,35 @@ def test_qwen35_linear_attention_uses_gated_rms_norm(monkeypatch): ) +@pytest.mark.parametrize("use_paged_attention,rows", [(False, "batch_size * sequence_length"), (True, "num_tokens")]) +def test_qwen35_moe_router_declares_the_token_layout(monkeypatch, use_paged_attention, rows): + """The router Reshape must declare the same row dim as the hidden states feeding the MoE op.""" + model = Qwen35MoETextModel.__new__(Qwen35MoETextModel) + model.io_dtype = ir.DataType.FLOAT16 + model.use_paged_attention = use_paged_attention + model.moe_attrs = {"num_experts": 4} + reshapes = [] + monkeypatch.setattr(model, "make_matmul", lambda *_args: "/model/layers.1/moe/router/MatMul") + monkeypatch.setattr( + model, "make_reshape", lambda name, inputs, dtype, shape: reshapes.append((name, inputs, dtype, shape)) + ) + + model.make_moe_router(1, types.SimpleNamespace(gate=object()), "hidden_states") + + assert reshapes == [ + ( + "/model/layers.1/moe/router/Reshape", + ["/model/layers.1/moe/router/MatMul/output_0", "/model/constants/INT64/[-1, 4]"], + ir.DataType.FLOAT16, + [rows, 4], + ) + ] + + def test_qwen35_moe_combines_shared_expert_with_gated_add(monkeypatch): model = Qwen35MoETextModel.__new__(Qwen35MoETextModel) model.io_dtype = ir.DataType.FLOAT16 + model.use_paged_attention = False model.hidden_size = 16 model.moe_intermediate_size = 4 model.moe_attrs = {