From 7835e7bc1d3686f14430410ea26406dd176358a9 Mon Sep 17 00:00:00 2001 From: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> Date: Tue, 5 May 2026 17:18:21 -0700 Subject: [PATCH 1/5] Onboard GPT_OSS with new modeling file approach Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> --- docs/source/models/supported-models.md | 3 +- .../cookbooks/gpt_oss_trtllm_cookbook.ipynb | 281 +++++++++ .../model_registry/configs/gpt_oss_120b.yaml | 21 + .../model_registry/configs/gpt_oss_20b.yaml | 21 + .../auto_deploy/model_registry/models.yaml | 8 +- .../custom_ops/attention/trtllm_attention.py | 16 +- .../auto_deploy/models/custom/__init__.py | 1 + .../models/custom/modeling_gpt_oss.py | 499 +++++++++++++++ .../models/patches/gptoss-mxfp4.py | 68 -- .../transform/library/mxfp4_moe.py | 20 + .../singlegpu/models/test_gpt_oss_modeling.py | 586 ++++++++++++++++++ 11 files changed, 1450 insertions(+), 74 deletions(-) create mode 100644 examples/auto_deploy/cookbooks/gpt_oss_trtllm_cookbook.ipynb create mode 100644 examples/auto_deploy/model_registry/configs/gpt_oss_120b.yaml create mode 100644 examples/auto_deploy/model_registry/configs/gpt_oss_20b.yaml create mode 100644 tensorrt_llm/_torch/auto_deploy/models/custom/modeling_gpt_oss.py delete mode 100644 tensorrt_llm/_torch/auto_deploy/models/patches/gptoss-mxfp4.py create mode 100644 tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_gpt_oss_modeling.py diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 07736cb1aa80..a8816cda9d5f 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -18,7 +18,7 @@ The following is a table of supported models for the PyTorch backend: | `Glm4MoeForCausalLM` | GLM-4.5, GLM-4.6, GLM-4.7 | `THUDM/GLM-4-100B-A10B` | | `Glm4MoeLiteForCausalLM` [^6] | GLM-4.7-Flash | `zai-org/GLM-4.7-Flash` | | `GlmMoeDsaForCausalLM` | GLM-5 | `zai-org/GLM-5` | -| `GptOssForCausalLM` | GPT-OSS | `openai/gpt-oss-120b` | +| `GptOssForCausalLM` [^10] | GPT-OSS | `openai/gpt-oss-20b`, `openai/gpt-oss-120b` | | `KimiK25ForConditionalGeneration` | Kimi-K2.5 | `moonshotai/Kimi-K2.5` | | `LlamaForCausalLM` | Llama 3.1, Llama 3, Llama 2, LLaMA | `meta-llama/Meta-Llama-3.1-70B` | | `Llama4ForConditionalGeneration` | Llama 4 | `meta-llama/Llama-4-Scout-17B-16E-Instruct` | @@ -65,6 +65,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl [^7]: Text-only support via the [AutoDeploy](../features/auto_deploy/auto-deploy.md) backend. See AD configs for [MoE](../../../examples/auto_deploy/model_registry/configs/gemma4_moe.yaml) and [dense](../../../examples/auto_deploy/model_registry/configs/gemma4_dense.yaml). [^8]: Text-only support via the [AutoDeploy](../features/auto_deploy/auto-deploy.md) backend. See [AD config](../../../examples/auto_deploy/model_registry/configs/gemma3n_e2b_it.yaml). [^9]: Supported via the [AutoDeploy](../features/auto_deploy/auto-deploy.md) backend. See [AD config](../../../examples/auto_deploy/model_registry/configs/minimax_m2.7.yaml). +[^10]: Also supported via the [AutoDeploy](../features/auto_deploy/auto-deploy.md) backend (native PyTorch backend support is unchanged). See AD configs for [20B](../../../examples/auto_deploy/model_registry/configs/gpt_oss_20b.yaml) and [120B](../../../examples/auto_deploy/model_registry/configs/gpt_oss_120b.yaml). # Multimodal Feature Support Matrix (PyTorch Backend) diff --git a/examples/auto_deploy/cookbooks/gpt_oss_trtllm_cookbook.ipynb b/examples/auto_deploy/cookbooks/gpt_oss_trtllm_cookbook.ipynb new file mode 100644 index 000000000000..0dcf571feb0b --- /dev/null +++ b/examples/auto_deploy/cookbooks/gpt_oss_trtllm_cookbook.ipynb @@ -0,0 +1,281 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Deploying GPT-OSS with TensorRT-LLM (AutoDeploy)\n", + "\n", + "This notebook walks you through deploying OpenAI's `openai/gpt-oss-20b` and `openai/gpt-oss-120b` models using TensorRT-LLM's AutoDeploy backend.\n", + "\n", + "[TensorRT-LLM](https://nvidia.github.io/TensorRT-LLM/) is NVIDIA's open-source library for accelerating and optimizing LLM inference on NVIDIA GPUs. AutoDeploy is a one-shot graph-transform pipeline that translates a HuggingFace model into a deploy-ready engine — see the [AutoDeploy guide](https://nvidia.github.io/TensorRT-LLM/torch/auto_deploy/auto-deploy.html) for details.\n", + "\n", + "**Model Resources:**\n", + "- [HuggingFace — gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b)\n", + "- [HuggingFace — gpt-oss-120b](https://huggingface.co/openai/gpt-oss-120b)\n", + "- [OpenAI — gpt-oss announcement](https://openai.com/index/introducing-gpt-oss/)\n", + "- [OpenAI — gpt-oss model card](https://openai.com/index/gpt-oss-model-card/)\n", + "\n", + "**Model Highlights:**\n", + "- Mixture-of-Experts (MoE) architecture released by OpenAI in 2025 under Apache 2.0\n", + "- `gpt-oss-20b`: 21B parameters total, 3.6B active, 24 layers, 32 experts, top-4 routing, ~16 GB MXFP4 weights\n", + "- `gpt-oss-120b`: 117B parameters total, 5.1B active, 36 layers, 128 experts, top-4 routing, ~63 GB MXFP4 weights\n", + "- Native MXFP4 quantization for MoE weights (handled by AutoDeploy's `quantize_mxfp4_moe` transform)\n", + "- 64 query heads / 8 key-value heads (GQA), `head_dim=64`, `hidden_size=2880`\n", + "- Per-head learnable attention sinks; alternating sliding (window=128) and full-attention layers\n", + "- 131,072 token context length (YaRN-scaled RoPE)\n", + "- Channel-based reasoning protocol (`analysis` → `final`)\n", + "\n", + "**Prerequisites:**\n", + "- NVIDIA GPU with CUDA 12.x and recent drivers\n", + " - `gpt-oss-20b`: ≥ 1×80 GB GPU (the cookbook below pairs it with 2 GPUs for parallelism)\n", + " - `gpt-oss-120b`: ≥ 4×80 GB or 8×80 GB GPUs (the cookbook below uses 8)\n", + "- Python 3.10+\n", + "- TensorRT-LLM ([container](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tensorrt-llm/containers/release) or pip install)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Prerequisites & Environment\n", + "\n", + "Set up a containerized environment for TensorRT-LLM by running the following command in a terminal:\n", + "\n", + "```shell\n", + "docker run --rm -it --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --gpus=all -p 8000:8000 nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc1\n", + "```\n", + "\n", + "You now have TensorRT-LLM set up!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# If pip not found\n", + "!python -m ensurepip --default-pip" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%pip install torch openai" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Verify GPU\n", + "\n", + "Check that CUDA is available and the GPU is detected correctly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Environment check\n", + "import sys\n", + "\n", + "import torch\n", + "\n", + "print(f\"Python: {sys.version}\")\n", + "print(f\"CUDA available: {torch.cuda.is_available()}\")\n", + "print(f\"Num GPUs: {torch.cuda.device_count()}\")\n", + "\n", + "if torch.cuda.is_available():\n", + " for i in range(torch.cuda.device_count()):\n", + " print(f\"GPU[{i}]: {torch.cuda.get_device_name(i)}\")" + ] + }, + { + "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." + ] + }, + { + "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." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Your server is now running!" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Use the API\n", + "\n", + "Use the OpenAI-compatible client to send requests to the TensorRT-LLM server." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from openai import OpenAI\n", + "\n", + "# Setup client\n", + "BASE_URL = \"http://0.0.0.0:8000/v1\"\n", + "API_KEY = \"null\"\n", + "client = OpenAI(base_url=BASE_URL, api_key=API_KEY)\n", + "\n", + "# Set this to whichever model you started the server with.\n", + "MODEL_ID = \"openai/gpt-oss-20b\" # or \"openai/gpt-oss-120b\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Basic chat completion\n", + "print(\"Chat Completion Example\")\n", + "print(\"=\" * 50)\n", + "\n", + "response = client.chat.completions.create(\n", + " model=MODEL_ID,\n", + " messages=[\n", + " {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n", + " {\"role\": \"user\", \"content\": \"Where is the capital of Iceland?\"},\n", + " ],\n", + " temperature=1.0,\n", + " top_p=1.0,\n", + " max_tokens=128,\n", + ")\n", + "\n", + "print(\"Response:\")\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Streaming chat completion\n", + "print(\"Streaming response:\")\n", + "print(\"=\" * 50)\n", + "\n", + "stream = client.chat.completions.create(\n", + " model=MODEL_ID,\n", + " messages=[\n", + " {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n", + " {\"role\": \"user\", \"content\": \"Write a Python function that checks if a number is prime.\"},\n", + " ],\n", + " temperature=1.0,\n", + " max_tokens=512,\n", + " stream=True,\n", + ")\n", + "\n", + "for chunk in stream:\n", + " if chunk.choices[0].delta.content:\n", + " print(chunk.choices[0].delta.content, end=\"\", flush=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Evaluation Parameters\n", + "\n", + "OpenAI's gpt-oss model card recommends the following defaults:\n", + "\n", + "- `temperature`: 1.0\n", + "- `top_p`: 1.0\n", + "- `max_tokens`: 131072 (model's context limit; trim for serving SLOs)\n", + "\n", + "The model uses a channel-based reasoning protocol (`analysis` → `final`); the `analysis` segment is the model's chain-of-thought and the `final` segment is the user-facing answer. Both are returned via the OpenAI-compatible API as a single response — the chat template handles the formatting." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Additional Resources\n", + "\n", + "- [TensorRT-LLM Documentation](https://nvidia.github.io/TensorRT-LLM/)\n", + "- [AutoDeploy Guide](https://nvidia.github.io/TensorRT-LLM/torch/auto_deploy/auto-deploy.html)\n", + "- [gpt-oss-20b on HuggingFace](https://huggingface.co/openai/gpt-oss-20b)\n", + "- [gpt-oss-120b on HuggingFace](https://huggingface.co/openai/gpt-oss-120b)\n", + "- [OpenAI gpt-oss announcement blog](https://openai.com/index/introducing-gpt-oss/)\n", + "- [OpenAI gpt-oss model card](https://openai.com/index/gpt-oss-model-card/)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/auto_deploy/model_registry/configs/gpt_oss_120b.yaml b/examples/auto_deploy/model_registry/configs/gpt_oss_120b.yaml new file mode 100644 index 000000000000..d6486ed041dc --- /dev/null +++ b/examples/auto_deploy/model_registry/configs/gpt_oss_120b.yaml @@ -0,0 +1,21 @@ +# 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: 8 +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 new file mode 100644 index 000000000000..212164d6c709 --- /dev/null +++ b/examples/auto_deploy/model_registry/configs/gpt_oss_20b.yaml @@ -0,0 +1,21 @@ +# 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: 2 +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 8095a9761a96..09e94e5e460e 100644 --- a/examples/auto_deploy/model_registry/models.yaml +++ b/examples/auto_deploy/model_registry/models.yaml @@ -155,8 +155,8 @@ models: config_id: default_ws_2 yaml_extra: ['dashboard_default.yaml', 'world_size_2.yaml'] - name: openai/gpt-oss-20b - config_id: default_ws_2 - yaml_extra: ['dashboard_default.yaml', 'world_size_2.yaml'] + config_id: gpt_oss_20b + yaml_extra: ['gpt_oss_20b.yaml'] - name: ibm-granite/granite-3.0-8b-instruct config_id: default_ws_2 yaml_extra: ['dashboard_default.yaml', 'world_size_2.yaml'] @@ -293,8 +293,8 @@ models: config_id: multimodal yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'multimodal.yaml'] - name: openai/gpt-oss-120b - config_id: num_hidden_layers_5 - yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'num_hidden_layers_5.yaml'] + config_id: gpt_oss_120b + yaml_extra: ['gpt_oss_120b.yaml'] - name: meta-llama/Llama-4-Scout-17B-16E-Instruct config_id: multimodal__llama4_scout yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'multimodal.yaml', 'llama4_scout.yaml'] 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 07cba4c51f06..fe99ba1eaa2f 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 @@ -536,6 +536,7 @@ def trtllm_mha_with_cache( rotary_cos_sin: Optional[torch.Tensor] = None, position_embedding_type: int = 0, rotary_embedding_dim: int = 0, + attention_sinks: Optional[torch.Tensor] = None, ) -> torch.Tensor: """TRT-LLM attention with paged KV cache for Auto-Deploy. @@ -588,6 +589,12 @@ def trtllm_mha_with_cache( qkv_fused = torch.cat([q_flat, k_flat, v_flat], dim=-1).contiguous() + # The thop.attention C++ kernel expects attention_sinks in float32 (see + # attentionOp.cpp:365). Models typically hold sinks at the same dtype as + # the rest of the parameters (e.g. bf16 for GPT-OSS), so cast here. + if attention_sinks is not None and attention_sinks.dtype != torch.float32: + attention_sinks = attention_sinks.to(torch.float32) + # Prepare output: if caller provided an `out` buffer, write directly into it total_padded_tokens = q_shape_og[0] * q_shape_og[1] # If out_scale is set, attention quantizes output to FP8. @@ -648,7 +655,7 @@ def trtllm_mha_with_cache( None, # latent_cache (MLA) None, # q_pe (MLA) None, # block_ids_per_seq - None, # attention_sinks + attention_sinks, # attention_sinks (per-head learnable scalar, e.g. GPT-OSS) True, # is_fused_qkv True, # update_kv_cache 1, # predicted_tokens_per_seq (always 1 except for MLA kernel) @@ -740,6 +747,7 @@ def trtllm_mha_with_cache_fake( rotary_cos_sin: Optional[torch.Tensor] = None, position_embedding_type: int = 0, rotary_embedding_dim: int = 0, + attention_sinks: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Fake implementation for torch.compile tracing.""" if out is not None: @@ -895,6 +903,11 @@ def get_constants(cls, source_attn_node: Node) -> List[Constant]: # Get sliding_window from source attention node sliding_window = extract_op_args(source_attn_node, "sliding_window")[0] + # Forward optional attention sinks (per-head learnable scalar) — used by + # GPT-OSS-style models. Stored as an FX get_attr Node when present and + # bound to the actual nn.Parameter at runtime. + sinks_node = extract_op_args(source_attn_node, "sinks")[0] + # Optional out_scale is injected by prepare_node_for_cache_insertion when available. out_scale = source_attn_node.meta.get(_TRTLLM_ATTN_OUT_SCALE_KEY) if not isinstance(out_scale, Node): @@ -921,4 +934,5 @@ def get_constants(cls, source_attn_node: Node) -> List[Constant]: rope_cos_sin, pos_emb_type, rot_emb_dim, + sinks_node, ] diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py b/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py index 85006863656f..fc126ba494f2 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py @@ -12,6 +12,7 @@ "modeling_gemma3n": ["Gemma3nForCausalLM", "Gemma3nForConditionalGeneration"], "modeling_gemma4": ["Gemma4ForCausalLM", "Gemma4ForConditionalGeneration"], "modeling_glm4_moe_lite": ["Glm4MoeLiteForCausalLM"], + "modeling_gpt_oss": ["GptOssForCausalLM"], "modeling_kimi_k2": ["KimiK2ForCausalLM", "KimiK25ForConditionalGeneration"], "modeling_llama4": ["Llama4ForCausalLM", "Llama4ForConditionalGeneration"], "modeling_minimax_m2": ["MiniMaxM2ForCausalLM"], 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 new file mode 100644 index 000000000000..2da8ecd52a6b --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_gpt_oss.py @@ -0,0 +1,499 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Slimmed-down PyTorch GPT-OSS model for AutoDeploy export (prefill only). + +Source: + https://huggingface.co/openai/gpt-oss-20b + https://huggingface.co/openai/gpt-oss-120b + +Both 20b and 120b share the same architecture (only num_hidden_layers and +num_local_experts differ), so this file covers both variants. + +Key architecture features: +* GQA: 64 Q heads / 8 KV heads, head_dim=64, hidden_size=2880 +* Attention sinks: per-head learnable scalar concatenated into softmax denominator +* Alternating sliding/full attention by layer (sliding_window=128) +* YaRN-scaled RoPE (factor=32, original_max=4096), Llama-style half-rotary +* MoE: 32 experts (20b) / 128 experts (120b), top-4 routing +* Stacked MoE weights with biases on both gate_up and down projections +* Custom GLU activation: ``(up + 1) * gate * sigmoid(gate * 1.702)`` with + ``gate.clamp(max=7)`` and ``up.clamp(-7, 7)`` +* MXFP4 quantized MoE weights handled by the AD ``quantize_mxfp4_moe`` transform + +Differences from the HF reference (modeling_gpt_oss.py): +* Stripped KV cache, training paths, dropout, mask construction, deprecated kwargs +* Uses AD canonical ops: + - ``torch_rmsnorm`` (normalization) + - ``torch_attention`` (with ``sinks=`` and ``sliding_window=``) + - ``torch_rope_with_explicit_cos_sin`` + - ``torch_moe_router`` (linear + topk + softmax + scatter) + - ``torch_moe_dense_mlp`` (dense bmm-based GPT-OSS expert math) +* 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`` +""" + +import math +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +import torch.nn as nn +from transformers.generation import GenerationMixin +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import ModelOutput + +from ..hf import AutoModelForCausalLMFactory + +# GPT-OSS hard-codes these in the HF reference (see modeling_gpt_oss.GptOssExperts). +# ``alpha`` controls the SwiGLU sigmoid scaling, ``limit`` clamps gate/up before the GLU. +_GPTOSS_GLU_ALPHA = 1.702 +_GPTOSS_GLU_LIMIT_FALLBACK = 7.0 + + +# --------------------------------------------------------------------------- +# Output dataclasses +# --------------------------------------------------------------------------- + + +@dataclass +class GptOssModelOutput(ModelOutput): + last_hidden_state: Optional[torch.FloatTensor] = None + + +@dataclass +class GptOssCausalLMOutput(ModelOutput): + logits: Optional[torch.FloatTensor] = None + + +# --------------------------------------------------------------------------- +# YaRN helpers (faithful copy of transformers._compute_yarn_parameters) +# --------------------------------------------------------------------------- + + +def _yarn_get_mscale(scale: float, mscale: float = 1.0) -> float: + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +def _yarn_find_correction_dim(num_rot: float, dim: int, base: float, max_pos: int) -> float: + return (dim * math.log(max_pos / (num_rot * 2 * math.pi))) / (2 * math.log(base)) + + +def _yarn_find_correction_range( + low_rot: float, high_rot: float, dim: int, base: float, max_pos: int, truncate: bool +) -> Tuple[float, float]: + low = _yarn_find_correction_dim(low_rot, dim, base, max_pos) + high = _yarn_find_correction_dim(high_rot, dim, base, max_pos) + if truncate: + low = math.floor(low) + high = math.ceil(high) + return max(low, 0), min(high, dim - 1) + + +def _yarn_linear_ramp_factor(min_v: float, max_v: float, dim: int) -> torch.Tensor: + if min_v == max_v: + max_v = max_v + 0.001 + factor = (torch.arange(dim, dtype=torch.float32) - min_v) / (max_v - min_v) + return torch.clamp(factor, 0.0, 1.0) + + +# --------------------------------------------------------------------------- +# RMSNorm (using AD canonical op) +# --------------------------------------------------------------------------- + + +class GptOssRMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.auto_deploy.torch_rmsnorm(x, self.weight, self.eps) + + +# --------------------------------------------------------------------------- +# Rotary Embedding (YaRN, pre-cached, sliced once per forward) +# --------------------------------------------------------------------------- + + +class GptOssRotaryEmbedding(nn.Module): + """YaRN-scaled rotary embedding for GPT-OSS. + + The HF reference applies RoPE via ``torch.chunk(x, 2, dim=-1)`` with cos/sin + of length ``head_dim/2``. This is mathematically identical to the standard + Llama RoPE (``rotate_half`` + ``cos = sin = cat(freqs, freqs)``), so we cache + a duplicated ``[max_pos, head_dim]`` table and feed it to the AD canonical + ``torch_rope_with_explicit_cos_sin`` op. + """ + + def __init__( + self, + head_dim: int, + max_position_embeddings: int, + rope_theta: float, + rope_scaling: Optional[dict] = None, + ): + super().__init__() + + attention_scaling = 1.0 + if rope_scaling is not None: + rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", "default")) + else: + rope_type = "default" + + if rope_type == "yarn": + factor = float(rope_scaling["factor"]) + beta_fast = float(rope_scaling.get("beta_fast", 32.0)) + beta_slow = float(rope_scaling.get("beta_slow", 1.0)) + mscale = rope_scaling.get("mscale", None) + mscale_all_dim = rope_scaling.get("mscale_all_dim", None) + attention_factor = rope_scaling.get("attention_factor", None) + original_max = int( + rope_scaling.get("original_max_position_embeddings") or max_position_embeddings + ) + truncate = bool(rope_scaling.get("truncate", True)) + + if attention_factor is None: + if mscale and mscale_all_dim: + attention_scaling = float( + _yarn_get_mscale(factor, float(mscale)) + / _yarn_get_mscale(factor, float(mscale_all_dim)) + ) + else: + attention_scaling = _yarn_get_mscale(factor) + else: + attention_scaling = float(attention_factor) + + pos_freqs = rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim) + inv_freq_extra = 1.0 / pos_freqs + inv_freq_inter = 1.0 / (factor * pos_freqs) + + low, high = _yarn_find_correction_range( + beta_fast, beta_slow, head_dim, rope_theta, original_max, truncate + ) + extra_factor = 1.0 - _yarn_linear_ramp_factor(low, high, head_dim // 2) + inv_freq = inv_freq_inter * (1.0 - extra_factor) + inv_freq_extra * extra_factor + else: + inv_freq = 1.0 / ( + rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim) + ) + + t = torch.arange(max_position_embeddings, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("_ad_cos_cached", emb.cos() * attention_scaling, persistent=False) + self.register_buffer("_ad_sin_cached", emb.sin() * attention_scaling, persistent=False) + + def forward( + self, x: torch.Tensor, position_ids: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + cos = self._ad_cos_cached[position_ids].to(dtype=x.dtype, device=x.device) + sin = self._ad_sin_cached[position_ids].to(dtype=x.dtype, device=x.device) + return cos, sin + + +# --------------------------------------------------------------------------- +# Router (replaces HF GptOssTopKRouter; eliminates the gptoss_topk_router patch) +# --------------------------------------------------------------------------- + + +class GptOssTopKRouter(nn.Module): + """Top-K router: linear projection + topk + softmax + scatter. + + Produces ``router_scores`` of shape ``[B*S, num_experts]`` with non-zero + entries only at the top-k expert positions, summing to 1 along dim=-1. + """ + + def __init__(self, config): + super().__init__() + self.top_k = int(config.num_experts_per_tok) + self.num_experts = int(config.num_local_experts) + self.weight = nn.Parameter(torch.empty(self.num_experts, config.hidden_size)) + self.bias = nn.Parameter(torch.empty(self.num_experts)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return torch.ops.auto_deploy.torch_moe_router( + hidden_states, self.weight, self.bias, self.top_k + ) + + +# --------------------------------------------------------------------------- +# Experts (stacked weights with biases; uses torch_moe_dense_mlp) +# --------------------------------------------------------------------------- + + +class GptOssExperts(nn.Module): + """GPT-OSS dense experts module. + + Holds the four stacked parameters that match the HF safetensors layout: + gate_up_proj : [E, H, 2I] (gate and up interleaved on the last dim) + gate_up_proj_bias : [E, 2I] + down_proj : [E, I, H] + down_proj_bias : [E, H] + + The forward delegates to ``torch_moe_dense_mlp``, which encodes GPT-OSS's + 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. + """ + + def __init__(self, config): + super().__init__() + self.num_experts = int(config.num_local_experts) + self.hidden_size = int(config.hidden_size) + self.expert_dim = int(config.intermediate_size) + self.alpha = _GPTOSS_GLU_ALPHA + # The HF safetensors / config carry ``swiglu_limit``; fall back to 7.0 + # for synthetic configs that omit it. + self.limit = float(getattr(config, "swiglu_limit", _GPTOSS_GLU_LIMIT_FALLBACK)) + + self.gate_up_proj = nn.Parameter( + torch.empty(self.num_experts, self.hidden_size, 2 * self.expert_dim) + ) + self.gate_up_proj_bias = nn.Parameter(torch.empty(self.num_experts, 2 * self.expert_dim)) + self.down_proj = nn.Parameter( + torch.empty(self.num_experts, self.expert_dim, self.hidden_size) + ) + self.down_proj_bias = nn.Parameter(torch.empty(self.num_experts, self.hidden_size)) + + def forward(self, hidden_states: torch.Tensor, routing_weights: torch.Tensor) -> torch.Tensor: + return torch.ops.auto_deploy.torch_moe_dense_mlp( + hidden_states, + routing_weights, + self.gate_up_proj, + self.gate_up_proj_bias, + self.down_proj, + self.down_proj_bias, + self.alpha, + self.limit, + ) + + +class GptOssMLP(nn.Module): + """Router + experts. Drop-in replacement for HF ``GptOssMLP`` in prefill.""" + + def __init__(self, config): + super().__init__() + self.router = GptOssTopKRouter(config) + self.experts = GptOssExperts(config) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + bsz, seq_len, hidden_dim = hidden_states.shape + routing_weights = self.router(hidden_states) # [B*S, E] + out = self.experts(hidden_states, routing_weights) + return out.view(bsz, seq_len, hidden_dim) + + +# --------------------------------------------------------------------------- +# Attention (GQA + sinks + per-layer sliding window) +# --------------------------------------------------------------------------- + + +class GptOssAttention(nn.Module): + """GPT-OSS attention with learnable per-head sinks and optional sliding window.""" + + def __init__(self, config, layer_idx: int): + super().__init__() + self.layer_idx = layer_idx + self.head_dim = int( + getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + ) + self.num_heads = int(config.num_attention_heads) + self.num_kv_heads = int(config.num_key_value_heads) + self.scaling = self.head_dim**-0.5 + self.attention_bias = bool(getattr(config, "attention_bias", True)) + + self.q_proj = nn.Linear( + config.hidden_size, self.num_heads * self.head_dim, bias=self.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, self.num_kv_heads * self.head_dim, bias=self.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, self.num_kv_heads * self.head_dim, bias=self.attention_bias + ) + self.o_proj = nn.Linear( + self.num_heads * self.head_dim, config.hidden_size, bias=self.attention_bias + ) + + self.sinks = nn.Parameter(torch.empty(self.num_heads)) + + # Per-layer sliding window: only enabled on layers tagged "sliding_attention". + layer_types = getattr(config, "layer_types", None) + is_sliding = layer_types is not None and layer_types[layer_idx] == "sliding_attention" + sliding_window = getattr(config, "sliding_window", None) + self.sliding_window = int(sliding_window) if (is_sliding and sliding_window) else None + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + ) -> 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) + + cos, sin = position_embeddings # [B, S, head_dim] + # Apply RoPE with unsqueeze_dim=2 for BSND layout. + 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 + # per-call kwargs. Causal mask is applied internally for prefill. + attn_output = torch.ops.auto_deploy.torch_attention( + q, + k, + v, + attn_mask=None, + dropout_p=0.0, + is_causal=True, + scale=self.scaling, + sinks=self.sinks, + 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) + + +# --------------------------------------------------------------------------- +# Decoder Layer +# --------------------------------------------------------------------------- + + +class GptOssDecoderLayer(nn.Module): + def __init__(self, config, layer_idx: int): + super().__init__() + self.self_attn = GptOssAttention(config, layer_idx) + self.mlp = GptOssMLP(config) + self.input_layernorm = GptOssRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = GptOssRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, position_embeddings=position_embeddings) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +# --------------------------------------------------------------------------- +# Model + CausalLM +# --------------------------------------------------------------------------- + + +class GptOssPreTrainedModel(PreTrainedModel): + base_model_prefix = "model" + _no_split_modules = ["GptOssDecoderLayer"] + supports_gradient_checkpointing = False + + +class GptOssModel(GptOssPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, getattr(config, "pad_token_id", None) + ) + self.layers = nn.ModuleList( + [GptOssDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = GptOssRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + head_dim = int( + getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + ) + 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_scaling=getattr(config, "rope_scaling", None), + ) + + self.post_init() + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> GptOssModelOutput: + assert position_ids is not None, "position_ids is required" + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + position_embeddings = self.rotary_emb(inputs_embeds, position_ids) + + hidden_states = inputs_embeds + for layer in self.layers: + hidden_states = layer(hidden_states, position_embeddings=position_embeddings) + hidden_states = self.norm(hidden_states) + return GptOssModelOutput(last_hidden_state=hidden_states) + + +class GptOssForCausalLM(GptOssPreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = GptOssModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, new_embeddings): + self.model.embed_tokens = new_embeddings + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> GptOssCausalLMOutput: + assert position_ids is not None, "position_ids is required" + outputs = self.model( + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + logits = self.lm_head(outputs.last_hidden_state) + return GptOssCausalLMOutput(logits=logits) + + +# --------------------------------------------------------------------------- +# Registration +# --------------------------------------------------------------------------- + +AutoModelForCausalLMFactory.register_custom_model_cls("GptOssConfig", GptOssForCausalLM) diff --git a/tensorrt_llm/_torch/auto_deploy/models/patches/gptoss-mxfp4.py b/tensorrt_llm/_torch/auto_deploy/models/patches/gptoss-mxfp4.py deleted file mode 100644 index a70c62bddfac..000000000000 --- a/tensorrt_llm/_torch/auto_deploy/models/patches/gptoss-mxfp4.py +++ /dev/null @@ -1,68 +0,0 @@ -from typing import Tuple - -import torch -from transformers.models.gpt_oss.modeling_gpt_oss import GptOssTopKRouter - -from ...export.interface import BaseExportPatch, ExportPatchRegistry - - -def _forward_router( - self: "GptOssTopKRouter", hidden_states: torch.Tensor -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Patched forward: - - Calls fused router op (returns only scores) - - Derives router_indices via topk(scores) to preserve original API - Returns: - router_scores: [T, E] - router_indices: [T, top_k] - """ - hs = hidden_states - # The custom op's fake kernel assumes 3D; ensure [B, S, H] if the caller gives [T, H] - if hs.dim() == 2: - hs = hs.unsqueeze(0) - - out = torch.ops.auto_deploy.torch_moe_router(hs, self.weight, self.bias, int(self.top_k)) - router_scores = out # [B*S, E] - - return router_scores, None - - -@ExportPatchRegistry.register("gptoss_topk_router") -class GptOssTopKRouterPatch(BaseExportPatch): - """Patch for GptOssTopKRouter to use the fused torch_moe_router custom op during export.""" - - def _apply_patch(self): - cls = self._resolve_router_class() - - # Keep original - key = f"{cls.__module__}.{cls.__qualname__}.forward" - self.original_values[key] = cls.forward - - # Apply patch - cls._original_forward = cls.forward - cls.forward = _forward_router - - def _revert_patch(self): - cls = self._resolve_router_class() - key = f"{cls.__module__}.{cls.__qualname__}.forward" - - # Restore original - cls.forward = self.original_values[key] - - # Cleanup - if hasattr(cls, "_original_forward"): - delattr(cls, "_original_forward") - - @staticmethod - def _resolve_router_class(): - """ - Resolve the GptOssTopKRouter class. If it's not in the current module scope, - import it from your project location instead. - """ - try: - return GptOssTopKRouter - except NameError: - from transformers.models.gpt_oss.modeling_gpt_oss import GptOssTopKRouter as _Cls - - return _Cls diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/mxfp4_moe.py b/tensorrt_llm/_torch/auto_deploy/transform/library/mxfp4_moe.py index 4123205c6cf1..3664bed0725a 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/mxfp4_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/mxfp4_moe.py @@ -189,6 +189,17 @@ def _register_mxfp4_expert_params( 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 ( @@ -301,6 +312,15 @@ def _apply( 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( diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_gpt_oss_modeling.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_gpt_oss_modeling.py new file mode 100644 index 000000000000..54a52bd939dd --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_gpt_oss_modeling.py @@ -0,0 +1,586 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Hierarchical equivalence tests for GPT-OSS AutoDeploy custom model. + +Reference is the HuggingFace GPT-OSS modeling code shipped in transformers +(``transformers.models.gpt_oss``). When that import is unavailable we skip +the affected tests rather than reproduce the math by hand. +""" + +from typing import Optional, Tuple + +import pytest +import torch +from torch.export import Dim +from transformers import GptOssConfig + +import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 — register canonical ops +from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.models.custom.modeling_gpt_oss import ( + GptOssAttention, + GptOssDecoderLayer, + GptOssExperts, + GptOssForCausalLM, + GptOssMLP, + GptOssRMSNorm, + GptOssRotaryEmbedding, + GptOssTopKRouter, +) + +# --------------------------------------------------------------------------- +# HF reference imports (skip tests if unavailable) +# --------------------------------------------------------------------------- + + +def _get_hf_classes(): + try: + from transformers.models.gpt_oss.modeling_gpt_oss import GptOssAttention as HFAttention + from transformers.models.gpt_oss.modeling_gpt_oss import GptOssDecoderLayer as HFLayer + from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts as HFExperts + from transformers.models.gpt_oss.modeling_gpt_oss import GptOssForCausalLM as HFForCausalLM + from transformers.models.gpt_oss.modeling_gpt_oss import GptOssMLP as HFMLP + from transformers.models.gpt_oss.modeling_gpt_oss import GptOssRMSNorm as HFRMSNorm + from transformers.models.gpt_oss.modeling_gpt_oss import GptOssRotaryEmbedding as HFRotary + from transformers.models.gpt_oss.modeling_gpt_oss import GptOssTopKRouter as HFRouter + + return { + "RMSNorm": HFRMSNorm, + "Rotary": HFRotary, + "Router": HFRouter, + "Experts": HFExperts, + "MLP": HFMLP, + "Attention": HFAttention, + "Layer": HFLayer, + "ForCausalLM": HFForCausalLM, + } + except ImportError: + return None + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def assert_rmse_close( + actual: torch.Tensor, expected: torch.Tensor, rmse_ratio_tol: float, msg: str = "" +) -> None: + diff = actual.float() - expected.float() + rmse_diff = torch.sqrt(torch.mean(diff**2)) + rmse_ref = torch.sqrt(torch.mean(expected.float() ** 2)) + ratio = (rmse_diff / rmse_ref.clamp(min=1e-12)).item() + assert ratio < rmse_ratio_tol, ( + f"{msg}RMSE ratio {ratio:.6f} exceeds tolerance {rmse_ratio_tol}. " + f"(rmse_diff={rmse_diff.item():.6f}, rmse_ref={rmse_ref.item():.6f})" + ) + + +def _device_and_dtype() -> Tuple[str, torch.dtype]: + if torch.cuda.is_available(): + return "cuda", torch.bfloat16 + return "cpu", torch.float32 + + +def _position_ids(batch: int, seq: int, device) -> torch.Tensor: + return torch.arange(seq, device=device).unsqueeze(0).expand(batch, -1) + + +def _small_config(num_layers: int = 3) -> GptOssConfig: + """Tiny but representative config covering both sliding and full attention. + + We use ``rope_type=default`` (no YaRN) here so equivalence comparisons + against HF do not have to thread through the rope scaling code path. + """ + return GptOssConfig( + num_hidden_layers=num_layers, + num_local_experts=4, + num_experts_per_tok=2, + vocab_size=1000, + hidden_size=64, + intermediate_size=32, + head_dim=16, + num_attention_heads=4, + num_key_value_heads=2, + sliding_window=4, + max_position_embeddings=64, + rope_scaling={"rope_type": "default"}, + rope_theta=10000.0, + rms_norm_eps=1e-5, + attention_bias=True, + attention_dropout=0.0, + initializer_range=0.02, + tie_word_embeddings=False, + pad_token_id=0, + ) + + +def _rope_position_embeddings( + config: GptOssConfig, B: int, S: int, device, dtype +) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: + """Return (ad_pos_emb, hf_pos_emb). + + AD pos_emb: (cos, sin) of shape ``[B, S, head_dim]`` (Llama-style duplicated). + HF pos_emb: (cos, sin) of shape ``[B, S, head_dim/2]`` (GPT-OSS native). + """ + head_dim = config.head_dim + # AD-side + rope = GptOssRotaryEmbedding( + head_dim=head_dim, + max_position_embeddings=config.max_position_embeddings, + rope_theta=config.rope_theta, + rope_scaling=config.rope_scaling, + ).to(device=device) + pos_ids = _position_ids(B, S, device) + dummy = torch.zeros(1, device=device, dtype=dtype) + ad_cos, ad_sin = rope(dummy, pos_ids) # [B, S, head_dim] + + # HF-side: cos/sin are half-size; equal to the front half of AD's duplicated cache. + hf_cos = ad_cos[..., : head_dim // 2] + hf_sin = ad_sin[..., : head_dim // 2] + return (ad_cos, ad_sin), (hf_cos, hf_sin) + + +# --------------------------------------------------------------------------- +# Block tests +# --------------------------------------------------------------------------- + + +def test_rmsnorm_equivalence(): + hf = _get_hf_classes() + if hf is None: + pytest.skip("transformers.models.gpt_oss not available") + device, dtype = _device_and_dtype() + H = 64 + ref = hf["RMSNorm"](H, eps=1e-5).to(device=device, dtype=dtype).eval() + ad = GptOssRMSNorm(H, eps=1e-5).to(device=device, dtype=dtype).eval() + ad.load_state_dict(ref.state_dict()) + + x = torch.randn(2, 8, H, device=device, dtype=dtype) + with torch.no_grad(): + torch.testing.assert_close(ad(x), ref(x), rtol=1e-3, atol=1e-3) + + +def test_rotary_embedding_equivalence(): + """Our pre-cached RoPE table matches HF's on-the-fly computation.""" + hf = _get_hf_classes() + if hf is None: + pytest.skip("transformers.models.gpt_oss not available") + device, dtype = _device_and_dtype() + config = _small_config() + B, S = 2, 8 + + ad_pos_emb, hf_pos_emb = _rope_position_embeddings(config, B, S, device, dtype) + ad_cos, ad_sin = ad_pos_emb + hf_cos_half, hf_sin_half = hf_pos_emb + + # Build HF rotary the way the HF model does and ask it for cos/sin. + hf_rope = hf["Rotary"](config=config).to(device=device).eval() + pos_ids = _position_ids(B, S, device) + dummy = torch.zeros(1, device=device, dtype=dtype) + with torch.no_grad(): + hf_cos_full, hf_sin_full = hf_rope(dummy, pos_ids) # [B, S, head_dim/2] + # The AD cache uses cat((freqs, freqs)); HF returns just freqs. Compare first half. + torch.testing.assert_close( + ad_cos[..., : config.head_dim // 2], hf_cos_full, rtol=1e-3, atol=1e-3 + ) + torch.testing.assert_close( + ad_sin[..., : config.head_dim // 2], hf_sin_full, rtol=1e-3, atol=1e-3 + ) + # Second half should equal the first half by construction (Llama-style duplication). + torch.testing.assert_close( + ad_cos[..., config.head_dim // 2 :], hf_cos_half, rtol=1e-3, atol=1e-3 + ) + torch.testing.assert_close( + ad_sin[..., config.head_dim // 2 :], hf_sin_half, rtol=1e-3, atol=1e-3 + ) + + +def test_router_equivalence(): + """torch_moe_router scatter output matches HF's router_scores.""" + hf = _get_hf_classes() + if hf is None: + pytest.skip("transformers.models.gpt_oss not available") + device, dtype = _device_and_dtype() + config = _small_config() + + ref = hf["Router"](config).to(device=device, dtype=dtype).eval() + # Initialise so weights aren't garbage from torch.empty. + torch.nn.init.normal_(ref.weight, std=0.02) + torch.nn.init.normal_(ref.bias, std=0.02) + ad = GptOssTopKRouter(config).to(device=device, dtype=dtype).eval() + ad.load_state_dict(ref.state_dict()) + + x = torch.randn(2, 8, config.hidden_size, device=device, dtype=dtype) + with torch.no_grad(): + ref_scores, _ = ref(x) # [B*S, E] + ad_scores = ad(x) # [B*S, E] + assert_rmse_close(ad_scores, ref_scores, rmse_ratio_tol=1e-3, msg="Router: ") + + +def _init_experts_(ref_experts, ad_experts): + """Initialise both experts modules with matching random weights.""" + torch.nn.init.normal_(ref_experts.gate_up_proj, std=0.02) + torch.nn.init.normal_(ref_experts.gate_up_proj_bias, std=0.02) + torch.nn.init.normal_(ref_experts.down_proj, std=0.02) + torch.nn.init.normal_(ref_experts.down_proj_bias, std=0.02) + ad_experts.load_state_dict(ref_experts.state_dict()) + + +def test_experts_equivalence(): + """torch_moe_dense_mlp matches HF GptOssExperts inference path.""" + hf = _get_hf_classes() + if hf is None: + pytest.skip("transformers.models.gpt_oss not available") + device, dtype = _device_and_dtype() + config = _small_config() + B, S = 2, 8 + T = B * S + E = config.num_local_experts + + ref = hf["Experts"](config).to(device=device, dtype=dtype).eval() + ad = GptOssExperts(config).to(device=device, dtype=dtype).eval() + _init_experts_(ref, ad) + + # Top-k routing weights with the same structure produced by GptOssTopKRouter. + router_logits = torch.randn(T, E, device=device, dtype=dtype) + top_v, top_idx = torch.topk(router_logits, config.num_experts_per_tok, dim=-1) + top_v = torch.softmax(top_v, dim=1, dtype=top_v.dtype) + routing_weights = torch.zeros_like(router_logits).scatter_(1, top_idx, top_v) + + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + with torch.no_grad(): + # HF GptOssExperts is in inference mode here; it picks the dense bmm branch on CUDA. + # On CPU it picks the sparse branch — both are mathematically equivalent. + ref_out = ref(x, router_indices=top_idx, routing_weights=routing_weights) + ad_out = ad(x, routing_weights) + ad_out = ad_out.view(B, S, -1) + assert_rmse_close(ad_out, ref_out, rmse_ratio_tol=0.02, msg="Experts: ") + + +def test_mlp_equivalence(): + """Full MoE block (router + experts) matches HF GptOssMLP.""" + hf = _get_hf_classes() + if hf is None: + pytest.skip("transformers.models.gpt_oss not available") + device, dtype = _device_and_dtype() + config = _small_config() + B, S = 2, 8 + + ref = hf["MLP"](config).to(device=device, dtype=dtype).eval() + ad = GptOssMLP(config).to(device=device, dtype=dtype).eval() + # Random init for the ref, then copy across. + torch.nn.init.normal_(ref.router.weight, std=0.02) + torch.nn.init.normal_(ref.router.bias, std=0.02) + torch.nn.init.normal_(ref.experts.gate_up_proj, std=0.02) + torch.nn.init.normal_(ref.experts.gate_up_proj_bias, std=0.02) + torch.nn.init.normal_(ref.experts.down_proj, std=0.02) + torch.nn.init.normal_(ref.experts.down_proj_bias, std=0.02) + ad.load_state_dict(ref.state_dict()) + + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + with torch.no_grad(): + ref_out, _ = ref(x) + ad_out = ad(x) + assert_rmse_close(ad_out, ref_out, rmse_ratio_tol=0.02, msg="MoE block: ") + + +# --------------------------------------------------------------------------- +# Attention test +# --------------------------------------------------------------------------- + + +def _init_attention_(ref: torch.nn.Module, ad: torch.nn.Module): + for proj in ("q_proj", "k_proj", "v_proj", "o_proj"): + torch.nn.init.normal_(getattr(ref, proj).weight, std=0.02) + if getattr(ref, proj).bias is not None: + torch.nn.init.normal_(getattr(ref, proj).bias, std=0.02) + torch.nn.init.normal_(ref.sinks, std=0.02) + ad.load_state_dict(ref.state_dict()) + + +def _hf_attention_forward( + hf_attn, + hidden_states: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> torch.Tensor: + """Run the HF GptOssAttention forward in eager mode without a mask. + + The HF attention requires an ``_attn_implementation == 'eager'`` attribute on + the config (we set it on the test config). We pass ``attention_mask=None`` — + the eager kernel applies its own causal handling via the sinks normaliser. + """ + # GPT-OSS HF eager kernel does not apply a causal mask itself when ``attention_mask`` is None. + # Build a [b,1,s,s] additive mask with -inf on the upper triangle so prefill is causal. + bsz, q_len, _ = hidden_states.shape + # We are only ever passing q_len tokens with no past, so s_q == s_k == q_len. + causal = torch.triu( + torch.full( + (q_len, q_len), float("-inf"), device=hidden_states.device, dtype=hidden_states.dtype + ), + diagonal=1, + ).view(1, 1, q_len, q_len) + out, _ = hf_attn( + hidden_states=hidden_states, + position_embeddings=(cos, sin), + attention_mask=causal, + ) + return out + + +def test_attention_equivalence_full(): + """GQA attention with sinks (full attention layer).""" + hf = _get_hf_classes() + if hf is None: + pytest.skip("transformers.models.gpt_oss not available") + device, dtype = _device_and_dtype() + config = _small_config() + config._attn_implementation = "eager" + # Use full-attention layer index (config alternates: 0=sliding, 1=full). + layer_idx = 1 + B, S = 2, 8 + + ref = hf["Attention"](config, layer_idx=layer_idx).to(device=device, dtype=dtype).eval() + ad = GptOssAttention(config, layer_idx=layer_idx).to(device=device, dtype=dtype).eval() + _init_attention_(ref, ad) + assert ad.sliding_window is None # this layer is "full_attention" + + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + ad_pos, hf_pos = _rope_position_embeddings(config, B, S, device, dtype) + + with torch.no_grad(): + ad_out = ad(x, position_embeddings=ad_pos) + ref_out = _hf_attention_forward(ref, x, hf_pos[0], hf_pos[1]) + assert_rmse_close(ad_out, ref_out, rmse_ratio_tol=0.10, msg="Attention (full): ") + + +def test_attention_equivalence_sliding(): + """GQA attention with sinks AND sliding window (sliding layer).""" + hf = _get_hf_classes() + if hf is None: + pytest.skip("transformers.models.gpt_oss not available") + device, dtype = _device_and_dtype() + config = _small_config() + config._attn_implementation = "eager" + layer_idx = 0 # sliding_attention + B, S = 2, 8 + + ref = hf["Attention"](config, layer_idx=layer_idx).to(device=device, dtype=dtype).eval() + ad = GptOssAttention(config, layer_idx=layer_idx).to(device=device, dtype=dtype).eval() + _init_attention_(ref, ad) + assert ad.sliding_window == config.sliding_window + + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + ad_pos, hf_pos = _rope_position_embeddings(config, B, S, device, dtype) + + # Build the HF sliding-window mask: causal AND within window. + q_pos = torch.arange(S, device=device).unsqueeze(1) + k_pos = torch.arange(S, device=device).unsqueeze(0) + diff = q_pos - k_pos + allowed = (diff >= 0) & (diff < config.sliding_window) + mask = torch.zeros((1, 1, S, S), device=device, dtype=dtype) + mask = mask.masked_fill(~allowed.view(1, 1, S, S), float("-inf")) + + with torch.no_grad(): + ad_out = ad(x, position_embeddings=ad_pos) + ref_out, _ = ref( + hidden_states=x, + position_embeddings=hf_pos, + attention_mask=mask, + ) + assert_rmse_close(ad_out, ref_out, rmse_ratio_tol=0.10, msg="Attention (sliding): ") + + +# --------------------------------------------------------------------------- +# Decoder layer tests +# --------------------------------------------------------------------------- + + +def _init_decoder_layer_(ref, ad): + """Init reference layer with random weights and copy to AD.""" + _init_attention_(ref.self_attn, ad.self_attn) + # HF GptOssMLP === GptOssTopKRouter + GptOssExperts (same names as ours). + torch.nn.init.normal_(ref.mlp.router.weight, std=0.02) + torch.nn.init.normal_(ref.mlp.router.bias, std=0.02) + torch.nn.init.normal_(ref.mlp.experts.gate_up_proj, std=0.02) + torch.nn.init.normal_(ref.mlp.experts.gate_up_proj_bias, std=0.02) + torch.nn.init.normal_(ref.mlp.experts.down_proj, std=0.02) + torch.nn.init.normal_(ref.mlp.experts.down_proj_bias, std=0.02) + torch.nn.init.normal_(ref.input_layernorm.weight, std=0.02) + torch.nn.init.normal_(ref.post_attention_layernorm.weight, std=0.02) + ad.load_state_dict(ref.state_dict()) + + +def _run_hf_decoder(hf_layer, x: torch.Tensor, hf_pos, sliding_window: Optional[int]): + """Run HF decoder with explicit causal (and optional sliding) mask.""" + bsz, q_len, _ = x.shape + if sliding_window is None: + causal = torch.triu( + torch.full((q_len, q_len), float("-inf"), device=x.device, dtype=x.dtype), + diagonal=1, + ).view(1, 1, q_len, q_len) + else: + q_pos = torch.arange(q_len, device=x.device).unsqueeze(1) + k_pos = torch.arange(q_len, device=x.device).unsqueeze(0) + diff = q_pos - k_pos + allowed = (diff >= 0) & (diff < sliding_window) + causal = torch.zeros((1, 1, q_len, q_len), device=x.device, dtype=x.dtype) + causal = causal.masked_fill(~allowed.view(1, 1, q_len, q_len), float("-inf")) + out = hf_layer( + hidden_states=x, + attention_mask=causal, + position_embeddings=hf_pos, + ) + if isinstance(out, tuple): + out = out[0] + return out + + +def test_decoder_layer_equivalence_full(): + hf = _get_hf_classes() + if hf is None: + pytest.skip("transformers.models.gpt_oss not available") + device, dtype = _device_and_dtype() + config = _small_config() + config._attn_implementation = "eager" + layer_idx = 1 # full_attention + B, S = 2, 8 + + ref = hf["Layer"](config, layer_idx=layer_idx).to(device=device, dtype=dtype).eval() + ad = GptOssDecoderLayer(config, layer_idx=layer_idx).to(device=device, dtype=dtype).eval() + _init_decoder_layer_(ref, ad) + + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + ad_pos, hf_pos = _rope_position_embeddings(config, B, S, device, dtype) + + with torch.no_grad(): + ad_out = ad(x, position_embeddings=ad_pos) + ref_out = _run_hf_decoder(ref, x, hf_pos, sliding_window=None) + assert_rmse_close(ad_out, ref_out, rmse_ratio_tol=0.05, msg="Decoder (full): ") + + +def test_decoder_layer_equivalence_sliding(): + hf = _get_hf_classes() + if hf is None: + pytest.skip("transformers.models.gpt_oss not available") + device, dtype = _device_and_dtype() + config = _small_config() + config._attn_implementation = "eager" + layer_idx = 0 # sliding_attention + B, S = 2, 8 + + ref = hf["Layer"](config, layer_idx=layer_idx).to(device=device, dtype=dtype).eval() + ad = GptOssDecoderLayer(config, layer_idx=layer_idx).to(device=device, dtype=dtype).eval() + _init_decoder_layer_(ref, ad) + + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + ad_pos, hf_pos = _rope_position_embeddings(config, B, S, device, dtype) + + with torch.no_grad(): + ad_out = ad(x, position_embeddings=ad_pos) + ref_out = _run_hf_decoder(ref, x, hf_pos, sliding_window=config.sliding_window) + assert_rmse_close(ad_out, ref_out, rmse_ratio_tol=0.05, msg="Decoder (sliding): ") + + +# --------------------------------------------------------------------------- +# Full model test +# --------------------------------------------------------------------------- + + +def _transfer_hf_to_ad_full_model(hf_model, ad_model: GptOssForCausalLM): + """HF and AD use the same parameter names; load_state_dict directly.""" + sd = hf_model.state_dict() + missing, unexpected = ad_model.load_state_dict(sd, strict=False) + # The AD-side rotary buffer (``model.rotary_emb._ad_*_cached``) is non-persistent + # so it's expected to be missing; HF's ``rotary_emb.inv_freq`` is non-persistent too. + assert not unexpected, f"Unexpected keys: {unexpected[:10]}" + + +def test_full_model_equivalence(): + hf = _get_hf_classes() + if hf is None: + pytest.skip("transformers.models.gpt_oss not available") + device, dtype = _device_and_dtype() + # Use 3 layers so we exercise both sliding and full attention plus an extra. + config = _small_config(num_layers=3) + config._attn_implementation = "eager" + + ref = hf["ForCausalLM"](config).to(device=device, dtype=dtype).eval() + # Random init (default _init_weights leaves stacked params with empty().normal_() + # which already runs in HFPreTrainedModel; just ensure deterministic seed). + torch.manual_seed(0) + for _, p in ref.named_parameters(): + if p.dim() > 0: + torch.nn.init.normal_(p, std=0.02) + + ad = GptOssForCausalLM(config).to(device=device, dtype=dtype).eval() + _transfer_hf_to_ad_full_model(ref, ad) + + B, S = 2, 8 + input_ids = torch.randint(0, config.vocab_size, (B, S), device=device) + pos_ids = _position_ids(B, S, device) + + with torch.no_grad(): + ref_out = ref(input_ids=input_ids, position_ids=pos_ids, use_cache=False) + ad_out = ad(input_ids=input_ids, position_ids=pos_ids) + + assert ad_out.logits.shape == (B, S, config.vocab_size) + assert torch.isfinite(ad_out.logits).all() + assert_rmse_close(ad_out.logits, ref_out.logits, rmse_ratio_tol=0.05, msg="Full model: ") + + +# --------------------------------------------------------------------------- +# Export test +# --------------------------------------------------------------------------- + + +def test_export(): + """Model can be exported with torch.export and produces correct output.""" + device = "cpu" + dtype = torch.float32 + + config = _small_config(num_layers=2) + model = GptOssForCausalLM(config).to(device=device, dtype=dtype).eval() + torch.manual_seed(0) + for _, p in model.named_parameters(): + if p.dim() > 0: + torch.nn.init.normal_(p, std=0.02) + + B, S = 2, 8 + input_ids = torch.randint(0, config.vocab_size, (B, S), device=device) + pos_ids = _position_ids(B, S, device) + + dynamic_shapes = { + "input_ids": {0: Dim.DYNAMIC, 1: Dim.DYNAMIC}, + "position_ids": {0: Dim.DYNAMIC, 1: Dim.DYNAMIC}, + } + + gm = torch_export_to_gm( + model, + args=(input_ids,), + kwargs={"position_ids": pos_ids}, + dynamic_shapes=dynamic_shapes, + ) + + with torch.no_grad(): + pre_export_out = model(input_ids=input_ids, position_ids=pos_ids) + exported_out = gm(input_ids, position_ids=pos_ids) + + logits = ( + exported_out[0] + if isinstance(exported_out, tuple) + else getattr(exported_out, "logits", exported_out) + ) + assert torch.isfinite(logits).all(), "Export produced non-finite values" + torch.testing.assert_close(logits, pre_export_out.logits, rtol=1e-3, atol=1e-3) + + B2, S2 = 1, 4 + ids2 = torch.randint(0, config.vocab_size, (B2, S2), device=device) + pos2 = _position_ids(B2, S2, device) + with torch.no_grad(): + out2 = gm(ids2, position_ids=pos2) + logits2 = out2[0] if isinstance(out2, tuple) else getattr(out2, "logits", out2) + assert logits2.shape == (B2, S2, config.vocab_size) + assert torch.isfinite(logits2).all() From a306871b86e51bd525d6116ef9c9a070b9f89880 Mon Sep 17 00:00:00 2001 From: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> Date: Tue, 5 May 2026 18:51:53 -0700 Subject: [PATCH 2/5] fix the B200 layout error Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> --- .../model_registry/configs/gpt_oss_120b.yaml | 2 +- .../custom_ops/fused_moe/mxfp4_moe.py | 12 +++++++- .../custom_ops/moe/test_mxfp4_moe_layout.py | 29 +++++++++++++++++++ 3 files changed, 41 insertions(+), 2 deletions(-) create mode 100644 tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py diff --git a/examples/auto_deploy/model_registry/configs/gpt_oss_120b.yaml b/examples/auto_deploy/model_registry/configs/gpt_oss_120b.yaml index d6486ed041dc..9ca974a5727e 100644 --- a/examples/auto_deploy/model_registry/configs/gpt_oss_120b.yaml +++ b/examples/auto_deploy/model_registry/configs/gpt_oss_120b.yaml @@ -9,7 +9,7 @@ model_factory: AutoModelForCausalLM attn_backend: trtllm compile_backend: torch-cudagraph skip_loading_weights: false -world_size: 8 +world_size: 4 max_batch_size: 128 max_seq_len: 4096 max_num_tokens: 8192 diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py index 6da2b61a8411..1391b4885498 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py @@ -31,6 +31,7 @@ ) from triton_kernels.numerics import InFlexData from triton_kernels.swiglu import swiglu_fn +from triton_kernels.target_info import cuda_capability_geq from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor from triton_kernels.tensor_details import layout from triton_kernels.tensor_details.layout import StridedLayout @@ -39,8 +40,17 @@ # copied from transformers.integrations.mxfp4::swizzle_mxfp4 with minor modification +def _mxfp4_value_layout(mx_axis: int): + # Blackwell's default value layout is only supported by the persistent TMA + # kernel. GPT-OSS MoE can select the non-persistent kernel for small shapes, + # where unswizzled values use the native MXFP4 dot_scaled path. + if cuda_capability_geq(10): + return StridedLayout, {} + return layout.make_default_matmul_mxfp4_w_layout(mx_axis=mx_axis) + + def _swizzle_mxfp4(w, w_scale): - value_layout, value_layout_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=1) + value_layout, value_layout_opts = _mxfp4_value_layout(mx_axis=1) w = convert_layout(wrap_torch_tensor(w, dtype=FP4), value_layout, **value_layout_opts) w_scale = convert_layout(wrap_torch_tensor(w_scale), StridedLayout) return w, w_scale diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py new file mode 100644 index 000000000000..fd076bf4bd9c --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from triton_kernels.tensor_details.layout import HopperMXValueLayout, StridedLayout + +from tensorrt_llm._torch.auto_deploy.custom_ops.fused_moe import mxfp4_moe + + +def test_mxfp4_value_layout_uses_strided_layout_on_blackwell(monkeypatch): + monkeypatch.setattr(mxfp4_moe, "cuda_capability_geq", lambda major, minor=0: major >= 10) + + value_layout, value_layout_opts = mxfp4_moe._mxfp4_value_layout(mx_axis=1) + + assert value_layout is StridedLayout + assert value_layout_opts == {} + + +def test_mxfp4_value_layout_keeps_default_layout_pre_blackwell(monkeypatch): + monkeypatch.setattr(mxfp4_moe, "cuda_capability_geq", lambda major, minor=0: False) + monkeypatch.setattr( + mxfp4_moe.layout, + "make_default_matmul_mxfp4_w_layout", + lambda mx_axis: (HopperMXValueLayout, {"mx_axis": mx_axis}), + ) + + value_layout, value_layout_opts = mxfp4_moe._mxfp4_value_layout(mx_axis=1) + + assert value_layout is HopperMXValueLayout + assert value_layout_opts == {"mx_axis": 1} From b360ff9942e323934ee41b9d1764b0c68efc2e0c Mon Sep 17 00:00:00 2001 From: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> Date: Tue, 5 May 2026 20:54:12 -0700 Subject: [PATCH 3/5] update the test Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> --- .../model_registry/configs/gpt_oss_20b.yaml | 2 +- .../defs/accuracy/references/gsm8k.yaml | 4 ++ .../defs/accuracy/test_llm_api_autodeploy.py | 60 +++++++++++++++++++ 3 files changed, 65 insertions(+), 1 deletion(-) diff --git a/examples/auto_deploy/model_registry/configs/gpt_oss_20b.yaml b/examples/auto_deploy/model_registry/configs/gpt_oss_20b.yaml index 212164d6c709..fee088fe9d4b 100644 --- a/examples/auto_deploy/model_registry/configs/gpt_oss_20b.yaml +++ b/examples/auto_deploy/model_registry/configs/gpt_oss_20b.yaml @@ -9,7 +9,7 @@ model_factory: AutoModelForCausalLM attn_backend: trtllm compile_backend: torch-cudagraph skip_loading_weights: false -world_size: 2 +world_size: 1 max_batch_size: 128 max_seq_len: 4096 max_num_tokens: 8192 diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index bf696094c64d..959e1740ffd1 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -303,6 +303,10 @@ microsoft/phi-4: accuracy: 90.64 mistralai/Codestral-22B-v0.1: - accuracy: 67.10 +openai/gpt-oss-120b: + - accuracy: 10.0 # TODO: update this when the perf is good. +openai/gpt-oss-20b: + - accuracy: 85.823 GPT-OSS/120B-MXFP4: - accuracy: 90.3 - spec_dec_algo: Eagle diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index b13f64d2e61d..7901ad5b5e68 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -58,6 +58,10 @@ def _get_registry_yaml_extra(model_name: str) -> tuple[list[str], int]: if "world_size_" in cfg_name and cfg_name.endswith(".yaml"): world_size = int( cfg_name.replace("world_size_", "").replace(".yaml", "")) + with open(config_dir / cfg) as config_file: + config = yaml.safe_load(config_file) or {} + if "world_size" in config: + world_size = int(config["world_size"]) return paths, world_size raise ValueError(f"Model '{model_name}' not found in model registry") @@ -1089,6 +1093,62 @@ def test_nvfp4(self, ep_size, attention_dp): task.evaluate(llm) +@skip_pre_hopper +@pytest.mark.skip_less_device_memory(80000) +class TestGPTOSS(LlmapiAccuracyTestHarness): + """GSM8K accuracy coverage for GPT-OSS via AutoDeploy.""" + + EXTRA_EVALUATOR_KWARGS = { + "fewshot_as_multiturn": True, + "apply_chat_template": True, + "chat_template_kwargs": { + "reasoning_effort": "low", + }, + } + GSM8K_MAX_OUTPUT_LEN = 512 + MODEL_PATHS = { + "20b": f"{llm_models_root()}/gpt_oss/gpt-oss-20b", + "120b": f"{llm_models_root()}/gpt_oss/gpt-oss-120b", + } + + MODEL_PARAMS = [ + pytest.param( + "20b", + "openai/gpt-oss-20b", + marks=pytest.mark.skip_less_device(2), + id="20b", + ), + pytest.param( + "120b", + "openai/gpt-oss-120b", + marks=pytest.mark.skip_less_device(4), + id="120b", + ), + ] + + @pytest.mark.parametrize("model_id,model_name", MODEL_PARAMS) + def test_mxfp4_gsm8k(self, model_id, model_name, 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: + pytest.skip("Not enough devices for world size, skipping test") + + model_path = self.MODEL_PATHS[model_id] + with AutoDeployLLM( + model=model_path, + tokenizer=model_path, + world_size=registry_world_size, + yaml_extra=yaml_paths, + max_seq_len=GSM8K.MAX_INPUT_LEN + self.GSM8K_MAX_OUTPUT_LEN, + ) as llm: + task = GSM8K(model_name) + task.evaluate(llm, + extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + + class TestGemma4MoE(LlmapiAccuracyTestHarness): """Bench-run coverage for Gemma4 MoE via AutoDeploy.""" From c1facf389f4f6db497d26db9d32312c01c895b97 Mon Sep 17 00:00:00 2001 From: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> Date: Wed, 6 May 2026 12:47:09 -0700 Subject: [PATCH 4/5] Cache MXFP4 prepared weights Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> --- .../custom_ops/fused_moe/mxfp4_moe.py | 133 +++++++++++++++++- .../custom_ops/moe/test_mxfp4_moe_layout.py | 106 ++++++++++++++ 2 files changed, 235 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py index 1391b4885498..3c471f3e3b80 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py @@ -15,7 +15,9 @@ # Triton-kernels-based MXFP4 MoE ops (GPT-OSS style) with routing, swizzling, and fused activation -from typing import Callable, Tuple +import weakref +from collections import OrderedDict +from collections.abc import Callable import torch import torch.nn.functional as F @@ -32,12 +34,32 @@ from triton_kernels.numerics import InFlexData from triton_kernels.swiglu import swiglu_fn from triton_kernels.target_info import cuda_capability_geq -from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor +from triton_kernels.tensor import FP4, Tensor, convert_layout, wrap_torch_tensor from triton_kernels.tensor_details import layout from triton_kernels.tensor_details.layout import StridedLayout from tensorrt_llm._torch.modules.fused_moe.fused_moe_triton import TritonEPRouter +PreparedWeights = tuple[Tensor, Tensor, Tensor, Tensor] +TensorCacheKey = tuple[ + str, + str, + int, + int, + int, + tuple[int, ...], + tuple[int, ...], + int | None, +] +WeightCacheKey = tuple[object, ...] + +_MXFP4_WEIGHT_CACHE_MAX_ENTRIES = 256 +# ``convert_layout`` swizzles the packed expert weights. Cache the result by +# underlying storage/view metadata so decode steps do not repeat that work. +_MXFP4_WEIGHT_CACHE: OrderedDict[WeightCacheKey, tuple[PreparedWeights, list[weakref.finalize]]] = ( + OrderedDict() +) + # copied from transformers.integrations.mxfp4::swizzle_mxfp4 with minor modification def _mxfp4_value_layout(mx_axis: int): @@ -56,7 +78,79 @@ def _swizzle_mxfp4(w, w_scale): return w, w_scale -RouteFn = Callable[[torch.Tensor], Tuple[RoutingData, GatherIndx, ScatterIndx]] +RouteFn = Callable[[torch.Tensor], tuple[RoutingData, GatherIndx, ScatterIndx]] + + +def _mxfp4_layout_cache_key() -> tuple[object, ...]: + value_layout, value_layout_opts = _mxfp4_value_layout(mx_axis=1) + return ( + value_layout.__module__, + value_layout.__qualname__, + tuple(sorted(value_layout_opts.items())), + ) + + +def _source_tensor(tensor: torch.Tensor) -> torch.Tensor: + base = getattr(tensor, "_base", None) + return base if isinstance(base, torch.Tensor) else tensor + + +def _tensor_cache_key(tensor: torch.Tensor) -> TensorCacheKey: + source = _source_tensor(tensor) + try: + version = source._version + except RuntimeError: + version = None + return ( + str(tensor.device), + str(tensor.dtype), + source.untyped_storage().data_ptr(), + tensor.data_ptr(), + tensor.storage_offset(), + tuple(tensor.shape), + tuple(tensor.stride()), + version, + ) + + +def _detach_finalizers(finalizers: list[weakref.finalize]) -> None: + for finalizer in finalizers: + finalizer.detach() + + +def _evict_mxfp4_weight_cache_entry(key: WeightCacheKey) -> None: + entry = _MXFP4_WEIGHT_CACHE.pop(key, None) + if entry is None: + return + _, finalizers = entry + _detach_finalizers(finalizers) + + +def _trim_mxfp4_weight_cache() -> None: + while len(_MXFP4_WEIGHT_CACHE) > _MXFP4_WEIGHT_CACHE_MAX_ENTRIES: + _, (_, finalizers) = _MXFP4_WEIGHT_CACHE.popitem(last=False) + _detach_finalizers(finalizers) + + +def _clear_mxfp4_weight_cache() -> None: + while _MXFP4_WEIGHT_CACHE: + _, (_, finalizers) = _MXFP4_WEIGHT_CACHE.popitem() + _detach_finalizers(finalizers) + + +def _register_cache_finalizers( + key: WeightCacheKey, tensors: tuple[torch.Tensor, ...] +) -> list[weakref.finalize]: + finalizers: list[weakref.finalize] = [] + seen_sources: set[int] = set() + for tensor in tensors: + source = _source_tensor(tensor) + source_id = id(source) + if source_id in seen_sources: + continue + seen_sources.add(source_id) + finalizers.append(weakref.finalize(source, _evict_mxfp4_weight_cache_entry, key)) + return finalizers def _prepare_weights_scales( @@ -90,6 +184,37 @@ def _prepare_weights_scales( ) +def _prepare_weights_scales_cached( + hidden_size: int, + gate_up_blocks: torch.Tensor, + gate_up_scales: torch.Tensor, + down_blocks: torch.Tensor, + down_scales: torch.Tensor, +) -> PreparedWeights: + raw_tensors = (gate_up_blocks, gate_up_scales, down_blocks, down_scales) + key: WeightCacheKey = ( + hidden_size, + _mxfp4_layout_cache_key(), + *(_tensor_cache_key(tensor) for tensor in raw_tensors), + ) + + entry = _MXFP4_WEIGHT_CACHE.get(key) + if entry is not None: + _MXFP4_WEIGHT_CACHE.move_to_end(key) + prepared_weights, _ = entry + return prepared_weights + + prepared_weights = _prepare_weights_scales( + hidden_size, gate_up_blocks, gate_up_scales, down_blocks, down_scales + ) + _MXFP4_WEIGHT_CACHE[key] = ( + prepared_weights, + _register_cache_finalizers(key, raw_tensors), + ) + _trim_mxfp4_weight_cache() + return prepared_weights + + def _run_mxfp4_mlp_core( hidden_states: torch.Tensor, # [B, S, H] or [B*S, H] router_weight: torch.Tensor, @@ -122,7 +247,7 @@ def _run_mxfp4_mlp_core( gate_up_w_scale_raw, triton_down_w, down_w_scale_raw, - ) = _prepare_weights_scales( + ) = _prepare_weights_scales_cached( hidden_size, gate_up_blocks, gate_up_scales, down_blocks, down_scales ) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py index fd076bf4bd9c..a689db818680 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/moe/test_mxfp4_moe_layout.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import torch from triton_kernels.tensor_details.layout import HopperMXValueLayout, StridedLayout from tensorrt_llm._torch.auto_deploy.custom_ops.fused_moe import mxfp4_moe @@ -27,3 +28,108 @@ def test_mxfp4_value_layout_keeps_default_layout_pre_blackwell(monkeypatch): assert value_layout is HopperMXValueLayout assert value_layout_opts == {"mx_axis": 1} + + +def test_mxfp4_weight_layout_cache_reuses_equivalent_views(monkeypatch): + mxfp4_moe._clear_mxfp4_weight_cache() + monkeypatch.setattr(mxfp4_moe, "_mxfp4_layout_cache_key", lambda: ("test-layout",)) + + calls = 0 + sentinel = (object(), object(), object(), object()) + + def fake_prepare(*_args): + nonlocal calls + calls += 1 + return sentinel + + monkeypatch.setattr(mxfp4_moe, "_prepare_weights_scales", fake_prepare) + + gate_up_blocks = torch.empty((4, 8, 1, 16), dtype=torch.uint8) + gate_up_scales = torch.empty((4, 8, 1), dtype=torch.uint8) + down_blocks = torch.empty((4, 32, 1, 16), dtype=torch.uint8) + down_scales = torch.empty((4, 32, 1), dtype=torch.uint8) + + try: + first_result = mxfp4_moe._prepare_weights_scales_cached( + 32, + gate_up_blocks[1:3], + gate_up_scales[1:3], + down_blocks[1:3], + down_scales[1:3], + ) + second_result = mxfp4_moe._prepare_weights_scales_cached( + 32, + gate_up_blocks[1:3], + gate_up_scales[1:3], + down_blocks[1:3], + down_scales[1:3], + ) + finally: + mxfp4_moe._clear_mxfp4_weight_cache() + + assert first_result is sentinel + assert second_result is sentinel + assert calls == 1 + + +def test_mxfp4_weight_layout_cache_invalidates_on_weight_update(monkeypatch): + mxfp4_moe._clear_mxfp4_weight_cache() + monkeypatch.setattr(mxfp4_moe, "_mxfp4_layout_cache_key", lambda: ("test-layout",)) + + calls = 0 + + def fake_prepare(*_args): + nonlocal calls + calls += 1 + return (object(), object(), object(), object()) + + monkeypatch.setattr(mxfp4_moe, "_prepare_weights_scales", fake_prepare) + + gate_up_blocks = torch.zeros((2, 8, 1, 16), dtype=torch.uint8) + gate_up_scales = torch.zeros((2, 8, 1), dtype=torch.uint8) + down_blocks = torch.zeros((2, 32, 1, 16), dtype=torch.uint8) + down_scales = torch.zeros((2, 32, 1), dtype=torch.uint8) + + try: + first_result = mxfp4_moe._prepare_weights_scales_cached( + 32, + gate_up_blocks, + gate_up_scales, + down_blocks, + down_scales, + ) + gate_up_blocks.add_(1) + second_result = mxfp4_moe._prepare_weights_scales_cached( + 32, + gate_up_blocks, + gate_up_scales, + down_blocks, + down_scales, + ) + finally: + mxfp4_moe._clear_mxfp4_weight_cache() + + assert first_result is not second_result + assert calls == 2 + + +def test_mxfp4_weight_layout_cache_accepts_inference_tensors(monkeypatch): + mxfp4_moe._clear_mxfp4_weight_cache() + monkeypatch.setattr(mxfp4_moe, "_mxfp4_layout_cache_key", lambda: ("test-layout",)) + + sentinel = (object(), object(), object(), object()) + monkeypatch.setattr(mxfp4_moe, "_prepare_weights_scales", lambda *_args: sentinel) + + try: + with torch.inference_mode(): + result = mxfp4_moe._prepare_weights_scales_cached( + 32, + torch.empty((2, 8, 1, 16), dtype=torch.uint8), + torch.empty((2, 8, 1), dtype=torch.uint8), + torch.empty((2, 32, 1, 16), dtype=torch.uint8), + torch.empty((2, 32, 1), dtype=torch.uint8), + ) + finally: + mxfp4_moe._clear_mxfp4_weight_cache() + + assert result is sentinel From aab928d1bdf8ae2e99a0ef1edea3bec5bc188944 Mon Sep 17 00:00:00 2001 From: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> Date: Tue, 12 May 2026 15:04:15 -0700 Subject: [PATCH 5/5] minor update Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com> --- .../_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py index 3c471f3e3b80..909e2eac91f9 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/mxfp4_moe.py @@ -53,7 +53,6 @@ ] WeightCacheKey = tuple[object, ...] -_MXFP4_WEIGHT_CACHE_MAX_ENTRIES = 256 # ``convert_layout`` swizzles the packed expert weights. Cache the result by # underlying storage/view metadata so decode steps do not repeat that work. _MXFP4_WEIGHT_CACHE: OrderedDict[WeightCacheKey, tuple[PreparedWeights, list[weakref.finalize]]] = ( @@ -127,7 +126,8 @@ def _evict_mxfp4_weight_cache_entry(key: WeightCacheKey) -> None: def _trim_mxfp4_weight_cache() -> None: - while len(_MXFP4_WEIGHT_CACHE) > _MXFP4_WEIGHT_CACHE_MAX_ENTRIES: + max_entries = 256 + while len(_MXFP4_WEIGHT_CACHE) > max_entries: _, (_, finalizers) = _MXFP4_WEIGHT_CACHE.popitem(last=False) _detach_finalizers(finalizers)