diff --git a/examples/auto_deploy/cookbooks/gpt_oss_trtllm_cookbook.ipynb b/examples/auto_deploy/cookbooks/gpt_oss_trtllm_cookbook.ipynb index 0dcf571feb0b..a99513ba25a0 100644 --- a/examples/auto_deploy/cookbooks/gpt_oss_trtllm_cookbook.ipynb +++ b/examples/auto_deploy/cookbooks/gpt_oss_trtllm_cookbook.ipynb @@ -100,48 +100,12 @@ { "cell_type": "markdown", "metadata": {}, - "source": [ - "## OpenAI-Compatible Server\n", - "\n", - "Start a local OpenAI-compatible server with TensorRT-LLM via the terminal, within the running docker container.\n", - "\n", - "Each gpt-oss size has its own AutoDeploy YAML under `examples/auto_deploy/model_registry/configs/`:\n", - "- `gpt_oss_20b.yaml` (world_size=2)\n", - "- `gpt_oss_120b.yaml` (world_size=8)\n", - "\n", - "Pick the YAML that matches the model size you want to deploy." - ] + "source": "## OpenAI-Compatible Server\n\nStart a local OpenAI-compatible server with TensorRT-LLM via the terminal, within the running docker container.\n\nBoth gpt-oss sizes share a single AutoDeploy YAML at `examples/auto_deploy/model_registry/configs/gpt_oss.yaml`. The same file is reused for 20B and 120B — only the HuggingFace model id changes between launches." }, { "cell_type": "markdown", "metadata": {}, - "source": [ - "### Load `gpt-oss-20b`\n", - "\n", - "Launch the TensorRT-LLM server on 2 GPUs:\n", - "\n", - "```shell\n", - "trtllm-serve \"openai/gpt-oss-20b\" \\\n", - " --host 0.0.0.0 \\\n", - " --port 8000 \\\n", - " --backend _autodeploy \\\n", - " --extra_llm_api_options examples/auto_deploy/model_registry/configs/gpt_oss_20b.yaml\n", - "```\n", - "\n", - "### Load `gpt-oss-120b`\n", - "\n", - "Launch the TensorRT-LLM server on 8 GPUs:\n", - "\n", - "```shell\n", - "trtllm-serve \"openai/gpt-oss-120b\" \\\n", - " --host 0.0.0.0 \\\n", - " --port 8000 \\\n", - " --backend _autodeploy \\\n", - " --extra_llm_api_options examples/auto_deploy/model_registry/configs/gpt_oss_120b.yaml\n", - "```\n", - "\n", - "Both YAMLs are self-contained — they include the compile backend, attention backend, world size, KV-cache settings and the CUDA-graph batch-size buckets needed for serving." - ] + "source": "### Load `gpt-oss-20b`\n\nLaunch the TensorRT-LLM server:\n\n```shell\ntrtllm-serve \"openai/gpt-oss-20b\" \\\n --host 0.0.0.0 \\\n --port 8000 \\\n --backend _autodeploy \\\n --extra_llm_api_options examples/auto_deploy/model_registry/configs/gpt_oss.yaml\n```\n\n### Load `gpt-oss-120b`\n\nLaunch the TensorRT-LLM server:\n\n```shell\ntrtllm-serve \"openai/gpt-oss-120b\" \\\n --host 0.0.0.0 \\\n --port 8000 \\\n --backend _autodeploy \\\n --extra_llm_api_options examples/auto_deploy/model_registry/configs/gpt_oss.yaml\n```\n\nThe shared YAML is self-contained — it includes the compile backend, attention backend, KV-cache settings and the CUDA-graph batch-size buckets needed for serving. `world_size` is supplied separately via the registry (e.g., `world_size_1.yaml`); pass it explicitly via `--extra_llm_api_options` when launching outside the registry." }, { "cell_type": "markdown", diff --git a/examples/auto_deploy/llmc/create_standalone_package.py b/examples/auto_deploy/llmc/create_standalone_package.py index c6aa87766abf..ff7c7295686e 100644 --- a/examples/auto_deploy/llmc/create_standalone_package.py +++ b/examples/auto_deploy/llmc/create_standalone_package.py @@ -185,6 +185,10 @@ # Imports utils.util.skip_pre_blackwell (not shipped in standalone) and exercises # fuse_finegrained_fp8_swiglu which depends on TRT-LLM runtime. "test_finegrained_fp8_swiglu.py", + # Exercise trtllm-gen MXFP4 MoE kernels (Blackwell-only) and import the + # prepare_trtllm_gen_moe_mxfp4_weights / utils.util helpers not in standalone. + "test_fuse_mxfp4_moe.py", + "test_trtllm_quant_mxfp4_trtllm_gen_moe.py", } # Import path rewrite: old -> new (applied to test files only). diff --git a/examples/auto_deploy/model_registry/configs/gpt_oss.yaml b/examples/auto_deploy/model_registry/configs/gpt_oss.yaml new file mode 100644 index 000000000000..f6078b77fdb8 --- /dev/null +++ b/examples/auto_deploy/model_registry/configs/gpt_oss.yaml @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# OpenAI GPT-OSS (20B / 120B, MXFP4 quantized) — shared AD serving config. +# - 20B: 24 layers, 32 experts, top-4 +# - 120B: 36 layers, 128 experts, top-4 +# Both share GQA (64 Q / 8 KV heads), head_dim=64, hidden=2880, and MXFP4 +# weights on HF that AD's `quantize_mxfp4_moe` transform handles. +# world_size is set via the registry's `world_size_N.yaml` overlay — not here. +model_factory: AutoModelForCausalLM +model_kwargs: + dtype: bfloat16 +attn_backend: trtllm +compile_backend: torch-cudagraph +skip_loading_weights: false +max_batch_size: 128 +max_seq_len: 4096 +max_num_tokens: 8192 +enable_chunked_prefill: true +cuda_graph_config: + batch_sizes: [1, 2, 4, 8, 16, 32, 64, 128] +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.8 +transforms: + detect_sharding: + enabled: false + sharding_transform_executor: + enabled: false + apply_sharding_hints: + enabled: true + requires_shape_prop: true + shard_layers: ["mha", "moe", "lm_head"] # V3: vocab-parallel lm_head (colwise + all_gather) + # TP2 trial: AUTO -> tunable_allreduce picks oneshot-lamport fused AR. + # NCCL has no fused residual+rmsnorm kernel -> RING_LL + separate rmsnorm. + allreduce_strategy: AUTO + quantize_mxfp4_moe: + backend: trtllm + trtllm_quant_act: mxfp8 + fuse_gemms_mixed_children: + enabled: true + fuse_gemms: + enabled: true + fuse_rope_into_trtllm_attention: + enabled: true + fuse_add_rms_norm: + enabled: true diff --git a/examples/auto_deploy/model_registry/configs/gpt_oss_120b.yaml b/examples/auto_deploy/model_registry/configs/gpt_oss_120b.yaml deleted file mode 100644 index 9ca974a5727e..000000000000 --- a/examples/auto_deploy/model_registry/configs/gpt_oss_120b.yaml +++ /dev/null @@ -1,21 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# OpenAI GPT-OSS-120B (128 experts, top-4, MXFP4 quantized) — standalone AD serving config. -# 36 layers (alternating sliding/full), GQA (64 Q / 8 KV heads), head_dim=64, hidden=2880. -# Weights are stored in MXFP4 on HF; AD's quantize_mxfp4_moe transform handles it. -runtime: trtllm -model_factory: AutoModelForCausalLM -attn_backend: trtllm -compile_backend: torch-cudagraph -skip_loading_weights: false -world_size: 4 -max_batch_size: 128 -max_seq_len: 4096 -max_num_tokens: 8192 -enable_chunked_prefill: true -cuda_graph_config: - batch_sizes: [1, 2, 4, 8, 16, 32, 64, 128] -kv_cache_config: - enable_block_reuse: false - free_gpu_memory_fraction: 0.8 diff --git a/examples/auto_deploy/model_registry/configs/gpt_oss_20b.yaml b/examples/auto_deploy/model_registry/configs/gpt_oss_20b.yaml deleted file mode 100644 index fee088fe9d4b..000000000000 --- a/examples/auto_deploy/model_registry/configs/gpt_oss_20b.yaml +++ /dev/null @@ -1,21 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# OpenAI GPT-OSS-20B (32 experts, top-4, MXFP4 quantized) — standalone AD serving config. -# 24 layers (alternating sliding/full), GQA (64 Q / 8 KV heads), head_dim=64, hidden=2880. -# Weights are stored in MXFP4 on HF; AD's quantize_mxfp4_moe transform handles it. -runtime: trtllm -model_factory: AutoModelForCausalLM -attn_backend: trtllm -compile_backend: torch-cudagraph -skip_loading_weights: false -world_size: 1 -max_batch_size: 128 -max_seq_len: 4096 -max_num_tokens: 8192 -enable_chunked_prefill: true -cuda_graph_config: - batch_sizes: [1, 2, 4, 8, 16, 32, 64, 128] -kv_cache_config: - enable_block_reuse: false - free_gpu_memory_fraction: 0.8 diff --git a/examples/auto_deploy/model_registry/models.yaml b/examples/auto_deploy/model_registry/models.yaml index 2864c1e2427e..a7cc7377cd1d 100644 --- a/examples/auto_deploy/model_registry/models.yaml +++ b/examples/auto_deploy/model_registry/models.yaml @@ -168,10 +168,9 @@ models: - name: nvidia/Mistral-NeMo-Minitron-8B-Base config_id: default_ws_2 yaml_extra: ['dashboard_default.yaml', 'world_size_2.yaml'] -# OOM during AutoDeploy run. -# - name: openai/gpt-oss-20b -# config_id: gpt_oss_20b -# yaml_extra: ['gpt_oss_20b.yaml'] +- name: openai/gpt-oss-20b + config_id: gpt_oss + yaml_extra: ['gpt_oss.yaml', 'world_size_1.yaml'] - name: ibm-granite/granite-3.0-8b-instruct config_id: default_ws_1 yaml_extra: ['dashboard_default.yaml', 'world_size_1.yaml'] @@ -333,10 +332,9 @@ models: # - name: meta-llama/Llama-3.2-90B-Vision-Instruct # config_id: multimodal # yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'multimodal.yaml'] -# torch.distributed.DistStoreError: Timed out after 601 seconds waiting for clients. 1/4 clients joined. -# - name: openai/gpt-oss-120b -# config_id: gpt_oss_120b -# yaml_extra: ['gpt_oss_120b.yaml'] +- name: openai/gpt-oss-120b + config_id: gpt_oss + yaml_extra: ['gpt_oss.yaml', 'world_size_2.yaml'] # [RANK 3] Error querying confidential compute state: Function Not Found # - name: meta-llama/Llama-4-Scout-17B-16E-Instruct # config_id: multimodal__llama4_scout diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index 831a9cc1e74e..9097aea0f57c 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -191,6 +191,9 @@ transforms: fuse_finegrained_fp8_linear: stage: post_load_fusion backend: trtllm + fuse_mxfp4_moe: + stage: post_load_fusion + expect_mem_change: true # adds padding for trtllm-gen kernel alignment during weight repack fuse_moe: stage: post_load_fusion expect_mem_change: true diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/prepare_trtllm_gen_moe_mxfp4_weights.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/prepare_trtllm_gen_moe_mxfp4_weights.py new file mode 100644 index 000000000000..cc1f75e9d863 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/prepare_trtllm_gen_moe_mxfp4_weights.py @@ -0,0 +1,648 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""MXFP4 weight prep for TRT-LLM-Gen ``bf16_mxe2m1_block_scale_moe_runner``. + +Transforms HF on-disk MXFP4 expert tensors (``gate_up_proj_*``, ``down_proj_*`` registered by +``quantize_mxfp4_moe``) into the kernel-ready stacked layout that +``auto_deploy::trtllm_quant_mxfp4_trtllm_gen_moe_fused`` expects: ``[E_local, 2I_pad, H_pad/2]`` weights, +``[E_local, 2I_pad]`` fp32 biases, etc., all run through ``torch.ops.trtllm.shuffle_matrix`` for +the TMA layout. + +Mirrors PT's ``MXFP4WeightTRTLLMGenFusedMoEMethod`` +(``tensorrt_llm/_torch/modules/fused_moe/quantization.py:4135``) — PT helpers +(``maybe_pad_for_mxfp4``, ``trtllmgen_maybe_get_cached_*``, ``_get_weight_alignment``) are imported +directly so the algorithm is byte-identical. +""" + +from dataclasses import dataclass +from typing import Dict, Tuple + +import torch + +from tensorrt_llm._torch.modules.fused_moe.quantization import ( + _get_weight_alignment, + maybe_pad_for_mxfp4, + trtllmgen_maybe_get_cached_w2_permute_indices, + trtllmgen_maybe_get_cached_w3_w1_permute_indices, +) +from tensorrt_llm.math_utils import pad_up + +# Cache permute indices to avoid recomputation across calls. +# Keyed by (shape, role, num_elts_per_sf) inside the PT helpers. +_PERMUTE_CACHE: Dict = {} + +# MXFP4 block size (UE8M0 scale per 32 elements). Matches HF gpt-oss layout. +_MXFP4_SCALING_VECTOR_SIZE: int = 32 + +# Kernel layout constants (mirror PT's MXFP4WeightTRTLLMGenFusedMoEMethod). +_INPUT_HIDDEN_ALIGNMENT: int = 512 +_WEIGHT_ALIGNMENT: int = 128 +_EPILOGUE_TILE_M: int = 128 + + +def _compute_padded_dims(per_rank_i: int, hidden_size: int) -> Tuple[int, int, int]: + """Returns ``(i_pad, h_w1_pad, h_w2_pad)`` for the trtllm-gen layout. + + ``i_pad`` / ``h_w2_pad`` align to 128 (TMA weight alignment); + ``h_w1_pad`` aligns to 512 (TMA input-hidden constraint on w1's K-axis). + """ + return ( + pad_up(per_rank_i, _WEIGHT_ALIGNMENT), + pad_up(hidden_size, _INPUT_HIDDEN_ALIGNMENT), + pad_up(hidden_size, _WEIGHT_ALIGNMENT), + ) + + +@dataclass(frozen=True) +class TRTLLMGenMXFP4MoEWeights: + """Output of :func:`prepare_trtllm_gen_moe_mxfp4_weights`.""" + + fc1_weights_mxfp4: torch.Tensor # [E, 2I_pad, H_pad/2] uint8 (shuffled) + fc1_weights_scale_ue8m0: torch.Tensor # [E, 2I_pad, H_pad/32] uint8 (shuffled) + fc1_bias_f32: torch.Tensor # [E, 2I_pad] float32 + fc2_weights_mxfp4: torch.Tensor # [E, H_pad, I_pad/2] uint8 (shuffled) + fc2_weights_scale_ue8m0: torch.Tensor # [E, H_pad, I_pad/32] uint8 (shuffled) + fc2_bias_f32: torch.Tensor # [E, H_pad] float32 (already /tp_size) + valid_hidden_size: int # original H + valid_intermediate_size: int # per-rank intermediate size (lean shape) + intermediate_size_padded: int # I_pad (per-rank, after pad) + hidden_size_padded: int # H_pad + + +@dataclass +class MXFP4PrepScratch: + """Reusable GPU scratch buffers for ``prepare_trtllm_gen_moe_mxfp4_weights``. + + Allocated once per build pass via :meth:`allocate` (sized for ONE MoE layer's per-rank shape — + gpt-oss guarantees H/I/E are constant across layers) and reused on every layer to avoid + per-layer pad/shuffle transients. ``trtllm.shuffle_matrix`` is not in-place, so we keep + separate ``_pad_buf`` (post-pad, pre-shuffle) and no-suffix (post-shuffle, kernel-ready) + buffers per tensor kind. + + Caller MUST ``copy_`` the prep result into final storage before the next call — the dataclass + holds VIEWS of these buffers and the next call overwrites them. The intended use is + :class:`FuseMXFP4Moe`: pre-allocate all layers' destination ``nn.Parameter`` storage first, + then per-layer prep + ``copy_`` from scratch. + """ + + # Shuffle outputs (= kernel-ready layout; what the prepared nn.Parameter + # will hold). + fc1_w_buf: torch.Tensor # [E_local, 2I_pad, H_w1_pad/2] uint8 + fc1_s_buf: torch.Tensor # [E_local, 2I_pad, H_w1_pad/32] uint8 + fc1_b_buf: torch.Tensor # [E_local, 2I_pad] fp32 + fc2_w_buf: torch.Tensor # [E_local, H_w2_pad, I_pad/2] uint8 + fc2_s_buf: torch.Tensor # [E_local, H_w2_pad, I_pad/32] uint8 + fc2_b_buf: torch.Tensor # [E_local, H_w2_pad] fp32 + + # Pad outputs (post pad, pre shuffle). Same shape as the corresponding + # shuffle output above. + fc1_w_pad_buf: torch.Tensor + fc1_s_pad_buf: torch.Tensor + fc1_b_pad_buf: torch.Tensor + fc2_w_pad_buf: torch.Tensor + fc2_s_pad_buf: torch.Tensor + fc2_b_pad_buf: torch.Tensor + + # Cached layout dimensions (for shape validation on subsequent calls). + e_local: int + hidden_size: int + per_rank_i: int + i_pad: int + h_w1_pad: int + h_w2_pad: int + + @classmethod + def allocate( + cls, + *, + e_local: int, + per_rank_i: int, + hidden_size: int, + device: torch.device | str, + ) -> "MXFP4PrepScratch": + """Allocate scratch for one MoE layer at the given per-rank shape. + + ``per_rank_i`` is the already-TP-sliced intermediate dim (= full ``I`` + when ``tp_size == 1``). + """ + i_pad, h_w1_pad, h_w2_pad = _compute_padded_dims(per_rank_i, hidden_size) + u8 = dict(dtype=torch.uint8, device=device) + f32 = dict(dtype=torch.float32, device=device) + return cls( + fc1_w_buf=torch.empty(e_local, 2 * i_pad, h_w1_pad // 2, **u8), + fc1_s_buf=torch.empty(e_local, 2 * i_pad, h_w1_pad // _MXFP4_SCALING_VECTOR_SIZE, **u8), + fc1_b_buf=torch.empty(e_local, 2 * i_pad, **f32), + fc2_w_buf=torch.empty(e_local, h_w2_pad, i_pad // 2, **u8), + fc2_s_buf=torch.empty(e_local, h_w2_pad, i_pad // _MXFP4_SCALING_VECTOR_SIZE, **u8), + fc2_b_buf=torch.empty(e_local, h_w2_pad, **f32), + fc1_w_pad_buf=torch.empty(e_local, 2 * i_pad, h_w1_pad // 2, **u8), + fc1_s_pad_buf=torch.empty( + e_local, 2 * i_pad, h_w1_pad // _MXFP4_SCALING_VECTOR_SIZE, **u8 + ), + fc1_b_pad_buf=torch.empty(e_local, 2 * i_pad, **f32), + fc2_w_pad_buf=torch.empty(e_local, h_w2_pad, i_pad // 2, **u8), + fc2_s_pad_buf=torch.empty(e_local, h_w2_pad, i_pad // _MXFP4_SCALING_VECTOR_SIZE, **u8), + fc2_b_pad_buf=torch.empty(e_local, h_w2_pad, **f32), + e_local=e_local, + hidden_size=hidden_size, + per_rank_i=per_rank_i, + i_pad=i_pad, + h_w1_pad=h_w1_pad, + h_w2_pad=h_w2_pad, + ) + + +def _flatten_block_dim(blocks_4d: torch.Tensor) -> torch.Tensor: + """Collapse ``[..., n_blocks, 16]`` -> ``[..., n_blocks * 16]`` (= H/2 or I/2).""" + if blocks_4d.dim() == 3: + return blocks_4d + if blocks_4d.dim() == 4: + return blocks_4d.reshape(*blocks_4d.shape[:-2], -1) + raise ValueError(f"Unexpected MXFP4 weight rank {blocks_4d.dim()}; expected 3 or 4.") + + +def _pad_per_expert_2d( + weight_3d: torch.Tensor, # [E, R, C] + col_alignment: int, + row_alignment: int, + *, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Pad each expert's 2-D matrix to the given row/col alignment. + + ``out=None``: build a fresh stacked tensor (one alloc). + ``out=...``: write each expert's padded matrix into ``out[i]`` in-place + so caller-provided storage (e.g. a scratch slice) is filled directly. + """ + e = weight_3d.size(0) + padded_per_expert = ( + maybe_pad_for_mxfp4(weight_3d[i], col_alignment, row_alignment) for i in range(e) + ) + if out is None: + return torch.stack(list(padded_per_expert), dim=0).contiguous() + assert out.shape[0] == e, f"out leading dim {out.shape[0]} != e {e}" + for i, padded in enumerate(padded_per_expert): + out[i].copy_(padded) + return out + + +def _shuffle_per_expert( + stacked: torch.Tensor, + permute_fn, + *, + num_elts_per_sf: int | None = None, + is_scale: bool = False, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Batched TMA-layout shuffle on a stacked ``[E, M, ...]`` tensor (weights, scales, biases). + + Derives the row permute ONCE on expert 0 (gpt-oss guarantees same per-expert shape, and + ``_PERMUTE_CACHE`` is keyed by shape so the per-expert loop would return the same index + every iteration anyway), then applies it to the whole stack via ``torch.index_select`` on + the M axis (= dim=1). When ``is_scale``, chains ``block_scale_interleave`` for the kernel's + scale layout. + + Biases use the SAME row permute as their weights so ``bias[i]`` aligns with + ``weight_row[i]`` post-shuffle — mismatch → kernel epilogue adds the wrong bias and MoE + output is garbage. + + ``out=None`` returns a fresh tensor; otherwise the result is ``copy_``-ed into the + caller-provided storage (used by :class:`MXFP4PrepScratch` to avoid per-layer transients). + """ + # Derive permute once on expert 0 — equivalent to per-expert calls because + # _PERMUTE_CACHE keys on shape and all experts share shape. + perm = permute_fn( + stacked[0], _PERMUTE_CACHE, _EPILOGUE_TILE_M, num_elts_per_sf=num_elts_per_sf + ).to(stacked.device) + + shuffled = torch.index_select(stacked, 1, perm) + if is_scale: + shuffled = torch.ops.trtllm.block_scale_interleave(shuffled).reshape(stacked.shape) + shuffled = shuffled.view(stacked.dtype) + + if out is None: + return shuffled.contiguous() + assert out.shape[0] == stacked.size(0) + out.copy_(shuffled) + return out + + +def _deinterleave_gate_up( + gu_3d: torch.Tensor, # [E, 2I, H/2] + gate_up_scales: torch.Tensor, # [E, 2I, H/32] + gate_up_bias: torch.Tensor, # [E, 2I] +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Split the HF interleaved 2I axis (gate at even rows, up at odd rows) into separate halves. + + Kept separate so downstream row-padding puts the zero-pad rows INSIDE each half before re-concat + as ``[up | gate]`` — matches PT's ``dst_w3 = up`` / ``dst_w1 = gate`` chunk layout. + """ + return ( + gu_3d[:, 0::2, :].contiguous(), # gate_rows_w + gu_3d[:, 1::2, :].contiguous(), # up_rows_w + gate_up_scales[:, 0::2, :].contiguous(), # gate_rows_s + gate_up_scales[:, 1::2, :].contiguous(), # up_rows_s + gate_up_bias[:, 0::2].contiguous(), # gate_b + gate_up_bias[:, 1::2].contiguous(), # up_b + ) + + +def _pad_and_slice_axis(t: torch.Tensor, dim: int, target: int, lo: int, hi: int) -> torch.Tensor: + """Pad ``t`` on ``dim`` to ``target`` then slice ``[lo:hi]`` on that dim.""" + cur = t.shape[dim] + if cur < target: + pad_amount = target - cur + # F.pad spec is reversed-axis order; build dynamically. + pad_spec = [0, 0] * (t.dim() - dim - 1) + [0, pad_amount] + [0, 0] * dim + t = torch.nn.functional.pad(t, pad_spec) + idx = [slice(None)] * t.dim() + idx[dim] = slice(lo, hi) + return t[tuple(idx)].contiguous() + + +def _tp_slice_intermediate_axis( + gate_rows_w: torch.Tensor, + up_rows_w: torch.Tensor, + gate_rows_s: torch.Tensor, + up_rows_s: torch.Tensor, + gate_b: torch.Tensor, + up_b: torch.Tensor, + dn_3d: torch.Tensor, + down_scales: torch.Tensor, + intermediate_size: int, + tp_size: int, + tp_rank: int, +): + """Pre-pad ``I`` to ``i_padded_tp`` then slice this rank's range (mirrors PT + ``quantization.py:4211-4234``). + + The alignment guarantees ``i_padded_tp / tp_size`` stays 128-aligned and that scaling-factor + blocks (32 elements) don't straddle rank boundaries. Example: gpt-oss I=2880 @ tp=8 → + ``alignment_tp=3072`` → ``per_rank_i=384``. + + No-op when ``tp_size == 1``: returns inputs unchanged with + ``per_rank_i = valid_intermediate = intermediate_size``. + + Returns the (possibly sliced) tensors plus ``(per_rank_i, valid_intermediate)``. + """ + if tp_size == 1: + return ( + gate_rows_w, + up_rows_w, + gate_rows_s, + up_rows_s, + gate_b, + up_b, + dn_3d, + down_scales, + intermediate_size, + intermediate_size, + ) + + alignment_tp = _get_weight_alignment( + _WEIGHT_ALIGNMENT, _MXFP4_SCALING_VECTOR_SIZE, tp_size, intermediate_size + ) + i_padded_tp = ((intermediate_size + alignment_tp - 1) // alignment_tp) * alignment_tp + per_rank_i = i_padded_tp // tp_size + slice_start = tp_rank * per_rank_i + slice_stop = (tp_rank + 1) * per_rank_i + valid_intermediate = max(0, min(intermediate_size, slice_stop) - slice_start) + + # Pad I axis (rows) of gate / up to i_padded_tp, then slice this rank's chunk. + def shard(t: torch.Tensor, dim: int, target: int, lo: int, hi: int) -> torch.Tensor: + return _pad_and_slice_axis(t, dim, target, lo, hi) + + sf_padded = i_padded_tp // _MXFP4_SCALING_VECTOR_SIZE + sf_start = slice_start // _MXFP4_SCALING_VECTOR_SIZE + sf_stop = slice_stop // _MXFP4_SCALING_VECTOR_SIZE + return ( + shard(gate_rows_w, 1, i_padded_tp, slice_start, slice_stop), + shard(up_rows_w, 1, i_padded_tp, slice_start, slice_stop), + shard(gate_rows_s, 1, i_padded_tp, slice_start, slice_stop), + shard(up_rows_s, 1, i_padded_tp, slice_start, slice_stop), + shard(gate_b, 1, i_padded_tp, slice_start, slice_stop), + shard(up_b, 1, i_padded_tp, slice_start, slice_stop), + shard(dn_3d, 2, i_padded_tp // 2, slice_start // 2, slice_stop // 2), + shard(down_scales, 2, sf_padded, sf_start, sf_stop), + per_rank_i, + valid_intermediate, + ) + + +def _pad_concat_fc1( + up_rows: torch.Tensor, + gate_rows: torch.Tensor, + col_alignment: int, + intermediate_size_pad: int, + *, + scratch_buf: torch.Tensor | None = None, +) -> torch.Tensor: + """Pad each half [E, I, X] then concat as ``[up | gate]`` on the 2I axis. + + Used for both weights (col_alignment = ``hidden_w1_pad // 2``) and scales + (col_alignment = ``hidden_w1_pad // _MXFP4_SCALING_VECTOR_SIZE``). When + ``scratch_buf`` is provided, the two halves are written directly into + ``scratch_buf[:, :i_pad]`` and ``scratch_buf[:, i_pad:]`` — no per-half + tensors or final concat alloc. + """ + if scratch_buf is None: + up_p = _pad_per_expert_2d(up_rows, col_alignment, intermediate_size_pad) + gate_p = _pad_per_expert_2d(gate_rows, col_alignment, intermediate_size_pad) + return torch.cat([up_p, gate_p], dim=1).contiguous() + i_pad = intermediate_size_pad + _pad_per_expert_2d(up_rows, col_alignment, intermediate_size_pad, out=scratch_buf[:, :i_pad, :]) + _pad_per_expert_2d( + gate_rows, col_alignment, intermediate_size_pad, out=scratch_buf[:, i_pad:, :] + ) + return scratch_buf + + +def _pad_fc2( + t: torch.Tensor, + col_alignment: int, + row_alignment: int, + *, + scratch_buf: torch.Tensor | None = None, +) -> torch.Tensor: + """Per-expert pad on the [E, H, X] down tensor. No concat (single half).""" + if scratch_buf is None: + return _pad_per_expert_2d(t, col_alignment, row_alignment) + _pad_per_expert_2d(t, col_alignment, row_alignment, out=scratch_buf) + return scratch_buf + + +def _shuffle_weights_and_scales( + gu_padded: torch.Tensor, + dn_padded: torch.Tensor, + gu_scale_padded: torch.Tensor, + dn_scale_padded: torch.Tensor, + *, + scratch: "MXFP4PrepScratch | None" = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Apply per-expert TMA-layout shuffle to weights + scales for both GEMMs.""" + w3w1 = trtllmgen_maybe_get_cached_w3_w1_permute_indices + w2 = trtllmgen_maybe_get_cached_w2_permute_indices + fc1_w_out = scratch.fc1_w_buf if scratch is not None else None + fc1_s_out = scratch.fc1_s_buf if scratch is not None else None + fc2_w_out = scratch.fc2_w_buf if scratch is not None else None + fc2_s_out = scratch.fc2_s_buf if scratch is not None else None + fc1_w = _shuffle_per_expert(gu_padded, w3w1, out=fc1_w_out) + fc1_s = _shuffle_per_expert( + gu_scale_padded, + w3w1, + num_elts_per_sf=_MXFP4_SCALING_VECTOR_SIZE, + is_scale=True, + out=fc1_s_out, + ) + fc2_w = _shuffle_per_expert(dn_padded, w2, out=fc2_w_out) + fc2_s = _shuffle_per_expert( + dn_scale_padded, + w2, + num_elts_per_sf=_MXFP4_SCALING_VECTOR_SIZE, + is_scale=True, + out=fc2_s_out, + ) + return fc1_w, fc1_s, fc2_w, fc2_s + + +def _prepare_fc1_bias( + up_b: torch.Tensor, + gate_b: torch.Tensor, + intermediate_size_pad: int, + *, + scratch_pad_buf: torch.Tensor | None = None, + scratch_out_buf: torch.Tensor | None = None, +) -> torch.Tensor: + """Pad each half [E, I] → [E, I_pad] (fp32), concat ``[up | gate]``, shuffle. + + The TMA-layout row shuffle on the bias is critical: PT applies the SAME + permute to bias rows as to weight rows so ``bias[i]`` aligns with + ``weight_row[i]`` after the shuffle. Skipping it → kernel's epilogue adds + the wrong bias to each output row → MoE output garbage (~2% GSM8K). + """ + if scratch_pad_buf is None: + up_p = ( + _pad_per_expert_2d(up_b.unsqueeze(-1), 1, intermediate_size_pad) + .squeeze(-1) + .float() + .contiguous() + ) + gate_p = ( + _pad_per_expert_2d(gate_b.unsqueeze(-1), 1, intermediate_size_pad) + .squeeze(-1) + .float() + .contiguous() + ) + fc1_bias_padded = torch.cat([up_p, gate_p], dim=1).contiguous() + else: + # ``_pad_per_expert_2d`` writes through ``copy_``; the scratch fp32 + # buffer absorbs bf16-padded values via copy_'s implicit cast. + i_pad = intermediate_size_pad + _pad_per_expert_2d( + up_b.unsqueeze(-1), + 1, + intermediate_size_pad, + out=scratch_pad_buf[:, :i_pad].unsqueeze(-1), + ) + _pad_per_expert_2d( + gate_b.unsqueeze(-1), + 1, + intermediate_size_pad, + out=scratch_pad_buf[:, i_pad:].unsqueeze(-1), + ) + fc1_bias_padded = scratch_pad_buf + return _shuffle_per_expert( + fc1_bias_padded, trtllmgen_maybe_get_cached_w3_w1_permute_indices, out=scratch_out_buf + ) + + +def _prepare_fc2_bias( + down_bias: torch.Tensor, + hidden_w2_pad: int, + tp_size: int, + *, + scratch_pad_buf: torch.Tensor | None = None, + scratch_out_buf: torch.Tensor | None = None, +) -> torch.Tensor: + """Pad ``[E, H] → [E, H_pad]`` (fp32), divide by ``tp_size``, shuffle. + + Scratch path enforces ``tp_size == 1`` upstream (TP slicing happens in + the load hook before this helper is reached), so no division is applied + when using scratch. + """ + if scratch_pad_buf is None: + fc2_b = ( + _pad_per_expert_2d(down_bias.unsqueeze(-1), 1, hidden_w2_pad) + .squeeze(-1) + .float() + .contiguous() + ) + if tp_size > 1: + fc2_b = fc2_b / tp_size + else: + _pad_per_expert_2d( + down_bias.unsqueeze(-1), + 1, + hidden_w2_pad, + out=scratch_pad_buf.unsqueeze(-1), + ) + fc2_b = scratch_pad_buf + return _shuffle_per_expert( + fc2_b, trtllmgen_maybe_get_cached_w2_permute_indices, out=scratch_out_buf + ) + + +def prepare_trtllm_gen_moe_mxfp4_weights( + gate_up_blocks: torch.Tensor, # [E, 2I, H/32, 16] or [E, 2I, H/2] uint8 + gate_up_scales: torch.Tensor, # [E, 2I, H/32] uint8 + gate_up_bias: torch.Tensor, # [E, 2I] bf16 + down_blocks: torch.Tensor, # [E, H, I/32, 16] or [E, H, I/2] uint8 + down_scales: torch.Tensor, # [E, H, I/32] uint8 + down_bias: torch.Tensor, # [E, H] bf16 + *, + hidden_size: int, + intermediate_size: int, + tp_size: int = 1, + tp_rank: int = 0, + scratch: MXFP4PrepScratch | None = None, +) -> TRTLLMGenMXFP4MoEWeights: + """Convert HF on-disk MXFP4 expert weights to the trtllm-gen kernel layout. + + Mirrors PT's ``MXFP4WeightTRTLLMGenFusedMoEMethod`` (``post_load_weights`` + + ``load_expert_w{3_w1,2}_weight{,_scale_mxfp4}``). + + Notes on optional args: + * ``tp_size > 1``: shard the intermediate dim before the kernel-layout pad+shuffle — see + :func:`_tp_slice_intermediate_axis` for the TP-aware pre-pad + slice math (mirrors PT + ``load_weight_shard``). + * ``scratch != None``: pad/shuffle outputs are written into the pre-allocated GPU buffers and + the returned dataclass holds VIEWS into that scratch, so caller must ``copy_`` results out + before the next call. Only supported at ``tp_size == 1`` (load hook does TP slicing first + — see :class:`MXFP4PrepScratch`). + + EP (expert-axis slicing) is NOT done here — the caller selects the expert subset before + invoking. + """ + if scratch is not None and tp_size != 1: + # Scratch path assumes inputs are already TP-sliced (load hook does + # that). Combining scratch with tp_size > 1 would double-slice. + raise ValueError( + "prepare_trtllm_gen_moe_mxfp4_weights: scratch is only supported with " + f"tp_size=1 (got tp_size={tp_size}). The caller is expected to do " + "TP slicing before this helper when using scratch." + ) + if tp_rank < 0 or tp_rank >= tp_size: + raise ValueError(f"tp_rank {tp_rank} out of range for tp_size {tp_size}") + + assert down_blocks.size(0) == gate_up_blocks.size(0) + + # 1. Flatten blocks ([..., 16] → flattened) and de-interleave gate/up halves. + gu_3d = _flatten_block_dim(gate_up_blocks) # [E, 2I, H/2] + dn_3d = _flatten_block_dim(down_blocks) # [E, H, I/2] + gate_rows_w, up_rows_w, gate_rows_s, up_rows_s, gate_b, up_b = _deinterleave_gate_up( + gu_3d, gate_up_scales, gate_up_bias + ) + + # 2. TP slicing on the intermediate axis (no-op for tp_size == 1). + ( + gate_rows_w, + up_rows_w, + gate_rows_s, + up_rows_s, + gate_b, + up_b, + dn_3d, + down_scales, + intermediate_size_for_local, + valid_intermediate, + ) = _tp_slice_intermediate_axis( + gate_rows_w, + up_rows_w, + gate_rows_s, + up_rows_s, + gate_b, + up_b, + dn_3d, + down_scales, + intermediate_size, + tp_size, + tp_rank, + ) + + # 3. Per-rank padded layout dims (mirrors PT + # ``MXFP4WeightTRTLLMGenFusedMoEMethod``: per-expert I + H padded + # BEFORE 2I row dim, so w1.2I = 2*I_pad — see ``_compute_padded_dims``). + intermediate_size_pad, hidden_w1_pad, hidden_w2_pad = _compute_padded_dims( + intermediate_size_for_local, hidden_size + ) + + # 4. Pad weights + scales (concat [up | gate] for fc1; single half for fc2). + sv = _MXFP4_SCALING_VECTOR_SIZE + gu_padded = _pad_concat_fc1( + up_rows_w, + gate_rows_w, + hidden_w1_pad // 2, + intermediate_size_pad, + scratch_buf=scratch.fc1_w_pad_buf if scratch is not None else None, + ) + dn_padded = _pad_fc2( + dn_3d, + intermediate_size_pad // 2, + hidden_w2_pad, + scratch_buf=scratch.fc2_w_pad_buf if scratch is not None else None, + ) + gu_scale_padded = _pad_concat_fc1( + up_rows_s, + gate_rows_s, + hidden_w1_pad // sv, + intermediate_size_pad, + scratch_buf=scratch.fc1_s_pad_buf if scratch is not None else None, + ) + dn_scale_padded = _pad_fc2( + down_scales, + intermediate_size_pad // sv, + hidden_w2_pad, + scratch_buf=scratch.fc2_s_pad_buf if scratch is not None else None, + ) + + # 5. Per-expert TMA-layout shuffle (weights + scales). + fc1_weights, fc1_weights_scale, fc2_weights, fc2_weights_scale = _shuffle_weights_and_scales( + gu_padded, dn_padded, gu_scale_padded, dn_scale_padded, scratch=scratch + ) + + # 6. Pad + shuffle biases (fp32, w2 bias divided by tp_size at tp>1). + fc1_bias_padded = _prepare_fc1_bias( + up_b, + gate_b, + intermediate_size_pad, + scratch_pad_buf=scratch.fc1_b_pad_buf if scratch is not None else None, + scratch_out_buf=scratch.fc1_b_buf if scratch is not None else None, + ) + fc2_bias_padded = _prepare_fc2_bias( + down_bias, + hidden_w2_pad, + tp_size, + scratch_pad_buf=scratch.fc2_b_pad_buf if scratch is not None else None, + scratch_out_buf=scratch.fc2_b_buf if scratch is not None else None, + ) + + intermediate_size_padded = fc1_weights.shape[1] // 2 # 2I_pad / 2 = I_pad + hidden_size_padded = fc1_weights.shape[-1] * 2 # (H_pad/2) * 2 = H_pad + return TRTLLMGenMXFP4MoEWeights( + fc1_weights_mxfp4=fc1_weights, + fc1_weights_scale_ue8m0=fc1_weights_scale, + fc1_bias_f32=fc1_bias_padded, + fc2_weights_mxfp4=fc2_weights, + fc2_weights_scale_ue8m0=fc2_weights_scale, + fc2_bias_f32=fc2_bias_padded, + valid_hidden_size=hidden_size, + valid_intermediate_size=valid_intermediate, + intermediate_size_padded=intermediate_size_padded, + hidden_size_padded=hidden_size_padded, + ) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py index 215653ffddd1..6e10bf5d1230 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py @@ -1299,3 +1299,214 @@ def trtllm_nvfp4_trtllm_gen_moe_fused_fake( batch_info_host: torch.Tensor | None = None, ) -> torch.Tensor: return torch.empty_like(x) + + +# ============================================================================= +# MXFP4 weights on TRT-LLM-Gen (W4A16 bf16-act or W4A8 mxfp8-act) +# ============================================================================= +# +# Same kernel path PT exercises for `gpt-oss-120b` on B200: +# * W4A16 (`act_dtype="bf16"`): `bf16_mxe2m1_block_scale_moe_runner` — +# PT mirror in fused_moe_trtllm_gen.py:652-712 (W4A16MXFP4TRTLLMGenFusedMoEMethod). +# * W4A8 (`act_dtype="mxfp8"`): `mxfp8_quantize` + `mxe4m3_mxe2m1_block_scale_moe_runner` — +# PT mirror in fused_moe_trtllm_gen.py:511 (W4A8MXFP4MXFP8TRTLLMGenFusedMoEMethod). +# The MXFP8 cubin family (median 9.1 µs/call vs 27 µs for bf16) unlocks bigger +# TileN candidates (up to 256 vs 64). +# +# Both paths use the SAME prepared weight layout (pad/shard/shuffle done by +# `prepare_trtllm_gen_moe_mxfp4_weights` in `prepare_trtllm_gen_moe_mxfp4_weights.py`); +# only the activation handling differs. At forward time we only pad activations +# to the kernel's expected hidden dim (and, for W4A8, also call `mxfp8_quantize`). +# +# Kernel-enforced weight layout: +# * Weights: uint8 packed (2 elements / byte), pre-padded + pre-shuffled +# * Scales: uint8 UE8M0 (block size 32) +# * Bias: float32 (kernel API) +# * input_hidden_alignment = 512 (TMA constraint, runner.cu:472) +# * weight_alignment = 128 (TMA 16U4 alignment) + + +@torch.library.custom_op("auto_deploy::trtllm_quant_mxfp4_trtllm_gen_moe_fused", mutates_args=()) +def trtllm_quant_mxfp4_trtllm_gen_moe_fused( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor, + top_k: int, + fc1_weights_mxfp4: torch.Tensor, + fc2_weights_mxfp4: torch.Tensor, + fc1_weights_scale_ue8m0: torch.Tensor, + fc2_weights_scale_ue8m0: torch.Tensor, + fc1_bias_f32: torch.Tensor, + fc2_bias_f32: torch.Tensor, + swiglu_alpha: torch.Tensor, + swiglu_beta: torch.Tensor, + swiglu_limit: torch.Tensor, + valid_hidden_size: int, + valid_intermediate_size: int, + act_dtype: str, + local_expert_offset: int = 0, + local_num_experts: int = -1, + routing_method_type: int = int(RoutingMethodType.Renormalize), +) -> torch.Tensor: + """TensorRT-LLM Gen MoE for MXFP4 weights with BF16 or MXFP8 activations. + + ``act_dtype`` selects the activation precision (see module header above for the runner + selection + cubin family details): + * ``"bf16"`` (W4A16) — bf16 hidden states fed directly to the bf16 MoE runner. + * ``"mxfp8"`` (W4A8) — bf16 hidden states pre-quantized to MXFP8 (E4M3 + UE8M0 + block scales via ``trtllm.mxfp8_quantize`` with alignment=512), then fed to + the MXFP8 MoE runner. + + The op takes the **raw router weight + bias** and computes top-k routing inside the + C++ runner via ``softmax(topk(F.linear(x, w, b)))`` (fused topk+softmax+cast). + + Args: + x: BF16 hidden states, shape ``(B, S, H)`` or ``(B*S, H)``. ``H`` may be smaller + than the kernel's expected (padded) hidden — the op zero-pads on entry and + slices the output back to ``valid_hidden_size``. + router_weight: ``[E_total, H]`` BF16 router projection. + router_bias: ``[E_total]`` BF16 router bias. + top_k: number of experts activated per token (4 for gpt-oss-120b). + fc1_weights_mxfp4: ``[E_local, 2*I_pad, H_pad/2]`` ``uint8`` (MXFP4 packed, + already pad+shard+shuffled; col-parallel along ``2*I``). + fc2_weights_mxfp4: ``[E_local, H_pad, I_pad/2]`` ``uint8`` (row-parallel along ``I``). + fc1_weights_scale_ue8m0: ``[E_local, 2*I_pad, H_pad/32]`` ``uint8`` UE8M0. + fc2_weights_scale_ue8m0: ``[E_local, H_pad, I_pad/32]`` ``uint8`` UE8M0. + fc1_bias_f32: ``[E_local, 2*I_pad]`` ``float32``. + fc2_bias_f32: ``[E_local, H_pad]`` ``float32`` (already divided by ``tp_size``). + swiglu_alpha / swiglu_beta / swiglu_limit: per-expert SwiGLU parameters, + ``[E_local]`` ``float32``. For gpt-oss: alpha=1.702, beta=1.0, limit=7.0. + valid_hidden_size: original (pre-pad) hidden size; output is sliced to this. + valid_intermediate_size: original per-rank intermediate size (kernel hint to skip + OOB MMA in padded regions). + act_dtype: ``"bf16"`` or ``"mxfp8"`` — selects W4A16 vs W4A8 cubin family. + local_expert_offset: ``slot_start`` for EP>1; ``0`` for EP=1. + local_num_experts: ``num_experts`` for EP=1, ``num_experts/ep_size`` for EP>1. + Pass ``-1`` to default to ``E_local`` inferred from ``fc1_weights_mxfp4``. + routing_method_type: integer from ``RoutingMethodType`` enum. Default + ``Renormalize`` (1) matches gpt-oss's ``RenormalizeMoeRoutingMethod``. + + Returns: + BF16 hidden states of shape ``(*x.shape[:-1], valid_hidden_size)``. + """ + x_shape = x.shape + x2d = x.view(-1, x_shape[-1]) + + # Routing: compute router logits and hand them to the C++ runner which performs + # fused topk + softmax + cast internally. routing_bias is None — the linear-layer + # bias was already folded into router_logits via F.linear. + router_logits = torch.nn.functional.linear(x2d, router_weight, router_bias) + + # Pad activations to the kernel's expected hidden (H_pad, multiple of 512). + expected_hidden = int(fc1_weights_mxfp4.shape[-1] * 2) + pad_size = expected_hidden - int(x2d.shape[-1]) + if pad_size > 0: + x2d = torch.nn.functional.pad(x2d, (0, pad_size)) + + num_experts_total = int(router_weight.shape[0]) + if local_num_experts < 0: + local_num_experts = int(fc1_weights_mxfp4.shape[0]) + intermediate_size_padded = int(fc1_weights_mxfp4.shape[1] // 2) + + if act_dtype == "mxfp8": + # Pre-quantize bf16 activation to MXFP8 (E4M3 elem + UE8M0 per-32-elem scale). + # Match PT's ``W4A8MXFP4MXFP8TRTLLMGenFusedMoEMethod.input_hidden_alignment = 512``. + # Keep ``x_scale`` 1D — the C++ runner asserts ``hidden_states_scale must be 1D``. + x_mxfp8, x_scale = torch.ops.trtllm.mxfp8_quantize( + x2d, + False, # is_sf_swizzled_layout + alignment=512, + ) + result = torch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runner( + router_logits, + None, # routing_bias + x_mxfp8, # hidden_states (E4M3-packed uint8) + x_scale, # hidden_states_scale (UE8M0 per-32-elem block scale) + fc1_weights_mxfp4, + fc1_weights_scale_ue8m0, + fc1_bias_f32, + swiglu_alpha, + swiglu_beta, + swiglu_limit, + fc2_weights_mxfp4, + fc2_weights_scale_ue8m0, + fc2_bias_f32, + num_experts_total, + int(top_k), + None, # n_group + None, # topk_group + intermediate_size_padded, + valid_hidden_size, + valid_intermediate_size, + local_expert_offset, + local_num_experts, + None, # routed_scaling_factor + routing_method_type, + 0, # act_type = SwiGlu + topk_weights=None, + topk_ids=None, + ) + elif act_dtype == "bf16": + result = torch.ops.trtllm.bf16_mxe2m1_block_scale_moe_runner( + router_logits, + None, # routing_bias + x2d, # hidden_states (bf16) + fc1_weights_mxfp4, + fc1_weights_scale_ue8m0, + fc1_bias_f32, + swiglu_alpha, + swiglu_beta, + swiglu_limit, + fc2_weights_mxfp4, + fc2_weights_scale_ue8m0, + fc2_bias_f32, + num_experts_total, + int(top_k), + None, # n_group + None, # topk_group + intermediate_size_padded, + valid_hidden_size, + valid_intermediate_size, + local_expert_offset, + local_num_experts, + None, # routed_scaling_factor + routing_method_type, + 0, # act_type = SwiGlu + # topk_weights/topk_ids omitted — kernel routes from router_logits. + ) + else: + raise ValueError( + f"trtllm_quant_mxfp4_trtllm_gen_moe_fused: act_dtype must be 'bf16' or 'mxfp8', " + f"got {act_dtype!r}." + ) + + if result.shape[-1] > valid_hidden_size: + result = result[..., :valid_hidden_size].contiguous() + return result.view(*x_shape[:-1], valid_hidden_size) + + +@trtllm_quant_mxfp4_trtllm_gen_moe_fused.register_fake +def trtllm_quant_mxfp4_trtllm_gen_moe_fused_fake( + x: torch.Tensor, + router_weight: torch.Tensor, + router_bias: torch.Tensor, + top_k: int, + fc1_weights_mxfp4: torch.Tensor, + fc2_weights_mxfp4: torch.Tensor, + fc1_weights_scale_ue8m0: torch.Tensor, + fc2_weights_scale_ue8m0: torch.Tensor, + fc1_bias_f32: torch.Tensor, + fc2_bias_f32: torch.Tensor, + swiglu_alpha: torch.Tensor, + swiglu_beta: torch.Tensor, + swiglu_limit: torch.Tensor, + valid_hidden_size: int, + valid_intermediate_size: int, + act_dtype: str, + local_expert_offset: int = 0, + local_num_experts: int = -1, + routing_method_type: int = int(RoutingMethodType.Renormalize), +) -> torch.Tensor: + out_shape = list(x.shape) + out_shape[-1] = valid_hidden_size + return x.new_empty(out_shape, dtype=x.dtype) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/linear.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/linear.py index 254410e6ba47..61bea3a3d6ae 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/linear.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/linear.py @@ -19,6 +19,8 @@ import torch +from ..._compat import get_sm_version + @torch.library.custom_op("auto_deploy::torch_linear_simple", mutates_args=()) def simple( @@ -65,6 +67,24 @@ def simple( Returns: Output tensor of shape ``(..., out_features)``. """ + # Blackwell (sm>=100) + bf16: route any bf16 linear to trtllm::cublas_mm. + # Selects single-pass cluster-mode cubins instead of cuBLAS-default + # split-K + reduce + zero-fill for small-M (decode) projection GEMMs. + # (Same trick PT introduced for GPT-OSS via use_custom_cublas_mm in + # modeling_gpt_oss.py; we apply it model-agnostically based on dtype + SM.) + if get_sm_version() >= 100 and input.dtype == torch.bfloat16 and weight.dtype == torch.bfloat16: + # cublas_mm requires 2D mat_a/mat_b. Flatten leading dims and unflatten on exit. + in_shape = input.shape + input_2d = input.reshape(-1, in_shape[-1]) + out_2d = torch.ops.trtllm.cublas_mm( + input_2d, + weight.t(), + bias, + None, # out_dtype + 0, # output_buffer_kind = DEFAULT + None, # group (no TP) + ) + return out_2d.view(*in_shape[:-1], out_2d.shape[-1]) return torch.ops.aten.linear(input, weight, bias) diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_gpt_oss.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_gpt_oss.py index 1c98b9378cec..b8589c1227cc 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_gpt_oss.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_gpt_oss.py @@ -36,6 +36,23 @@ * No ``repeat_kv`` (``torch_attention`` handles GQA natively) * RoPE cos/sin is computed once per forward and pre-sliced by ``position_ids`` * The HF config class ``GptOssConfig`` is reused directly from ``transformers`` + +Sharding-IR convention: every attention Linear is expressed via ``torch.ops.auto_deploy.torch_linear_simple`` +with sharding hint kwargs (``tp_mode``, ``tp_min_local_shape``, ``layer_type``), +and the post-attention all-reduce uses the ``auto_deploy.all_reduce`` placeholder. +The exported graph is a self-contained spec of how attention should be TP-sharded; +``apply_sharding_hints`` reads those hints + a runtime ``DistConfig`` to produce +deterministic, node-local sharding. + + * Attention q/k/v/o: ``torch_linear_simple`` (q/k/v colwise + + ``tp_min_local_shape=head_dim`` for GQA, o rowwise) + trailing ``all_reduce``. + * q/k/v/attn_out views use ``auto_deploy.view`` with ``tp_scaled_dim=2`` so the + head-count dimension scales with TP. + * MoE router + experts stay replicated under sharding-IR; EP/TP-MoE for the + trtllm-gen path is applied later by a separate ``ShardableNode``. + * ``lm_head`` stays as a plain ``nn.Linear`` — no canonical sharding-IR pattern + for col-parallel-linear-then-all-gather, and the gain is marginal + (~80 us/token at TP=4 for gpt-oss-120b). """ import math @@ -48,6 +65,8 @@ from transformers.modeling_utils import PreTrainedModel from transformers.utils import ModelOutput +from tensorrt_llm._utils import get_hf_rope_theta + from ..hf import AutoModelForCausalLMFactory # GPT-OSS hard-codes these in the HF reference (see modeling_gpt_oss.GptOssExperts). @@ -243,10 +262,9 @@ class GptOssExperts(nn.Module): custom GLU: ``(up + 1) * gate * sigmoid(alpha * gate)`` with clamps on gate (max=limit) and up (-limit, limit). - The MXFP4 quantization path replaces this op (and the upstream router op) - with ``triton_mxfp4_moe`` in the AD ``quantize_mxfp4_moe`` graph transform; - the ``_blocks`` / ``_scales`` parameters are registered there at transform - time so we do not declare them here. + Quantization (MXFP4 → Triton / TRT-LLM-Gen) is handled by the + ``quantize_mxfp4_moe`` transform, which rewrites the FX graph and swaps + parameters at PATTERN_MATCHER time (see :mod:`...transform.library.fused_moe_mxfp4`). """ def __init__(self, config): @@ -297,12 +315,21 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # --------------------------------------------------------------------------- -# Attention (GQA + sinks + per-layer sliding window) +# Attention (GQA + sinks + per-layer sliding window) -- sharding-IR variant # --------------------------------------------------------------------------- class GptOssAttention(nn.Module): - """GPT-OSS attention with learnable per-head sinks and optional sliding window.""" + """GPT-OSS attention with sharding hints (see module docstring for the + sharding-IR convention). + + Sharding strategy: + q_proj -> colwise (+ tp_min_local_shape=head_dim for GQA) + k_proj -> colwise (+ tp_min_local_shape=head_dim for GQA) + v_proj -> colwise (+ tp_min_local_shape=head_dim for GQA) + view -> tp_scaled_dim=2 (head-count dim shrinks with TP) + o_proj -> rowwise + auto_deploy.all_reduce + """ def __init__(self, config, layer_idx: int): super().__init__() @@ -343,16 +370,59 @@ def forward( ) -> torch.Tensor: bsz, q_len, _ = hidden_states.size() - # Project Q/K/V and reshape to [B, S, N, head_dim] (BSND layout). - q = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim) - k = self.k_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim) - v = self.v_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim) + # Project Q/K/V via torch_linear_simple with colwise sharding hints + # (tp_min_local_shape=head_dim guards GQA where num_kv_heads < tp_size). + q = torch.ops.auto_deploy.torch_linear_simple( + hidden_states, + self.q_proj.weight, + self.q_proj.bias, + tp_mode="colwise", + tp_min_local_shape=self.head_dim, + layer_type="mha", + ) + k = torch.ops.auto_deploy.torch_linear_simple( + hidden_states, + self.k_proj.weight, + self.k_proj.bias, + tp_mode="colwise", + tp_min_local_shape=self.head_dim, + layer_type="mha", + ) + v = torch.ops.auto_deploy.torch_linear_simple( + hidden_states, + self.v_proj.weight, + self.v_proj.bias, + tp_mode="colwise", + tp_min_local_shape=self.head_dim, + layer_type="mha", + ) + + # Reshape to [B, S, N, head_dim] (BSND layout). ``tp_scaled_dim=2`` lets the + # head-count axis shrink with TP after apply_sharding_hints rewrites the view. + q = torch.ops.auto_deploy.view( + q, + [bsz, q_len, self.num_heads, self.head_dim], + tp_scaled_dim=2, + layer_type="mha", + ) + k = torch.ops.auto_deploy.view( + k, + [bsz, q_len, self.num_kv_heads, self.head_dim], + tp_scaled_dim=2, + layer_type="mha", + ) + v = torch.ops.auto_deploy.view( + v, + [bsz, q_len, self.num_kv_heads, self.head_dim], + tp_scaled_dim=2, + layer_type="mha", + ) - cos, sin = position_embeddings # [B, S, head_dim] # Apply RoPE with unsqueeze_dim=2 for BSND layout. + cos, sin = position_embeddings q, k = torch.ops.auto_deploy.torch_rope_with_explicit_cos_sin(q, k, cos, sin, 2) - # ``torch_attention`` handles GQA natively; sinks/sliding_window are + # ``torch_attention`` handles GQA natively; sinks / sliding_window are # per-call kwargs. Causal mask is applied internally for prefill. attn_output = torch.ops.auto_deploy.torch_attention( q, @@ -366,9 +436,25 @@ def forward( sliding_window=self.sliding_window, layout="bsnd", ) + # [B, S, N, D] -> [B, S, N*D] - attn_output = attn_output.reshape(bsz, q_len, -1) - return self.o_proj(attn_output) + attn_output = torch.ops.auto_deploy.view( + attn_output, + [bsz, q_len, self.num_heads * self.head_dim], + tp_scaled_dim=2, + layer_type="mha", + ) + + # o_proj is rowwise; ``apply_sharding_hints`` adds the trailing all_reduce. + attn_output = torch.ops.auto_deploy.torch_linear_simple( + attn_output, + self.o_proj.weight, + self.o_proj.bias, + tp_mode="rowwise", + layer_type="mha", + ) + attn_output = torch.ops.auto_deploy.all_reduce(attn_output, layer_type="mha") + return attn_output # --------------------------------------------------------------------------- @@ -429,7 +515,7 @@ def __init__(self, config): self.rotary_emb = GptOssRotaryEmbedding( head_dim=head_dim, max_position_embeddings=config.max_position_embeddings, - rope_theta=float(getattr(config, "rope_theta", 10000.0)), + rope_theta=get_hf_rope_theta(config, 10000.0), rope_scaling=getattr(config, "rope_scaling", None), ) @@ -462,6 +548,7 @@ class GptOssForCausalLM(GptOssPreTrainedModel, GenerationMixin): def __init__(self, config): super().__init__(config) self.model = GptOssModel(config) + # lm_head stays as plain nn.Linear; see module docstring for rationale. self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.post_init() diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe_mxfp4.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe_mxfp4.py new file mode 100644 index 000000000000..e5782caaf53c --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe_mxfp4.py @@ -0,0 +1,1303 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Literal, Optional, Tuple, Type + +import torch +import torch.nn as nn +from pydantic import Field +from torch.fx import GraphModule, Node + +from ..._compat import get_sm_version +from ...utils.logger import ad_logger +from ...utils.module import get_submodule_of_param +from ...utils.node_utils import is_op +from ...utils.pattern_matcher import ADPatternMatcherPass, register_ad_pattern +from ..interface import BaseTransform, TransformConfig, TransformInfo, TransformRegistry + +# MXFP4 layout constants (mirror the on-disk HF format the trtllm-gen kernel +# consumes). Used by both the load hook below and the TP-aware pre-pad math +# in ``QuantizeMXFP4MOE._apply_trtllm``. +_MXFP4_SCALING_VECTOR_SIZE = 32 +_WEIGHT_ALIGNMENT = 128 + +# Backend selection for MXFP4 MoE quantization. +# - "triton": use the triton_mxfp4_moe kernel (Ampere/Hopper compatible). +# - "trtllm": use the trtllm-gen MXFP4 MoE kernel (Blackwell SM>=100 only). +# When ``backend`` is left unset (``None``) on the transform config, the +# default is auto-resolved from the current SM: ``trtllm`` on SM>=100, +# ``triton`` otherwise. ``backend="trtllm"`` on SM<100 falls back to +# ``triton`` with a warning (silent fallback, not an error). +MxFP4Backend = Literal["triton", "trtllm"] + + +def _moe_dense_mlp_pattern( + hidden_states: torch.Tensor, + routing_weights: torch.Tensor, + gate_up_w: torch.Tensor, + gate_up_b: torch.Tensor, + down_w: torch.Tensor, + down_b: torch.Tensor, + alpha: float = 1.0, + limit: float = 10.0, + minus_limit: float = -10.0, +) -> torch.Tensor: + batch_size = hidden_states.shape[0] + hidden_size = hidden_states.shape[2] + hidden_states = hidden_states.reshape(-1, hidden_size) # (num_tokens, hidden_size) + num_experts = routing_weights.shape[1] + + hidden_states = hidden_states.repeat(num_experts, 1) + hidden_states = hidden_states.view(num_experts, -1, hidden_size) + gate_up = torch.bmm(hidden_states, gate_up_w) + gate_up_b.unsqueeze(-2) + gate, up = gate_up[..., ::2], gate_up[..., 1::2] + gate = gate.clamp(min=None, max=limit) + up = up.clamp(min=minus_limit, max=limit) + glu = gate * torch.sigmoid(gate * alpha) + next_states = torch.bmm(((up + 1) * glu), down_w) + next_states = next_states + down_b.unsqueeze(-2) + next_states = next_states.view(num_experts, batch_size, -1, hidden_size) + next_states = ( + next_states * routing_weights.transpose(0, 1).view(num_experts, batch_size, -1)[..., None] + ) + next_states = next_states.sum(dim=0) # [B, S, H] + return next_states + + +def _moe_dense_mlp_repl( + hidden_states: torch.Tensor, + routing_weights: torch.Tensor, + gate_up_w: torch.Tensor, + gate_up_b: torch.Tensor, + down_w: torch.Tensor, + down_b: torch.Tensor, + alpha: float, + limit: float, + minus_limit: float, +) -> torch.Tensor: + return torch.ops.auto_deploy.torch_moe_dense_mlp( + hidden_states, routing_weights, gate_up_w, gate_up_b, down_w, down_b, alpha, limit + ) + + +@TransformRegistry.register("match_dense_moe_pattern") +class MatchMXFP4MoePattern(BaseTransform): + def _apply( + self, + gm: GraphModule, + cm, + factory, + shared_config, + ) -> Tuple[GraphModule, TransformInfo]: + graph = gm.graph + patterns = ADPatternMatcherPass() + + B, S, H = 2, 4, 8 # batch, seq, hidden + E, In = 3, 16 # experts, intermediate (I); gate_up has 2I + T = B * S + + dummy_args = [ + torch.randn(B, S, H, device="meta", dtype=torch.float16), # hidden_states + torch.randn(T, E, device="meta", dtype=torch.float16), # routing_weights + torch.randn(E, H, 2 * In, device="meta", dtype=torch.float16), # gate_up_w [E,H,2I] + torch.randn(E, 2 * In, device="meta", dtype=torch.float16), # gate_up_b [E,2I] + torch.randn(E, In, H, device="meta", dtype=torch.float16), # down_w [E,I,H] + torch.randn(E, H, device="meta", dtype=torch.float16), # down_b [E,H] + 1.07, + 10.1, + -10.1, + ] + + op_ignore_types = { + torch.ops.aten.view.default: (int,), + torch.ops.aten.reshape.default: (int,), + torch.ops.auto_deploy.view.default: (int,), + torch.ops.aten.repeat.default: (int,), + torch.ops.aten.slice.Tensor: (int,), + torch.ops.aten.unsqueeze.default: (int,), + torch.ops.aten.transpose.int: (int,), + } + + scalar_workaround = {"alpha": 1.07, "limit": 10.1, "minus_limit": -10.1} + + register_ad_pattern( + search_fn=_moe_dense_mlp_pattern, + replace_fn=_moe_dense_mlp_repl, + patterns=patterns, + dummy_args=dummy_args, + op_ignore_types=op_ignore_types, + scalar_workaround=scalar_workaround, + ) + + num_matches = patterns.apply(graph) + info = TransformInfo( + skipped=False, + num_matches=num_matches, + is_clean=num_matches == 0, + has_valid_shapes=num_matches == 0, + ) + return gm, info + + +def _get_alpha_limit_from_dense(node: Node) -> Tuple[float, float]: + # torch_moe_dense_mlp(hidden, routing, gu_w, gu_b, dn_w, dn_b, alpha, limit) + # alpha/limit may be in args or kwargs + alpha = node.kwargs.get("alpha", None) + limit = node.kwargs.get("limit", None) + if alpha is None: + alpha = float(node.args[6]) if len(node.args) >= 7 else 1.0 + if limit is None: + limit = float(node.args[7]) if len(node.args) >= 8 else 10.0 + return float(alpha), float(limit) + + +def _get_topk_from_router(node: Node) -> int: + # torch_moe_router(hidden, weight, bias, top_k=2) + if "top_k" in node.kwargs: + return int(node.kwargs["top_k"]) + return int(node.args[3]) if len(node.args) >= 4 else 2 + + +def _register_mxfp4_expert_params( + gm: GraphModule, + gate_up_w_name: str, + gate_up_b_name: str, + down_w_name: str, + down_b_name: str, +) -> Tuple[str, str, str, str]: + """Create (if missing) the four MXFP4 params under the experts module and return their full names. + + Returns: + (gu_blocks_name, gu_scales_name, dn_blocks_name, dn_scales_name) + """ + # Shapes from existing params + gu_b = gm.get_parameter(gate_up_b_name) # [E, 2I] + gu_w = gm.get_parameter(gate_up_w_name) # [E, 2I, H] + dn_b = gm.get_parameter(down_b_name) # [E, H] + + E = int(gu_b.shape[0]) + I2 = int(gu_b.shape[1]) # 2I + In = I2 // 2 + + # infer H from gu_w shape + assert gu_w.dim() == 3, "gate_up_w must be rank-3" + if gu_w.shape[1] == I2: + H = int(gu_w.shape[2]) + else: + # Fallback: use down bias last dim + H = int(dn_b.shape[1]) + + # Compute block dims (assume divisible; zero-init anyway) + H_blk = max(1, H // 32) + I_blk = max(1, In // 32) + + experts_mod, experts_path, _ = get_submodule_of_param(gm, gate_up_w_name) + + # New param names under experts module + gu_blocks_name = "gate_up_proj_blocks" + gu_scales_name = "gate_up_proj_scales" + dn_blocks_name = "down_proj_blocks" + dn_scales_name = "down_proj_scales" + + # Uninitialized placeholders — names match HF safetensors so the standard + # state_dict load path overwrites them. Reuse the existing param's device + # (meta in the normal meta-device build) so we don't materialize giant CPU + # buffers before load. + param_device = gu_w.device + gu_blocks = torch.empty((E, 2 * In, H_blk, 16), dtype=torch.uint8, device=param_device) + gu_scales = torch.empty((E, 2 * In, H_blk), dtype=torch.uint8, device=param_device) + dn_blocks = torch.empty((E, H, I_blk, 16), dtype=torch.uint8, device=param_device) + dn_scales = torch.empty((E, H, I_blk), dtype=torch.uint8, device=param_device) + + experts_mod.register_parameter(gu_blocks_name, nn.Parameter(gu_blocks, requires_grad=False)) + experts_mod.register_parameter(gu_scales_name, nn.Parameter(gu_scales, requires_grad=False)) + experts_mod.register_parameter(dn_blocks_name, nn.Parameter(dn_blocks, requires_grad=False)) + experts_mod.register_parameter(dn_scales_name, nn.Parameter(dn_scales, requires_grad=False)) + + # Free the now-unused bf16 stacked weight params; the biases are still + # consumed by ``triton_mxfp4_moe`` and must remain. + gu_w_local = gate_up_w_name.split(".")[-1] + dn_w_local = down_w_name.split(".")[-1] + for local_name in (gu_w_local, dn_w_local): + if local_name in experts_mod._parameters: + del experts_mod._parameters[local_name] + + # Full GM attribute paths for new params + prefix = (experts_path + ".") if experts_path else "" + return ( + prefix + gu_blocks_name, + prefix + gu_scales_name, + prefix + dn_blocks_name, + prefix + dn_scales_name, + ) + + +# EP+TP load hook — slices raw HF MXFP4 state_dict tensors on CPU before copy +# so per-rank module params (registered at the EP-sliced shape) accept them. +# Kernel-layout prep is deferred to FuseMXFP4Moe on GPU. + + +def make_mxfp4_sharding_load_hook( + *, + num_layers: int, + num_experts: int, + intermediate_size: int, + moe_ep_size: int, + moe_ep_rank: int, + moe_tp_size: int, + moe_tp_rank: int, + layer_prefix: str = "model.layers", + experts_subpath: str = "mlp.experts", +): + """Build a ``load_state_dict`` pre-hook that EP+TP-shards raw HF MXFP4 keys. + + For each layer's six raw HF MXFP4 keys + (``gate_up_proj_{blocks,scales,bias}``, ``down_proj_{blocks,scales,bias}``) + the hook slices on CPU before copy so per-rank GPU memory only holds this + rank's shard. Kernel-layout prep (H-pad, TMA shuffle, bias dtype/scale) is + deferred to ``FuseMXFP4Moe`` on GPU. + + Slicing axes: + + 1. **EP (leading expert axis)** — ``t[ep_start:ep_stop]`` where + ``experts_per_rank = num_experts / moe_ep_size``. No-op when + ``moe_ep_size == 1``. + + 2. **TP (intermediate axis)** — only when ``moe_tp_size > 1``. ``I`` is + padded to ``i_padded_tp`` so that ``per_rank_i = i_padded_tp / + moe_tp_size`` is a multiple of 128 (TMA weight alignment), then: + + * ``gate_up_proj_*`` — axis 1 of the interleaved 2I layout, + ``[2*tp_start : 2*tp_stop]``. Alternating gate(k)/up(k) means a + contiguous slice covers ``gate(k:k+m) ∪ up(k:k+m)``. + * ``down_proj_{blocks,scales}`` — axis 2 (``I_blk = I/32``), + ``[tp_start/32 : tp_stop/32]``. ``per_rank_i`` is a multiple of 32. + * ``down_proj_bias`` ``[E, H]`` is left intact (H not TP-split); + ``FuseMXFP4Moe`` divides it by ``moe_tp_size`` after dtype convert. + + Args: + num_layers: number of decoder layers to scan. + num_experts: total expert count on disk. + intermediate_size: per-expert intermediate dim ``I`` on disk + (before any padding/slicing). + moe_ep_size / moe_ep_rank: expert-parallel group size + this rank. + moe_tp_size / moe_tp_rank: MoE tensor-parallel group size + this rank + (intermediate-axis split). + layer_prefix: where layers live, default ``"model.layers"``. + experts_subpath: where the experts module sits within each layer, + default ``"mlp.experts"``. + + Returns: + A hook with the standard ``(state_dict, prefix, ...)`` signature. + """ + if num_experts % moe_ep_size != 0: + raise ValueError( + f"num_experts ({num_experts}) must be divisible by moe_ep_size ({moe_ep_size})" + ) + experts_per_rank = num_experts // moe_ep_size + ep_start = moe_ep_rank * experts_per_rank + ep_stop = ep_start + experts_per_rank + + # TP-aware pre-pad/slice math (only used when moe_tp_size > 1). + if moe_tp_size > 1: + # Lazy import: TRT-LLM-only helper. Keeps this module importable in + # standalone (no tensorrt_llm) so its transforms still register. + from tensorrt_llm._torch.modules.fused_moe.quantization import _get_weight_alignment + + alignment_tp = _get_weight_alignment( + _WEIGHT_ALIGNMENT, _MXFP4_SCALING_VECTOR_SIZE, moe_tp_size, intermediate_size + ) + i_padded_tp = ((intermediate_size + alignment_tp - 1) // alignment_tp) * alignment_tp + per_rank_i = i_padded_tp // moe_tp_size + tp_start = moe_tp_rank * per_rank_i + tp_stop = (moe_tp_rank + 1) * per_rank_i + if per_rank_i % _MXFP4_SCALING_VECTOR_SIZE != 0: + raise ValueError( + f"per_rank_i ({per_rank_i}) must be divisible by " + f"_MXFP4_SCALING_VECTOR_SIZE ({_MXFP4_SCALING_VECTOR_SIZE}); " + f"check _get_weight_alignment output." + ) + # Block-axis bounds for down_proj's I_blk = I / 32 axis. + blk_pad = i_padded_tp // _MXFP4_SCALING_VECTOR_SIZE + blk_start = tp_start // _MXFP4_SCALING_VECTOR_SIZE + blk_stop = tp_stop // _MXFP4_SCALING_VECTOR_SIZE + else: + i_padded_tp = intermediate_size + per_rank_i = intermediate_size + tp_start = 0 + tp_stop = intermediate_size + blk_pad = intermediate_size // _MXFP4_SCALING_VECTOR_SIZE + blk_start = 0 + blk_stop = blk_pad + + def _pad_axis(t: torch.Tensor, dim: int, target: int) -> torch.Tensor: + cur = t.shape[dim] + if cur >= target: + return t + pad_amount = target - cur + # F.pad spec is (pad_lastdim_left, pad_lastdim_right, ..., pad_dim_left, pad_dim_right) + pad = [0, 0] * (t.dim() - dim - 1) + [0, pad_amount] + [0, 0] * dim + return torch.nn.functional.pad(t, pad) + + def hook(state_dict, prefix, *args, local_metadata=None, **kwargs): + do_ep = moe_ep_size > 1 + do_tp = moe_tp_size > 1 + if not (do_ep or do_tp): + # Nothing to slice — leave state_dict alone. + return + for layer_idx in range(num_layers): + base = f"{prefix}{layer_prefix}.{layer_idx}.{experts_subpath}." + + # ---- EP slice (leading expert axis) ---- + if do_ep: + for s in ( + "gate_up_proj_blocks", + "gate_up_proj_scales", + "gate_up_proj_bias", + "down_proj_blocks", + "down_proj_scales", + "down_proj_bias", + ): + k = base + s + t = state_dict.get(k) + if t is None: + continue + state_dict[k] = t[ep_start:ep_stop].contiguous() + + # ---- TP-aware pre-pad + slice (intermediate axis) ---- + if do_tp: + # gate_up_*: axis 1 (the 2I interleaved axis); pad to + # 2*i_padded_tp, then slice [2*tp_start : 2*tp_stop]. + for s in ("gate_up_proj_blocks", "gate_up_proj_scales", "gate_up_proj_bias"): + k = base + s + t = state_dict.get(k) + if t is None: + continue + t = _pad_axis(t, 1, 2 * i_padded_tp) + state_dict[k] = t[:, 2 * tp_start : 2 * tp_stop].contiguous() + + # down_proj_blocks / scales: axis 2 (I_blk = I / 32); pad to + # blk_pad, then slice [blk_start : blk_stop]. Inner 16 axis + # (blocks only) is untouched. + for s in ("down_proj_blocks", "down_proj_scales"): + k = base + s + t = state_dict.get(k) + if t is None: + continue + t = _pad_axis(t, 2, blk_pad) + state_dict[k] = t[:, :, blk_start:blk_stop].contiguous() + # down_proj_bias [E, H]: H axis is not TP-split. Leave as-is + # and let FuseMXFP4Moe divide by moe_tp_size after dtype + # conversion (matches the prep helper's tp-aware bias path). + + return hook + + +def make_swiglu_param_tensors( + num_local_experts: int, + *, + alpha: float = 1.702, + beta: float = 1.0, + limit: float = 7.0, + device: torch.device | str | None = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the per-expert SwiGLU-bias parameter triple expected by the kernel. + + These constants are NOT in HF safetensors (only in the model config), so + the transform constructs them here and registers them as ``nn.Parameter`` + on the experts module. For gpt-oss-120b: alpha=1.702, beta=1.0, limit=7.0. + """ + dev = torch.device(device) if device is not None else None + a = torch.full((num_local_experts,), alpha, dtype=torch.float32, device=dev) + b = torch.full((num_local_experts,), beta, dtype=torch.float32, device=dev) + c = torch.full((num_local_experts,), limit, dtype=torch.float32, device=dev) + return a, b, c + + +class QuantizeMXFP4MOEConfig(TransformConfig): + """Configuration for ``quantize_mxfp4_moe``.""" + + backend: Optional[MxFP4Backend] = Field( + default=None, + description=( + "MXFP4 MoE kernel backend selection. When unset (``None``), the " + "default is SM-based: ``trtllm`` on SM>=100 (Blackwell), ``triton`` " + "otherwise. Explicit ``triton`` or ``trtllm`` overrides the default. " + "``trtllm`` on SM<100 silently falls back to ``triton`` with a warning." + ), + ) + trtllm_quant_act: Literal["bf16", "mxfp8"] = Field( + default="mxfp8", + description=( + "Only used when ``backend='trtllm'``. Activation precision for the trtllm-gen " + "MoE GEMM, passed as ``act_dtype`` to " + "``trtllm_quant_mxfp4_trtllm_gen_moe_fused``: ``bf16`` dispatches to the bf16 " + "MoE runner (W4A16), ``mxfp8`` pre-quantizes the activation to MXFP8 and " + "dispatches to the MXFP8 MoE runner (W4A8, faster cubin family). " + "Default ``mxfp8`` matches the modeling-side default." + ), + ) + + +@TransformRegistry.register("quantize_mxfp4_moe") +class QuantizeMXFP4MOE(BaseTransform): + """Quantize MXFP4 MoE: dispatch to triton or trtllm-gen backend. + + Replaces ``(torch_moe_router -> torch_moe_dense_mlp)`` with a single fused + MoE op. The chosen backend determines the destination op and the parameter + layout registered on the experts module: + + * ``backend="triton"`` → ``auto_deploy::triton_mxfp4_moe`` with raw HF + MXFP4 layout (``_blocks`` / ``_scales`` / ``_bias``). Lazy weight + swizzling happens inside the Triton kernel on first forward. + * ``backend="trtllm"`` → ``auto_deploy::trtllm_quant_mxfp4_*_moe_fused`` with + trtllm-gen prepared layout (``fc1_w_trtllm`` / ``fc1_w_scale_trtllm`` / + ...). Weight preparation (shuffle + interleave) is done on CPU inside + a state-dict pre-hook registered by this transform, so the raw HF + tensors are converted before being moved to GPU. + """ + + algo_name: str = "mxfp4" + config: QuantizeMXFP4MOEConfig + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return QuantizeMXFP4MOEConfig + + def _resolve_backend(self) -> MxFP4Backend: + """Resolve the effective backend from config + runtime SM. + + - ``config.backend is None`` → SM-based default + * SM>=100 → ``trtllm`` + * SM<100 → ``triton`` + - ``config.backend="trtllm"`` + SM<100 → warn + fallback to ``triton`` + - Otherwise honour the explicit config value. + """ + requested = self.config.backend + sm = get_sm_version() + if requested is None: + return "trtllm" if sm >= 100 else "triton" + if requested == "trtllm" and sm < 100: + ad_logger.warning( + f"quantize_mxfp4_moe: backend='trtllm' requires SM>=100 (Blackwell), " + f"but current SM={sm}. Falling back to backend='triton'." + ) + return "triton" + return requested + + def _apply( + self, + gm: GraphModule, + cm, + factory, + shared_config, + ) -> Tuple[GraphModule, TransformInfo]: + """Dispatcher: pick a backend and delegate to the corresponding method. + + The actual graph rewrite + parameter swap lives in + :meth:`_apply_triton` / :meth:`_apply_trtllm`. This method only: + 1. Skips if quant_method != "mxfp4". + 2. Resolves the backend (``triton`` | ``trtllm``) and dispatches. + """ + qcfg = factory.get_quant_config() + if not qcfg or qcfg.get("quant_method", "") != self.algo_name: + return gm, TransformInfo( + skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True + ) + + backend = self._resolve_backend() + ad_logger.info(f"quantize_mxfp4_moe: dispatching to backend={backend!r}") + + if backend == "triton": + return self._apply_triton(gm, cm, factory, shared_config) + elif backend == "trtllm": + return self._apply_trtllm(gm, cm, factory, shared_config) + else: + # _resolve_backend should only return "triton" or "trtllm". + raise ValueError(f"Unexpected backend resolved: {backend!r}") + + def _apply_triton( + self, + gm: GraphModule, + cm, + factory, + shared_config, + ) -> Tuple[GraphModule, TransformInfo]: + """Triton backend: graph rewrite to ``triton_mxfp4_moe``. + + Replaces ``(torch_moe_router -> torch_moe_dense_mlp)`` with a single + ``auto_deploy::triton_mxfp4_moe`` op and registers raw HF-layout MXFP4 params + (``_blocks`` / ``_scales``) on the experts module via :func:`_register_mxfp4_expert_params`. + The bf16 placeholders (``gate_up_proj`` / ``down_proj``) are deleted; biases are kept. + + Weight swizzling for the Triton kernel happens lazily inside the kernel on first forward + (see ``_prepare_weights_scales_cached`` in ``custom_ops/fused_moe/mxfp4_moe.py``) -- no + load hook needed because the HF state-dict keys already match the registered param names + (``gate_up_proj_blocks``, ``gate_up_proj_scales``, etc.). + """ + num_matches = 0 + + for n in list(gm.graph.nodes): + if not is_op(n, torch.ops.auto_deploy.torch_moe_dense_mlp): + continue + + # Expect: torch_moe_dense_mlp(hidden, routing, gu_w, gu_b, dn_w, dn_b, alpha, limit) + if len(n.args) < 6: + continue + + hidden_node = n.args[0] + routing_node = n.args[1] + gate_up_w_node = n.args[2] + gate_up_b_node = n.args[3] + down_w_node = n.args[4] + down_b_node = n.args[5] + + if not isinstance(routing_node, Node) or not is_op( + routing_node, torch.ops.auto_deploy.torch_moe_router + ): + continue + + # Router params: weight, bias, top_k + router_weight_node = routing_node.args[1] + router_bias_node = routing_node.args[2] + top_k = _get_topk_from_router(routing_node) + + # Resolve parameter names so we can find the experts module + if gate_up_w_node.op != "get_attr" or gate_up_b_node.op != "get_attr": + continue + if down_w_node.op != "get_attr" or down_b_node.op != "get_attr": + continue + + gu_w_name = gate_up_w_node.target + gu_b_name = gate_up_b_node.target + dn_w_name = down_w_node.target + dn_b_name = down_b_node.target + + # Register MXFP4 params on experts + gu_blocks_name, gu_scales_name, dn_blocks_name, dn_scales_name = ( + _register_mxfp4_expert_params(gm, gu_w_name, gu_b_name, dn_w_name, dn_b_name) + ) + + # Alpha/limit (from dense call) + alpha, limit = _get_alpha_limit_from_dense(n) + + # Insert the new get_attr nodes for MXFP4 params + with gm.graph.inserting_before(n): + gu_blocks_attr = gm.graph.create_node("get_attr", gu_blocks_name) + gu_scales_attr = gm.graph.create_node("get_attr", gu_scales_name) + dn_blocks_attr = gm.graph.create_node("get_attr", dn_blocks_name) + dn_scales_attr = gm.graph.create_node("get_attr", dn_scales_name) + + n.target = torch.ops.auto_deploy.triton_mxfp4_moe.default + n.kwargs = {} + + # triton_mxfp4_moe( + # hidden_states, + # router_weight, router_bias, top_k, + # gate_up_blocks, gate_up_bias, gate_up_scales, alpha, limit, + # down_blocks, down_bias, down_scales) + new_args = ( + hidden_node, + router_weight_node, + router_bias_node, + top_k, + gu_blocks_attr, + gate_up_b_node, + gu_scales_attr, + float(alpha), + float(limit), + dn_blocks_attr, + down_b_node, + dn_scales_attr, + ) + n.args = new_args + + # Remove the now-unneeded router node if nobody else uses it + if len(routing_node.users) == 0: + gm.graph.erase_node(routing_node) + + # Erase the old get_attr nodes for gate_up_proj and down_proj. + # _register_mxfp4_expert_params deleted those attributes from the + # experts module, so these nodes now reference non-existent attrs. + # They have no users after the args replacement above, so it is + # safe to erase them directly. + for stale_node in (gate_up_w_node, down_w_node): + if len(stale_node.users) == 0: + gm.graph.erase_node(stale_node) + + num_matches += 1 + + info = TransformInfo( + skipped=(num_matches == 0), + num_matches=num_matches, + is_clean=num_matches == 0, + has_valid_shapes=num_matches == 0, + ) + return gm, info + + def _apply_trtllm( + self, + gm: GraphModule, + cm, + factory, + shared_config, + ) -> Tuple[GraphModule, TransformInfo]: + """TRT-LLM-Gen backend: graph rewrite + raw HF param registration. + + Per MoE node: + + 1. Find ``torch_moe_dense_mlp`` + its upstream ``torch_moe_router``. + 2. Look up the experts module that owns the bf16 placeholder params. + 3. Delete the bf16 placeholders (``gate_up_proj`` / ``down_proj`` / biases). + 4. Register **raw HF MXFP4 params** at the EP-sliced shape + (``E_local = E_full / moe_ep_size``) on the experts module: + ``gate_up_proj_{blocks,scales,bias}`` and ``down_proj_{blocks,scales,bias}``. Names match + HF safetensors so the standard ``load_state_dict`` path can populate them (after the + slim EP-slice hook below trims the leading expert axis when ``moe_ep_size > 1``). + 5. Also register the per-expert SwiGLU constants (``swiglu_alpha_trtllm`` / beta / limit) — + these are not in HF safetensors so they are populated with their numeric defaults at + registration time. + 6. Rewrite the ``torch_moe_dense_mlp`` node to + ``trtllm_quant_mxfp4_trtllm_gen_moe_fused`` (with ``act_dtype`` set from + ``config.trtllm_quant_act``) with args pointing at the **raw** params for + now. The downstream :class:`FuseMXFP4Moe` POST_LOAD_FUSION transform will run + :func:`prepare_trtllm_gen_moe_mxfp4_weights` on the actually-loaded GPU tensors, + register prepared-shape params, and re-point the op args. The op call is therefore not + runnable between PATTERN_MATCHER and POST_LOAD_FUSION, but no forward pass happens in + that window. + 7. If ``tp_size > 1`` insert an ``auto_deploy.all_reduce`` node after the downstream view + (covers both MoE-TP and MoE-EP). + + Then once for the whole module: + + 8. Register a top-level ``load_state_dict`` pre-hook + (:func:`make_mxfp4_sharding_load_hook`) that slices raw HF MXFP4 tensors on the expert + axis when ``moe_ep_size > 1``. The hook does **not** run any kernel-layout prep — + that runs on GPU in :class:`FuseMXFP4Moe` after the weights are loaded. + """ + import re + + # MoE topology comes from the build-time ``DistConfig`` on + # ``shared_config``; passed directly into the sharding load hook + # below. + dc = getattr(shared_config, "dist_config", None) + moe_tp_size = int(getattr(dc, "moe_tp_size", 1)) if dc is not None else 1 + moe_tp_rank = int(getattr(dc, "moe_tp_rank", 0)) if dc is not None else 0 + moe_ep_size = int(getattr(dc, "moe_ep_size", 1)) if dc is not None else 1 + moe_ep_rank = int(getattr(dc, "moe_ep_rank", 0)) if dc is not None else 0 + # Cover MoE-EP as well: any distributed case (tp_size>1) needs the + # configured strategy. ``moe_tp_size > 1`` alone would miss EP-only. + _tp_size = int(getattr(dc, "tp_size", 1)) if dc is not None else 1 + allreduce_strategy = ( + str(dc.allreduce_strategy) if dc is not None and _tp_size > 1 else "NCCL" + ) + + # Single op handles both activation precisions via the ``act_dtype`` arg: + # ``"bf16"`` → W4A16 (bf16 MoE runner), ``"mxfp8"`` → W4A8 (mxfp8_quantize + + # MXFP8 MoE runner). + target_op = torch.ops.auto_deploy.trtllm_quant_mxfp4_trtllm_gen_moe_fused.default + quant_act = self.config.trtllm_quant_act + + # Module-level info needed once for the load hook factory. + hidden_size_global: Optional[int] = None + intermediate_size_global: Optional[int] = None + num_experts_global: Optional[int] = None + layer_indices: list = [] + + layer_re = re.compile(r"\.layers\.(\d+)\.") + num_matches = 0 + + for n in list(gm.graph.nodes): + if not is_op(n, torch.ops.auto_deploy.torch_moe_dense_mlp): + continue + # Expect: torch_moe_dense_mlp(hidden, routing, gu_w, gu_b, dn_w, dn_b, alpha, limit) + if len(n.args) < 6: + continue + + hidden_node = n.args[0] + routing_node = n.args[1] + gate_up_w_node = n.args[2] + gate_up_b_node = n.args[3] + down_w_node = n.args[4] + down_b_node = n.args[5] + + if not isinstance(routing_node, Node) or not is_op( + routing_node, torch.ops.auto_deploy.torch_moe_router + ): + continue + if ( + gate_up_w_node.op != "get_attr" + or gate_up_b_node.op != "get_attr" + or down_w_node.op != "get_attr" + or down_b_node.op != "get_attr" + ): + continue + + router_weight_node = routing_node.args[1] + router_bias_node = routing_node.args[2] + top_k = _get_topk_from_router(routing_node) + + gu_w_name = gate_up_w_node.target + dn_w_name = down_w_node.target + + # Shapes from the bf16 placeholders (meta is fine — only .shape is read). + # gu_w shape: [E, H, 2I]; dn_w shape: [E, I, H] (we infer I from gu_w). + gu_w_t = gm.get_parameter(gu_w_name) + E_full = int(gu_w_t.shape[0]) + H = int(gu_w_t.shape[1]) + two_I = int(gu_w_t.shape[2]) + i_size = two_I // 2 + + # Cross-layer consistency check (the load hook is registered once + # for the whole module, so all layers must share these). + if hidden_size_global is None: + hidden_size_global = H + intermediate_size_global = i_size + num_experts_global = E_full + else: + if (H, i_size, E_full) != ( + hidden_size_global, + intermediate_size_global, + num_experts_global, + ): + raise ValueError( + f"quantize_mxfp4_moe(backend=trtllm): inconsistent MoE shapes " + f"across layers (got H={H}, I={i_size}, E={E_full}; previously " + f"H={hidden_size_global}, I={intermediate_size_global}, " + f"E={num_experts_global}). All MoE layers must share shape." + ) + + if E_full % moe_ep_size != 0: + raise ValueError( + f"num_experts ({E_full}) must be divisible by moe_ep_size ({moe_ep_size})" + ) + e_local = E_full // moe_ep_size + + # Locate the experts module via the gate_up param path. + experts_mod, experts_path, _ = get_submodule_of_param(gm, gu_w_name) + + # Per-rank dims after EP+TP slicing. The kernel-layout work + # (H-axis pad, TMA shuffle, dtype convert) is deferred to + # ``FuseMXFP4Moe`` at POST_LOAD_FUSION on GPU. EP+TP sharding is + # done on CPU inside the load hook (see + # :func:`make_mxfp4_sharding_load_hook`). + h_blk = max(1, H // 32) + + # TP-aware pre-pad math (mirrors the hook). The hook pads the raw + # intermediate axis to ``i_padded_tp`` then slices ``per_rank_i`` + # rows; ``per_rank_i`` is guaranteed to be a multiple of 128 by + # ``_get_weight_alignment``, so it's also the per-rank kernel + # weight-alignment size that the trtllm-gen runner expects. + if moe_tp_size > 1: + # Lazy import: TRT-LLM-only helper (see module-level note above). + from tensorrt_llm._torch.modules.fused_moe.quantization import _get_weight_alignment + + alignment_tp = _get_weight_alignment( + _WEIGHT_ALIGNMENT, _MXFP4_SCALING_VECTOR_SIZE, moe_tp_size, i_size + ) + i_padded_tp = ((i_size + alignment_tp - 1) // alignment_tp) * alignment_tp + per_rank_i = i_padded_tp // moe_tp_size + slice_start = moe_tp_rank * per_rank_i + slice_stop = (moe_tp_rank + 1) * per_rank_i + # ``valid_intermediate_size`` reports the unpadded portion of + # this rank's slice — used by the kernel to mask OOB MMA in + # padded regions. + valid_intermediate_size = max(0, min(i_size, slice_stop) - slice_start) + else: + per_rank_i = i_size + valid_intermediate_size = i_size + + # Local I block-count for down_proj after TP slicing. + i_blk_local = max(1, per_rank_i // 32) + two_i_local = 2 * per_rank_i # gate_up's 2I axis is per-rank too + + num_local_experts = e_local + local_expert_offset = moe_ep_rank * e_local + valid_hidden_size = H + + # Register RAW HF MXFP4 params at the EP+TP-sliced shape — names + # match HF safetensors so the standard load path populates them + # after the sharding hook does the leading-axis (EP) + intermediate + # (TP) slice on the state-dict tensors. + raw_specs = [ + ("gate_up_proj_blocks", (e_local, two_i_local, h_blk, 16), torch.uint8), + ("gate_up_proj_scales", (e_local, two_i_local, h_blk), torch.uint8), + ("gate_up_proj_bias", (e_local, two_i_local), torch.bfloat16), + ("down_proj_blocks", (e_local, H, i_blk_local, 16), torch.uint8), + ("down_proj_scales", (e_local, H, i_blk_local), torch.uint8), + ("down_proj_bias", (e_local, H), torch.bfloat16), + ] + # Reuse the existing placeholder's device (meta in the normal + # meta-device build) so we don't materialize giant CPU buffers + # before load_weights runs. Safe because names match HF + # safetensors and the load path overwrites the bytes. + param_device = gu_w_t.device + for name, shape, dtype in raw_specs: + experts_mod.register_parameter( + name, + nn.Parameter( + torch.empty(shape, dtype=dtype, device=param_device), + requires_grad=False, + ), + ) + + # SwiGLU constants. These are NOT in HF safetensors, so we set + # them with their numeric defaults here (matches gpt-oss config: + # alpha=1.702, beta=1.0, limit=7.0). The kernel expects fp32 + # tensors of length ``num_local_experts``. + # IMPORTANT: do NOT pass ``device=param_device`` here — on a + # meta-device build that would create meta tensors and the + # constants (1.702 / 1.0 / 7.0) would be lost; nothing later + # re-injects them, which silently breaks SwiGLU and tanks + # accuracy. + a, b, c = make_swiglu_param_tensors(num_local_experts) + experts_mod.register_parameter( + "swiglu_alpha_trtllm", nn.Parameter(a, requires_grad=False) + ) + experts_mod.register_parameter( + "swiglu_beta_trtllm", nn.Parameter(b, requires_grad=False) + ) + experts_mod.register_parameter( + "swiglu_limit_trtllm", nn.Parameter(c, requires_grad=False) + ) + + # Track layer index so the load hook iterates the right range. + m = layer_re.search(experts_path or "") + if m: + layer_indices.append(int(m.group(1))) + + # Build get_attr nodes for the RAW params (will be replaced by + # ``FuseMXFP4Moe`` once GPU-side prep produces the kernel layout). + prefix_path = (experts_path + ".") if experts_path else "" + with gm.graph.inserting_before(n): + gu_blocks_attr = gm.graph.create_node( + "get_attr", prefix_path + "gate_up_proj_blocks" + ) + gu_scales_attr = gm.graph.create_node( + "get_attr", prefix_path + "gate_up_proj_scales" + ) + gu_bias_attr = gm.graph.create_node("get_attr", prefix_path + "gate_up_proj_bias") + dn_blocks_attr = gm.graph.create_node("get_attr", prefix_path + "down_proj_blocks") + dn_scales_attr = gm.graph.create_node("get_attr", prefix_path + "down_proj_scales") + dn_bias_attr = gm.graph.create_node("get_attr", prefix_path + "down_proj_bias") + sa_attr = gm.graph.create_node("get_attr", prefix_path + "swiglu_alpha_trtllm") + sb_attr = gm.graph.create_node("get_attr", prefix_path + "swiglu_beta_trtllm") + sl_attr = gm.graph.create_node("get_attr", prefix_path + "swiglu_limit_trtllm") + + # Rewrite the op call. Op target is chosen by ``trtllm_quant_act``. + # The op args point at RAW HF MXFP4 buffers for now — the op is + # NOT runnable until ``FuseMXFP4Moe`` (POST_LOAD_FUSION) swaps in + # the prepared layout. That is safe because no forward pass runs + # between PATTERN_MATCHER and POST_LOAD_FUSION. + # Single op trtllm_quant_mxfp4_trtllm_gen_moe_fused; ``act_dtype`` arg + # selects W4A16 (bf16) vs W4A8 (mxfp8) cubin family at runtime. + n.target = target_op + n.kwargs = {} + n.args = ( + hidden_node, + router_weight_node, + router_bias_node, + int(top_k), + gu_blocks_attr, # fc1_weights_mxfp4 (raw uint8; FuseMXFP4Moe replaces) + dn_blocks_attr, # fc2_weights_mxfp4 (raw uint8) + gu_scales_attr, # fc1_weights_scale_ue8m0 (raw uint8) + dn_scales_attr, # fc2_weights_scale_ue8m0 (raw uint8) + gu_bias_attr, # fc1_bias_f32 (raw bf16; FuseMXFP4Moe converts/pads/shuffles) + dn_bias_attr, # fc2_bias_f32 (raw bf16) + sa_attr, + sb_attr, + sl_attr, + valid_hidden_size, + valid_intermediate_size, + quant_act, # act_dtype: "bf16" (W4A16) or "mxfp8" (W4A8) + local_expert_offset, + num_local_experts, + 1, # routing_method_type = RoutingMethodType.Renormalize + ) + + # Distributed MoE: insert an all_reduce after the downstream view so + # the ``MoE -> view -> AR -> add -> norm`` ordering matches + # ``fuse_allreduce_residual_rmsnorm`` (see legacy transform's + # rationale for the same placement). + # + # Both MoE-TP and MoE-EP need an AR after the local MoE op: + # - MoE-TP: each rank computes partial inner-product (summed + # by AR to reconstruct the full intermediate-dim contraction). + # - MoE-EP: each rank computes outputs only for its local + # expert range (zero contribution from other experts); + # AR sums per-token outputs across ranks. + # Use ``tp_size > 1`` (= ``moe_tp_size * moe_ep_size * + # moe_cluster_size > 1``) so the AR fires for any distributed + # configuration. Matches taylor's pre-refactor modeling code + # which emitted an unconditional AR placeholder at this exact + # spot (commit bad1871004 + 93f78e962c, validated EP=2 GSM8K + # 88.02%). + tp_size = int(getattr(dc, "tp_size", 1)) if dc is not None else 1 + if tp_size > 1: + from .sharding import _get_dist_ops + + _, all_reduce_op = _get_dist_ops("auto") + view_node = next( + ( + u + for u in n.users.keys() + if u.op == "call_function" and u.target == torch.ops.aten.view.default + ), + None, + ) + anchor = view_node if view_node is not None else n + with gm.graph.inserting_after(anchor): + red = gm.graph.call_function( + all_reduce_op, + args=(anchor, allreduce_strategy), + ) + anchor.replace_all_uses_with(red) + red.replace_input_with(red, anchor) + + # Erase old router node + stale bf16 get_attr nodes if unused. + if len(routing_node.users) == 0: + gm.graph.erase_node(routing_node) + for stale_node in ( + gate_up_w_node, + gate_up_b_node, + down_w_node, + down_b_node, + ): + if len(stale_node.users) == 0: + gm.graph.erase_node(stale_node) + + # Free bf16 placeholders from the experts module so they don't + # linger as orphaned attributes (and don't get loaded from HF + # via the standard load_state_dict path). Skip the *bias* names + # because we re-registered them with the SAME names as new raw + # HF MXFP4 params (``gate_up_proj_bias`` / ``down_proj_bias``); + # those are the ones we want to keep, not delete. Only the + # ``gate_up_proj`` / ``down_proj`` weight tensors (which don't + # collide with any raw param name) need to be cleaned up here. + for stale_name in (gu_w_name, dn_w_name): + owner_mod, _path, attr_short = get_submodule_of_param(gm, stale_name) + _delete_module_attr(owner_mod, attr_short) + + num_matches += 1 + + # Register top-level EP+TP sharding load hook whenever there is any + # actual sharding on the MoE axes. The hook only does *sharding* + # (EP leading-axis slice + TP-aware pre-pad / intermediate-axis + # slice) on the raw HF MXFP4 state-dict entries; the kernel-layout + # work (H-axis pad, TMA shuffle, dtype convert, bias / tp_size) is + # deferred to :class:`FuseMXFP4Moe` on GPU at POST_LOAD_FUSION. + if num_matches > 0 and (moe_ep_size > 1 or moe_tp_size > 1): + assert num_experts_global is not None # for type checker + assert intermediate_size_global is not None + num_layers = (max(layer_indices) + 1) if layer_indices else num_matches + gm._register_load_state_dict_pre_hook( + make_mxfp4_sharding_load_hook( + num_layers=num_layers, + num_experts=num_experts_global, + intermediate_size=intermediate_size_global, + moe_ep_size=moe_ep_size, + moe_ep_rank=moe_ep_rank, + moe_tp_size=moe_tp_size, + moe_tp_rank=moe_tp_rank, + ) + ) + ad_logger.info( + f"quantize_mxfp4_moe (backend=trtllm, quant_act={quant_act}): " + f"rewrote {num_matches} MoE node(s); registered load hook for " + f"{num_layers} layer slots." + ) + + info = TransformInfo( + skipped=(num_matches == 0), + num_matches=num_matches, + is_clean=(num_matches == 0), + has_valid_shapes=(num_matches == 0), + ) + return gm, info + + +def _delete_module_attr(module: nn.Module, name: str) -> None: + """Remove a parameter/buffer/attr from a Module if present.""" + if name in module._parameters: + del module._parameters[name] + elif name in module._buffers: + del module._buffers[name] + elif hasattr(module, name): + delattr(module, name) + + +class FuseMXFP4MoeConfig(TransformConfig): + """Configuration for ``fuse_mxfp4_moe`` (POST_LOAD_FUSION).""" + + +@TransformRegistry.register("fuse_mxfp4_moe") +class FuseMXFP4Moe(BaseTransform): + """POST_LOAD_FUSION transform: GPU-side MXFP4 MoE weight prep for the trtllm-gen backend. + + Runs after ``QuantizeMXFP4MOE`` registered raw HF MXFP4 buffers and the EP-slice load hook + populated them. For each ``trtllm_quant_mxfp4_trtllm_gen_moe_fused`` node, calls + :func:`prepare_trtllm_gen_moe_mxfp4_weights` on the loaded GPU tensors to produce the kernel + layout, swaps the op args to the prepared params, and deletes the raw buffers. + + Skipped when the op already references prepared params (idempotent). + """ + + config: FuseMXFP4MoeConfig + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return FuseMXFP4MoeConfig + + def _apply( + self, + gm: GraphModule, + cm, + factory, + shared_config, + ) -> Tuple[GraphModule, TransformInfo]: + """Two-pass GPU prep with shared scratch + contiguous prepared blocks. + + Pass 1 (``_collect_moe_nodes``): walk the graph, find every + ``trtllm_quant_mxfp4_trtllm_gen_moe_fused`` op whose weight args still reference + raw HF buffers, record the per-layer info (experts module path, raw + ``get_attr`` nodes, shapes). Cross-layer consistency is asserted + (gpt-oss guarantees same H/I/E across all MoE layers). + + Pass 2: allocate ``MXFP4PrepScratch`` once for the per-rank shape. + Reused for every layer's pad + shuffle work. + + Pass 3: pre-allocate the SIX prepared ``nn.Parameter`` storages on + every experts module *before* any layer's prep runs. This is the + fragmentation-prevention step — all prepared blocks for all layers + come from the allocator's frontier in one back-to-back run, so no + transient alloc/free cycle from the prep work can interleave them. + + Pass 4: per layer, run ``prepare_trtllm_gen_moe_mxfp4_weights`` with + ``scratch=`` (pad + shuffle outputs land in scratch buffers, no + per-layer transient allocations of the big intermediates). Then + ``data.copy_`` scratch outputs into the pre-allocated prepared + params, re-point the op args, delete raw params + raw get_attrs. + """ + from ...custom_ops.fused_moe.prepare_trtllm_gen_moe_mxfp4_weights import ( + MXFP4PrepScratch, + prepare_trtllm_gen_moe_mxfp4_weights, + ) + + # Resolve runtime topology — used to divide ``fc2_bias`` by + # ``moe_tp_size`` (the prep helper's tp_size > 1 branch is skipped in + # the scratch path, so we do the division ourselves after). + dc = getattr(shared_config, "dist_config", None) + moe_tp_size = int(getattr(dc, "moe_tp_size", 1)) if dc is not None else 1 + + # Single MXFP4 trtllm-gen MoE op (act_dtype="bf16" or "mxfp8"). + target_op = torch.ops.auto_deploy.trtllm_quant_mxfp4_trtllm_gen_moe_fused.default + + # ---- Pass 1: collect MoE node info, validate consistent shape ---- + # Arg index layout from ``_apply_trtllm`` (kept in sync; comment + # block there documents the full slot list). + ARG_FC1_W, ARG_FC2_W, ARG_FC1_S, ARG_FC2_S, ARG_FC1_B, ARG_FC2_B = 4, 5, 6, 7, 8, 9 + + layer_infos: list = [] + e_local_g: Optional[int] = None + per_rank_i_g: Optional[int] = None + H_g: Optional[int] = None + device_g: Optional[torch.device] = None + for n in list(gm.graph.nodes): + if n.op != "call_function" or n.target is not target_op: + continue + if len(n.args) < 13: + continue + + raw_get_attrs = ( + n.args[ARG_FC1_W], # gate_up_proj_blocks + n.args[ARG_FC2_W], # down_proj_blocks + n.args[ARG_FC1_S], # gate_up_proj_scales + n.args[ARG_FC2_S], # down_proj_scales + n.args[ARG_FC1_B], # gate_up_proj_bias + n.args[ARG_FC2_B], # down_proj_bias + ) + if not all(isinstance(a, Node) and a.op == "get_attr" for a in raw_get_attrs): + continue + if not str(raw_get_attrs[0].target).endswith("gate_up_proj_blocks"): + # Already prepped or unexpected layout — skip. + continue + + gu_blocks_name = raw_get_attrs[0].target + experts_mod, experts_path, _ = get_submodule_of_param(gm, gu_blocks_name) + gu_blocks = gm.get_parameter(gu_blocks_name).data + dn_blocks = gm.get_parameter(raw_get_attrs[1].target).data + + e_local = int(gu_blocks.shape[0]) + two_i_local = int(gu_blocks.shape[1]) + per_rank_i = two_i_local // 2 + H = int(dn_blocks.shape[1]) + device = gu_blocks.device + + if e_local_g is None: + e_local_g, per_rank_i_g, H_g, device_g = e_local, per_rank_i, H, device + else: + if (e_local, per_rank_i, H) != (e_local_g, per_rank_i_g, H_g): + raise ValueError( + f"fuse_mxfp4_moe: cross-layer shape mismatch — layer " + f"got (E={e_local}, I={per_rank_i}, H={H}) but previous " + f"layers had (E={e_local_g}, I={per_rank_i_g}, H={H_g})." + ) + + layer_infos.append( + { + "node": n, + "experts_mod": experts_mod, + "experts_path": experts_path, + "raw_get_attrs": raw_get_attrs, + "raw_names": tuple(a.target for a in raw_get_attrs), + } + ) + + num_matches = len(layer_infos) + if num_matches == 0: + info = TransformInfo(skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True) + return gm, info + + # ---- Pass 2: allocate scratch ONCE ---- + scratch = MXFP4PrepScratch.allocate( + e_local=e_local_g, + per_rank_i=per_rank_i_g, + hidden_size=H_g, + device=device_g, + ) + + # ---- Pass 3: pre-allocate ALL prepared params (no data yet) ---- + # The six prepared kinds (fc1/fc2 × {w, s, b}). Shape + dtype mirror + # scratch fields; allocating them now (before any per-layer prep + # work) places them at the allocator frontier in one contiguous run, + # with no per-layer transient alloc/free in between. + prepared_kinds = ( + ("fc1_w_trtllm", scratch.fc1_w_buf.shape, scratch.fc1_w_buf.dtype), + ("fc1_w_scale_trtllm", scratch.fc1_s_buf.shape, scratch.fc1_s_buf.dtype), + ("fc1_bias_trtllm", scratch.fc1_b_buf.shape, scratch.fc1_b_buf.dtype), + ("fc2_w_trtllm", scratch.fc2_w_buf.shape, scratch.fc2_w_buf.dtype), + ("fc2_w_scale_trtllm", scratch.fc2_s_buf.shape, scratch.fc2_s_buf.dtype), + ("fc2_bias_trtllm", scratch.fc2_b_buf.shape, scratch.fc2_b_buf.dtype), + ) + for info_dict in layer_infos: + experts_mod = info_dict["experts_mod"] + for name, shape, dtype in prepared_kinds: + experts_mod.register_parameter( + name, + nn.Parameter( + torch.empty(shape, dtype=dtype, device=device_g), + requires_grad=False, + ), + ) + + # ---- Pass 4: per-layer prep into scratch + copy into prepared ---- + for info_dict in layer_infos: + n = info_dict["node"] + experts_mod = info_dict["experts_mod"] + experts_path = info_dict["experts_path"] + raw_get_attrs = info_dict["raw_get_attrs"] + raw_names = info_dict["raw_names"] + + gu_blocks = gm.get_parameter(raw_names[0]).data + dn_blocks_t = gm.get_parameter(raw_names[1]).data + gu_scales = gm.get_parameter(raw_names[2]).data + dn_scales = gm.get_parameter(raw_names[3]).data + gu_bias = gm.get_parameter(raw_names[4]).data + dn_bias = gm.get_parameter(raw_names[5]).data + + # Run prep with shared scratch — outputs are views into scratch, + # we copy_ them into the pre-allocated prepared params below. + prep = prepare_trtllm_gen_moe_mxfp4_weights( + gu_blocks, + gu_scales, + gu_bias, + dn_blocks_t, + dn_scales, + dn_bias, + hidden_size=H_g, + intermediate_size=per_rank_i_g, + tp_size=1, + tp_rank=0, + scratch=scratch, + ) + + # Copy scratch outputs into the pre-allocated prepared params. + # ``fc2_bias`` gets divided by ``moe_tp_size`` so the post-AR sum + # reproduces the unsharded bias (mirrors the prep helper's + # ``tp_size > 1`` branch which we skip in the scratch path). + getp = experts_mod.get_parameter + getp("fc1_w_trtllm").data.copy_(prep.fc1_weights_mxfp4) + getp("fc1_w_scale_trtllm").data.copy_(prep.fc1_weights_scale_ue8m0) + getp("fc1_bias_trtllm").data.copy_(prep.fc1_bias_f32) + getp("fc2_w_trtllm").data.copy_(prep.fc2_weights_mxfp4) + getp("fc2_w_scale_trtllm").data.copy_(prep.fc2_weights_scale_ue8m0) + if moe_tp_size > 1: + getp("fc2_bias_trtllm").data.copy_(prep.fc2_bias_f32 / moe_tp_size) + else: + getp("fc2_bias_trtllm").data.copy_(prep.fc2_bias_f32) + + # Build prepared get_attr nodes inserted right before the op call, + # then re-point the op's weight args to the prepared get_attrs. + prefix_path = (experts_path + ".") if experts_path else "" + with gm.graph.inserting_before(n): + fc1_w_attr = gm.graph.create_node("get_attr", prefix_path + "fc1_w_trtllm") + fc2_w_attr = gm.graph.create_node("get_attr", prefix_path + "fc2_w_trtllm") + fc1_s_attr = gm.graph.create_node("get_attr", prefix_path + "fc1_w_scale_trtllm") + fc2_s_attr = gm.graph.create_node("get_attr", prefix_path + "fc2_w_scale_trtllm") + fc1_b_attr = gm.graph.create_node("get_attr", prefix_path + "fc1_bias_trtllm") + fc2_b_attr = gm.graph.create_node("get_attr", prefix_path + "fc2_bias_trtllm") + + new_args = list(n.args) + new_args[ARG_FC1_W] = fc1_w_attr + new_args[ARG_FC2_W] = fc2_w_attr + new_args[ARG_FC1_S] = fc1_s_attr + new_args[ARG_FC2_S] = fc2_s_attr + new_args[ARG_FC1_B] = fc1_b_attr + new_args[ARG_FC2_B] = fc2_b_attr + n.args = tuple(new_args) + + # Erase raw get_attr nodes if no other consumer. + for stale_node in raw_get_attrs: + if len(stale_node.users) == 0: + gm.graph.erase_node(stale_node) + + # Delete raw module params now that prepared replaces them. + for raw_name in ( + "gate_up_proj_blocks", + "gate_up_proj_scales", + "gate_up_proj_bias", + "down_proj_blocks", + "down_proj_scales", + "down_proj_bias", + ): + _delete_module_attr(experts_mod, raw_name) + + # Scratch goes out of scope here → CUDA caching allocator reclaims + # the scratch region. The persistent prepared blocks remain + # contiguous (allocated before scratch was freed and after raw was + # being deleted layer by layer). + del scratch + + ad_logger.info( + f"fuse_mxfp4_moe: GPU-prepped {num_matches} MoE node(s) " + f"with shared scratch (E={e_local_g}, I={per_rank_i_g}, H={H_g})" + ) + + info = TransformInfo( + skipped=False, + num_matches=num_matches, + is_clean=False, + has_valid_shapes=True, + ) + return gm, info diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fusion.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fusion.py index de83dab56df9..62624b7ab6dc 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/fusion.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fusion.py @@ -66,6 +66,11 @@ def _insert_fused_gemm( y = x @ w.T y1 = y.narrow(-1, 0, out1).contiguous() # contiguous copy y2 = y.narrow(-1, out1, out2).contiguous() # contiguous copy + + Bias handling: + All children must have uniform bias state (all with bias or none). + Each bias must be 1D per-channel matching its weight's out_features. + Stacked bias is the dim=0 concatenation, mirroring weight stacking. """ keys_unfused = [extract_weight_name(n) for n in linear_nodes] params_unfused = [gm.get_parameter(k) for k in keys_unfused] @@ -77,13 +82,55 @@ def _insert_fused_gemm( return False weight_dtype = dtypes.pop() + # --- Bias fusibility check (all-or-none + 1D per-channel + size match) --- + bias_args = [n.args[2] for n in linear_nodes] + bias_present = [b is not None for b in bias_args] + if any(bias_present) and not all(bias_present): + # Mixed bias state — would require padding with zeros; bail out. + return False + has_bias = bias_present[0] + bias_params: List[torch.Tensor] = [] + if has_bias: + for n, w_param in zip(linear_nodes, params_unfused): + bnode = n.args[2] + # Only fuse statically known biases (get_attr nodes). + if bnode.op != "get_attr": + ad_logger.warning( + f"Skipping GEMM fusion for {keys_unfused}: bias is not a get_attr node" + ) + return False + bp = gm.get_parameter(bnode.target) + # Reject anything other than per-channel 1D bias matching out_features. + if bp.dim() != 1 or bp.size(0) != w_param.size(0): + ad_logger.warning( + f"Skipping GEMM fusion for {keys_unfused}: non per-channel bias " + f"(weight out={w_param.size(0)}, bias shape={tuple(bp.shape)})" + ) + return False + bias_params.append(bp) + bias_dtypes = {p.dtype for p in bias_params} + if len(bias_dtypes) != 1: + ad_logger.warning( + f"Skipping GEMM fusion for {keys_unfused}: mixed bias dtypes {bias_dtypes}" + ) + return False + key_fused = f"fused_weight_{idx}" fused_weight = torch.cat(params_unfused, dim=0).to(weight_dtype) param_fused = nn.Parameter(fused_weight, requires_grad=False) setattr(gm, key_fused, param_fused) + bias_key_fused = None + if has_bias: + bias_key_fused = f"fused_bias_{idx}" + bias_dtype = bias_params[0].dtype + fused_bias = torch.cat(bias_params, dim=0).to(bias_dtype) + bias_param_fused = nn.Parameter(fused_bias, requires_grad=False) + setattr(gm, bias_key_fused, bias_param_fused) + ad_logger.warning( - f"Fusing {len(linear_nodes)} GEMMs ({keys_unfused}) into {key_fused} (dtype={weight_dtype})" + f"Fusing {len(linear_nodes)} GEMMs ({keys_unfused}) into {key_fused} " + f"(dtype={weight_dtype}, bias={'yes' if has_bias else 'no'})" ) fused_kwargs = dict(linear_nodes[0].kwargs) @@ -91,11 +138,12 @@ def _insert_fused_gemm( with gm.graph.inserting_before(linear_nodes[0]): get_param_node = gm.graph.get_attr(key_fused, torch.Tensor) + get_bias_node = gm.graph.get_attr(bias_key_fused, torch.Tensor) if has_bias else None with gm.graph.inserting_before(linear_nodes[0]): fused_linear_node = gm.graph.call_function( linear_nodes[0].target, - args=(parent_node, get_param_node, None), + args=(parent_node, get_param_node, get_bias_node), kwargs=fused_kwargs, ) if ref_val is not None: @@ -325,18 +373,20 @@ def _apply( factory: ModelFactory, shared_config: SharedConfig, ) -> Tuple[GraphModule, TransformInfo]: - # sort linear nodes by parent node + # sort linear nodes by (parent, has_bias). Bias and no-bias siblings + # can't co-fuse (would need zero-padding), so bucket them separately + # to preserve partial fusion when a subset is bias-uniform. linear_nodes = defaultdict(list) for node in gm.graph.nodes: - # TODO: we don't handle bias for now... - if is_linear_op(node) and node.args[2] is None: - linear_nodes[node.args[0]].append(node) + if is_linear_op(node): + has_bias = node.args[2] is not None + linear_nodes[(node.args[0], has_bias)].append(node) # fuse linear nodes idx = -1 num_matches = 0 with cuda_memory_tracker(): - for parent_node, lin_children in linear_nodes.items(): + for (parent_node, _has_bias), lin_children in linear_nodes.items(): if len(lin_children) < 2: continue if not check_same_children(parent_node, is_linear_op): diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/mxfp4_moe.py b/tensorrt_llm/_torch/auto_deploy/transform/library/mxfp4_moe.py deleted file mode 100644 index ae6343b7aa5b..000000000000 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/mxfp4_moe.py +++ /dev/null @@ -1,346 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -from typing import Tuple - -import torch -import torch.nn as nn -from torch.fx import GraphModule, Node - -from ...utils.module import get_submodule_of_param -from ...utils.node_utils import is_op -from ...utils.pattern_matcher import ADPatternMatcherPass, register_ad_pattern -from ..interface import BaseTransform, TransformInfo, TransformRegistry - - -def _moe_dense_mlp_pattern( - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - gate_up_w: torch.Tensor, - gate_up_b: torch.Tensor, - down_w: torch.Tensor, - down_b: torch.Tensor, - alpha: float = 1.0, - limit: float = 10.0, - minus_limit: float = -10.0, -) -> torch.Tensor: - batch_size = hidden_states.shape[0] - hidden_size = hidden_states.shape[2] - hidden_states = hidden_states.reshape(-1, hidden_size) # (num_tokens, hidden_size) - num_experts = routing_weights.shape[1] - - hidden_states = hidden_states.repeat(num_experts, 1) - hidden_states = hidden_states.view(num_experts, -1, hidden_size) - gate_up = torch.bmm(hidden_states, gate_up_w) + gate_up_b.unsqueeze(-2) - gate, up = gate_up[..., ::2], gate_up[..., 1::2] - gate = gate.clamp(min=None, max=limit) - up = up.clamp(min=minus_limit, max=limit) - glu = gate * torch.sigmoid(gate * alpha) - next_states = torch.bmm(((up + 1) * glu), down_w) - next_states = next_states + down_b.unsqueeze(-2) - next_states = next_states.view(num_experts, batch_size, -1, hidden_size) - next_states = ( - next_states * routing_weights.transpose(0, 1).view(num_experts, batch_size, -1)[..., None] - ) - next_states = next_states.sum(dim=0) # [B, S, H] - return next_states - - -def _moe_dense_mlp_repl( - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - gate_up_w: torch.Tensor, - gate_up_b: torch.Tensor, - down_w: torch.Tensor, - down_b: torch.Tensor, - alpha: float, - limit: float, - minus_limit: float, -) -> torch.Tensor: - return torch.ops.auto_deploy.torch_moe_dense_mlp( - hidden_states, routing_weights, gate_up_w, gate_up_b, down_w, down_b, alpha, limit - ) - - -@TransformRegistry.register("match_dense_moe_pattern") -class MatchMOEDenseMLP(BaseTransform): - def _apply( - self, - gm: GraphModule, - cm, - factory, - shared_config, - ) -> Tuple[GraphModule, TransformInfo]: - graph = gm.graph - patterns = ADPatternMatcherPass() - - B, S, H = 2, 4, 8 # batch, seq, hidden - E, In = 3, 16 # experts, intermediate (I); gate_up has 2I - T = B * S - - dummy_args = [ - torch.randn(B, S, H, device="meta", dtype=torch.float16), # hidden_states - torch.randn(T, E, device="meta", dtype=torch.float16), # routing_weights - torch.randn(E, H, 2 * In, device="meta", dtype=torch.float16), # gate_up_w [E,H,2I] - torch.randn(E, 2 * In, device="meta", dtype=torch.float16), # gate_up_b [E,2I] - torch.randn(E, In, H, device="meta", dtype=torch.float16), # down_w [E,I,H] - torch.randn(E, H, device="meta", dtype=torch.float16), # down_b [E,H] - 1.07, - 10.1, - -10.1, - ] - - op_ignore_types = { - torch.ops.aten.view.default: (int,), - torch.ops.aten.reshape.default: (int,), - torch.ops.auto_deploy.view.default: (int,), - torch.ops.aten.repeat.default: (int,), - torch.ops.aten.slice.Tensor: (int,), - torch.ops.aten.unsqueeze.default: (int,), - torch.ops.aten.transpose.int: (int,), - } - - scalar_workaround = {"alpha": 1.07, "limit": 10.1, "minus_limit": -10.1} - - register_ad_pattern( - search_fn=_moe_dense_mlp_pattern, - replace_fn=_moe_dense_mlp_repl, - patterns=patterns, - dummy_args=dummy_args, - op_ignore_types=op_ignore_types, - scalar_workaround=scalar_workaround, - ) - - num_matches = patterns.apply(graph) - info = TransformInfo( - skipped=False, - num_matches=num_matches, - is_clean=num_matches == 0, - has_valid_shapes=num_matches == 0, - ) - return gm, info - - -def _get_alpha_limit_from_dense(node: Node) -> Tuple[float, float]: - # torch_moe_dense_mlp(hidden, routing, gu_w, gu_b, dn_w, dn_b, alpha, limit) - # alpha/limit may be in args or kwargs - alpha = node.kwargs.get("alpha", None) - limit = node.kwargs.get("limit", None) - if alpha is None: - alpha = float(node.args[6]) if len(node.args) >= 7 else 1.0 - if limit is None: - limit = float(node.args[7]) if len(node.args) >= 8 else 10.0 - return float(alpha), float(limit) - - -def _get_topk_from_router(node: Node) -> int: - # torch_moe_router(hidden, weight, bias, top_k=2) - if "top_k" in node.kwargs: - return int(node.kwargs["top_k"]) - return int(node.args[3]) if len(node.args) >= 4 else 2 - - -def _register_mxfp4_expert_params( - gm: GraphModule, - gate_up_w_name: str, - gate_up_b_name: str, - down_w_name: str, - down_b_name: str, -) -> Tuple[str, str, str, str]: - """Create (if missing) the four MXFP4 params under the experts module and return their full names. - - Returns: - (gu_blocks_name, gu_scales_name, dn_blocks_name, dn_scales_name) - """ - # Shapes from existing params - gu_b = gm.get_parameter(gate_up_b_name) # [E, 2I] - gu_w = gm.get_parameter(gate_up_w_name) # [E, 2I, H] - dn_b = gm.get_parameter(down_b_name) # [E, H] - - E = int(gu_b.shape[0]) - I2 = int(gu_b.shape[1]) # 2I - In = I2 // 2 - - # infer H from gu_w shape - assert gu_w.dim() == 3, "gate_up_w must be rank-3" - if gu_w.shape[1] == I2: - H = int(gu_w.shape[2]) - else: - # Fallback: use down bias last dim - H = int(dn_b.shape[1]) - - # Compute block dims (assume divisible; zero-init anyway) - H_blk = max(1, H // 32) - I_blk = max(1, In // 32) - - experts_mod, experts_path, _ = get_submodule_of_param(gm, gate_up_w_name) - - # New param names under experts module - gu_blocks_name = "gate_up_proj_blocks" - gu_scales_name = "gate_up_proj_scales" - dn_blocks_name = "down_proj_blocks" - dn_scales_name = "down_proj_scales" - - # Zero-init tensors (uint8 for blocks/scales) - gu_blocks = torch.zeros((E, 2 * In, H_blk, 16), dtype=torch.uint8) - gu_scales = torch.zeros((E, 2 * In, H_blk), dtype=torch.uint8) - dn_blocks = torch.zeros((E, H, I_blk, 16), dtype=torch.uint8) - dn_scales = torch.zeros((E, H, I_blk), dtype=torch.uint8) - - experts_mod.register_parameter(gu_blocks_name, nn.Parameter(gu_blocks, requires_grad=False)) - experts_mod.register_parameter(gu_scales_name, nn.Parameter(gu_scales, requires_grad=False)) - experts_mod.register_parameter(dn_blocks_name, nn.Parameter(dn_blocks, requires_grad=False)) - experts_mod.register_parameter(dn_scales_name, nn.Parameter(dn_scales, requires_grad=False)) - - # Free the now-unused bf16 stacked weight params (`gate_up_proj`, `down_proj`). - # The biases (`gate_up_proj_bias`, `down_proj_bias`) are still consumed by - # ``triton_mxfp4_moe`` and must remain. For models like GPT-OSS-120B - # (128 experts × 36 layers × ~33 MB per layer of bf16 placeholder) freeing - # these saves ~150 GB per rank. - gu_w_local = gate_up_w_name.split(".")[-1] - dn_w_local = down_w_name.split(".")[-1] - for local_name in (gu_w_local, dn_w_local): - if local_name in experts_mod._parameters: - del experts_mod._parameters[local_name] - - # Full GM attribute paths for new params - prefix = (experts_path + ".") if experts_path else "" - return ( - prefix + gu_blocks_name, - prefix + gu_scales_name, - prefix + dn_blocks_name, - prefix + dn_scales_name, - ) - - -@TransformRegistry.register("quantize_mxfp4_moe") -class InsertMXFP4MLP(BaseTransform): - """ - Replace (torch_moe_router -> torch_moe_dense_mlp) with a single auto_deploy::triton_mxfp4_moe op, - and register MXFP4 expert params (blocks + scales) on the experts module. - """ - - algo_name: str = "mxfp4" - - def _apply( - self, - gm: GraphModule, - cm, - factory, - shared_config, - ) -> Tuple[GraphModule, TransformInfo]: - qcfg = factory.get_quant_config() - if not qcfg or qcfg.get("quant_method", "") != self.algo_name: - return gm, TransformInfo( - skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True - ) - num_matches = 0 - - for n in list(gm.graph.nodes): - if not is_op(n, torch.ops.auto_deploy.torch_moe_dense_mlp): - continue - - # Expect: torch_moe_dense_mlp(hidden, routing, gu_w, gu_b, dn_w, dn_b, alpha, limit) - if len(n.args) < 6: - continue - - hidden_node = n.args[0] - routing_node = n.args[1] - gate_up_w_node = n.args[2] - gate_up_b_node = n.args[3] - down_w_node = n.args[4] - down_b_node = n.args[5] - - if not isinstance(routing_node, Node) or not is_op( - routing_node, torch.ops.auto_deploy.torch_moe_router - ): - continue - - # Router params: weight, bias, top_k - router_weight_node = routing_node.args[1] - router_bias_node = routing_node.args[2] - top_k = _get_topk_from_router(routing_node) - - # Resolve parameter names so we can find the experts module - if gate_up_w_node.op != "get_attr" or gate_up_b_node.op != "get_attr": - continue - if down_w_node.op != "get_attr" or down_b_node.op != "get_attr": - continue - - gu_w_name = gate_up_w_node.target - gu_b_name = gate_up_b_node.target - dn_w_name = down_w_node.target - dn_b_name = down_b_node.target - - # Register MXFP4 params on experts - gu_blocks_name, gu_scales_name, dn_blocks_name, dn_scales_name = ( - _register_mxfp4_expert_params(gm, gu_w_name, gu_b_name, dn_w_name, dn_b_name) - ) - - # Alpha/limit (from dense call) - alpha, limit = _get_alpha_limit_from_dense(n) - - # Insert the new get_attr nodes for MXFP4 params - with gm.graph.inserting_before(n): - gu_blocks_attr = gm.graph.create_node("get_attr", gu_blocks_name) - gu_scales_attr = gm.graph.create_node("get_attr", gu_scales_name) - dn_blocks_attr = gm.graph.create_node("get_attr", dn_blocks_name) - dn_scales_attr = gm.graph.create_node("get_attr", dn_scales_name) - - n.target = torch.ops.auto_deploy.triton_mxfp4_moe.default - n.kwargs = {} - - # triton_mxfp4_moe( - # hidden_states, - # router_weight, router_bias, top_k, - # gate_up_blocks, gate_up_bias, gate_up_scales, alpha, limit, - # down_blocks, down_bias, down_scales) - new_args = ( - hidden_node, - router_weight_node, - router_bias_node, - top_k, - gu_blocks_attr, - gate_up_b_node, - gu_scales_attr, - float(alpha), - float(limit), - dn_blocks_attr, - down_b_node, - dn_scales_attr, - ) - n.args = new_args - - # Remove the now-unneeded router node if nobody else uses it - if len(routing_node.users) == 0: - gm.graph.erase_node(routing_node) - - # Erase the old get_attr nodes for gate_up_proj and down_proj. - # _register_mxfp4_expert_params deleted those attributes from the - # experts module, so these nodes now reference non-existent attrs. - # They have no users after the args replacement above, so it is - # safe to erase them directly. - for stale_node in (gate_up_w_node, down_w_node): - if len(stale_node.users) == 0: - gm.graph.erase_node(stale_node) - - num_matches += 1 - - info = TransformInfo( - skipped=(num_matches == 0), - num_matches=num_matches, - is_clean=num_matches == 0, - has_valid_shapes=num_matches == 0, - ) - return gm, info diff --git a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py index a6df9e7c3bcd..0f9eb4a5e618 100644 --- a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py @@ -1902,10 +1902,10 @@ def get_valid_tactics(self, inputs: List[torch.Tensor], self.top_k, hidden_size, self.intermediate_size, - self.valid_hidden_size or hidden_size, - self.valid_intermediate_size or self.intermediate_size, self.local_num_experts, num_tokens, + self.valid_hidden_size or hidden_size, + self.valid_intermediate_size or self.intermediate_size, ) return tactics diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 37fc0553e9af..55506016720c 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -330,7 +330,7 @@ microsoft/phi-4: mistralai/Codestral-22B-v0.1: - accuracy: 67.10 openai/gpt-oss-120b: - - accuracy: 10.0 # TODO: update this when the perf is good. + - accuracy: 90.3 openai/gpt-oss-20b: - accuracy: 85.823 GPT-OSS/120B-MXFP4: diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index 366da2f629e7..4dd331be76d5 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -1264,7 +1264,7 @@ class TestGPTOSS(LlmapiAccuracyTestHarness): "reasoning_effort": "low", }, } - GSM8K_MAX_OUTPUT_LEN = 512 + GSM8K_MAX_OUTPUT_LEN = 8192 MODEL_PATHS = { "20b": f"{llm_models_root()}/gpt_oss/gpt-oss-20b", "120b": f"{llm_models_root()}/gpt_oss/gpt-oss-120b", @@ -1274,34 +1274,84 @@ class TestGPTOSS(LlmapiAccuracyTestHarness): pytest.param( "20b", "openai/gpt-oss-20b", - marks=pytest.mark.skip_less_device(2), + None, + None, id="20b", ), pytest.param( "120b", "openai/gpt-oss-120b", - marks=pytest.mark.skip_less_device(4), + None, + None, id="120b", ), + pytest.param( + "120b", + "openai/gpt-oss-120b", + 2, + "tp", + id="120b-tp2", + ), + pytest.param( + "120b", + "openai/gpt-oss-120b", + 2, + "ep", + id="120b-ep2", + ), ] - @pytest.mark.parametrize("model_id,model_name", MODEL_PARAMS) - def test_mxfp4_gsm8k(self, model_id, model_name, mocker): + @skip_pre_blackwell + @pytest.mark.skip_less_device(2) + @pytest.mark.parametrize( + "model_id,model_name,world_size_override,moe_topology", MODEL_PARAMS) + def test_mxfp4_gsm8k(self, model_id, model_name, world_size_override, + moe_topology, mocker): mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", self.GSM8K_MAX_OUTPUT_LEN) mocker.patch.dict(GSM8K.EVALUATE_KWARGS, {"scores_filter": "exact_match,flexible-extract"}) yaml_paths, registry_world_size = _get_registry_yaml_extra(model_name) - if get_device_count() < registry_world_size: + # world_size: yaml-driven; ``world_size_override`` is only used for + # MoE-TP / MoE-EP cases that exercise sharding on top of the same yaml. + world_size = (world_size_override if world_size_override is not None + else registry_world_size) + if get_device_count() < world_size: pytest.skip("Not enough devices for world size, skipping test") + # Override the default MoE topology via `apply_sharding_hints.dist_mapping`. + # The sharding invariants (`enabled`, `shard_layers: ["mha", "moe"]`, and + # `detect_sharding`/`sharding_transform_executor` disable) live in + # `gpt_oss.yaml`; here we only set the per-parametrize TP/EP mapping. + # `shard_layers=["mha","moe"]` (from yaml) lets the sharding pass wire up + # the MoE all_reduce inserted by ``QuantizeMXFP4MOE._apply_trtllm`` when + # tp_size > 1. + extra_kwargs = {} + if moe_topology is not None and world_size > 1: + if moe_topology == "tp": + moe_tp, moe_ep = world_size, 1 + elif moe_topology == "ep": + moe_tp, moe_ep = 1, world_size + else: + raise ValueError(f"unknown moe_topology={moe_topology!r}") + extra_kwargs["transforms"] = { + "apply_sharding_hints": { + "dist_mapping": { + "tp": world_size, + "moe_tp": moe_tp, + "moe_ep": moe_ep, + }, + }, + } + model_path = self.MODEL_PATHS[model_id] with AutoDeployLLM( model=model_path, tokenizer=model_path, - world_size=registry_world_size, + world_size=world_size, yaml_extra=yaml_paths, max_seq_len=GSM8K.MAX_INPUT_LEN + self.GSM8K_MAX_OUTPUT_LEN, + **extra_kwargs, ) as llm: task = GSM8K(model_name) task.evaluate(llm, diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 2818cf78b433..1c9d36e9bd6e 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -365,6 +365,9 @@ l0_dgx_b200: - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[fp8-4-attn_dp_off-trtllm] - accuracy/test_llm_api_autodeploy.py::TestNemotronUltraV3::test_accuracy[nvfp4-4] - accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[nvidia_Llama-3.1-8B-Instruct-NVFP4-True] + - accuracy/test_llm_api_autodeploy.py::TestGPTOSS::test_mxfp4_gsm8k[120b-tp2] + # ------------- AutoDeploy Perf Sanity --------------- + - perf/test_perf_sanity.py::test_e2e[aggr_upload-super_ad_blackwell-super_ad_ws4_1k1k] TIMEOUT (120) - condition: ranges: system_gpu_count: @@ -388,6 +391,9 @@ l0_dgx_b200: - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_accuracy[nvfp4-4-attn_dp_on-trtllm] - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-flashinfer] - accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[bf16_ws4_180gb-trtllm] + - accuracy/test_llm_api_autodeploy.py::TestGPTOSS::test_mxfp4_gsm8k[20b] + - accuracy/test_llm_api_autodeploy.py::TestGPTOSS::test_mxfp4_gsm8k[120b] + - accuracy/test_llm_api_autodeploy.py::TestGPTOSS::test_mxfp4_gsm8k[120b-ep2] # ------------- AutoDeploy Perf Sanity --------------- - perf/test_perf_sanity.py::test_e2e[aggr_upload-super_ad_blackwell-super_ad_ws4_1k1k] TIMEOUT (120) - perf/test_perf_sanity.py::test_e2e[aggr_upload-super_mtp_ad_blackwell-super_mtp_ad_ws4_1k1k] TIMEOUT (120) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_quant_mxfp4_trtllm_gen_moe.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_quant_mxfp4_trtllm_gen_moe.py new file mode 100644 index 000000000000..5c63535aed15 --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_quant_mxfp4_trtllm_gen_moe.py @@ -0,0 +1,451 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the trtllm-gen MXFP4 MoE custom-op family. + +Covers two layers of the same code path: + +1. **Weight preparation** (``prepare_trtllm_gen_moe_mxfp4_weights``) — pins the + kernel-layout invariants the trtllm-gen ``bf16_mxe2m1_block_scale_moe_runner`` + relies on: + + * fc1 / fc2 biases must go through the SAME row permutation as fc1 / fc2 + weights (gated-act-gemm interleave + epilogue-tile reorder for w3/w1; + only the epilogue-tile reorder for w2). Without this the kernel adds the + wrong bias to each post-shuffle output row and MoE output is garbage + (~2% on gpt-oss-120b GSM8K instead of ~90%). + * Byte-identical match against PT's MXFP4 reference loader. + +2. **Unified op dispatch** (``trtllm_quant_mxfp4_trtllm_gen_moe_fused``) — + the op dispatches to either the bf16 (W4A16) or the MXFP8 (W4A8) + trtllm-gen runner depending on the ``act_dtype`` arg. The focus is the + dispatch contract: both branches run end-to-end, invalid values raise. +""" + +import pytest +import torch +from utils.util import skip_pre_blackwell + +import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 (op registration) +from tensorrt_llm._torch.auto_deploy.custom_ops.fused_moe.prepare_trtllm_gen_moe_mxfp4_weights import ( + prepare_trtllm_gen_moe_mxfp4_weights, +) + +# Both the prep helper and the op rely on ``torch.ops.trtllm.shuffle_matrix``, +# which is CUDA-only. The op-dispatch tests additionally require Blackwell+ +# (decorated individually). +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="trtllm-gen MXFP4 MoE prep + op rely on torch.ops.trtllm.shuffle_matrix (CUDA only)", +) + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + +# gpt-oss-120b layout (H=2880, I=2880) for the prep tests (kernel-layout +# invariants depend on the real padded shapes); op-dispatch tests use a +# smaller H/I to keep runtime down. +GPTOSS_HIDDEN_SIZE = 2880 +GPTOSS_INTERMEDIATE_SIZE = 2880 +NUM_EXPERTS = 4 + + +def _build_synthetic_mxfp4_inputs( + e: int = NUM_EXPERTS, + h: int = GPTOSS_HIDDEN_SIZE, + i: int = GPTOSS_INTERMEDIATE_SIZE, + device: str = "cuda", +): + """Build deterministic uint8/bf16 expert tensors in the HF on-disk layout.""" + g = torch.Generator(device="cpu").manual_seed(0) + gu_blocks = torch.randint(0, 256, (e, 2 * i, h // 32, 16), dtype=torch.uint8, generator=g).to( + device + ) + gu_scales = torch.randint(0, 256, (e, 2 * i, h // 32), dtype=torch.uint8, generator=g).to( + device + ) + gu_bias = torch.randn(e, 2 * i, dtype=torch.bfloat16, generator=g).to(device) + dn_blocks = torch.randint(0, 256, (e, h, i // 32, 16), dtype=torch.uint8, generator=g).to( + device + ) + dn_scales = torch.randint(0, 256, (e, h, i // 32), dtype=torch.uint8, generator=g).to(device) + dn_bias = torch.randn(e, h, dtype=torch.bfloat16, generator=g).to(device) + return gu_blocks, gu_scales, gu_bias, dn_blocks, dn_scales, dn_bias + + +def _make_random_mxfp4_inputs( + *, + num_experts: int, + hidden_size: int, + intermediate_size: int, + device: str = "cuda", +): + """Build a minimal set of raw HF-layout MXFP4 weights + scales + biases for op-dispatch tests. + + Values are random but consistent in shape with HF MXFP4 safetensors: + * blocks ``[E, 2I, H/32, 16]`` uint8 (packed nibbles) + * scales ``[E, 2I, H/32]`` uint8 (UE8M0) + * biases ``[E, 2I]`` bf16 + and mirrored for ``down_*`` with H and I swapped. + """ + H = hidden_size + I = intermediate_size # noqa: E741 + E = num_experts + assert H % 32 == 0 and I % 32 == 0, "H and I must be multiples of 32 (MXFP4 block size)." + + gate_up_blocks = torch.randint( + 0, 256, (E, 2 * I, H // 32, 16), dtype=torch.uint8, device=device + ) + gate_up_scales = torch.randint(126, 130, (E, 2 * I, H // 32), dtype=torch.uint8, device=device) + gate_up_bias = torch.randn(E, 2 * I, dtype=torch.bfloat16, device=device) * 0.01 + + down_blocks = torch.randint(0, 256, (E, H, I // 32, 16), dtype=torch.uint8, device=device) + down_scales = torch.randint(126, 130, (E, H, I // 32), dtype=torch.uint8, device=device) + down_bias = torch.randn(E, H, dtype=torch.bfloat16, device=device) * 0.01 + + return ( + gate_up_blocks, + gate_up_scales, + gate_up_bias, + down_blocks, + down_scales, + down_bias, + ) + + +def _call_op(prep, *, act_dtype, x, router_weight, router_bias, top_k=2): + return torch.ops.auto_deploy.trtllm_quant_mxfp4_trtllm_gen_moe_fused( + x, + router_weight, + router_bias, + top_k, + prep.fc1_weights_mxfp4, + prep.fc2_weights_mxfp4, + prep.fc1_weights_scale_ue8m0, + prep.fc2_weights_scale_ue8m0, + prep.fc1_bias_f32, + prep.fc2_bias_f32, + # SwiGLU constants (gpt-oss defaults). + torch.full((prep.fc1_bias_f32.shape[0],), 1.702, dtype=torch.float32, device=x.device), + torch.full((prep.fc1_bias_f32.shape[0],), 1.0, dtype=torch.float32, device=x.device), + torch.full((prep.fc1_bias_f32.shape[0],), 7.0, dtype=torch.float32, device=x.device), + prep.valid_hidden_size, + prep.valid_intermediate_size, + act_dtype, + ) + + +# --------------------------------------------------------------------------- +# Section 1: weight-preparation invariants (CPU/CUDA, SM-agnostic shuffle ops) +# --------------------------------------------------------------------------- + + +def test_fc1_bias_is_shuffled_with_same_row_permutation_as_fc1_weights(): + """Regression: fc1 bias must follow the gated-act-gemm + TMA row permute. + + Reproduces the gpt-oss-120b MoE/GEMM accuracy bug where bias was only + padded (not shuffled), causing the trtllm-gen kernel to add the wrong + bias to each output row. + """ + from tensorrt_llm._torch.modules.fused_moe.quantization import ( + trtllmgen_maybe_get_cached_w3_w1_permute_indices, + ) + + device = "cuda" + gu_blocks, gu_scales, gu_bias, dn_blocks, dn_scales, dn_bias = _build_synthetic_mxfp4_inputs( + device=device + ) + e, two_i_pad = 4, 5888 # I_pad = 2944 (= ceil(2880/128)*128); 2*I_pad = 5888 + + # Reconstruct the pre-shuffle bias the prep helper builds (after pad + + # de-interleave + cat([up | gate])). Then derive the expected shuffled + # bias by reusing PT's permute helpers and compare against the actual + # output of ``prepare_trtllm_gen_moe_mxfp4_weights``. + gate_b = gu_bias[:, 0::2].contiguous() # [E, I] + up_b = gu_bias[:, 1::2].contiguous() # [E, I] + pad_amount = (128 - GPTOSS_INTERMEDIATE_SIZE % 128) % 128 + up_b_padded = torch.nn.functional.pad(up_b, (0, pad_amount)).float() + gate_b_padded = torch.nn.functional.pad(gate_b, (0, pad_amount)).float() + pre_shuffle_fc1_bias = torch.cat([up_b_padded, gate_b_padded], dim=1).contiguous() + assert pre_shuffle_fc1_bias.shape == (e, two_i_pad) + + cache: dict = {} + expected_fc1_bias_per_expert = [] + for k in range(e): + slc = pre_shuffle_fc1_bias[k].contiguous() + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices(slc, cache, 128) + expected_fc1_bias_per_expert.append(torch.index_select(slc, 0, perm.to(slc.device))) + expected_fc1_bias = torch.stack(expected_fc1_bias_per_expert, dim=0).contiguous() + + prep = prepare_trtllm_gen_moe_mxfp4_weights( + gu_blocks, + gu_scales, + gu_bias, + dn_blocks, + dn_scales, + dn_bias, + hidden_size=GPTOSS_HIDDEN_SIZE, + intermediate_size=GPTOSS_INTERMEDIATE_SIZE, + tp_size=1, + tp_rank=0, + ) + + assert prep.fc1_bias_f32.shape == (e, two_i_pad) + torch.testing.assert_close(prep.fc1_bias_f32, expected_fc1_bias, atol=0, rtol=0) + + # And it should NOT equal the unshuffled (just-padded) baseline — that's + # the buggy state we are guarding against. + assert not torch.equal(prep.fc1_bias_f32, pre_shuffle_fc1_bias), ( + "fc1 bias was not shuffled — this is the regression that wrecks gpt-oss-120b accuracy." + ) + + +def test_fc2_bias_is_shuffled_with_same_row_permutation_as_fc2_weights(): + """Regression: fc2 bias must follow the (non-gated) TMA row permute used by w2.""" + from tensorrt_llm._torch.modules.fused_moe.quantization import ( + trtllmgen_maybe_get_cached_w2_permute_indices, + ) + + device = "cuda" + gu_blocks, gu_scales, gu_bias, dn_blocks, dn_scales, dn_bias = _build_synthetic_mxfp4_inputs( + device=device + ) + e, h_pad = 4, 2944 # H_pad = ceil(2880/128)*128 = 2944 + + pad_amount = (128 - GPTOSS_HIDDEN_SIZE % 128) % 128 + pre_shuffle_fc2_bias = torch.nn.functional.pad(dn_bias, (0, pad_amount)).float() + assert pre_shuffle_fc2_bias.shape == (e, h_pad) + + cache: dict = {} + expected_fc2_bias_per_expert = [] + for k in range(e): + slc = pre_shuffle_fc2_bias[k].contiguous() + perm = trtllmgen_maybe_get_cached_w2_permute_indices(slc, cache, 128) + expected_fc2_bias_per_expert.append(torch.index_select(slc, 0, perm.to(slc.device))) + expected_fc2_bias = torch.stack(expected_fc2_bias_per_expert, dim=0).contiguous() + + prep = prepare_trtllm_gen_moe_mxfp4_weights( + gu_blocks, + gu_scales, + gu_bias, + dn_blocks, + dn_scales, + dn_bias, + hidden_size=GPTOSS_HIDDEN_SIZE, + intermediate_size=GPTOSS_INTERMEDIATE_SIZE, + tp_size=1, + tp_rank=0, + ) + + assert prep.fc2_bias_f32.shape == (e, h_pad) + torch.testing.assert_close(prep.fc2_bias_f32, expected_fc2_bias, atol=0, rtol=0) + + assert not torch.equal(prep.fc2_bias_f32, pre_shuffle_fc2_bias), ( + "fc2 bias was not shuffled — this is the regression that wrecks gpt-oss-120b accuracy." + ) + + +def test_prep_against_pt_reference_loader_byte_identical(): + """Compare the AD prep output byte-for-byte against PT's MXFP4 loader. + + PT's ``MXFP4WeightTRTLLMGenFusedMoEMethod.{load_expert_w3_w1_weight, + load_expert_w2_weight, load_expert_w3_w1_weight_scale_mxfp4, + load_expert_w2_weight_scale_mxfp4}`` is the gold standard the AD prep + helper must mirror. Any divergence here is a kernel-layout bug. + """ + from tensorrt_llm._torch.modules.fused_moe.quantization import ( + _get_weight_alignment, + maybe_pad_for_mxfp4, + trtllmgen_maybe_get_cached_w2_permute_indices, + trtllmgen_maybe_get_cached_w3_w1_permute_indices, + ) + + device = "cuda" + gu_blocks, gu_scales, gu_bias, dn_blocks, dn_scales, dn_bias = _build_synthetic_mxfp4_inputs( + device=device + ) + e, h, i = NUM_EXPERTS, GPTOSS_HIDDEN_SIZE, GPTOSS_INTERMEDIATE_SIZE + weight_alignment = 128 + input_hidden_alignment = 512 + scaling_vector_size = 32 + epilogue_tile_m = 128 + + gu_blocks_3d = gu_blocks.contiguous().view(e, 2 * i, h // 2) # [E, 2I, H/2] + dn_blocks_3d = dn_blocks.contiguous().view(e, h, i // 2) # [E, H, I/2] + + # PT-style per-expert reference for fc1 weight + bias. + # Step 1: deinterleave gate/up halves so dst gets [up | gate] in the row dim. + gate_w = gu_blocks_3d[:, 0::2, :].contiguous() # [E, I, H/2] + up_w = gu_blocks_3d[:, 1::2, :].contiguous() # [E, I, H/2] + gate_s = gu_scales[:, 0::2, :].contiguous() + up_s = gu_scales[:, 1::2, :].contiguous() + gate_b = gu_bias[:, 0::2].contiguous() + up_b = gu_bias[:, 1::2].contiguous() + + cache: dict = {} + + fc1_weight_ref = [] + fc1_scale_ref = [] + fc1_bias_ref = [] + fc2_weight_ref = [] + fc2_scale_ref = [] + fc2_bias_ref = [] + + for k in range(e): + # ---- fc1 weight ---- + alignment_w = _get_weight_alignment(weight_alignment, scaling_vector_size, 1, i) + u = maybe_pad_for_mxfp4(up_w[k], input_hidden_alignment // 2, alignment_w) + gp = maybe_pad_for_mxfp4(gate_w[k], input_hidden_alignment // 2, alignment_w) + dst = torch.cat([u, gp], dim=0).contiguous() # [2*I_pad, H_pad/2] + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices(dst, cache, epilogue_tile_m) + fc1_weight_ref.append(torch.index_select(dst, 0, perm.to(dst.device))) + + # ---- fc1 weight scale ---- + u_s = maybe_pad_for_mxfp4( + up_s[k], input_hidden_alignment // scaling_vector_size, alignment_w + ) + gp_s = maybe_pad_for_mxfp4( + gate_s[k], input_hidden_alignment // scaling_vector_size, alignment_w + ) + dst_s = torch.cat([u_s, gp_s], dim=0).contiguous() # [2*I_pad, H_pad/32] + perm_s = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + dst_s, cache, epilogue_tile_m, num_elts_per_sf=scaling_vector_size + ) + shuffled_s = torch.index_select(dst_s, 0, perm_s.to(dst_s.device)) + fc1_scale_ref.append( + torch.ops.trtllm.block_scale_interleave(shuffled_s).reshape(dst_s.shape) + ) + + # ---- fc1 bias ---- + ub = maybe_pad_for_mxfp4(up_b[k], alignment_w).float() + gb = maybe_pad_for_mxfp4(gate_b[k], alignment_w).float() + dst_b = torch.cat([ub, gb], dim=0).contiguous() # [2*I_pad] + perm_b = trtllmgen_maybe_get_cached_w3_w1_permute_indices(dst_b, cache, epilogue_tile_m) + fc1_bias_ref.append(torch.index_select(dst_b, 0, perm_b.to(dst_b.device))) + + # ---- fc2 weight ---- + alignment_w2 = _get_weight_alignment(weight_alignment, scaling_vector_size, 1, i) + d = maybe_pad_for_mxfp4(dn_blocks_3d[k], alignment_w2 // 2, weight_alignment) + perm_w2 = trtllmgen_maybe_get_cached_w2_permute_indices(d, cache, epilogue_tile_m) + fc2_weight_ref.append(torch.index_select(d, 0, perm_w2.to(d.device))) + + # ---- fc2 weight scale ---- + alignment_w2_s = _get_weight_alignment( + weight_alignment, scaling_vector_size, 1, dn_scales[k].shape[-1] + ) + d_s = maybe_pad_for_mxfp4( + dn_scales[k], alignment_w2_s // scaling_vector_size, weight_alignment + ) + perm_w2_s = trtllmgen_maybe_get_cached_w2_permute_indices( + d_s, cache, epilogue_tile_m, num_elts_per_sf=scaling_vector_size + ) + shuffled_s2 = torch.index_select(d_s, 0, perm_w2_s.to(d_s.device)) + fc2_scale_ref.append( + torch.ops.trtllm.block_scale_interleave(shuffled_s2).reshape(d_s.shape) + ) + + # ---- fc2 bias ---- + db = maybe_pad_for_mxfp4(dn_bias[k], weight_alignment).float() + perm_b2 = trtllmgen_maybe_get_cached_w2_permute_indices(db, cache, epilogue_tile_m) + fc2_bias_ref.append(torch.index_select(db, 0, perm_b2.to(db.device))) + + fc1_weight_ref_t = torch.stack(fc1_weight_ref, dim=0).contiguous() + fc1_scale_ref_t = torch.stack(fc1_scale_ref, dim=0).contiguous() + fc1_bias_ref_t = torch.stack(fc1_bias_ref, dim=0).contiguous() + fc2_weight_ref_t = torch.stack(fc2_weight_ref, dim=0).contiguous() + fc2_scale_ref_t = torch.stack(fc2_scale_ref, dim=0).contiguous() + fc2_bias_ref_t = torch.stack(fc2_bias_ref, dim=0).contiguous() + + prep = prepare_trtllm_gen_moe_mxfp4_weights( + gu_blocks, + gu_scales, + gu_bias, + dn_blocks, + dn_scales, + dn_bias, + hidden_size=h, + intermediate_size=i, + tp_size=1, + tp_rank=0, + ) + + assert torch.equal(prep.fc1_weights_mxfp4, fc1_weight_ref_t) + assert torch.equal(prep.fc1_weights_scale_ue8m0, fc1_scale_ref_t) + torch.testing.assert_close(prep.fc1_bias_f32, fc1_bias_ref_t, atol=0, rtol=0) + assert torch.equal(prep.fc2_weights_mxfp4, fc2_weight_ref_t) + assert torch.equal(prep.fc2_weights_scale_ue8m0, fc2_scale_ref_t) + torch.testing.assert_close(prep.fc2_bias_f32, fc2_bias_ref_t, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# Section 2: unified op act_dtype dispatch (Blackwell+ only) +# --------------------------------------------------------------------------- + + +@skip_pre_blackwell +@pytest.mark.parametrize("act_dtype", ["bf16", "mxfp8"]) +def test_trtllm_quant_mxfp4_trtllm_gen_moe_fused_act_dtype_dispatch(act_dtype): + """Both ``act_dtype`` branches run end-to-end on Blackwell. + + Verifies that each branch yields a finite output of the expected shape/dtype. + """ + torch.manual_seed(0) + device = "cuda" + E, H, I_, top_k = 4, 256, 64, 2 + B = 8 + + inputs = _make_random_mxfp4_inputs( + num_experts=E, hidden_size=H, intermediate_size=I_, device=device + ) + prep = prepare_trtllm_gen_moe_mxfp4_weights( + *inputs, hidden_size=H, intermediate_size=I_, tp_size=1, tp_rank=0 + ) + + x = torch.randn(B, H, dtype=torch.bfloat16, device=device) + router_weight = torch.randn(E, H, dtype=torch.bfloat16, device=device) * 0.02 + router_bias = torch.zeros(E, dtype=torch.bfloat16, device=device) + + y = _call_op( + prep, + act_dtype=act_dtype, + x=x, + router_weight=router_weight, + router_bias=router_bias, + top_k=top_k, + ) + + assert y.shape == (B, H), f"unexpected output shape {tuple(y.shape)}, want {(B, H)}" + assert y.dtype == torch.bfloat16, f"unexpected output dtype {y.dtype}" + assert torch.isfinite(y).all(), f"non-finite output for act_dtype={act_dtype!r}" + + +@skip_pre_blackwell +def test_trtllm_quant_mxfp4_trtllm_gen_moe_fused_invalid_act_dtype(): + """Invalid ``act_dtype`` raises ``ValueError``. + + Loud failure, no silent dispatch to one of the two real branches. + """ + torch.manual_seed(0) + device = "cuda" + E, H, I_ = 4, 256, 64 + + inputs = _make_random_mxfp4_inputs( + num_experts=E, hidden_size=H, intermediate_size=I_, device=device + ) + prep = prepare_trtllm_gen_moe_mxfp4_weights( + *inputs, hidden_size=H, intermediate_size=I_, tp_size=1, tp_rank=0 + ) + + x = torch.randn(8, H, dtype=torch.bfloat16, device=device) + router_weight = torch.randn(E, H, dtype=torch.bfloat16, device=device) * 0.02 + router_bias = torch.zeros(E, dtype=torch.bfloat16, device=device) + + with pytest.raises(ValueError, match="act_dtype"): + _call_op( + prep, + act_dtype="fp16_invalid", + x=x, + router_weight=router_weight, + router_bias=router_bias, + ) diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_mxfp4_moe.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_mxfp4_moe.py new file mode 100644 index 000000000000..4d513e390e26 --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_mxfp4_moe.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the ``fuse_mxfp4_moe`` POST_LOAD_FUSION transform. + +Pins the post-quantize / post-load contract of :class:`FuseMXFP4Moe`: + +* After running, raw HF MXFP4 buffers on the experts module + (``gate_up_proj_{blocks,scales,bias}`` / ``down_proj_{blocks,scales,bias}``) + are deleted and replaced by the six kernel-layout ``*_trtllm`` params + produced by :func:`prepare_trtllm_gen_moe_mxfp4_weights`. +* The ``trtllm_quant_mxfp4_trtllm_gen_moe_fused`` op's weight/bias arg slots + are re-pointed at the new prepared get_attr nodes — the op is again + runnable. +* When ``moe_tp_size > 1``, the prepared fc2 bias is divided by + ``moe_tp_size`` so that the post-AR sum reproduces the unsharded bias. +""" + +from typing import Tuple + +import pytest +import torch +import torch.nn as nn + +import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 (op registration) +import tensorrt_llm._torch.auto_deploy.transform.library.fused_moe_mxfp4 # noqa: F401 +from tensorrt_llm._torch.auto_deploy.custom_ops.fused_moe.prepare_trtllm_gen_moe_mxfp4_weights import ( + prepare_trtllm_gen_moe_mxfp4_weights, +) +from tensorrt_llm._torch.auto_deploy.transform.interface import SharedConfig, TransformRegistry +from tensorrt_llm._torch.auto_deploy.utils.dist_config import DistConfig + +# The transform calls ``prepare_trtllm_gen_moe_mxfp4_weights`` which itself +# invokes ``torch.ops.trtllm.shuffle_matrix`` — registered CUDA-only. +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="fuse_mxfp4_moe runs prepare_trtllm_gen_moe_mxfp4_weights which is CUDA-only", +) + + +# Small shapes that still respect the MXFP4 block size (32) and the kernel +# weight alignment (128) so prep runs without padding surprises. +E = 4 +H = 128 +I = 128 # noqa: E741 + + +def _make_raw_mxfp4_tensors(device: str = "cuda") -> Tuple[torch.Tensor, ...]: + """Build a deterministic raw-HF-layout MXFP4 expert set on ``device``.""" + g = torch.Generator(device="cpu").manual_seed(0) + gu_blocks = torch.randint(0, 256, (E, 2 * I, H // 32, 16), dtype=torch.uint8, generator=g).to( + device + ) + gu_scales = torch.randint(126, 130, (E, 2 * I, H // 32), dtype=torch.uint8, generator=g).to( + device + ) + gu_bias = (torch.randn(E, 2 * I, dtype=torch.bfloat16, generator=g) * 0.01).to(device) + dn_blocks = torch.randint(0, 256, (E, H, I // 32, 16), dtype=torch.uint8, generator=g).to( + device + ) + dn_scales = torch.randint(126, 130, (E, H, I // 32), dtype=torch.uint8, generator=g).to(device) + dn_bias = (torch.randn(E, H, dtype=torch.bfloat16, generator=g) * 0.01).to(device) + return gu_blocks, gu_scales, gu_bias, dn_blocks, dn_scales, dn_bias + + +def _build_pre_fuse_gm(raw_tensors: Tuple[torch.Tensor, ...]) -> torch.fx.GraphModule: + """Build a tiny GM in the exact pre-``FuseMXFP4Moe`` shape ``QuantizeMXFP4MOE`` leaves. + + Shape mirrors ``_apply_trtllm``'s output: + * root has an ``experts`` submodule with the six raw HF MXFP4 params + and the three SwiGLU constant params. + * Graph: ``(hidden, router_w, router_b) -> trtllm_quant_mxfp4_trtllm_gen_moe_fused`` + whose weight/bias args are get_attrs pointing at the raw experts params. + """ + gu_blocks, gu_scales, gu_bias, dn_blocks, dn_scales, dn_bias = raw_tensors + + root = nn.Module() + root.experts = nn.Module() + raw_specs = [ + ("gate_up_proj_blocks", gu_blocks), + ("gate_up_proj_scales", gu_scales), + ("gate_up_proj_bias", gu_bias), + ("down_proj_blocks", dn_blocks), + ("down_proj_scales", dn_scales), + ("down_proj_bias", dn_bias), + ] + for name, t in raw_specs: + root.experts.register_parameter(name, nn.Parameter(t.clone(), requires_grad=False)) + + # SwiGLU constants (gpt-oss defaults). Must live on the experts module + # because their get_attr nodes are inserted with that path. + a = torch.full((E,), 1.702, dtype=torch.float32, device=gu_blocks.device) + b = torch.full((E,), 1.0, dtype=torch.float32, device=gu_blocks.device) + c = torch.full((E,), 7.0, dtype=torch.float32, device=gu_blocks.device) + root.experts.register_parameter("swiglu_alpha_trtllm", nn.Parameter(a, requires_grad=False)) + root.experts.register_parameter("swiglu_beta_trtllm", nn.Parameter(b, requires_grad=False)) + root.experts.register_parameter("swiglu_limit_trtllm", nn.Parameter(c, requires_grad=False)) + + graph = torch.fx.Graph() + hidden = graph.placeholder("hidden") + router_w = graph.placeholder("router_w") + router_b = graph.placeholder("router_b") + + gu_blocks_n = graph.get_attr("experts.gate_up_proj_blocks") + dn_blocks_n = graph.get_attr("experts.down_proj_blocks") + gu_scales_n = graph.get_attr("experts.gate_up_proj_scales") + dn_scales_n = graph.get_attr("experts.down_proj_scales") + gu_bias_n = graph.get_attr("experts.gate_up_proj_bias") + dn_bias_n = graph.get_attr("experts.down_proj_bias") + sa_n = graph.get_attr("experts.swiglu_alpha_trtllm") + sb_n = graph.get_attr("experts.swiglu_beta_trtllm") + sl_n = graph.get_attr("experts.swiglu_limit_trtllm") + + moe = graph.call_function( + torch.ops.auto_deploy.trtllm_quant_mxfp4_trtllm_gen_moe_fused.default, + args=( + hidden, + router_w, + router_b, + 2, # top_k + gu_blocks_n, + dn_blocks_n, + gu_scales_n, + dn_scales_n, + gu_bias_n, + dn_bias_n, + sa_n, + sb_n, + sl_n, + H, # valid_hidden_size + I, # valid_intermediate_size + "mxfp8", # act_dtype + 0, # local_expert_offset + E, # num_local_experts + 1, # routing_method_type = Renormalize + ), + ) + graph.output(moe) + + return torch.fx.GraphModule(root, graph) + + +def _run_fuse(gm: torch.fx.GraphModule, dist_config: DistConfig): + """Apply just ``FuseMXFP4Moe`` with the given ``dist_config``.""" + shared_config = SharedConfig( + local_rank=dist_config.rank, + world_size=dist_config.world_size, + dist_config=dist_config, + ) + config_cls = TransformRegistry.get_config_class("fuse_mxfp4_moe") + config = config_cls(stage="post_load_fusion") + transform = TransformRegistry.get("fuse_mxfp4_moe")(config) + return transform._apply(gm, cm=None, factory=None, shared_config=shared_config) + + +def _moe_node(gm: torch.fx.GraphModule) -> torch.fx.Node: + target_op = torch.ops.auto_deploy.trtllm_quant_mxfp4_trtllm_gen_moe_fused.default + nodes = [n for n in gm.graph.nodes if n.op == "call_function" and n.target is target_op] + assert len(nodes) == 1, f"expected exactly one MoE op node, found {len(nodes)}" + return nodes[0] + + +# --------------------------------------------------------------------------- +# TP=1 — single-rank: raw → prepared swap, no bias /= moe_tp_size +# --------------------------------------------------------------------------- + + +def test_fuse_mxfp4_moe_tp1_raw_to_prepared_swap(): + """Single-rank: every raw HF buffer becomes a prepared ``*_trtllm`` buffer; arg slots re-pointed.""" + device = "cuda" + raw = _make_raw_mxfp4_tensors(device=device) + gm = _build_pre_fuse_gm(raw) + + dc = DistConfig(world_size=1, rank=0, tp_size=1, moe_tp_size=1, moe_ep_size=1) + _, info = _run_fuse(gm, dc) + + # TransformInfo: exactly one MoE node was prepped, not idempotent-skip. + assert info.skipped is False + assert info.num_matches == 1 + + # Raw HF params are gone. + raw_names = ( + "gate_up_proj_blocks", + "gate_up_proj_scales", + "gate_up_proj_bias", + "down_proj_blocks", + "down_proj_scales", + "down_proj_bias", + ) + for name in raw_names: + assert not hasattr(gm.experts, name) or getattr(gm.experts, name, None) is None, ( + f"raw param {name!r} should have been removed" + ) + + # Prepared params are registered (six kinds). + prepared_names = ( + "fc1_w_trtllm", + "fc1_w_scale_trtllm", + "fc1_bias_trtllm", + "fc2_w_trtllm", + "fc2_w_scale_trtllm", + "fc2_bias_trtllm", + ) + for name in prepared_names: + assert hasattr(gm.experts, name), f"prepared param {name!r} missing" + + # Op args 4..9 (fc1_w, fc2_w, fc1_s, fc2_s, fc1_b, fc2_b) point at prepared get_attrs. + n = _moe_node(gm) + ARG_FC1_W, ARG_FC2_W, ARG_FC1_S, ARG_FC2_S, ARG_FC1_B, ARG_FC2_B = 4, 5, 6, 7, 8, 9 + expected_targets = { + ARG_FC1_W: "experts.fc1_w_trtllm", + ARG_FC2_W: "experts.fc2_w_trtllm", + ARG_FC1_S: "experts.fc1_w_scale_trtllm", + ARG_FC2_S: "experts.fc2_w_scale_trtllm", + ARG_FC1_B: "experts.fc1_bias_trtllm", + ARG_FC2_B: "experts.fc2_bias_trtllm", + } + for slot, want in expected_targets.items(): + arg = n.args[slot] + assert isinstance(arg, torch.fx.Node) and arg.op == "get_attr", ( + f"arg slot {slot} is not a get_attr Node (got {arg!r})" + ) + assert arg.target == want, f"arg slot {slot} target = {arg.target!r}, want {want!r}" + + # TP=1: fc2_bias matches the raw prep output exactly (no /= moe_tp_size division). + prep = prepare_trtllm_gen_moe_mxfp4_weights( + *raw, hidden_size=H, intermediate_size=I, tp_size=1, tp_rank=0 + ) + torch.testing.assert_close(gm.experts.fc2_bias_trtllm.data, prep.fc2_bias_f32, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# TP=2 — fc2 bias must be divided by moe_tp_size so post-AR sum reproduces the unsharded bias +# --------------------------------------------------------------------------- + + +def test_fuse_mxfp4_moe_tp2_divides_fc2_bias_by_moe_tp_size(): + """``moe_tp_size > 1`` divides only the prepared ``fc2_bias`` by ``moe_tp_size``. + + Other prepared tensors (fc1/fc2 weights, fc1/fc2 scales, fc1 bias) must + match the TP=1 prep output 1:1 — the transform leaves them alone in the + scratch path; only ``fc2_bias`` is scaled. + """ + device = "cuda" + raw = _make_raw_mxfp4_tensors(device=device) + gm = _build_pre_fuse_gm(raw) + + moe_tp_size = 2 + dc = DistConfig( + world_size=2, + rank=0, + tp_size=moe_tp_size, + moe_tp_size=moe_tp_size, + moe_ep_size=1, + ) + _, info = _run_fuse(gm, dc) + assert info.num_matches == 1 + + # Golden: run prep on the SAME raw tensors at tp=1 (the transform path with + # scratch skips the helper's tp_size > 1 branch and does the division itself). + prep = prepare_trtllm_gen_moe_mxfp4_weights( + *raw, hidden_size=H, intermediate_size=I, tp_size=1, tp_rank=0 + ) + + # fc2_bias was divided by moe_tp_size; everything else matches 1:1. + torch.testing.assert_close( + gm.experts.fc2_bias_trtllm.data, prep.fc2_bias_f32 / moe_tp_size, atol=0, rtol=0 + ) + torch.testing.assert_close(gm.experts.fc1_bias_trtllm.data, prep.fc1_bias_f32, atol=0, rtol=0) + assert torch.equal(gm.experts.fc1_w_trtllm.data, prep.fc1_weights_mxfp4) + assert torch.equal(gm.experts.fc2_w_trtllm.data, prep.fc2_weights_mxfp4) + assert torch.equal(gm.experts.fc1_w_scale_trtllm.data, prep.fc1_weights_scale_ue8m0) + assert torch.equal(gm.experts.fc2_w_scale_trtllm.data, prep.fc2_weights_scale_ue8m0) + + +# --------------------------------------------------------------------------- +# Idempotency: re-running on an already-prepped graph is a no-op +# --------------------------------------------------------------------------- + + +def test_fuse_mxfp4_moe_idempotent_on_already_prepped_graph(): + """Re-running ``FuseMXFP4Moe`` on its own output skips (no double-prep).""" + device = "cuda" + raw = _make_raw_mxfp4_tensors(device=device) + gm = _build_pre_fuse_gm(raw) + + dc = DistConfig(world_size=1, rank=0, tp_size=1, moe_tp_size=1, moe_ep_size=1) + _, info1 = _run_fuse(gm, dc) + assert info1.num_matches == 1 + + _, info2 = _run_fuse(gm, dc) + assert info2.skipped is True, "second run should skip — no raw HF buffers left to prep" + assert info2.num_matches == 0 diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_gemm_fusion.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_gemm_fusion.py index 5e62540e4a23..421dd53e086a 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_gemm_fusion.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_gemm_fusion.py @@ -133,13 +133,17 @@ def __init__(self, **kwargs): class FusableModel3(FusableModel): - """Same as FusableModel1 except one GEMM is not fusable due to missing bias support.""" + """Same as FusableModel1 plus a bias=True sibling. + + Bias / no-bias linears are bucketed separately by ``FuseGemms``, so fc1+fc2 + fuse while fc3 stays alone (single in its bias=True bucket). + """ def __init__(self, cls=nn.Linear, **kwargs): super().__init__(**kwargs) self.fc1 = cls(self.in_features, self.out_features, bias=False) self.fc2 = cls(self.in_features, self.out_features, bias=False) - self.fc3 = cls(self.in_features, self.out_features, bias=True) # no bias support yet + self.fc3 = cls(self.in_features, self.out_features, bias=True) def forward(self, x): y1 = self.fc1(x) @@ -188,6 +192,49 @@ def __init__(self, **kwargs): super().__init__(**{"cls": FakeFP8Linear, **kwargs}) +class FusableModelAllBias(FusableModel): + """Three GEMMs sharing the same input, *all* with bias=True. + + FuseGemms should fuse weight + bias into one stacked GEMM with concat(bias) + along dim=0. + """ + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.fc1 = nn.Linear(self.in_features, self.out_features, bias=True) + self.fc2 = nn.Linear(self.in_features, self.out_features, bias=True) + self.fc3 = nn.Linear(self.in_features, self.out_features, bias=True) + + def forward(self, x): + return self.fc1(x) * self.fc2(x) + self.fc3(x) + + @property + def num_gemms_after_fusion(self) -> int: + return 1 + + +class FusableModelSplitByBias(FusableModel): + """Two bias=True and two bias=False linears sharing the same input. + + FuseGemms must bucket them separately and emit one fused bias group + one + fused no-bias group (= 2 GEMMs total). + """ + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.fc1 = nn.Linear(self.in_features, self.out_features, bias=False) + self.fc2 = nn.Linear(self.in_features, self.out_features, bias=False) + self.fc3 = nn.Linear(self.in_features, self.out_features, bias=True) + self.fc4 = nn.Linear(self.in_features, self.out_features, bias=True) + + def forward(self, x): + return self.fc1(x) * self.fc2(x) + self.fc3(x) + self.fc4(x) + + @property + def num_gemms_after_fusion(self) -> int: + return 2 + + # TODO: consider adding test cases for classic GQA and MLP layers @pytest.mark.parametrize( "get_model,dtype", @@ -218,6 +265,12 @@ def __init__(self, **kwargs): (FusableModel2, "bfloat16"), (FusableModel3, "bfloat16"), (FusableModel4, "bfloat16"), + # Bias-fusion (all siblings share bias state) + split bucketing (bias / + # no-bias go into separate fusion groups). + (FusableModelAllBias, "float16"), + (FusableModelAllBias, "bfloat16"), + (FusableModelSplitByBias, "float16"), + (FusableModelSplitByBias, "bfloat16"), pytest.param( FusableModel1_M_FP8, "fp8", @@ -312,6 +365,80 @@ def test_fusion(get_model: Callable[[], TestModel], dtype: str): assert not all_close(y_model, y_random) +@pytest.mark.parametrize("dtype", ["float16", "bfloat16"]) +@torch.inference_mode() +def test_fuse_gemms_bias_fusion_structure(dtype: str): + """All-bias siblings: verify the fused linear has a non-None bias. + + Also checks the stacked bias tensor matches dim=0 concat of the originals. + """ + torch_dtype = getattr(torch, dtype) + model = FusableModelAllBias().to(device="cuda", dtype=torch_dtype) + x = model.get_input(device="cuda", dtype=torch_dtype) + y_ref = model(x) + + # Capture original biases (per-fc) before export. + orig_biases = [ + model.fc1.bias.detach().clone(), + model.fc2.bias.detach().clone(), + model.fc3.bias.detach().clone(), + ] + expected_fused_bias = torch.cat(orig_biases, dim=0) + + gm = torch_export_to_gm(model, args=(x,), clone=True) + gm_transformed = InferenceOptimizer(None, {"fuse_gemms": {"stage": "post_load_fusion"}})( + None, gm + ) + + linear_nodes = [n for n in gm_transformed.graph.nodes if is_linear_op(n)] + assert len(linear_nodes) == 1, f"expected 1 fused linear, got {len(linear_nodes)}" + + # bias arg of the fused linear must be a get_attr node (not None). + bias_node = linear_nodes[0].args[2] + assert bias_node is not None, "fused linear has None bias — bias fusion did not run" + assert bias_node.op == "get_attr", f"bias arg is not get_attr: {bias_node.op}" + + fused_bias = gm_transformed.get_parameter(bias_node.target).detach().cpu() + assert fused_bias.dim() == 1, f"fused bias must be 1D, got shape {fused_bias.shape}" + assert fused_bias.numel() == sum(b.numel() for b in orig_biases), ( + f"fused bias numel {fused_bias.numel()} != sum of original biases" + ) + # Concat preserves values bit-for-bit. + torch.testing.assert_close(fused_bias, expected_fused_bias.cpu(), atol=0.0, rtol=0.0) + + # Sanity: forward still matches. + y_fused = gm_transformed.to("cuda")(x) + torch.testing.assert_close(y_ref, y_fused, atol=1e-3, rtol=1e-3) + + +@torch.inference_mode() +def test_fuse_gemms_split_by_bias_bucketing(): + """Mixed bias / no-bias siblings on the same parent. + + FuseGemms must keep them in separate buckets, producing one fused linear + *with* bias and one *without* bias (= 2 GEMMs total). + """ + model = FusableModelSplitByBias().to(device="cuda", dtype=torch.float16) + x = model.get_input(device="cuda", dtype=torch.float16) + y_ref = model(x) + + gm = torch_export_to_gm(model, args=(x,), clone=True) + gm_transformed = InferenceOptimizer(None, {"fuse_gemms": {"stage": "post_load_fusion"}})( + None, gm + ) + + linear_nodes = [n for n in gm_transformed.graph.nodes if is_linear_op(n)] + assert len(linear_nodes) == 2, f"expected 2 fused linears, got {len(linear_nodes)}" + + bias_states = {n.args[2] is not None for n in linear_nodes} + assert bias_states == {True, False}, ( + f"expected one bias=True and one bias=False fused linear, got {bias_states}" + ) + + y_fused = gm_transformed.to("cuda")(x) + torch.testing.assert_close(y_ref, y_fused, atol=1e-3, rtol=1e-3) + + # =========================================================================== # Tests for fuse_gemms_mixed_children (relaxed fusion with narrow / zero-copy views) # ===========================================================================