@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d9105ea800

[Qwen3.8-Flash-Next] Remove torch.compile for NVIDIA implementation (#55272)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

Thien Tran committed Sep 7, 2026 at 12:57 UTC d9105ea8001e0a6d77a96327d17515bb5791fb36

13 files changed +165 -185

tests/models/qwen4_exp/test_config.py

+37 -14

@@ -147,9 +147,9 @@ def test_qwen4_exp_model_state_prepares_ngram_context() -> None:

147 model_state.uses_ngram_embedding = True

148 model_state.ngram_context_len = 3

149 model_state.ngram_eos_token_id = 99

150 - model_state.ngram_context = torch.empty((4, 3), dtype=torch.int32)

150 + model_state.ngram_context = torch.empty((8, 3), dtype=torch.int32)

151 model_state.ngram_context_offsets = torch.arange(-3, 0, dtype=torch.int64)

152 - model_state.ple_query_start_loc = torch.empty(5, dtype=torch.int32)

152 + model_state.ple_query_start_loc = torch.empty(9, dtype=torch.int32)

153

154 input_batch = SimpleNamespace(

155 num_reqs=2,

@@ -167,34 +167,57 @@ def test_qwen4_exp_model_state_prepares_ngram_context() -> None:

167 with patch.object(MambaHybridModelState, "prepare_inputs", return_value={}):

168 model_inputs = model_state.prepare_inputs(input_batch, req_states)

169

170 + expected_query_start_loc = torch.full((9,), 3, dtype=torch.int32)

171 + expected_query_start_loc[0] = 0

172 + expected_query_start_loc[1] = 2

173 torch.testing.assert_close(

171 - model_inputs["query_start_loc"],

172 - torch.tensor([0, 2, 3, 3], dtype=torch.int32),

174 + model_inputs["query_start_loc"], expected_query_start_loc

175 )

176 + expected_context = torch.full((8, 3), 99, dtype=torch.int32)

177 + expected_context[:2] = torch.tensor([[99, 99, 20], [1, 2, 3]])

178 + torch.testing.assert_close(model_inputs["ngram_context"], expected_context)

179 +

180 + # Retain the views to detect reallocations as the request layout changes.

181 + query_start_loc = model_inputs["query_start_loc"]

182 + ngram_context = model_inputs["ngram_context"]

183 + input_batch.num_reqs = 1

184 + input_batch.num_reqs_after_padding = 1

185 + input_batch.idx_mapping = torch.tensor([0])

186 + input_batch.query_start_loc = torch.tensor([0, 3], dtype=torch.int32)

187 + with patch.object(MambaHybridModelState, "prepare_inputs", return_value={}):

188 + model_inputs = model_state.prepare_inputs(input_batch, req_states)

189 +

190 + expected_query_start_loc.fill_(3)

191 + expected_query_start_loc[0] = 0

192 torch.testing.assert_close(

175 - model_inputs["ngram_context"],

176 - torch.tensor([[99, 99, 20], [1, 2, 3], [99, 99, 99]], dtype=torch.int32),

193 + model_inputs["query_start_loc"], expected_query_start_loc

194 )

195 + expected_context.fill_(99)

196 + expected_context[0] = torch.tensor([1, 2, 3])

197 + torch.testing.assert_close(model_inputs["ngram_context"], expected_context)

198 + assert model_inputs["query_start_loc"].data_ptr() == query_start_loc.data_ptr()

199 + assert model_inputs["ngram_context"].data_ptr() == ngram_context.data_ptr()

200

201

202 def test_qwen4_exp_model_state_prepares_stable_dummy_ngram_inputs() -> None:

203 model_state = object.__new__(Qwen4ExpModelState)

204 model_state.uses_ngram_embedding = True

205 model_state.ngram_eos_token_id = 99

184 - model_state.ngram_context = torch.empty((4, 3), dtype=torch.int32)

185 - model_state.ple_query_start_loc = torch.empty(5, dtype=torch.int32)

206 + model_state.ngram_context = torch.empty((8, 3), dtype=torch.int32)

207 + model_state.ple_query_start_loc = torch.empty(9, dtype=torch.int32)

208

209 with patch.object(MambaHybridModelState, "prepare_dummy_inputs", return_value={}):

188 - first = model_state.prepare_dummy_inputs(num_reqs=3, num_tokens=8)

210 + first = model_state.prepare_dummy_inputs(num_reqs=3, num_tokens=4)

211 + # Dummy runs establish the addresses used during CUDA graph capture.

212 query_start_loc_ptr = first["query_start_loc"].data_ptr()

213 ngram_context_ptr = first["ngram_context"].data_ptr()

191 - second = model_state.prepare_dummy_inputs(num_reqs=3, num_tokens=8)

214 + second = model_state.prepare_dummy_inputs(num_reqs=3, num_tokens=4)

215

216 + expected_query_start_loc = torch.full((9,), 4, dtype=torch.int32)

217 + expected_query_start_loc[:4] = torch.tensor([0, 1, 2, 4], dtype=torch.int32)

218 + torch.testing.assert_close(second["query_start_loc"], expected_query_start_loc)

219 torch.testing.assert_close(

194 - second["query_start_loc"], torch.tensor([0, 2, 5, 8], dtype=torch.int32)

195 - )

196 - torch.testing.assert_close(

197 - second["ngram_context"], torch.full((3, 3), 99, dtype=torch.int32)

220 + second["ngram_context"], torch.full((8, 3), 99, dtype=torch.int32)

221 )

222 assert second["query_start_loc"].data_ptr() == query_start_loc_ptr

223 assert second["ngram_context"].data_ptr() == ngram_context_ptr

tests/models/qwen4_exp/test_ple.py

+12

@@ -498,6 +498,17 @@ def _ngram_hash_params(device: torch.device, context_len: int) -> dict:

498 ([3], [1], [[11]], 20),

499 ([4, 4], [0, 7], [[11, 12], [13, 14]], 20),

500 ([4, 0, 3], [], [[11, 12], [13, 14], [15, 16]], 20),

501 + (

502 + [3, 2, 0, 0],

503 + [],

504 + [

505 + [11, 12],

506 + [13, 14],

507 + [_NGRAM_EOS_TOKEN_ID, _NGRAM_EOS_TOKEN_ID],

508 + [_NGRAM_EOS_TOKEN_ID, _NGRAM_EOS_TOKEN_ID],

509 + ],

510 + 20,

511 + ),

512 (

513 [1, 33, 2],

514 [5, 32],

@@ -528,6 +539,7 @@ def _ngram_hash_params(device: torch.device, context_len: int) -> dict:

539 "bigram-only",

540 "power-of-two",

541 "empty-request",

542 + "trailing-padded-requests",

543 "three-requests",

544 "six-requests",

545 "four-gram",

tests/models/qwen4_exp/test_qsa_reference.py

+66 -2

@@ -45,7 +45,6 @@ def test_qsa_mtp_index_share_updates_cache_but_skips_selection(

45 indexer = SimpleNamespace(

46 skip_topk=True,

47 _metadata=lambda: (raw_metadata, compressed_metadata),

48 - index_qk_proj=lambda hidden: (torch.zeros(2, 2), None),

48 index_n_heads=1,

49 index_kv_heads=1,

50 index_head_dim=1,

@@ -80,7 +79,7 @@ def test_qsa_mtp_index_share_updates_cache_but_skips_selection(

80 actual = indexer_qsa.QSAIndexer.forward(

81 indexer,

83 - torch.zeros(2, 4),

82 + torch.zeros(2, 2),

83 torch.tensor([7, 8]),

84 rows,

85 )

@@ -446,6 +445,71 @@ def test_qsa_compressed_metadata_keeps_dummy_slots_inert() -> None:

445 assert metadata.k_work_metadata.tolist() == [[0, 0], [2, 0], [2, 1], [-1, -1]]

446

447

448 +@requires_qsa_kernels

449 +@pytest.mark.usefixtures("default_vllm_config")

450 +def test_qsa_unfused_cache_update_ignores_padded_qk() -> None:

451 + """Padded projected Q/K rows must not affect either side cache."""

452 + from vllm.model_executor.layers.rotary_embedding import get_rope

453 +

454 + device = torch.device("cuda")

455 + # Five tokens complete one compressed group and retain four keys in the ring.

456 + raw_metadata = SimpleNamespace(

457 + num_actual_tokens=5,

458 + slot_mapping=torch.tensor([-1, 1, 2, 3, 0], device=device),

459 + block_table=torch.zeros((1, 1), dtype=torch.int32, device=device),

460 + token_to_req=torch.zeros(5, dtype=torch.int32, device=device),

461 + query_start_loc=torch.tensor([0, 5], dtype=torch.int32, device=device),

462 + logical_positions=torch.arange(5, device=device),

463 + )

464 + compressed_metadata = SimpleNamespace(

465 + slot_mapping=torch.tensor([-1, -1, -1, 0, -1], device=device),

466 + )

467 + with torch.device(device):

468 + rope = get_rope(

469 + head_size=128,

470 + max_position=32,

471 + rope_parameters={"rope_type": "default", "partial_rotary_factor": 0.5},

472 + dtype=torch.bfloat16,

473 + )

474 + raw_cache = torch.zeros((1, 4, 1, 64), dtype=torch.bfloat16, device=device)

475 + compressed_cache = torch.zeros((1, 2, 1, 64), dtype=torch.bfloat16, device=device)

476 + norm = SimpleNamespace(

477 + weight=torch.zeros(64, dtype=torch.bfloat16, device=device),

478 + variance_epsilon=1e-6,

479 + )

480 + indexer = SimpleNamespace(

481 + _metadata=lambda: (raw_metadata, compressed_metadata),

482 + skip_topk=True,

483 + index_kv_heads=1,

484 + use_fused_pre_indexer=False,

485 + index_n_heads=1,

486 + index_head_dim=64,

487 + q_layernorm=norm,

488 + k_layernorm=norm,

489 + rotary_emb=rope,

490 + compress_ratio=4,

491 + raw_key_cache=SimpleNamespace(

492 + kv_cache=raw_cache, key_cache=raw_cache, rope_position_cache=None

493 + ),

494 + compressed_key_cache=SimpleNamespace(kv_cache=compressed_cache),

495 + )

496 + keys = torch.arange(1, 6, dtype=torch.bfloat16, device=device)[:, None].expand(

497 + 5, 64

498 + )

499 + padded_keys = torch.full((8, 64), torch.nan, dtype=torch.bfloat16, device=device)

500 + padded_keys[:5].copy_(keys)

501 + indexer_qsa.QSAIndexer.forward(

502 + indexer,

503 + torch.cat((torch.ones_like(padded_keys), padded_keys), dim=-1),

504 + torch.zeros(8, dtype=torch.long, device=device),

505 + torch.full((5, 5), -1, dtype=torch.int32, device=device),

506 + )

507 + torch.testing.assert_close(raw_cache[0, :, 0], keys[[4, 1, 2, 3]])

508 + expected_compressed = torch.zeros_like(compressed_cache)

509 + expected_compressed[0, 0] = 1

510 + torch.testing.assert_close(compressed_cache, expected_compressed)

511 +

512 +

513 @requires_qsa_kernels

514 @pytest.mark.parametrize("compress_ratio", [1, 4])

515 @pytest.mark.parametrize("num_reqs", [2, 3, 4, 7, 8, 9])

tests/test_config.py

+7 -1

@@ -338,9 +338,15 @@ def test_dsa_models_default_to_mrv2_and_breakable_cudagraph(

338 ("DeepseekV32MTPModel", True, False),

339 ("GlmMoeDsaForCausalLM", False, True),

340 ("GlmMoeDsaForCausalLM", True, False),

341 + ("Qwen4ExpForCausalLM", False, True),

342 + ("Qwen4ExpForCausalLM", True, False),

343 + ("Qwen4ExpForConditionalGeneration", False, True),

344 + ("Qwen4ExpForConditionalGeneration", True, False),

345 + ("Qwen4ExpMTP", False, True),

346 + ("Qwen4ExpMTP", True, False),

347 ],

348 )

343 -def test_dsa_breakable_cudagraph_platform_default(

349 +def test_breakable_cudagraph_platform_default(

350 monkeypatch, architecture, is_rocm, expected

351 ):

352 from vllm.config.vllm import default_breakable_cudagraph_architectures

vllm/config/compilation.py

+1 -1

@@ -768,7 +768,7 @@ class CompilationConfig:

768 "vllm::mamba_mixer2",

769 "vllm::mamba_mixer",

770 "vllm::short_conv",

771 - "vllm::qwen4_exp_compute_ple_ngram_ids",

771 + # Qwen4Exp's AMD backend still uses these splitting ops.

772 "vllm::qwen4_exp_ple_short_conv",

773 "vllm::qwen4_exp_qsa_with_output",

774 "vllm::linear_attention",

vllm/config/vllm.py

@@ -94,6 +94,9 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset(

94 "KimiLinearForCausalLM",

95 "MiniMaxM3SparseForCausalLM",

96 "MiniMaxM3SparseForConditionalGeneration",

97 + "Qwen4ExpForCausalLM",

98 + "Qwen4ExpForConditionalGeneration",

99 + "Qwen4ExpMTP",

100 }

101 )

102

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

@@ -12,6 +12,7 @@ from torch import nn

12 from vllm import _custom_ops as ops

13 from vllm import envs

14 from vllm._aiter_ops import rocm_aiter_ops

15 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture

16 from vllm.config import (

17 VllmConfig,

18 get_current_vllm_config,

@@ -1909,6 +1910,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):

1910 )

1911

1912

1913 +@eager_break_during_capture

1914 def qwen_gdn_attention_core(

1915 qkv_or_qkvz: torch.Tensor,

1916 b_or_ba: torch.Tensor,

@@ -1956,6 +1958,7 @@ direct_register_custom_op(

1958 )

1959

1960

1961 +@eager_break_during_capture

1962 def qwen_gdn_attention_core_fused_norm_packed(

1963 mixed_qkvz: torch.Tensor,

1964 ba: torch.Tensor,

vllm/models/qwen4_exp/nvidia/indexer_qsa.py

+6 -8

@@ -88,7 +88,7 @@ def _supports_fused_pre_indexer(

90 class QSAIndexer(nn.Module):

91 - """Replicated Q/K projection plus paged, weight-free QSA selection.

91 + """QSA projection weights, side caches, and paged, weight-free selection.

93 ``prefix`` must be the checkpoint's indexer prefix, normally

94 ``model.layers.N.self_attn.indexer``. Consequently the trainable names are

@@ -218,11 +218,11 @@ class QSAIndexer(nn.Module):

218

219 def forward(

220 self,

221 - hidden_states: torch.Tensor,

221 + projected_qk: torch.Tensor,

222 positions: torch.Tensor,

223 out: torch.Tensor | None = None,

224 ) -> torch.Tensor:

225 - """Select each query row's token indices.

225 + """Update side caches and select token indices from pre-projected Q/K.

226

227 Returns the packed buffer of shape [num_tokens, output_width + 1]:

228 the leading ``output_width`` columns are ``-1``-padded

@@ -237,10 +237,10 @@ class QSAIndexer(nn.Module):

237 if self.skip_topk and out is not None:

238 return out

239 result = torch.full(

240 - (hidden_states.shape[0], self.packed_output_width),

240 + (projected_qk.shape[0], self.packed_output_width),

241 -1,

242 dtype=torch.int32,

243 - device=hidden_states.device,

243 + device=projected_qk.device,

244 )

245 # Inert rows carry a zero valid count (empty loop bound), not -1.

246 result[:, -1] = 0

@@ -258,11 +258,9 @@ class QSAIndexer(nn.Module):

258

259 raw_metadata, compressed_metadata = metadata

260 num_tokens = raw_metadata.num_actual_tokens

261 - hidden_states = hidden_states[:num_tokens]

261 + projected_qk = projected_qk[:num_tokens]

262 positions = positions[..., :num_tokens]

263

264 - # Q/K projection

265 - projected_qk, _ = self.index_qk_proj(hidden_states)

264 projected_q, raw_keys = projected_qk.split(

265 (

266 self.index_n_heads * self.index_head_dim,

vllm/models/qwen4_exp/nvidia/model.py

-12

@@ -8,7 +8,6 @@ from itertools import islice

8 import torch

9 from torch import nn

11 -from vllm.compilation.decorators import support_torch_compile

11 from vllm.config import VllmConfig

12 from vllm.distributed import get_pp_group

13 from vllm.model_executor.layers.fused_moe.utils import (

@@ -381,17 +380,6 @@ class Qwen4ExpMixtureOfExperts(MixtureOfExperts):

380 moe.experts.update_expert_map()

381

382

384 -@support_torch_compile(

385 - dynamic_arg_dims={

386 - "input_ids": 0,

387 - "positions": -1,

388 - "intermediate_tensors": 0,

389 - "inputs_embeds": 0,

390 - "query_start_loc": 0,

391 - "ngram_context": 0,

392 - "deepstack_input_embeds": 0,

393 - }

394 -)

383 class Qwen4ExpModel(nn.Module):

384 hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER

385

vllm/models/qwen4_exp/nvidia/model_state.py

+11 -7

@@ -44,6 +44,8 @@ class Qwen4ExpModelState(MambaHybridModelState):

44 if self.ngram_context_len <= 0:

45 raise ValueError("N-gram embedding requires context length >= 1.")

46 self.ngram_eos_token_id = int(config.eos_token_id)

47 + # PLE runs inside captured regions, so these buffers keep a fixed shape

48 + # and address as the active request count changes between replays.

49 self.ngram_context = torch.full(

50 (self.max_num_reqs, self.ngram_context_len),

51 self.ngram_eos_token_id,

@@ -68,8 +70,7 @@ class Qwen4ExpModelState(MambaHybridModelState):

70 req_states: RequestState,

71 ) -> torch.Tensor:

72 num_reqs = input_batch.num_reqs

71 - num_reqs_padded = input_batch.num_reqs_after_padding

72 - context = self.ngram_context[:num_reqs_padded]

73 + context = self.ngram_context

74 context.fill_(self.ngram_eos_token_id)

75 if num_reqs == 0:

76 return context

@@ -101,8 +102,10 @@ class Qwen4ExpModelState(MambaHybridModelState):

102 return model_inputs

103

104 num_reqs_padded = input_batch.num_reqs_after_padding

104 - query_start_loc = self.ple_query_start_loc[: num_reqs_padded + 1]

105 - query_start_loc.copy_(input_batch.query_start_loc[: num_reqs_padded + 1])

105 + query_start_loc = self.ple_query_start_loc

106 + query_start_loc[: num_reqs_padded + 1].copy_(input_batch.query_start_loc)

107 + # Represent unused capacity as trailing zero-length requests.

108 + query_start_loc[num_reqs_padded + 1 :].copy_(input_batch.query_start_loc[-1])

109 model_inputs.update(

110 query_start_loc=query_start_loc,

111 ngram_context=self._prepare_ngram_context(input_batch, req_states),

@@ -118,7 +121,7 @@ class Qwen4ExpModelState(MambaHybridModelState):

121 if not self.uses_ngram_embedding:

122 return model_inputs

123

121 - query_start_loc = self.ple_query_start_loc[: num_reqs + 1]

124 + query_start_loc = self.ple_query_start_loc

125 query_start_loc[0] = 0

126 tokens_per_req, num_extra_tokens = divmod(num_tokens, num_reqs)

127 query_lens = torch.full(

@@ -129,9 +132,10 @@ class Qwen4ExpModelState(MambaHybridModelState):

132 )

133 if num_extra_tokens > 0:

134 query_lens[-num_extra_tokens:] += 1

132 - torch.cumsum(query_lens, dim=0, out=query_start_loc[1:])

135 + torch.cumsum(query_lens, dim=0, out=query_start_loc[1 : num_reqs + 1])

136 + query_start_loc[num_reqs + 1 :].fill_(num_tokens)

137

134 - ngram_context = self.ngram_context[:num_reqs]

138 + ngram_context = self.ngram_context

139 ngram_context.fill_(self.ngram_eos_token_id)

140 model_inputs.update(

141 query_start_loc=query_start_loc,

vllm/models/qwen4_exp/nvidia/mtp.py

-19

@@ -19,7 +19,6 @@ import regex as re

19 import torch

20 from torch import nn

22 -from vllm.compilation.decorators import support_torch_compile

22 from vllm.config import VllmConfig, replace, set_current_vllm_config

23 from vllm.distributed import get_pp_group

24 from vllm.model_executor.layers.fused_moe.utils import (

@@ -147,15 +146,6 @@ def _make_draft_vllm_config(

146 return draft_vllm_config

147

148

150 -@support_torch_compile(

151 - dynamic_arg_dims={

152 - "input_ids": 0,

153 - "positions": -1,

154 - "intermediate_tensors": 0,

155 - "inputs_embeds": 0,

156 - "hidden_states": 0,

157 - }

158 -)

149 class Qwen4ExpMultiTokenPredictor(nn.Module):

150 hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER

151

@@ -355,15 +345,6 @@ class Qwen4ExpMultiTokenPredictor(nn.Module):

345 return loader.load_weights(weights, mapper=mapper)

346

347

358 -@support_torch_compile(

359 - dynamic_arg_dims={

360 - "input_ids": 0,

361 - "positions": -1,

362 - "intermediate_tensors": 0,

363 - "inputs_embeds": 0,

364 - "hidden_states": 0,

365 - }

366 -)

348 class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts):

349 packed_modules_mapping = {

350 "qkv_proj": ["q_proj", "k_proj", "v_proj"],

vllm/models/qwen4_exp/nvidia/ple_layer.py

+5 -61

@@ -8,6 +8,7 @@ import torch

8 import torch.nn.functional as F

9 from torch import nn

11 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture

12 from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config

13 from vllm.forward_context import get_forward_context

14 from vllm.model_executor.layers.linear import MergedColumnParallelLinear

@@ -38,7 +39,6 @@ from vllm.model_executor.parameter import PerTensorScaleParameter

39 from vllm.transformers_utils.configs.qwen4_exp import (

40 Qwen4ExpTextConfig,

41 )

41 -from vllm.utils.torch_utils import direct_register_custom_op

42 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum

43 from vllm.v1.attention.backends.short_conv_attn import (

44 PleShortConvAttentionBackend,

@@ -264,12 +264,10 @@ class Qwen4ExpNGramEmbedding(nn.Module):

264 embedding_dim: int,

265 ple_dense_layer_id: int,

266 prefix: str,

267 - layer_name: str,

267 quant_config: QuantizationConfig | None = None,

268 params_dtype: torch.dtype | None = None,

269 ) -> None:

270 super().__init__()

272 - self.layer_name = layer_name

271 self.embedding_dim = embedding_dim

272 self.ngram_size = int(config.ngram_size)

273 self.heads_per_ngram = int(config.heads_per_ngram)

@@ -445,20 +443,7 @@ class Qwen4ExpNGramEmbedding(nn.Module):

443 query_start_loc: torch.Tensor,

444 ngram_context: torch.Tensor,

445 ) -> torch.Tensor:

448 - # Keep num_reqs-dependent ID generation outside PIECEWISE CUDA graphs,

449 - # which dispatch only on the padded token count.

450 - # torch.compile requires the splitting op to write graph-owned storage.

451 - # Once compilation is removed, the op can return the IDs directly.

452 - ngram_ids = input_ids.new_empty(

453 - (input_ids.numel(), self.ngram_heads), dtype=torch.long

454 - )

455 - torch.ops.vllm.qwen4_exp_compute_ple_ngram_ids(

456 - input_ids,

457 - query_start_loc,

458 - ngram_context,

459 - ngram_ids,

460 - self.layer_name,

461 - )

446 + ngram_ids = self.compute_ngram_ids(input_ids, query_start_loc, ngram_context)

447 return self.ngram_embedding(ngram_ids).flatten(-2)

448

449 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

@@ -563,7 +548,6 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):

548 int(config.ple_embed_dim),

549 self.ple_dense_layer_id,

550 prefix=f"{prefix}.ple_embedding",

566 - layer_name=prefix,

551 quant_config=quant_config,

552 params_dtype=model_config.dtype,

553 )

@@ -779,6 +763,8 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):

763 token_indices=non_spec_token_indices,

764 )

765

766 + # State routing consumes the current request metadata on every replay.

767 + @eager_break_during_capture

768 def _short_conv(self, inputs: torch.Tensor, residual: torch.Tensor) -> None:

769 forward_context = get_forward_context()

770 attn_metadata = forward_context.attn_metadata

@@ -854,52 +840,10 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):

840 self.norm_conv.weight,

841 self.norm_key.eps,

842 )

857 - # State routing depends on runtime request metadata and remains outside

858 - # the piecewise graph; short convolution accumulates into gated_output.

859 - torch.ops.vllm.qwen4_exp_ple_short_conv(conv_input, gated_output, self.prefix)

843 + self._short_conv(conv_input, gated_output)

844 return gated_output

845

846

863 -def qwen4_exp_compute_ple_ngram_ids(

864 - input_ids: torch.Tensor,

865 - query_start_loc: torch.Tensor,

866 - ngram_context: torch.Tensor,

867 - output: torch.Tensor,

868 - layer_name: str,

869 -) -> None:

870 - """Compute request-dependent PLE n-gram IDs outside piecewise graphs."""

871 - layer = get_forward_context().no_compile_layers[layer_name]

872 - layer.ple_embedding.compute_ngram_ids(

873 - input_ids,

874 - query_start_loc,

875 - ngram_context,

876 - output,

877 - )

878 -

879 -

880 -def qwen4_exp_ple_short_conv(

881 - inputs: torch.Tensor,

882 - residual_output: torch.Tensor,

883 - layer_name: str,

884 -) -> None:

885 - layer = get_forward_context().no_compile_layers[layer_name]

886 - layer._short_conv(inputs, residual_output)

887 -

888 -

889 -direct_register_custom_op(

890 - op_name="qwen4_exp_compute_ple_ngram_ids",

891 - op_func=qwen4_exp_compute_ple_ngram_ids,

892 - mutates_args=["output"],

893 -)

894 -

895 -

896 -direct_register_custom_op(

897 - op_name="qwen4_exp_ple_short_conv",

898 - op_func=qwen4_exp_ple_short_conv,

899 - mutates_args=["residual_output"],

900 -)

901 -

902 -

847 __all__ = [

848 "Qwen4ExpNGramEmbedding",

849 "Qwen4ExpPLEGroupedNorm",

vllm/models/qwen4_exp/nvidia/qsa.py

+14 -60

@@ -9,6 +9,7 @@ from typing import ClassVar, cast

9 import torch

10 from torch import nn

12 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture

13 from vllm.config import VllmConfig

14 from vllm.config.cache import CacheDType

15 from vllm.distributed import get_tensor_model_parallel_world_size

@@ -27,10 +28,6 @@ from vllm.transformers_utils.configs.qwen4_exp import (

28 Qwen4ExpTextConfig,

29 )

30 from vllm.utils.torch_utils import (

30 - LayerNameType,

31 - _encode_layer_name,

32 - _resolve_layer_name,

33 - direct_register_custom_op,

31 kv_cache_dtype_str_to_dtype,

32 )

33 from vllm.v1.attention.backend import (

@@ -347,9 +344,10 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):

344 kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),

345 )

346

347 + @eager_break_during_capture

348 def _run_qsa(

349 self,

352 - hidden_states: torch.Tensor,

350 + projected_qk: torch.Tensor,

351 positions: torch.Tensor,

352 query: torch.Tensor,

353 key: torch.Tensor,

@@ -374,7 +372,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):

372 if side_metadata.num_actual_tokens != num_tokens:

373 raise RuntimeError("QSA main and side metadata token counts disagree")

374 selected = self.indexer(

377 - hidden_states,

375 + projected_qk,

376 positions,

377 self.topk_indices_buffer[:num_tokens],

378 )

@@ -412,27 +410,16 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):

410 key = k.view(num_tokens, self.num_kv_heads, self.head_dim)

411 value = v.view(num_tokens, self.num_kv_heads, self.head_dim)

412 attn_output = torch.empty_like(query)

415 - encoded_layer_name = _encode_layer_name(self.layer_name)

416 - if current_platform.opaque_attention_op():

417 - torch.ops.vllm.qwen4_exp_qsa_with_output(

418 - hidden_states,

419 - positions,

420 - query,

421 - key,

422 - value,

423 - attn_output,

424 - encoded_layer_name,

425 - )

426 - else:

427 - qwen4_exp_qsa_with_output(

428 - hidden_states,

429 - positions,

430 - query,

431 - key,

432 - value,

433 - attn_output,

434 - encoded_layer_name,

435 - )

413 + # Keep the index projection outside the eager break.

414 + projected_qk, _ = self.indexer.index_qk_proj(hidden_states)

415 + self._run_qsa(

416 + projected_qk,

417 + positions,

418 + query,

419 + key,

420 + value,

421 + attn_output,

422 + )

423 flat_output = attn_output.view(num_tokens, -1)

424 if gate is not None:

425 flat_output = flat_output * torch.sigmoid(gate)

@@ -440,42 +427,9 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):

427 return output

428

429

443 -def qwen4_exp_qsa_with_output(

444 - hidden_states: torch.Tensor,

445 - positions: torch.Tensor,

446 - query: torch.Tensor,

447 - key: torch.Tensor,

448 - value: torch.Tensor,

449 - output: torch.Tensor,

450 - layer_name: LayerNameType,

451 -) -> None:

452 - """Run the complete QSA state/update/attend transaction."""

453 -

454 - layer_name = _resolve_layer_name(layer_name)

455 - layer = get_forward_context().no_compile_layers[layer_name]

456 - if not isinstance(layer, Qwen4ExpQSAAttention):

457 - raise TypeError(f"{layer_name} is not a Qwen4Exp QSA owner")

458 - layer._run_qsa(

459 - hidden_states,

460 - positions,

461 - query,

462 - key,

463 - value,

464 - output,

465 - )

466 -

467 -

468 -direct_register_custom_op(

469 - op_name="qwen4_exp_qsa_with_output",

470 - op_func=qwen4_exp_qsa_with_output,

471 - mutates_args=["output"],

472 -)

473 -

474 -

430 __all__ = [

431 "QSAIndexer",

432 "Qwen4ExpQSAAttention",

433 "Qwen4ExpQSAFlashAttentionBackend",

434 "Qwen4ExpQSAFlashAttentionImpl",

480 - "qwen4_exp_qsa_with_output",

435 ]