[None][feat] Support breakable CUDA graph for Qwen3.8 Flash-Next

docs/source/features/torch_compile_and_piecewise_cuda_graph.md

@@ -59,14 +59,18 @@ prefill_cuda_graph_backend: breakable
 prefill_capture_num_tokens: [128, 256, 512]
 ```
 
-The breakable backend is experimental. The integration coverage in this change
-includes BF16 Qwen3.5 on one GPU and NVFP4 DeepSeek models on multiple GPUs,
-with both context-only and mixed context/decode batches using the KV cache.
+The breakable backend is experimental. Integration coverage includes BF16
+Qwen3.5 on one GPU, NVFP4 DeepSeek models on multiple GPUs, and Qwen3.8
+Flash-Next NVFP4 with TP4/EP4, attention DP, MTP3, and PLE host offload. The
+Qwen3.8 test covers 2,048- and 8,192-token graph buckets, including native QSA
+sparse-attention dispatch for a long prompt; the other tests cover both
+context-only and mixed context/decode batches using the KV cache.
 
 The following restrictions are enforced:
 
-- `torch_compile_config`, LoRA, and multimodal models are rejected during
-  engine initialization.
+- `torch_compile_config` and LoRA are rejected during engine initialization.
+  Multimodal wrappers are also rejected unless the multimodal encoder is
+  explicitly disabled and the wrapper exposes a decoder-only text model.
 - Speculative decoding is supported.
 - Context-logit requests run eagerly instead of replaying a breakable CUDA
   graph.

tensorrt_llm/_torch/attention/backends/sparse/qsa/custom_ops.py

@@ -0,0 +1,109 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""QSA custom operations for breakable CUDA graph capture."""
+
+from typing import Optional
+
+import torch
+
+from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import eager_on_graph
+
+from ...interface import PredefinedAttentionMask
+from .metadata import QSAAttentionMetadata
+from .module import QSASparseHooks
+
+
+def _extract_qsa_extra_attrs(layer_idx: str):
+    from tensorrt_llm._torch.attention.attention import Attention, extract_extra_attrs
+
+    metadata, attention = extract_extra_attrs(layer_idx, "attn")
+    if not isinstance(metadata, QSAAttentionMetadata):
+        raise TypeError("QSA breakable CUDA graph received incompatible metadata")
+    if not isinstance(attention, Attention):
+        raise TypeError("QSA breakable CUDA graph requires an Attention layer")
+    if not isinstance(attention.sparse_attn_hooks, QSASparseHooks):
+        raise TypeError("QSA breakable CUDA graph requires QSA sparse hooks")
+    return metadata, attention
+
+
+@torch.library.custom_op("trtllm::qsa_attn_inplace", mutates_args=("output",))
+def qsa_attn_inplace(
+    q: torch.Tensor,
+    k: Optional[torch.Tensor],
+    v: Optional[torch.Tensor],
+    q_index: torch.Tensor,
+    token_k: torch.Tensor,
+    position_coordinates: torch.Tensor,
+    mrope_rotary_cos_sin: Optional[torch.Tensor],
+    mrope_position_deltas: Optional[torch.Tensor],
+    output_gate: Optional[torch.Tensor],
+    layer_idx: str,
+    output: torch.Tensor,
+) -> None:
+    """Run live QSA dispatch and write the fixed physical output in place."""
+    metadata, attention = _extract_qsa_extra_attrs(layer_idx)
+    output.zero_()
+    sparse_output = attention.sparse_attn_hooks.forward(
+        attention,
+        q,
+        k,
+        v,
+        metadata,
+        PredefinedAttentionMask.CAUSAL,
+        None,
+        None,
+        None,
+        None,
+        None,
+        0,
+        False,
+        output_gate,
+        qsa_index_projection=(q_index, token_k, position_coordinates),
+    )
+    if sparse_output is not None:
+        output[: sparse_output.shape[0]].copy_(sparse_output)
+        return
+
+    dense_output, output_sf = attention._attn_impl(
+        q,
+        k,
+        v,
+        metadata,
+        PredefinedAttentionMask.CAUSAL,
+        mrope_rotary_cos_sin,
+        mrope_position_deltas,
+        None,
+        None,
+        output=output,
+    )
+    if output_sf is not None:
+        raise RuntimeError("QSA output gating requires an unquantized attention output")
+    if output_gate is not None:
+        dense_output = attention.apply_output_gate(
+            dense_output, output_gate[: dense_output.shape[0]]
+        )
+    output[: dense_output.shape[0]].copy_(dense_output)
+
+
+@qsa_attn_inplace.register_fake
+def _qsa_attn_inplace_fake(
+    q: torch.Tensor,
+    k: Optional[torch.Tensor],
+    v: Optional[torch.Tensor],
+    q_index: torch.Tensor,
+    token_k: torch.Tensor,
+    position_coordinates: torch.Tensor,
+    mrope_rotary_cos_sin: Optional[torch.Tensor],
+    mrope_position_deltas: Optional[torch.Tensor],
+    output_gate: Optional[torch.Tensor],
+    layer_idx: str,
+    output: torch.Tensor,
+) -> None:
+    """Model the in-place output mutation during fake-tensor propagation."""
+
+
+maybe_bcg_qsa_attn_inplace = eager_on_graph(qsa_attn_inplace)
+
+
+__all__ = ["maybe_bcg_qsa_attn_inplace", "qsa_attn_inplace"]

tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py

@@ -8,6 +8,7 @@
 import torch
 
 from tensorrt_llm._torch.modules.top_k import TopK
+from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph
 from tensorrt_llm.logger import logger
 
 from ...interface import AttentionMask, PredefinedAttentionMask
@@ -403,6 +404,18 @@ def select_qsa_paged_tokens(
 class QSASparseHooks(AttentionSparseHooks):
     """Keep the QSA side cache current and replace dense attention when useful."""
 
+    def __init__(self) -> None:
+        self._num_sparse_prefill_dispatches = 0
+
+    @property
+    def num_sparse_prefill_dispatches(self) -> int:
+        """Return exact sparse prefill calls observed by this layer.
+
+        The count includes graph capture and warmup calls. Runtime diagnostics
+        should compare snapshots when attributing calls to live requests.
+        """
+        return self._num_sparse_prefill_dispatches
+
     def forward(
         self,
         attention: "Attention",
@@ -440,6 +453,11 @@ def forward(
             raise NotImplementedError("QSA sparse attention does not support LoRA")
         # The model's QKNormRoPEAttention preprocesses q/k before this hook;
         # mrope_config is therefore backend state, not work for the QSA kernel.
+        mrope_rotary_cos_sin = None
+        mrope_position_deltas = None
+        if mrope_config is not None:
+            mrope_rotary_cos_sin = mrope_config.get("mrope_rotary_cos_sin")
+            mrope_position_deltas = mrope_config.get("mrope_position_deltas")
         del mrope_config
         if not isinstance(attn_metadata, QSAAttentionMetadata):
             raise TypeError("QSA sparse attention received incompatible metadata")
@@ -450,9 +468,40 @@ def forward(
             return None
         hidden_states = kwargs.get("qsa_index_hidden_states")
         position_ids = kwargs.get("qsa_position_ids")
-        if hidden_states is None or position_ids is None:
+        index_projection = kwargs.get("qsa_index_projection")
+        if index_projection is None and (hidden_states is None or position_ids is None):
             raise ValueError("QSA sparse attention requires hidden states and position IDs")
 
+        needs_fixed_output = is_in_breakable_cuda_graph() or q.shape[0] > num_tokens
+        if needs_fixed_output and index_projection is None:
+            # Keep the token-wise index projection captured. Cache mutation and
+            # live batch dispatch execute in the eager bridge below. The same
+            # bridge preserves physical padding during graph warmup, which runs
+            # before capture and must still return a bucket-sized output.
+            q_index, token_k, position_coordinates = attention.indexer.project(
+                hidden_states, position_ids
+            )
+            output = q.new_empty((q.shape[0], attention.num_heads * attention.head_dim))
+            from . import custom_ops
+
+            custom_ops.maybe_bcg_qsa_attn_inplace(
+                q,
+                k,
+                v,
+                q_index,
+                token_k,
+                position_coordinates,
+                mrope_rotary_cos_sin,
+                mrope_position_deltas,
+                output_gate,
+                attention.layer_idx_str,
+                output,
+            )
+            return output
+
+        if output_gate is not None:
+            output_gate = output_gate[:num_tokens]
+
         # The standard TRT-LLM backend owns quantized formats with auxiliary
         # scale pages, such as NVFP4. Use it until QSA kernels consume the data
         # and scale pages together; the cache manager itself remains generic.
@@ -478,13 +527,23 @@ def forward(
             )
             return None
 
-        hidden_states = hidden_states[:num_tokens]
-        q_index = attention.indexer.project_and_update_cache(
-            hidden_states,
-            position_ids,
-            attention.layer_idx,
-            attn_metadata,
-        )
+        if index_projection is None:
+            hidden_states = hidden_states[:num_tokens]
+            q_index = attention.indexer.project_and_update_cache(
+                hidden_states,
+                position_ids,
+                attention.layer_idx,
+                attn_metadata,
+            )
+        else:
+            q_index, token_k, position_coordinates = index_projection
+            q_index = q_index[:num_tokens]
+            attention.indexer.update_cache_and_compress(
+                attention.layer_idx,
+                token_k[:num_tokens],
+                position_coordinates[:num_tokens],
+                attn_metadata,
+            )
 
         params = attention.sparse_params
         max_kv_len = int(attn_metadata.kv_lens_runtime[: attn_metadata.num_seqs].max())
@@ -613,6 +672,7 @@ def forward(
                 output = attention.apply_output_gate(output, output_gate)
             return output
 
+        self._num_sparse_prefill_dispatches += 1
         index_cache = attn_metadata.kv_cache_manager.get_index_k_buffer(attention.layer_idx)
         if index_cache is None:
             raise RuntimeError(f"QSA index cache is unavailable for layer {attention.layer_idx}")