[#14828][feat] AutoDeploy: enable trtllm multi KV cache pool (gpt-oss branch)

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

@@ -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).

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

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

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

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
 

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]: