[Qwen3.8-Flash-Next] Remove torch.compile for NVIDIA implementation

tests/models/qwen4_exp/test_config.py

@@ -147,9 +147,9 @@ def test_qwen4_exp_model_state_prepares_ngram_context() -> None:
     model_state.uses_ngram_embedding = True
     model_state.ngram_context_len = 3
     model_state.ngram_eos_token_id = 99
-    model_state.ngram_context = torch.empty((4, 3), dtype=torch.int32)
+    model_state.ngram_context = torch.empty((8, 3), dtype=torch.int32)
     model_state.ngram_context_offsets = torch.arange(-3, 0, dtype=torch.int64)
-    model_state.ple_query_start_loc = torch.empty(5, dtype=torch.int32)
+    model_state.ple_query_start_loc = torch.empty(9, dtype=torch.int32)
 
     input_batch = SimpleNamespace(
         num_reqs=2,
@@ -167,34 +167,57 @@ def test_qwen4_exp_model_state_prepares_ngram_context() -> None:
     with patch.object(MambaHybridModelState, "prepare_inputs", return_value={}):
         model_inputs = model_state.prepare_inputs(input_batch, req_states)
 
+    expected_query_start_loc = torch.full((9,), 3, dtype=torch.int32)
+    expected_query_start_loc[0] = 0
+    expected_query_start_loc[1] = 2
     torch.testing.assert_close(
-        model_inputs["query_start_loc"],
-        torch.tensor([0, 2, 3, 3], dtype=torch.int32),
+        model_inputs["query_start_loc"], expected_query_start_loc
     )
+    expected_context = torch.full((8, 3), 99, dtype=torch.int32)
+    expected_context[:2] = torch.tensor([[99, 99, 20], [1, 2, 3]])
+    torch.testing.assert_close(model_inputs["ngram_context"], expected_context)
+
+    # Retain the views to detect reallocations as the request layout changes.
+    query_start_loc = model_inputs["query_start_loc"]
+    ngram_context = model_inputs["ngram_context"]
+    input_batch.num_reqs = 1
+    input_batch.num_reqs_after_padding = 1
+    input_batch.idx_mapping = torch.tensor([0])
+    input_batch.query_start_loc = torch.tensor([0, 3], dtype=torch.int32)
+    with patch.object(MambaHybridModelState, "prepare_inputs", return_value={}):
+        model_inputs = model_state.prepare_inputs(input_batch, req_states)
+
+    expected_query_start_loc.fill_(3)
+    expected_query_start_loc[0] = 0
     torch.testing.assert_close(
-        model_inputs["ngram_context"],
-        torch.tensor([[99, 99, 20], [1, 2, 3], [99, 99, 99]], dtype=torch.int32),
+        model_inputs["query_start_loc"], expected_query_start_loc
     )
+    expected_context.fill_(99)
+    expected_context[0] = torch.tensor([1, 2, 3])
+    torch.testing.assert_close(model_inputs["ngram_context"], expected_context)
+    assert model_inputs["query_start_loc"].data_ptr() == query_start_loc.data_ptr()
+    assert model_inputs["ngram_context"].data_ptr() == ngram_context.data_ptr()
 
 
 def test_qwen4_exp_model_state_prepares_stable_dummy_ngram_inputs() -> None:
     model_state = object.__new__(Qwen4ExpModelState)
     model_state.uses_ngram_embedding = True
     model_state.ngram_eos_token_id = 99
-    model_state.ngram_context = torch.empty((4, 3), dtype=torch.int32)
-    model_state.ple_query_start_loc = torch.empty(5, dtype=torch.int32)
+    model_state.ngram_context = torch.empty((8, 3), dtype=torch.int32)
+    model_state.ple_query_start_loc = torch.empty(9, dtype=torch.int32)
 
     with patch.object(MambaHybridModelState, "prepare_dummy_inputs", return_value={}):
-        first = model_state.prepare_dummy_inputs(num_reqs=3, num_tokens=8)
+        first = model_state.prepare_dummy_inputs(num_reqs=3, num_tokens=4)
+        # Dummy runs establish the addresses used during CUDA graph capture.
         query_start_loc_ptr = first["query_start_loc"].data_ptr()
         ngram_context_ptr = first["ngram_context"].data_ptr()
-        second = model_state.prepare_dummy_inputs(num_reqs=3, num_tokens=8)
+        second = model_state.prepare_dummy_inputs(num_reqs=3, num_tokens=4)
 
+    expected_query_start_loc = torch.full((9,), 4, dtype=torch.int32)
+    expected_query_start_loc[:4] = torch.tensor([0, 1, 2, 4], dtype=torch.int32)
+    torch.testing.assert_close(second["query_start_loc"], expected_query_start_loc)
     torch.testing.assert_close(
-        second["query_start_loc"], torch.tensor([0, 2, 5, 8], dtype=torch.int32)
-    )
-    torch.testing.assert_close(
-        second["ngram_context"], torch.full((3, 3), 99, dtype=torch.int32)
+        second["ngram_context"], torch.full((8, 3), 99, dtype=torch.int32)
     )
     assert second["query_start_loc"].data_ptr() == query_start_loc_ptr
     assert second["ngram_context"].data_ptr() == ngram_context_ptr

tests/models/qwen4_exp/test_ple.py

@@ -498,6 +498,17 @@ def _ngram_hash_params(device: torch.device, context_len: int) -> dict:
         ([3], [1], [[11]], 20),
         ([4, 4], [0, 7], [[11, 12], [13, 14]], 20),
         ([4, 0, 3], [], [[11, 12], [13, 14], [15, 16]], 20),
+        (
+            [3, 2, 0, 0],
+            [],
+            [
+                [11, 12],
+                [13, 14],
+                [_NGRAM_EOS_TOKEN_ID, _NGRAM_EOS_TOKEN_ID],
+                [_NGRAM_EOS_TOKEN_ID, _NGRAM_EOS_TOKEN_ID],
+            ],
+            20,
+        ),
         (
             [1, 33, 2],
             [5, 32],
@@ -528,6 +539,7 @@ def _ngram_hash_params(device: torch.device, context_len: int) -> dict:
         "bigram-only",
         "power-of-two",
         "empty-request",
+        "trailing-padded-requests",
         "three-requests",
         "six-requests",
         "four-gram",

tests/models/qwen4_exp/test_qsa_reference.py

@@ -45,7 +45,6 @@ def test_qsa_mtp_index_share_updates_cache_but_skips_selection(
     indexer = SimpleNamespace(
         skip_topk=True,
         _metadata=lambda: (raw_metadata, compressed_metadata),
-        index_qk_proj=lambda hidden: (torch.zeros(2, 2), None),
         index_n_heads=1,
         index_kv_heads=1,
         index_head_dim=1,
@@ -80,7 +79,7 @@ def test_qsa_mtp_index_share_updates_cache_but_skips_selection(
 
     actual = indexer_qsa.QSAIndexer.forward(
         indexer,
-        torch.zeros(2, 4),
+        torch.zeros(2, 2),
         torch.tensor([7, 8]),
         rows,
     )
@@ -446,6 +445,71 @@ def test_qsa_compressed_metadata_keeps_dummy_slots_inert() -> None:
     assert metadata.k_work_metadata.tolist() == [[0, 0], [2, 0], [2, 1], [-1, -1]]
 
 
+@requires_qsa_kernels
+@pytest.mark.usefixtures("default_vllm_config")
+def test_qsa_unfused_cache_update_ignores_padded_qk() -> None:
+    """Padded projected Q/K rows must not affect either side cache."""
+    from vllm.model_executor.layers.rotary_embedding import get_rope
+
+    device = torch.device("cuda")
+    # Five tokens complete one compressed group and retain four keys in the ring.
+    raw_metadata = SimpleNamespace(
+        num_actual_tokens=5,
+        slot_mapping=torch.tensor([-1, 1, 2, 3, 0], device=device),
+        block_table=torch.zeros((1, 1), dtype=torch.int32, device=device),
+        token_to_req=torch.zeros(5, dtype=torch.int32, device=device),
+        query_start_loc=torch.tensor([0, 5], dtype=torch.int32, device=device),
+        logical_positions=torch.arange(5, device=device),
+    )
+    compressed_metadata = SimpleNamespace(
+        slot_mapping=torch.tensor([-1, -1, -1, 0, -1], device=device),
+    )
+    with torch.device(device):
+        rope = get_rope(
+            head_size=128,
+            max_position=32,
+            rope_parameters={"rope_type": "default", "partial_rotary_factor": 0.5},
+            dtype=torch.bfloat16,
+        )
+    raw_cache = torch.zeros((1, 4, 1, 64), dtype=torch.bfloat16, device=device)
+    compressed_cache = torch.zeros((1, 2, 1, 64), dtype=torch.bfloat16, device=device)
+    norm = SimpleNamespace(
+        weight=torch.zeros(64, dtype=torch.bfloat16, device=device),
+        variance_epsilon=1e-6,
+    )
+    indexer = SimpleNamespace(
+        _metadata=lambda: (raw_metadata, compressed_metadata),
+        skip_topk=True,
+        index_kv_heads=1,
+        use_fused_pre_indexer=False,
+        index_n_heads=1,
+        index_head_dim=64,
+        q_layernorm=norm,
+        k_layernorm=norm,
+        rotary_emb=rope,
+        compress_ratio=4,
+        raw_key_cache=SimpleNamespace(
+            kv_cache=raw_cache, key_cache=raw_cache, rope_position_cache=None
+        ),
+        compressed_key_cache=SimpleNamespace(kv_cache=compressed_cache),
+    )
+    keys = torch.arange(1, 6, dtype=torch.bfloat16, device=device)[:, None].expand(
+        5, 64
+    )
+    padded_keys = torch.full((8, 64), torch.nan, dtype=torch.bfloat16, device=device)
+    padded_keys[:5].copy_(keys)
+    indexer_qsa.QSAIndexer.forward(
+        indexer,
+        torch.cat((torch.ones_like(padded_keys), padded_keys), dim=-1),
+        torch.zeros(8, dtype=torch.long, device=device),
+        torch.full((5, 5), -1, dtype=torch.int32, device=device),
+    )
+    torch.testing.assert_close(raw_cache[0, :, 0], keys[[4, 1, 2, 3]])
+    expected_compressed = torch.zeros_like(compressed_cache)
+    expected_compressed[0, 0] = 1
+    torch.testing.assert_close(compressed_cache, expected_compressed)
+
+
 @requires_qsa_kernels
 @pytest.mark.parametrize("compress_ratio", [1, 4])
 @pytest.mark.parametrize("num_reqs", [2, 3, 4, 7, 8, 9])

tests/test_config.py

@@ -338,9 +338,15 @@ def test_dsa_models_default_to_mrv2_and_breakable_cudagraph(
         ("DeepseekV32MTPModel", True, False),
         ("GlmMoeDsaForCausalLM", False, True),
         ("GlmMoeDsaForCausalLM", True, False),
+        ("Qwen4ExpForCausalLM", False, True),
+        ("Qwen4ExpForCausalLM", True, False),
+        ("Qwen4ExpForConditionalGeneration", False, True),
+        ("Qwen4ExpForConditionalGeneration", True, False),
+        ("Qwen4ExpMTP", False, True),
+        ("Qwen4ExpMTP", True, False),
     ],
 )
-def test_dsa_breakable_cudagraph_platform_default(
+def test_breakable_cudagraph_platform_default(
     monkeypatch, architecture, is_rocm, expected
 ):
     from vllm.config.vllm import default_breakable_cudagraph_architectures

vllm/config/compilation.py

@@ -768,7 +768,7 @@ class CompilationConfig:
         "vllm::mamba_mixer2",
         "vllm::mamba_mixer",
         "vllm::short_conv",
-        "vllm::qwen4_exp_compute_ple_ngram_ids",
+        # Qwen4Exp's AMD backend still uses these splitting ops.
         "vllm::qwen4_exp_ple_short_conv",
         "vllm::qwen4_exp_qsa_with_output",
         "vllm::linear_attention",

vllm/config/vllm.py

@@ -94,6 +94,9 @@
         "KimiLinearForCausalLM",
         "MiniMaxM3SparseForCausalLM",
         "MiniMaxM3SparseForConditionalGeneration",
+        "Qwen4ExpForCausalLM",
+        "Qwen4ExpForConditionalGeneration",
+        "Qwen4ExpMTP",
     }
 )
 

vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py

@@ -12,6 +12,7 @@
 from vllm import _custom_ops as ops
 from vllm import envs
 from vllm._aiter_ops import rocm_aiter_ops
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
 from vllm.config import (
     VllmConfig,
     get_current_vllm_config,
@@ -1909,6 +1910,7 @@ def _forward_core_fused_norm(
         )
 
 
+@eager_break_during_capture
 def qwen_gdn_attention_core(
     qkv_or_qkvz: torch.Tensor,
     b_or_ba: torch.Tensor,
@@ -1956,6 +1958,7 @@ def qwen_gdn_attention_core(
 )
 
 
+@eager_break_during_capture
 def qwen_gdn_attention_core_fused_norm_packed(
     mixed_qkvz: torch.Tensor,
     ba: torch.Tensor,

vllm/models/qwen4_exp/nvidia/indexer_qsa.py

@@ -88,7 +88,7 @@ def _supports_fused_pre_indexer(
 
 
 class QSAIndexer(nn.Module):
-    """Replicated Q/K projection plus paged, weight-free QSA selection.
+    """QSA projection weights, side caches, and paged, weight-free selection.
 
     ``prefix`` must be the checkpoint's indexer prefix, normally
     ``model.layers.N.self_attn.indexer``.  Consequently the trainable names are
@@ -218,11 +218,11 @@ def _metadata(
 
     def forward(
         self,
-        hidden_states: torch.Tensor,
+        projected_qk: torch.Tensor,
         positions: torch.Tensor,
         out: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        """Select each query row's token indices.
+        """Update side caches and select token indices from pre-projected Q/K.
 
         Returns the packed buffer of shape [num_tokens, output_width + 1]:
         the leading ``output_width`` columns are ``-1``-padded
@@ -237,10 +237,10 @@ def forward(
             if self.skip_topk and out is not None:
                 return out
             result = torch.full(
-                (hidden_states.shape[0], self.packed_output_width),
+                (projected_qk.shape[0], self.packed_output_width),
                 -1,
                 dtype=torch.int32,
-                device=hidden_states.device,
+                device=projected_qk.device,
             )
             # Inert rows carry a zero valid count (empty loop bound), not -1.
             result[:, -1] = 0
@@ -258,11 +258,9 @@ def forward(
 
         raw_metadata, compressed_metadata = metadata
         num_tokens = raw_metadata.num_actual_tokens
-        hidden_states = hidden_states[:num_tokens]
+        projected_qk = projected_qk[:num_tokens]
         positions = positions[..., :num_tokens]
 
-        # Q/K projection
-        projected_qk, _ = self.index_qk_proj(hidden_states)
         projected_q, raw_keys = projected_qk.split(
             (
                 self.index_n_heads * self.index_head_dim,

vllm/models/qwen4_exp/nvidia/model.py

@@ -8,7 +8,6 @@
 import torch
 from torch import nn
 
-from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
 from vllm.distributed import get_pp_group
 from vllm.model_executor.layers.fused_moe.utils import (
@@ -381,17 +380,6 @@ def update_physical_experts_metadata(
             moe.experts.update_expert_map()
 
 
-@support_torch_compile(
-    dynamic_arg_dims={
-        "input_ids": 0,
-        "positions": -1,
-        "intermediate_tensors": 0,
-        "inputs_embeds": 0,
-        "query_start_loc": 0,
-        "ngram_context": 0,
-        "deepstack_input_embeds": 0,
-    }
-)
 class Qwen4ExpModel(nn.Module):
     hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER
 

vllm/models/qwen4_exp/nvidia/model_state.py

@@ -44,6 +44,8 @@ def __init__(
         if self.ngram_context_len <= 0:
             raise ValueError("N-gram embedding requires context length >= 1.")
         self.ngram_eos_token_id = int(config.eos_token_id)
+        # PLE runs inside captured regions, so these buffers keep a fixed shape
+        # and address as the active request count changes between replays.
         self.ngram_context = torch.full(
             (self.max_num_reqs, self.ngram_context_len),
             self.ngram_eos_token_id,
@@ -68,8 +70,7 @@ def _prepare_ngram_context(
         req_states: RequestState,
     ) -> torch.Tensor:
         num_reqs = input_batch.num_reqs
-        num_reqs_padded = input_batch.num_reqs_after_padding
-        context = self.ngram_context[:num_reqs_padded]
+        context = self.ngram_context
         context.fill_(self.ngram_eos_token_id)
         if num_reqs == 0:
             return context
@@ -101,8 +102,10 @@ def prepare_inputs(
             return model_inputs
 
         num_reqs_padded = input_batch.num_reqs_after_padding
-        query_start_loc = self.ple_query_start_loc[: num_reqs_padded + 1]
-        query_start_loc.copy_(input_batch.query_start_loc[: num_reqs_padded + 1])
+        query_start_loc = self.ple_query_start_loc
+        query_start_loc[: num_reqs_padded + 1].copy_(input_batch.query_start_loc)
+        # Represent unused capacity as trailing zero-length requests.
+        query_start_loc[num_reqs_padded + 1 :].copy_(input_batch.query_start_loc[-1])
         model_inputs.update(
             query_start_loc=query_start_loc,
             ngram_context=self._prepare_ngram_context(input_batch, req_states),
@@ -118,7 +121,7 @@ def prepare_dummy_inputs(
         if not self.uses_ngram_embedding:
             return model_inputs
 
-        query_start_loc = self.ple_query_start_loc[: num_reqs + 1]
+        query_start_loc = self.ple_query_start_loc
         query_start_loc[0] = 0
         tokens_per_req, num_extra_tokens = divmod(num_tokens, num_reqs)
         query_lens = torch.full(
@@ -129,9 +132,10 @@ def prepare_dummy_inputs(
         )
         if num_extra_tokens > 0:
             query_lens[-num_extra_tokens:] += 1
-        torch.cumsum(query_lens, dim=0, out=query_start_loc[1:])
+        torch.cumsum(query_lens, dim=0, out=query_start_loc[1 : num_reqs + 1])
+        query_start_loc[num_reqs + 1 :].fill_(num_tokens)
 
-        ngram_context = self.ngram_context[:num_reqs]
+        ngram_context = self.ngram_context
         ngram_context.fill_(self.ngram_eos_token_id)
         model_inputs.update(
             query_start_loc=query_start_loc,