[#13816][feat] AutoDeploy: Optimize gpt-oss-120b perf

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",

examples/auto_deploy/llmc/create_standalone_package.py

@@ -185,6 +185,10 @@
     # Imports utils.util.skip_pre_blackwell (not shipped in standalone) and exercises
     # fuse_finegrained_fp8_swiglu which depends on TRT-LLM runtime.
     "test_finegrained_fp8_swiglu.py",
+    # Exercise trtllm-gen MXFP4 MoE kernels (Blackwell-only) and import the
+    # prepare_trtllm_gen_moe_mxfp4_weights / utils.util helpers not in standalone.
+    "test_fuse_mxfp4_moe.py",
+    "test_trtllm_quant_mxfp4_trtllm_gen_moe.py",
 }
 
 # Import path rewrite: old -> new (applied to test files only).

examples/auto_deploy/model_registry/configs/gpt_oss.yaml

@@ -0,0 +1,47 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# OpenAI GPT-OSS (20B / 120B, MXFP4 quantized) — shared AD serving config.
+# - 20B: 24 layers, 32 experts, top-4
+# - 120B: 36 layers, 128 experts, top-4
+# Both share GQA (64 Q / 8 KV heads), head_dim=64, hidden=2880, and MXFP4
+# weights on HF that AD's `quantize_mxfp4_moe` transform handles.
+# world_size is set via the registry's `world_size_N.yaml` overlay — not here.
+model_factory: AutoModelForCausalLM
+model_kwargs:
+  dtype: bfloat16
+attn_backend: trtllm
+compile_backend: torch-cudagraph
+skip_loading_weights: false
+max_batch_size: 128
+max_seq_len: 4096
+max_num_tokens: 8192
+enable_chunked_prefill: true
+cuda_graph_config:
+  batch_sizes: [1, 2, 4, 8, 16, 32, 64, 128]
+kv_cache_config:
+  enable_block_reuse: false
+  free_gpu_memory_fraction: 0.8
+transforms:
+  detect_sharding:
+    enabled: false
+  sharding_transform_executor:
+    enabled: false
+  apply_sharding_hints:
+    enabled: true
+    requires_shape_prop: true
+    shard_layers: ["mha", "moe", "lm_head"]  # V3: vocab-parallel lm_head (colwise + all_gather)
+    # TP2 trial: AUTO -> tunable_allreduce picks oneshot-lamport fused AR.
+    # NCCL has no fused residual+rmsnorm kernel -> RING_LL + separate rmsnorm.
+    allreduce_strategy: AUTO
+  quantize_mxfp4_moe:
+    backend: trtllm
+    trtllm_quant_act: mxfp8
+  fuse_gemms_mixed_children:
+    enabled: true
+  fuse_gemms:
+    enabled: true
+  fuse_rope_into_trtllm_attention:
+    enabled: true
+  fuse_add_rms_norm:
+    enabled: true

examples/auto_deploy/model_registry/models.yaml

@@ -168,10 +168,9 @@ models:
 - name: nvidia/Mistral-NeMo-Minitron-8B-Base
   config_id: default_ws_2
   yaml_extra: ['dashboard_default.yaml', 'world_size_2.yaml']
-# OOM during AutoDeploy run.
-# - name: openai/gpt-oss-20b
-#   config_id: gpt_oss_20b
-#   yaml_extra: ['gpt_oss_20b.yaml']
+- name: openai/gpt-oss-20b
+  config_id: gpt_oss
+  yaml_extra: ['gpt_oss.yaml', 'world_size_1.yaml']
 - name: ibm-granite/granite-3.0-8b-instruct
   config_id: default_ws_1
   yaml_extra: ['dashboard_default.yaml', 'world_size_1.yaml']
@@ -333,10 +332,9 @@ models:
 # - name: meta-llama/Llama-3.2-90B-Vision-Instruct
 #   config_id: multimodal
 #   yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'multimodal.yaml']
-# torch.distributed.DistStoreError: Timed out after 601 seconds waiting for clients. 1/4 clients joined.
-# - name: openai/gpt-oss-120b
-#   config_id: gpt_oss_120b
-#   yaml_extra: ['gpt_oss_120b.yaml']
+- name: openai/gpt-oss-120b
+  config_id: gpt_oss
+  yaml_extra: ['gpt_oss.yaml', 'world_size_2.yaml']
 # [RANK 3] Error querying confidential compute state: Function Not Found
 # - name: meta-llama/Llama-4-Scout-17B-16E-Instruct
 #   config_id: multimodal__llama4_scout

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