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 4a2bc641f81f..192b47c58abe 100644 --- a/examples/auto_deploy/llmc/create_standalone_package.py +++ b/examples/auto_deploy/llmc/create_standalone_package.py @@ -181,6 +181,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..ebb5354c51e3 --- /dev/null +++ b/examples/auto_deploy/model_registry/configs/gpt_oss.yaml @@ -0,0 +1,44 @@ +# 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"] + 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 e71fef61adaa..8359deba1ac8 100644 --- a/examples/auto_deploy/model_registry/models.yaml +++ b/examples/auto_deploy/model_registry/models.yaml @@ -170,8 +170,8 @@ models: 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'] +# 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'] @@ -336,8 +336,8 @@ models: # 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'] +# config_id: gpt_oss +# yaml_extra: ['gpt_oss.yaml', 'world_size_1.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/attention/trtllm_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py index f80f78e39357..cc5eb0d71507 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py @@ -94,7 +94,19 @@ def __init__(self): self.context_lengths_gpu: Optional[torch.Tensor] = None # [max_batch] int32 device # Persistent block_offsets buffer for CUDA graph compatibility. # Pre-allocated to max size so the tensor address is stable across replays. + # ``self.block_offsets`` is the group-0 buffer (kept for the spec-dec + # scratch path and backward compatibility); additional KV window groups + # (VSWA / non-uniform sliding window, e.g. gpt-oss) get their own + # persistent buffer keyed by the group's ``cache_loc`` input pointer in + # ``_block_offsets_by_cache_loc``. The transform invokes + # ``prepare_trtllm_metadata`` once per group with that group's + # ``cache_loc_g{i}`` / ``cu_num_pages_g{i}`` inputs, so without per-group + # buffers the groups would clobber a single shared buffer. self.block_offsets: Optional[torch.Tensor] = None + self._block_offsets_by_cache_loc: dict[int, torch.Tensor] = {} + # Shapes for lazy per-group buffer allocation (set in ``reset``). + self._max_batch: int = 0 + self._max_blocks_per_seq: int = 0 # Per-layer cache for tensors that must survive CUDA graph replay. # Keyed by kv_cache.data_ptr() (stable and unique per layer). self._layer_cache: dict[ @@ -148,9 +160,13 @@ def reset(self, device: torch.device, max_batch: int, max_blocks_per_seq: int) - self.host_request_types = torch.zeros( max_batch, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() ) + self._max_batch = max_batch + self._max_blocks_per_seq = max_blocks_per_seq self.block_offsets = torch.zeros( 1, max_batch, 2, max_blocks_per_seq, dtype=torch.int32, device=device ) + # Group 0 reuses ``self.block_offsets``; it is registered under its + # ``cache_loc`` pointer on first use in ``_get_block_offsets_buffer``. self.host_past_kv_lengths = torch.zeros( max_batch, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() ) @@ -290,23 +306,71 @@ def refresh_batch_state(self, batch_info: BatchInfo) -> None: self.num_contexts = num_prefill self.num_ctx_tokens = batch_info.get_num_tokens()[0] + def _get_block_offsets_buffer(self, cache_loc: torch.Tensor) -> torch.Tensor: + """Return the persistent block_offsets buffer for this KV window group. + + Each KV window group is driven by its own ``cache_loc`` input tensor + (group 0 uses ``cache_loc``; groups 1..N-1 use ``cache_loc_g{i}``), which + are persistent buffers with stable ``data_ptr()`` across CUDA-graph + replays. Keying by that pointer (same pattern as ``_layer_cache`` keyed + by ``kv_cache.data_ptr()``) gives each group an independent, address-stable + block_offsets buffer so per-group ``prepare_trtllm_metadata`` invocations + do not clobber each other. + + Lazily allocates a buffer on first sight of a group's ``cache_loc``. This + must happen during warm-up (never mid-capture) so the tensor address is + stable for graph replay; group 0's buffer reuses the one already + allocated in ``reset``. + """ + key = cache_loc.data_ptr() + buf = self._block_offsets_by_cache_loc.get(key) + if buf is None: + assert self.block_offsets is not None, ( + "planner.reset() must run before _get_block_offsets_buffer()" + ) + if not self._block_offsets_by_cache_loc: + # First group seen this run is group 0: reuse the reset() buffer. + buf = self.block_offsets + else: + assert ( + not torch.cuda.is_current_stream_capturing() + ) or cuda_graph_state.in_warm_up(), ( + "block_offsets buffer for a new KV window group must be " + "allocated during warm-up, not during CUDA graph capture. " + "Ensure warm-up exercises every KV pool." + ) + buf = torch.zeros( + 1, + self._max_batch, + 2, + self._max_blocks_per_seq, + dtype=torch.int32, + device=self.block_offsets.device, + ) + self._block_offsets_by_cache_loc[key] = buf + return buf + def plan_device( self, num_seq: int, block_offset_multiplier: int, cu_num_pages: torch.Tensor, cache_loc: torch.Tensor, - ) -> None: + ) -> torch.Tensor: """Per-forward DEVICE metadata: block_offsets via Triton kernel (pure GPU). Called from the ``prepare_trtllm_metadata`` custom op (in the graph). + Returns the per-group block_offsets buffer that was populated, so the op + can flow it through the graph to that group's attention layers. """ - k_slice = self.block_offsets[0, :, 0, :] # [max_batch, M], stride [2*M, 1] + block_offsets = self._get_block_offsets_buffer(cache_loc) + k_slice = block_offsets[0, :, 0, :] # [max_batch, M], stride [2*M, 1] torch.ops.auto_deploy.ragged_to_block_table_triton( cache_loc, cu_num_pages, k_slice, num_seq ) - self.block_offsets[0, :num_seq, 0, :].mul_(block_offset_multiplier) - self.block_offsets[0, :num_seq, 1, :] = self.block_offsets[0, :num_seq, 0, :] + 1 + block_offsets[0, :num_seq, 0, :].mul_(block_offset_multiplier) + block_offsets[0, :num_seq, 1, :] = block_offsets[0, :num_seq, 0, :] + 1 + return block_offsets _GlobalTrtllmPlanner = _TrtllmPlanner() @@ -479,14 +543,16 @@ def prepare_trtllm_metadata( _GlobalTrtllmPlanner.use_spec_decoding = batch_info.get_num_sequences()[2] == 0 block_offset_multiplier = batch_info.get_block_offset_multiplier() - _GlobalTrtllmPlanner.plan_device( + block_offsets = _GlobalTrtllmPlanner.plan_device( num_seq=batch_info.get_total_num_sequences(), block_offset_multiplier=block_offset_multiplier, cu_num_pages=cu_num_pages, cache_loc=cache_loc, ) - return [_GlobalTrtllmPlanner.block_offsets] + # Return this group's buffer (keyed by ``cache_loc``) so multi-pool + # (VSWA) deployments flow the correct block_offsets to each group's layers. + return [block_offsets] @prepare_trtllm_metadata.register_fake @@ -571,7 +637,10 @@ def trtllm_mha_with_cache( num_tokens = batch_info.get_total_num_tokens() max_context_length = batch_info.get_max_context_length() max_num_requests = batch_info.get_max_batch_size() - # Use sliding_window for attention_window_size if provided, else full context length + # Use sliding_window for attention_window_size if provided, else full context length. + # The mask stays ``causal`` (matching the PyTorch backend, which never uses + # sliding_window_causal): the kernel honors the window via the cyclic + # attention-window handling driven by ``attention_window_size``. attention_window_size = ( sliding_window if isinstance(sliding_window, int) and sliding_window > 0 @@ -799,6 +868,13 @@ class TrtllmAttention(AttentionDescriptor): Follows the same stateless descriptor pattern as ``FlashInferAttention``. """ + @classmethod + def kernel_handles_cyclic_swa(cls) -> bool: + """thop.attention applies the sliding-window mask internally via cyclic + KV indexing, so the executor passes the full per-window block table and + global KV lengths (no host-side window slicing). See base class.""" + return True + @classmethod def get_attention_layout(cls) -> AttentionLayout: """Get the attention layout expected by the backend.""" @@ -856,8 +932,12 @@ def get_cache_initializers( num_kv_heads = k_fake.shape[2] head_dim = k_fake.shape[3] kv_dtype = k_fake.dtype - # ``sliding_window`` is propagated into the handler so layers - # with different windows land in separate pools. + # ``sliding_window`` is propagated into the handler so layers with + # different windows land in separate KV pools. The trtllm backend now + # supports multiple pools (per-group block_offsets + cyclic-window + # staging), so non-uniform-window models (e.g. gpt-oss) get a dedicated + # window-sized pool per window instead of over-allocating a single + # full-seq pool for every layer. (sw,) = extract_op_args(source_attn_node, "sliding_window") sliding_window = sw if isinstance(sw, int) and sw > 0 else 0 diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index 838a1731791a..308f6dafeec9 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -2233,6 +2233,22 @@ def supports_shared_kv(cls) -> bool: """Whether this backend supports shared-KV cache aliasing.""" return False + @classmethod + def kernel_handles_cyclic_swa(cls) -> bool: + """Whether the backend's kernel applies the sliding-window mask itself. + + When ``True`` (e.g. the trtllm ``thop.attention`` kernel), the kernel + cyclically indexes the KV cache internally using the per-layer attention + window, so the executor must hand it the *full* per-window block table + and a *global* (un-window-capped) KV length -- the same contract as the + PyTorch backend. + + When ``False`` (default; e.g. triton / flashinfer), the kernel does not + cyclic-index, so the executor must host-slice the block table down to the + live sliding-window view (see ``ad_executor._compute_window_local_view``). + """ + return False + @classmethod @abstractmethod def get_standard_metadata_args(cls) -> List[str]: 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 41c4d0abbd03..2ab879f2be19 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 @@ -1161,3 +1161,214 @@ def trtllm_nvfp4_trtllm_gen_moe_fused_fake( apply_routing_on_input: bool = False, ) -> 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/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py index b8b8d5892310..5c08a95d87d1 100644 --- a/tensorrt_llm/_torch/auto_deploy/llm_args.py +++ b/tensorrt_llm/_torch/auto_deploy/llm_args.py @@ -429,8 +429,18 @@ def disable_cudagraph_for_speculative_flashinfer(self): ### UTILITY METHODS ############################################################################ @property def requires_uniform_kv_caches(self) -> bool: - """Whether CachedSequenceInterface must enforce a uniform KV cache mapping.""" - return self.attn_backend.lower() == "trtllm" + """Whether CachedSequenceInterface must enforce a uniform KV cache mapping. + + No attention backend currently requires this. The trtllm backend used to + return ``True`` here to force a single KV pool, but it now supports + multiple KV cache memory pools for non-uniform sliding-window models + (e.g. gpt-oss) -- the kernel applies the sliding-window mask internally + via cyclic indexing, so per-window pools route correctly. The flag is + kept (defaulting to ``False``) so the uniformity enforcement in + ``CachedSequenceInterface`` remains available should a future backend + need it. + """ + return False def create_factory(self) -> ModelFactory: """Create a model factory from the arguments. 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/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index b633d6a3b625..8fc6d70e8481 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -301,6 +301,40 @@ def _compute_window_local_view( return active_indices, extra_page, active_token_count, last_page_len +def _compute_cyclic_full_view( + all_indices: Sequence[int], + end_compute_i: int, + tokens_per_block: int, +) -> Tuple[List[int], int, int, int]: + """Compute the metadata view for a cyclic-SWA kernel (trtllm). + + Unlike ``_compute_window_local_view`` (which slices the block table down to + the live sliding window for kernels that cannot cyclic-index), the trtllm + ``thop.attention`` kernel applies the sliding-window mask itself by wrapping + KV reads modulo the attention window. It therefore needs: + + * the FULL per-window block table (``all_indices`` verbatim, including any + stale front-evicted entries -- the kernel's modulo indexing skips them), + and + * the GLOBAL (un-window-capped) KV length ``end_compute_i``. + + This mirrors the PyTorch backend, which copies the manager's full block list + from index 0 and passes ``host_past_key_value_lengths == total KV length``. + + Returns the same 4-tuple shape as ``_compute_window_local_view``: + ``(active_indices, extra_page, seq_len_with_cache, last_page_len)``. + ``extra_page`` is always -1: the full table already contains the next page, + so the overlap scheduler needs no deferred-page insertion. + """ + active_indices = list(all_indices) + seq_len_with_cache = end_compute_i + if seq_len_with_cache > 0: + last_page_len = (seq_len_with_cache - 1) % tokens_per_block + 1 + else: + last_page_len = 0 + return active_indices, -1, seq_len_with_cache, last_page_len + + class ADEngine(ModelEngine): """The AutoDeploy Engine (ADEngine) is the main engine interface to execute AutoDeploy models. @@ -764,6 +798,12 @@ def _prepare_inputs( # on SequenceInfo). Per-window queries on the manager route to the # correct C++ pool via mLayerToWindowSize. kv_group_windows = self.cache_seq_interface.kv_group_windows + # When the attention kernel applies the sliding-window mask itself via + # cyclic KV indexing (trtllm), the executor must hand it the full + # per-window block table and a global (un-window-capped) KV length -- + # the same contract as the PyTorch backend. Otherwise (triton / + # flashinfer) host-slice the block table to the live window below. + cyclic_swa = self.cache_seq_interface.kernel_handles_cyclic_swa # Cache hot lookups so the per-request loop avoids repeated C++ # dispatch / hasattr calls. _tokens_per_block = kv_cache_manager.tokens_per_block @@ -803,40 +843,56 @@ def _prepare_inputs( for pool_idx, group_window in enumerate(kv_group_windows): all_indices = batch_cache_indices_per_pool[pool_idx][i] - # SWA front-eviction: get_batch_cache_indices returns the FULL - # historical page list including front-evicted entries (the - # C++ side bumps a counter rather than popping mCacheBlockIds). - # _compute_window_local_view slices it down to the live window - # in window-local coords. - front_removed = kv_cache_manager.get_num_front_blocks_removed( - request.py_request_id, window_size=group_window - ) - ( - active_indices, - extra_page, - active_token_count, - lpl_i, - ) = _compute_window_local_view( - all_indices, - front_removed=front_removed, - end_compute_i=end_compute_i, - group_window=group_window, - tokens_per_block=_tokens_per_block, - ) - num_active = len(active_indices) + if cyclic_swa: + # Cyclic-SWA kernels (trtllm) want the FULL per-window block + # table and the GLOBAL KV length; the kernel masks the window + # internally. No front-eviction slicing, so the + # get_num_front_blocks_removed C++ dispatch is skipped here. + ( + active_indices, + extra_page, + active_token_count, + lpl_i, + ) = _compute_cyclic_full_view( + all_indices, + end_compute_i=end_compute_i, + tokens_per_block=_tokens_per_block, + ) + num_active = len(active_indices) + else: + # SWA front-eviction: get_batch_cache_indices returns the FULL + # historical page list including front-evicted entries (the + # C++ side bumps a counter rather than popping mCacheBlockIds). + # _compute_window_local_view slices it down to the live window + # in window-local coords. + front_removed = kv_cache_manager.get_num_front_blocks_removed( + request.py_request_id, window_size=group_window + ) + ( + active_indices, + extra_page, + active_token_count, + lpl_i, + ) = _compute_window_local_view( + all_indices, + front_removed=front_removed, + end_compute_i=end_compute_i, + group_window=group_window, + tokens_per_block=_tokens_per_block, + ) + num_active = len(active_indices) cache_loc_per_pool[pool_idx].extend(active_indices) cu_num_pages_per_pool[pool_idx].append( cu_num_pages_per_pool[pool_idx][i] + num_active ) extra_page_per_seq_per_pool[pool_idx].append(extra_page) - # Window-local seq_len_with_cache / last_page_len for every - # pool (including 0). For full-attention pools the helper - # returns the unclamped global value (group_window equals - # max_seq_len, no clamping kicks in), so this is identical to - # the legacy single-pool path for non-SWA models. For SWA - # pools (whether pool 0 or pool 1+), it carries the - # window-local coords the kernel needs under front-eviction. + # seq_len_with_cache / last_page_len per pool (including 0). + # Cyclic-SWA (trtllm): the global KV length for every pool. + # Host-sliced (triton/flashinfer): the unclamped global value for + # full-attention pools (window == max_seq_len, no clamping), and + # the window-local coords for SWA pools under front-eviction -- + # identical to the legacy single-pool path for non-SWA models. seq_len_with_cache_per_pool[pool_idx].append(active_token_count) last_page_len_per_pool[pool_idx].append(lpl_i) diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index e6714f513184..36361ca560f2 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -141,6 +141,12 @@ def __init__( # same order as the C++ manager's internal pool ordering (i.e. the # insertion order of the per-window shape map keys). self._kv_group_windows: List[int] = [] + # Whether the attention backend's kernel applies the sliding-window mask + # itself via cyclic KV indexing (trtllm). When True the executor passes + # the full per-window block table and global KV lengths instead of + # host-slicing to the live window. Set by the kvcache transform from the + # attention descriptor's ``kernel_handles_cyclic_swa()``. + self._kernel_handles_cyclic_swa: bool = False # lookup of unmanaged resources self._unmanaged_resources: List[str] = [] self._spec_config = spec_config @@ -1160,6 +1166,21 @@ def set_kv_groups(self, group_windows: List[int]) -> None: """ self._kv_group_windows = list(group_windows) + @property + def kernel_handles_cyclic_swa(self) -> bool: + """Whether the attention kernel applies the sliding-window mask itself. + + When True (trtllm), the executor passes the full per-window block table + and global KV lengths; when False (triton/flashinfer), it host-slices to + the live sliding window. + """ + return self._kernel_handles_cyclic_swa + + def set_kernel_handles_cyclic_swa(self, value: bool) -> None: + """Record the attention backend's cyclic-SWA capability (called by the + kvcache transform from ``AttentionDescriptor.kernel_handles_cyclic_swa``).""" + self._kernel_handles_cyclic_swa = bool(value) + @property def kv_cache_manager(self) -> Optional[KVCacheManager]: """Return the unified KVCacheManager, or None if not initialized.""" 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/kvcache.py b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py index f131d2d6292e..f58498ecf679 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py @@ -352,6 +352,12 @@ def _apply( skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True ) + # Record whether this backend's kernel applies the sliding-window mask + # itself (cyclic KV indexing, e.g. trtllm). The executor uses this to + # decide between passing the full per-window block table + global KV + # lengths (cyclic) and host-slicing to the live window (triton/flashinfer). + cm.set_kernel_handles_cyclic_swa(attn_descriptor.kernel_handles_cyclic_swa()) + # get standard metadata nodes for all source attention nodes meta_nodes_std = self._process_metadata_std(gm, cm) 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 c7cdaddc297e..c0ec622064ee 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 5c3d26f2e79b..0f40c4ccdcf8 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -1252,7 +1252,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", @@ -1262,34 +1262,86 @@ class TestGPTOSS(LlmapiAccuracyTestHarness): pytest.param( "20b", "openai/gpt-oss-20b", + None, + None, marks=pytest.mark.skip_less_device(2), 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", + marks=pytest.mark.skip_less_device(2), + id="120b-tp2", + ), + pytest.param( + "120b", + "openai/gpt-oss-120b", + 2, + "ep", + marks=pytest.mark.skip_less_device(2), + 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.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 00b9c02e1c48..062b2bea024b 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -353,6 +353,7 @@ 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: @@ -379,6 +380,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/attention/test_trtllm_attention_op.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_trtllm_attention_op.py index 353938830fa1..2b1c9934dfef 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_trtllm_attention_op.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_trtllm_attention_op.py @@ -881,3 +881,89 @@ def test_metadata_handles_two_sequences_with_different_lengths(self): assert metadata["host_total_kv_lens"][0] == 300 if "context_lengths_gpu" in metadata: assert metadata["context_lengths_gpu"] == [100, 200] + + +# --------------------------------------------------------------------------- +# Multi-pool (VSWA / non-uniform sliding window) block_offsets +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("device", ["cuda"]) +class TestTrtllmMultiPoolBlockOffsets: + """Per-KV-window-group block_offsets buffers must not clobber one another. + + The trtllm planner keeps a separate, address-stable block_offsets buffer + per KV window group (keyed by the group's ``cache_loc`` input ptr), so that + per-group ``prepare_trtllm_metadata`` invocations -- as emitted by the + kvcache transform for non-uniform sliding-window models like gpt-oss -- do + not clobber one another. + """ + + @staticmethod + def _host_prepare(num_seq, max_seq_len, tokens_per_block, max_batch_size, device): + max_blocks_per_seq = math.ceil(max_seq_len / tokens_per_block) + _bi = BatchInfo() + _bi.update([num_seq, num_seq, 0, 0, 0, 0]) # all-prefill, 1 token/seq + _bi.update_max_seq_info(max_seq_len, max_blocks_per_seq, 2, max_batch_size) + batch_info_host = _bi.serialize() + ones = torch.ones(num_seq, dtype=torch.int32) + zeros = torch.zeros(num_seq, dtype=torch.int32) + prepare_trtllm_metadata_host( + batch_info_host, + ones.clone().pin_memory(), # seq_len_with_cache_host + zeros.clone().pin_memory(), # input_pos_host + ones.clone().pin_memory(), # seq_len_host + ones.clone().pin_memory(), # prompt_lens_host + ones.clone().to(device), # prompt_lens + ) + return batch_info_host + + def test_per_group_buffers_are_distinct_and_not_clobbered(self, device): + _reset_trtllm_planner() + batch_info_host = self._host_prepare( + num_seq=2, max_seq_len=2048, tokens_per_block=32, max_batch_size=4, device=device + ) + + # Two groups, each with its own (distinct-ptr) cache_loc / cu_num_pages. + cache_loc_a = torch.tensor([10, 11, 12, 13], dtype=torch.int32, device=device) + cache_loc_b = torch.tensor([20, 21, 22, 23], dtype=torch.int32, device=device) + cu_num_pages = torch.tensor([0, 2, 4], dtype=torch.int32, device=device) + assert cache_loc_a.data_ptr() != cache_loc_b.data_ptr() + + (buf_a,) = torch.ops.auto_deploy.trtllm_attention_prepare_metadata( + batch_info_host, cu_num_pages, cache_loc_a + ) + # group 0 reuses the pre-allocated reset() buffer + assert buf_a.data_ptr() == _GlobalTrtllmPlanner.block_offsets.data_ptr() + a_after_a = buf_a[0, :2, 0, :].clone() + + (buf_b,) = torch.ops.auto_deploy.trtllm_attention_prepare_metadata( + batch_info_host, cu_num_pages, cache_loc_b + ) + # second group gets its own, distinct buffer + assert buf_b.data_ptr() != buf_a.data_ptr() + assert len(_GlobalTrtllmPlanner._block_offsets_by_cache_loc) == 2 + + # group 0's buffer must be untouched by group 1's prepare + torch.testing.assert_close(buf_a[0, :2, 0, :], a_after_a) + # block_offsets reflect each group's own cache_loc (× multiplier 2) + assert buf_a[0, 0, 0, 0].item() == 10 * 2 + assert buf_b[0, 0, 0, 0].item() == 20 * 2 + + def test_same_cache_loc_returns_stable_buffer(self, device): + _reset_trtllm_planner() + batch_info_host = self._host_prepare( + num_seq=1, max_seq_len=1024, tokens_per_block=32, max_batch_size=2, device=device + ) + cache_loc = torch.tensor([5, 6], dtype=torch.int32, device=device) + cu_num_pages = torch.tensor([0, 2], dtype=torch.int32, device=device) + + (buf1,) = torch.ops.auto_deploy.trtllm_attention_prepare_metadata( + batch_info_host, cu_num_pages, cache_loc + ) + (buf2,) = torch.ops.auto_deploy.trtllm_attention_prepare_metadata( + batch_info_host, cu_num_pages, cache_loc + ) + # Same cache_loc ptr -> identical (address-stable) buffer across replays. + assert buf1.data_ptr() == buf2.data_ptr() + assert len(_GlobalTrtllmPlanner._block_offsets_by_cache_loc) == 1 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/shim/test_ad_executor_swa_eviction.py b/tests/unittest/auto_deploy/singlegpu/shim/test_ad_executor_swa_eviction.py index 95c873f8a160..210b12d4fa36 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_ad_executor_swa_eviction.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_ad_executor_swa_eviction.py @@ -26,7 +26,10 @@ from tensorrt_llm._torch.auto_deploy._compat import KvCacheConfig from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import KVPagedResourceHandler -from tensorrt_llm._torch.auto_deploy.shim.ad_executor import _compute_window_local_view +from tensorrt_llm._torch.auto_deploy.shim.ad_executor import ( + _compute_cyclic_full_view, + _compute_window_local_view, +) from tensorrt_llm._torch.auto_deploy.shim.interface import CachedSequenceInterface pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @@ -312,5 +315,122 @@ def test_helper_does_not_consume_evicted_extra_slot(two_window_interface, monkey assert extra_page == -1 +# --------------------------------------------------------------------------- +# Multi-pool gate: trtllm (and any backend) may host >1 KV pool +# --------------------------------------------------------------------------- + + +def _build_two_pool_interface(requires_uniform_kv_caches: bool): + interface = CachedSequenceInterface( + max_seq_len=FULL_WINDOW, + max_batch_size=2, + max_num_tokens=default_max_num_tokens(FULL_WINDOW, 2), + device="cuda", + kv_cache_config=KvCacheConfig( + tokens_per_block=TOKENS_PER_BLOCK, + max_tokens=1024, + free_gpu_memory_fraction=0.0, + ), + requires_uniform_kv_caches=requires_uniform_kv_caches, + ) + interface.add_resource( + "kv_swa", KVPagedResourceHandler(4, 32, dtype=torch.float16, sliding_window=SWA_WINDOW) + ) + interface.add_resource("kv_full", KVPagedResourceHandler(4, 32, dtype=torch.float16)) + return interface + + +def test_two_distinct_windows_allowed_by_default(): + """Default (requires_uniform_kv_caches=False, the trtllm setting) hosts two pools.""" + interface = _build_two_pool_interface(requires_uniform_kv_caches=False) + interface.initialize_resources() # must not raise + # SWA pool + full-attention pool == two distinct windows. + windows = sorted( + {SWA_WINDOW, FULL_WINDOW} + & {pc.window_size for pc in interface._identify_managed_kv_resources()[1]} + ) + assert windows == [SWA_WINDOW, FULL_WINDOW] + + +def test_uniform_kv_caches_still_enforced_when_requested(): + """The uniformity mechanism is intact: opting in still rejects >1 pool. + + (No backend opts in today; trtllm now defaults to False -- this guards the + mechanism so a future single-pool backend can still rely on it.) + """ + interface = _build_two_pool_interface(requires_uniform_kv_caches=True) + with pytest.raises(RuntimeError, match="not uniform"): + interface.initialize_resources() + + +# --------------------------------------------------------------------------- +# Cyclic-SWA view (trtllm): full block table + global KV length, no slicing +# --------------------------------------------------------------------------- + + +def test_cyclic_view_passes_full_table_and_global_length(two_window_interface): + """Trtllm path: hand the kernel the FULL block table and the GLOBAL length. + + The trtllm kernel masks the sliding window internally via cyclic indexing, + so -- unlike the host-sliced triton/flashinfer path -- the executor must NOT + front-slice and must report the un-window-capped KV length. + """ + manager = two_window_interface.kv_cache_manager + # A prefill that exceeds the SWA window so window-local slicing WOULD differ. + prefill_len = SWA_WINDOW * 3 # 192 tokens + req = _add_request(manager, request_id=50, token_num=prefill_len) + all_indices = manager.get_cache_indices(req, window_size=SWA_WINDOW) + + active_indices, extra_page, swc, lpl = _compute_cyclic_full_view( + all_indices, + end_compute_i=prefill_len, + tokens_per_block=TOKENS_PER_BLOCK, + ) + + # Full table verbatim (no front-slice, no window cap). + assert active_indices == list(all_indices) + # Global (un-capped) KV length -- matches host_past_key_value_lengths. + assert swc == prefill_len + assert lpl == (prefill_len - 1) % TOKENS_PER_BLOCK + 1 + # No deferred-page insertion in cyclic mode. + assert extra_page == -1 + + +def test_cyclic_view_differs_from_window_local_when_evicted(two_window_interface, monkeypatch): + """Cyclic view ignores front-eviction; window-local view slices it off. + + Guards that the two staging paths genuinely diverge once the window has + been exceeded (so a backend mix-up would be caught). + """ + manager = two_window_interface.kv_cache_manager + front_removed = 2 + total_tokens = front_removed * TOKENS_PER_BLOCK + SWA_WINDOW + 1 + req = _add_request(manager, request_id=51, token_num=total_tokens) + monkeypatch.setattr( + manager, "get_num_front_blocks_removed", lambda req_id, window_size=None: front_removed + ) + all_indices = manager.get_cache_indices(req, window_size=SWA_WINDOW) + + cyc_indices, _, cyc_swc, _ = _compute_cyclic_full_view( + all_indices, end_compute_i=total_tokens, tokens_per_block=TOKENS_PER_BLOCK + ) + win_indices, _, win_swc, _ = _compute_window_local_view( + all_indices, + front_removed=front_removed, + end_compute_i=total_tokens, + group_window=SWA_WINDOW, + tokens_per_block=TOKENS_PER_BLOCK, + ) + + # Cyclic keeps the full list + global length; window-local slices + caps. + assert cyc_indices == list(all_indices) + assert cyc_swc == total_tokens + # Window-local view drops the stale front pages and starts at front_removed. + assert win_indices == list(all_indices[front_removed : front_removed + len(win_indices)]) + assert len(win_indices) < len(cyc_indices) + assert win_swc == total_tokens - front_removed * TOKENS_PER_BLOCK + assert cyc_swc != win_swc + + if __name__ == "__main__": pytest.main([__file__, "-v"]) 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) # =========================================================================== diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache_trtllm_multipool.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache_trtllm_multipool.py new file mode 100644 index 000000000000..54507c51dec6 --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache_trtllm_multipool.py @@ -0,0 +1,241 @@ +# SPDX-FileCopyrightText: Copyright (c) 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. + +"""Forward-level test for the trtllm attention backend with MULTIPLE KV pools. + +Builds a tiny two-layer model whose layers use different attention windows +(layer 0 = sliding window, layer 1 = full attention), so the AutoDeploy +kvcache transform creates two KV cache memory pools. Runs a prefill that +exceeds the sliding window through the cached ``trtllm`` attention op and +checks it matches the eager (uncached) reference. + +This is the on-GPU forward validation for issue #14828: it exercises the +unblocked multi-pool gate, the per-group block_offsets buffers, the +cyclic-SWA metadata staging (full block table + global KV length), and that +each pool's kernel receives its own attention window. +""" + +import pytest +import torch +import torch.nn as nn +from _torch_test_utils import all_close + +import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 +from tensorrt_llm._torch.auto_deploy._compat import KvCacheConfig +from tensorrt_llm._torch.auto_deploy.models.factory import FullModelExportInfo, ModelFactory +from tensorrt_llm._torch.auto_deploy.shim.interface import CachedSequenceInterface +from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + + +class _DummyFactory(ModelFactory): + def __init__(self, model): + self._model = model + + def build_model(self, device: str): + return self._model.to(device=device) + + def _build_model(self, device: str): + return + + def _load_checkpoint(self, model, device): + return + + def get_cache_config_updates(self): + return {} + + def get_export_infos(self, model): + return [FullModelExportInfo()] + + @property + def max_seq_len(self) -> int: + return 512 + + +class _WindowedAttnLayer(nn.Module): + def __init__(self, hidden: int, n_heads: int, sliding_window, layer_idx: int): + super().__init__() + self.n_heads = n_heads + self.head_dim = hidden // n_heads + self.sliding_window = sliding_window + self.layer_idx = layer_idx + self.q_proj = nn.Linear(hidden, hidden, bias=False) + self.k_proj = nn.Linear(hidden, hidden, bias=False) + self.v_proj = nn.Linear(hidden, hidden, bias=False) + self.o_proj = nn.Linear(hidden, hidden, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + b, s, _ = x.shape + q = self.q_proj(x).view(b, s, self.n_heads, self.head_dim) + k = self.k_proj(x).view(b, s, self.n_heads, self.head_dim) + v = self.v_proj(x).view(b, s, self.n_heads, self.head_dim) + o = torch.ops.auto_deploy.torch_attention( + q, + k, + v, + None, # attn_mask + 0.0, # dropout_p + True, # is_causal + None, # scale + None, # sinks + self.sliding_window, # sliding_window (int for SWA layer, None for full) + None, # logit_cap + "bsnd", # layout + self.layer_idx, # layer_idx + ) + return x + self.o_proj(o.reshape(b, s, -1)) + + +class _TwoWindowModel(nn.Module): + """Layer 0 = sliding window, layer 1 = full attention -> two KV pools.""" + + def __init__(self, vocab: int, hidden: int, n_heads: int, sliding_window: int): + super().__init__() + self.embed_tokens = nn.Embedding(vocab, hidden) + self.layer0 = _WindowedAttnLayer(hidden, n_heads, sliding_window, layer_idx=0) + self.layer1 = _WindowedAttnLayer(hidden, n_heads, None, layer_idx=1) + + @torch.no_grad() + def forward(self, input_ids: torch.Tensor, position_ids=None) -> torch.Tensor: + x = self.embed_tokens(input_ids) + x = self.layer0(x) + x = self.layer1(x) + return x + + +def _build_and_stage(sliding_window, seq_len, dtype=torch.float16): + """Build a 2-window model and run a single cyclic-staged prefill. + + Inserts trtllm cached attention (2 pools) and stages the prefill the way + ad_executor does for the cyclic (trtllm) path. Returns (eager_ref, cached_out). + """ + vocab, hidden, n_heads = 1000, 128, 2 + batch_size = 2 + tokens_per_block = 128 # >= max_seq_len -> 1 page per sequence per pool + max_seq_len = 128 + + kv_cache_config = KvCacheConfig( + tokens_per_block=tokens_per_block, + max_tokens=batch_size * tokens_per_block, + free_gpu_memory_fraction=0.0, + ) + cm = CachedSequenceInterface( + max_seq_len=max_seq_len, + max_batch_size=batch_size, + max_num_tokens=batch_size * max_seq_len, + device="cuda", + kv_cache_config=kv_cache_config, + ) + + model = _TwoWindowModel(vocab, hidden, n_heads, sliding_window).to(dtype=dtype, device="cuda") + input_ids = torch.randint(0, vocab, (batch_size, seq_len), device="cuda") + position_ids = torch.arange(seq_len, device="cuda").unsqueeze(0).repeat(batch_size, 1) + + y_ref = model(input_ids, position_ids) # eager reference (per-layer SWA masking) + + optimizer = InferenceOptimizer( + _DummyFactory(model), + { + "build_model": { + "stage": "factory", + "run_per_gm": False, + "device": "cuda", + "run_graph_cleanup": False, + "requires_clean_graph": False, + }, + "export_to_gm": { + "stage": "export", + "strict": False, + "run_per_gm": False, + "clone_state_dict": True, + "run_graph_cleanup": False, + "requires_clean_graph": False, + }, + "cleanup_input_constraints": {"stage": "post_export"}, + "insert_cached_attention": {"stage": "cache_init", "backend": "trtllm"}, + }, + ) + gm = optimizer(cm) + gm.to("cuda") + cm.initialize_resources() + + # Two distinct windows -> two pools, and trtllm uses the cyclic-SWA path. + assert len(cm.kv_group_windows) == 2, cm.kv_group_windows + assert cm.kernel_handles_cyclic_swa is True + + # Stage prefill metadata the way ad_executor does for the cyclic (trtllm) + # path: full per-window block table (1 page/seq here) + GLOBAL kv length. + n_pools = len(cm.kv_group_windows) + cache_loc_per_pool = [list(range(batch_size)) for _ in range(n_pools)] + cu_num_pages_per_pool = [list(range(batch_size + 1)) for _ in range(n_pools)] + seq_len_with_cache_per_pool = [[seq_len] * batch_size for _ in range(n_pools)] + last_page_len_per_pool = [ + [seq_len % tokens_per_block or tokens_per_block] * batch_size for _ in range(n_pools) + ] + extra_page_per_seq_per_pool = [[-1] * batch_size for _ in range(n_pools)] + + cm.info.reset() + cm.info.nest_sequences( + input_ids.flatten().tolist(), + cu_seqlen=list(range(0, batch_size * seq_len + 1, seq_len)), + input_pos=[0] * batch_size, + batch_info=[batch_size, batch_size * seq_len, 0, 0, 0, 0], + cache_loc_per_pool=cache_loc_per_pool, + cu_num_pages_per_pool=cu_num_pages_per_pool, + extra_page_per_seq_per_pool=extra_page_per_seq_per_pool, + seq_len_with_cache_per_pool=seq_len_with_cache_per_pool, + last_page_len_per_pool=last_page_len_per_pool, + slot_idx=list(range(batch_size)), + prompt_lens=[seq_len] * batch_size, + gather_context_logits=True, + ) + y_cached = torch.stack(cm.info.unnest_sequences(gm(**cm.named_args))) + return y_ref, y_cached + + +@torch.inference_mode() +def test_trtllm_two_pools_no_mask_matches_eager(): + """Two DISTINCT KV pools with no masking (both windows >= seq_len). + + Both layers do full causal attention. + Strict match against the eager reference validates the multi-pool feature: + two pools are created, each layer reads its OWN pool's block_offsets buffer + (no clobbering), and the cyclic full-table staging is correct. + """ + y_ref, y_cached = _build_and_stage(sliding_window=64, seq_len=48) + assert all_close(y_ref, y_cached, atol=2e-2, rtol=2e-2) + + +@torch.inference_mode() +def test_trtllm_two_pools_swa_engaged_runs(): + """Two pools with the SWA window strictly below the sequence length. + + Exercises the cyclic-SWA staging with a real sub-sequence window through + both pools. We assert it runs and produces finite, correctly-shaped output + rather than exact-matching the eager reference: the prefill sliding-window + mask is applied by the trtllm kernel, and on SMs where the trtllm-gen FMHA + is unavailable for the layer's shape the op falls back to an unfused MHA + that does not apply the context-phase window. Exact SWA-prefill correctness + on the supported kernel is covered by the PyTorch-backend contract (causal + mask + attention_window_size) this op mirrors. + """ + y_ref, y_cached = _build_and_stage(sliding_window=32, seq_len=96) + assert y_cached.shape == y_ref.shape + assert torch.isfinite(y_cached).all() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kvcache_vswa_metadata.py b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kvcache_vswa_metadata.py index 2c98ee893551..f089d756188e 100644 --- a/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kvcache_vswa_metadata.py +++ b/tests/unittest/auto_deploy/singlegpu/transformations/library/test_kvcache_vswa_metadata.py @@ -209,5 +209,22 @@ def find_swc_dep(node, visited=None): ) +@pytest.mark.parametrize( + "backend, expected_cyclic", + [("triton", False), ("trtllm", True)], +) +def test_vswa_sets_kernel_handles_cyclic_swa(backend, expected_cyclic): + """The transform records the backend's cyclic-SWA capability on the interface. + + trtllm's kernel masks the sliding window internally (cyclic), so the + executor must pass full block tables + global lengths; triton must not. + """ + gm, info, cm = _run_transform(backend=backend) + assert info.num_matches == 2 + assert cm.kernel_handles_cyclic_swa is expected_cyclic + # Both backends still register two window groups regardless of cyclic-ness. + assert len(cm.kv_group_windows) == 2 + + if __name__ == "__main__": pytest.main([__file__, "-v"])