@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3f55ad2f07

[Qwen3.8-Flash-Next] Fuse PLE residual and QSA output gate (#55309)

Signed-off-by: Canlin <canlinguosdu@gmail.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Kimi Code <noreply@moonshot.cn>

Canlin Guo committed Sep 14, 2026 at 20:16 UTC 3f55ad2f073a7d1f573db5a5a3fe3dc9c688ddee

7 files changed +111 -12

tests/models/qwen4_exp/test_ple.py

+12 -3

@@ -1610,6 +1610,12 @@ def test_fused_conv_correctness(

1610 dtype=torch.bfloat16,

1611 generator=rng,

1612 )

1613 + outer_residual = torch.randn(

1614 + inputs.shape,

1615 + device=device,

1616 + dtype=torch.bfloat16,

1617 + generator=rng,

1618 + )

1619 null_state = conv_state[NULL_BLOCK_ID].clone()

1620 residual_kernel = residual.clone()

1621 residual_reference = residual.clone()

@@ -1617,6 +1623,7 @@ def test_fused_conv_correctness(

1623 module._short_conv_dilated_dispatch(

1624 inputs=inputs,

1625 residual=residual_kernel,

1626 + outer_residual=outer_residual,

1627 metadata=metadata,

1628 conv_state=conv_state,

1629 conv_weights=weights,

@@ -1630,15 +1637,17 @@ def test_fused_conv_correctness(

1637 conv_state_len=module.conv_state_len,

1638 dilation=module.short_conv_dilation,

1639 )

1640 + residual_reference = outer_residual + residual_reference

1641

1634 - torch.testing.assert_close(

1635 - residual_kernel.float(), residual_reference.float(), atol=3e-2, rtol=3e-2

1642 + assert torch.equal(

1643 + residual_kernel[:num_real_tokens], residual_reference[:num_real_tokens]

1644 )

1645 assert torch.equal(conv_state, state_reference)

1646 assert torch.equal(conv_state[NULL_BLOCK_ID], null_state)

1647 if case.graph_padding:

1648 assert torch.equal(

1641 - residual_kernel[num_real_tokens:], residual[num_real_tokens:]

1649 + residual_kernel[num_real_tokens:],

1650 + (outer_residual + residual)[num_real_tokens:],

1651 )

1652

1653

tests/models/qwen4_exp/test_qsa_reference.py

@@ -944,6 +944,7 @@ def test_qsa_sparse_paged_attention_correctness(

944 q = torch.randn(

945 num_rows, num_query_heads, head_dim, device="cuda", dtype=torch.bfloat16

946 )

947 + output_gate = torch.randn_like(q)

948 kv_cache = torch.randn(

949 num_cache_blocks,

950 page_size,

@@ -1023,6 +1024,7 @@ def test_qsa_sparse_paged_attention_correctness(

1024 block_table,

1025 token_to_req,

1026 use_prefill_config=use_prefill_config,

1027 + output_gate=output_gate,

1028 )

1029 expected = _qsa_sparse_paged_attention_reference(

1030 q,

@@ -1033,6 +1035,7 @@ def test_qsa_sparse_paged_attention_correctness(

1035 token_to_req,

1036 scale,

1037 )

1038 + expected = expected * torch.sigmoid(output_gate)

1039

1040 torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)

1041

vllm/models/qwen4_exp/nvidia/model.py

+1 -1

@@ -297,7 +297,7 @@ class Qwen4ExpDecoderLayer(nn.Module):

297

298 if input_ids is None or query_start_loc is None or ngram_context is None:

299 raise RuntimeError("PLE inputs were not prepared")

300 - hidden_states = hidden_states + self.ple(

300 + hidden_states = self.ple(

301 hidden_states,

302 input_ids,

303 query_start_loc,

vllm/models/qwen4_exp/nvidia/ops/ple.py

+15 -2

@@ -322,6 +322,7 @@ def _ple_conv_kernel(

322 state_ptr,

323 w_ptr,

324 residual_ptr,

325 + outer_residual_ptr,

326 state_idx_ptr,

327 qsl_ptr,

328 num_acc_ptr,

@@ -442,11 +443,21 @@ def _ple_conv_kernel(

443 mask=c_mask,

444 other=0.0,

445 )

446 + # Preserve the original eager operation boundaries: short convolution is

447 + # first accumulated into the BF16/FP16 PLE output, then the outer residual

448 + # is added to that rounded value.

449 + ple_output = (residual + conv_output).to(residual_ptr.dtype.element_ty)

450 + outer_residual = tl.load(

451 + outer_residual_ptr + output_t * C + c_offs,

452 + mask=c_mask,

453 + other=0.0,

454 + )

455 + ple_output = outer_residual.to(tl.float32) + ple_output.to(tl.float32)

456 if launch_pdl:

457 tl.extra.cuda.gdc_launch_dependents()

458 tl.store(

459 residual_ptr + output_t * C + c_offs,

449 - residual + conv_output,

460 + ple_output,

461 mask=c_mask,

462 )

463

@@ -564,6 +575,7 @@ def ple_conv(

575 conv_state: torch.Tensor,

576 conv_weights: torch.Tensor,

577 state_indices: torch.Tensor,

578 + outer_residual: torch.Tensor,

579 *,

580 mode: Literal["decode", "spec", "prefill"],

581 dilation: int,

@@ -573,7 +585,7 @@ def ple_conv(

585 spec_query_len: int = 1,

586 token_indices: torch.Tensor | None = None,

587 ) -> None:

576 - """Add short-convolution output to ``residual`` and update its state."""

588 + """Add short convolution and the outer residual; update state."""

589 BLOCK_C = 512

590 kernel_spec_query_len = spec_query_len if mode == "spec" else 1

591 T, C = inputs.shape

@@ -624,6 +636,7 @@ def ple_conv(

636 conv_state,

637 conv_weights,

638 residual,

639 + outer_residual,

640 state_indices,

641 query_start_loc,

642 num_accepted_tokens,

vllm/models/qwen4_exp/nvidia/ops/qsa.py

+48

@@ -24,6 +24,7 @@ def _qsa_sparse_paged_gqa_splitk_kernel(

24 partial_output_ptr,

25 partial_lse_ptr,

26 output_ptr,

27 + output_gate_ptr,

28 stride_q_row,

29 stride_q_head,

30 stride_k_block,

@@ -36,6 +37,8 @@ def _qsa_sparse_paged_gqa_splitk_kernel(

37 stride_table_req,

38 stride_output_row,

39 stride_output_head,

40 + stride_output_gate_row,

41 + stride_output_gate_head,

42 num_rows,

43 num_cache_blocks,

44 num_requests,

@@ -151,6 +154,18 @@ def _qsa_sparse_paged_gqa_splitk_kernel(

154 )

155 output_mask = head_offsets[:, None] < GROUP_SIZE

156 if NUM_SPLITS == 1:

157 + # Preserve the unfused path's BF16 attention-output rounding before

158 + # applying the gate in FP32.

159 + normalized_output = normalized_output.to(output_ptr.dtype.element_ty)

160 + output_gate = tl.load(

161 + output_gate_ptr

162 + + row * stride_output_gate_row

163 + + (first_head + head_offsets[:, None]) * stride_output_gate_head

164 + + dim_offsets[None, :],

165 + mask=output_mask,

166 + other=0.0,

167 + ).to(tl.float32)

168 + normalized_output = normalized_output.to(tl.float32) * tl.sigmoid(output_gate)

169 tl.store(

170 output_ptr

171 + row * stride_output_row

@@ -186,8 +201,11 @@ def _qsa_merge_splitk_kernel(

201 partial_output_ptr,

202 partial_lse_ptr,

203 output_ptr,

204 + output_gate_ptr,

205 stride_output_row,

206 stride_output_head,

207 + stride_output_gate_row,

208 + stride_output_gate_head,

209 num_rows,

210 HEAD_DIM: tl.constexpr,

211 NUM_QUERY_HEADS: tl.constexpr,

@@ -219,6 +237,16 @@ def _qsa_merge_splitk_kernel(

237 )

238 merged = tl.sum(partial_output * weights[:, None], axis=0)

239 merged = tl.where(denominator > 0, merged / denominator, 0.0)

240 + # Preserve the unfused path's BF16 attention-output rounding before

241 + # applying the gate in FP32.

242 + merged = merged.to(output_ptr.dtype.element_ty)

243 + output_gate = tl.load(

244 + output_gate_ptr

245 + + row * stride_output_gate_row

246 + + head * stride_output_gate_head

247 + + dim_offsets

248 + ).to(tl.float32)

249 + merged = merged.to(tl.float32) * tl.sigmoid(output_gate)

250 tl.store(

251 output_ptr + row * stride_output_row + head * stride_output_head + dim_offsets,

252 merged,

@@ -455,6 +483,8 @@ def qsa_sparse_paged_attention(

483 token_to_req: torch.Tensor,

484 use_prefill_config: bool,

485 out: torch.Tensor | None = None,

486 + *,

487 + output_gate: torch.Tensor,

488 ) -> torch.Tensor:

489 """Run sparse GQA directly over paged BF16 K/V caches.

490

@@ -496,6 +526,8 @@ def qsa_sparse_paged_attention(

526 raise ValueError("QSA sparse output must match its query")

527 assert out.dtype == q.dtype and out.device == q.device

528 assert out.stride(2) == 1

529 + assert output_gate.is_contiguous()

530 + output_gate_view = output_gate.view_as(q)

531 if not q.shape[0]:

532 return out

533

@@ -533,6 +565,7 @@ def qsa_sparse_paged_attention(

565 partial_output,

566 partial_lse,

567 out,

568 + output_gate_view,

569 q.stride(0),

570 q.stride(1),

571 k_cache.stride(0),

@@ -545,6 +578,8 @@ def qsa_sparse_paged_attention(

578 block_table.stride(0),

579 out.stride(0),

580 out.stride(1),

581 + output_gate_view.stride(0),

582 + output_gate_view.stride(1),

583 q.shape[0],

584 k_cache.shape[0],

585 block_table.shape[0],

@@ -568,8 +603,11 @@ def qsa_sparse_paged_attention(

603 partial_output,

604 partial_lse,

605 out,

606 + output_gate_view,

607 out.stride(0),

608 out.stride(1),

609 + output_gate_view.stride(0),

610 + output_gate_view.stride(1),

611 q.shape[0],

612 HEAD_DIM=q.shape[2],

613 NUM_QUERY_HEADS=q.shape[1],

@@ -632,6 +670,10 @@ def warmup_qsa_sparse_paged_attention(

670 output_ptr = TritonWarmupTensor(

671 torch.bfloat16, shape=(num_rows, num_query_heads, head_dim)

672 )

673 + # The output gate is mandatory at runtime; warm the gated specialization.

674 + output_gate_ptr = TritonWarmupTensor(

675 + torch.bfloat16, shape=(num_rows, num_query_heads, head_dim)

676 + )

677 head_stride = head_dim

678 row_stride = num_query_heads * head_dim

679 num_cache_blocks = triton_scalar_specialization_rep(kv_cache.shape[0])

@@ -659,6 +701,7 @@ def warmup_qsa_sparse_paged_attention(

701 partial_output_ptr,

702 partial_lse_ptr,

703 output_ptr,

704 + output_gate_ptr,

705 row_stride,

706 head_stride,

707 key_cache.stride(0),

@@ -671,6 +714,8 @@ def warmup_qsa_sparse_paged_attention(

714 block_table.stride(0),

715 row_stride,

716 head_stride,

717 + row_stride,

718 + head_stride,

719 num_rows,

720 num_cache_blocks,

721 num_requests,

@@ -693,6 +738,9 @@ def warmup_qsa_sparse_paged_attention(

738 partial_output_ptr,

739 partial_lse_ptr,

740 output_ptr,

741 + output_gate_ptr,

742 + row_stride,

743 + head_stride,

744 row_stride,

745 head_stride,

746 num_rows,

vllm/models/qwen4_exp/nvidia/ple_layer.py

+26 -4

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

195 self,

196 inputs: torch.Tensor,

197 residual: torch.Tensor,

198 + outer_residual: torch.Tensor,

199 metadata: PleShortConvAttentionMetadata,

200 conv_state: torch.Tensor,

201 conv_weights: torch.Tensor,

@@ -209,6 +210,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):

210 has_non_spec = has_prefill or has_decode

211 inputs = inputs[: metadata.num_actual_tokens]

212 residual = residual[: metadata.num_actual_tokens]

213 + outer_residual = outer_residual[: metadata.num_actual_tokens]

214

215 spec_token_indices = None

216 non_spec_token_indices = None

@@ -235,6 +237,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):

237 conv_state=conv_state,

238 conv_weights=conv_weights,

239 state_indices=spec_state_indices,

240 + outer_residual=outer_residual,

241 mode="spec",

242 dilation=self.short_conv_dilation,

243 query_start_loc=query_start_loc,

@@ -259,11 +262,17 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):

262 residual_d, residual_p = torch.split(

263 residual, [num_decode_tokens, num_prefill_tokens], dim=0

264 )

265 + outer_residual_d, outer_residual_p = torch.split(

266 + outer_residual,

267 + [num_decode_tokens, num_prefill_tokens],

268 + dim=0,

269 + )

270 token_indices_d = None

271 token_indices_p = None

272 else:

273 inputs_d = inputs_p = inputs

274 residual_d = residual_p = residual

275 + outer_residual_d = outer_residual_p = outer_residual

276 token_indices_d, token_indices_p = torch.split(

277 non_spec_token_indices,

278 [num_decode_tokens, num_prefill_tokens],

@@ -277,6 +286,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):

286 conv_state=conv_state,

287 conv_weights=conv_weights,

288 state_indices=state_indices_d,

289 + outer_residual=outer_residual_d,

290 mode="decode",

291 dilation=self.short_conv_dilation,

292 has_initial_states=metadata.has_initial_states_d,

@@ -298,6 +308,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):

308 conv_state=conv_state,

309 conv_weights=conv_weights,

310 state_indices=state_indices_p,

311 + outer_residual=outer_residual_p,

312 mode="prefill",

313 dilation=self.short_conv_dilation,

314 query_start_loc=query_start_loc,

@@ -316,6 +327,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):

327 conv_state=conv_state,

328 conv_weights=conv_weights,

329 state_indices=state_indices[:num_decode_rows],

330 + outer_residual=outer_residual,

331 mode="decode",

332 dilation=self.short_conv_dilation,

333 has_initial_states=metadata.has_initial_states_d,

@@ -324,12 +336,18 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):

336

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

338 @eager_break_during_capture

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

339 + def _short_conv(

340 + self,

341 + inputs: torch.Tensor,

342 + residual: torch.Tensor,

343 + outer_residual: torch.Tensor,

344 + ) -> None:

345 forward_context = get_forward_context()

346 attn_metadata = forward_context.attn_metadata

330 - # Profiling omits all metadata or this Mamba entry. The residual

331 - # already contains the gated output, so short convolution is a no-op.

347 + # Profiling omits all metadata or this Mamba entry. Short convolution

348 + # is a no-op there, but preserve the outer residual addition.

349 if attn_metadata is None:

350 + residual.add_(outer_residual)

351 return

352

353 if not isinstance(attn_metadata, dict):

@@ -340,6 +358,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):

358

359 layer_attn_metadata = attn_metadata.get(self.prefix)

360 if layer_attn_metadata is None:

361 + residual.add_(outer_residual)

362 return

363 if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata):

364 raise TypeError(

@@ -367,6 +386,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):

386 self._short_conv_dilated_dispatch(

387 inputs=inputs,

388 residual=residual,

389 + outer_residual=outer_residual,

390 metadata=layer_attn_metadata,

391 conv_state=conv_state,

392 conv_weights=conv_weights.to(dtype=inputs.dtype),

@@ -404,7 +424,9 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):

424 self.norm_conv.weight,

425 self.norm_key.eps,

426 )

407 - self._short_conv(conv_input, gated_output)

427 + # State routing depends on runtime request metadata and remains outside

428 + # the graph; short convolution accumulates into gated_output.

429 + self._short_conv(conv_input, gated_output, hidden_states)

430 return gated_output

431

432

vllm/models/qwen4_exp/nvidia/qsa.py

+6 -2

@@ -121,6 +121,7 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):

121 output: torch.Tensor,

122 token_to_req: torch.Tensor,

123 use_prefill_config: bool,

124 + output_gate: torch.Tensor,

125 output_scale: torch.Tensor | None = None,

126 output_block_scale: torch.Tensor | None = None,

127 ) -> torch.Tensor:

@@ -157,6 +158,7 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):

158 token_to_req,

159 use_prefill_config,

160 output[:num_tokens],

161 + output_gate=output_gate[:num_tokens],

162 )

163 return output

164

@@ -353,6 +355,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):

355 key: torch.Tensor,

356 value: torch.Tensor,

357 output: torch.Tensor,

358 + output_gate: torch.Tensor,

359 ) -> None:

360 metadata = get_forward_context().attn_metadata

361 if isinstance(metadata, list):

@@ -396,6 +399,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):

399 output,

400 token_to_req=side_metadata.token_to_req,

401 use_prefill_config=main_metadata.max_query_len > self._max_decode_query_len,

402 + output_gate=output_gate,

403 )

404

405 def forward(

@@ -405,6 +409,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):

409 ) -> torch.Tensor:

410 qkv, _ = self.qkv_proj(hidden_states)

411 q, k, v, gate = self._project_qkv_gate(qkv, positions)

412 + assert gate is not None

413 num_tokens = hidden_states.shape[0]

414 query = q.view(num_tokens, self.num_heads, self.head_dim)

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

@@ -419,10 +424,9 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):

424 key,

425 value,

426 attn_output,

427 + gate,

428 )

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

424 - if gate is not None:

425 - flat_output = flat_output * torch.sigmoid(gate)

430 output, _ = self.o_proj(flat_output)

431 return output

432