diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 020e1cc8fc..e7d0e72d8a 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -1610,6 +1610,12 @@ def test_fused_conv_correctness( dtype=torch.bfloat16, generator=rng, ) + outer_residual = torch.randn( + inputs.shape, + device=device, + dtype=torch.bfloat16, + generator=rng, + ) null_state = conv_state[NULL_BLOCK_ID].clone() residual_kernel = residual.clone() residual_reference = residual.clone() @@ -1617,6 +1623,7 @@ def test_fused_conv_correctness( module._short_conv_dilated_dispatch( inputs=inputs, residual=residual_kernel, + outer_residual=outer_residual, metadata=metadata, conv_state=conv_state, conv_weights=weights, @@ -1630,15 +1637,17 @@ def test_fused_conv_correctness( conv_state_len=module.conv_state_len, dilation=module.short_conv_dilation, ) + residual_reference = outer_residual + residual_reference - torch.testing.assert_close( - residual_kernel.float(), residual_reference.float(), atol=3e-2, rtol=3e-2 + assert torch.equal( + residual_kernel[:num_real_tokens], residual_reference[:num_real_tokens] ) assert torch.equal(conv_state, state_reference) assert torch.equal(conv_state[NULL_BLOCK_ID], null_state) if case.graph_padding: assert torch.equal( - residual_kernel[num_real_tokens:], residual[num_real_tokens:] + residual_kernel[num_real_tokens:], + (outer_residual + residual)[num_real_tokens:], ) diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py index 9259972ac9..16294a30d7 100644 --- a/tests/models/qwen4_exp/test_qsa_reference.py +++ b/tests/models/qwen4_exp/test_qsa_reference.py @@ -944,6 +944,7 @@ def test_qsa_sparse_paged_attention_correctness( q = torch.randn( num_rows, num_query_heads, head_dim, device="cuda", dtype=torch.bfloat16 ) + output_gate = torch.randn_like(q) kv_cache = torch.randn( num_cache_blocks, page_size, @@ -1023,6 +1024,7 @@ def test_qsa_sparse_paged_attention_correctness( block_table, token_to_req, use_prefill_config=use_prefill_config, + output_gate=output_gate, ) expected = _qsa_sparse_paged_attention_reference( q, @@ -1033,6 +1035,7 @@ def test_qsa_sparse_paged_attention_correctness( token_to_req, scale, ) + expected = expected * torch.sigmoid(output_gate) torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 28c3ea69d1..54ee6f02b4 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -297,7 +297,7 @@ class Qwen4ExpDecoderLayer(nn.Module): if input_ids is None or query_start_loc is None or ngram_context is None: raise RuntimeError("PLE inputs were not prepared") - hidden_states = hidden_states + self.ple( + hidden_states = self.ple( hidden_states, input_ids, query_start_loc, diff --git a/vllm/models/qwen4_exp/nvidia/ops/ple.py b/vllm/models/qwen4_exp/nvidia/ops/ple.py index d76e2705a7..bc4c99957d 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/ple.py +++ b/vllm/models/qwen4_exp/nvidia/ops/ple.py @@ -322,6 +322,7 @@ def _ple_conv_kernel( state_ptr, w_ptr, residual_ptr, + outer_residual_ptr, state_idx_ptr, qsl_ptr, num_acc_ptr, @@ -442,11 +443,21 @@ def _ple_conv_kernel( mask=c_mask, other=0.0, ) + # Preserve the original eager operation boundaries: short convolution is + # first accumulated into the BF16/FP16 PLE output, then the outer residual + # is added to that rounded value. + ple_output = (residual + conv_output).to(residual_ptr.dtype.element_ty) + outer_residual = tl.load( + outer_residual_ptr + output_t * C + c_offs, + mask=c_mask, + other=0.0, + ) + ple_output = outer_residual.to(tl.float32) + ple_output.to(tl.float32) if launch_pdl: tl.extra.cuda.gdc_launch_dependents() tl.store( residual_ptr + output_t * C + c_offs, - residual + conv_output, + ple_output, mask=c_mask, ) @@ -564,6 +575,7 @@ def ple_conv( conv_state: torch.Tensor, conv_weights: torch.Tensor, state_indices: torch.Tensor, + outer_residual: torch.Tensor, *, mode: Literal["decode", "spec", "prefill"], dilation: int, @@ -573,7 +585,7 @@ def ple_conv( spec_query_len: int = 1, token_indices: torch.Tensor | None = None, ) -> None: - """Add short-convolution output to ``residual`` and update its state.""" + """Add short convolution and the outer residual; update state.""" BLOCK_C = 512 kernel_spec_query_len = spec_query_len if mode == "spec" else 1 T, C = inputs.shape @@ -624,6 +636,7 @@ def ple_conv( conv_state, conv_weights, residual, + outer_residual, state_indices, query_start_loc, num_accepted_tokens, diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index 709acbde00..d3986ba38a 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -24,6 +24,7 @@ def _qsa_sparse_paged_gqa_splitk_kernel( partial_output_ptr, partial_lse_ptr, output_ptr, + output_gate_ptr, stride_q_row, stride_q_head, stride_k_block, @@ -36,6 +37,8 @@ def _qsa_sparse_paged_gqa_splitk_kernel( stride_table_req, stride_output_row, stride_output_head, + stride_output_gate_row, + stride_output_gate_head, num_rows, num_cache_blocks, num_requests, @@ -151,6 +154,18 @@ def _qsa_sparse_paged_gqa_splitk_kernel( ) output_mask = head_offsets[:, None] < GROUP_SIZE if NUM_SPLITS == 1: + # Preserve the unfused path's BF16 attention-output rounding before + # applying the gate in FP32. + normalized_output = normalized_output.to(output_ptr.dtype.element_ty) + output_gate = tl.load( + output_gate_ptr + + row * stride_output_gate_row + + (first_head + head_offsets[:, None]) * stride_output_gate_head + + dim_offsets[None, :], + mask=output_mask, + other=0.0, + ).to(tl.float32) + normalized_output = normalized_output.to(tl.float32) * tl.sigmoid(output_gate) tl.store( output_ptr + row * stride_output_row @@ -186,8 +201,11 @@ def _qsa_merge_splitk_kernel( partial_output_ptr, partial_lse_ptr, output_ptr, + output_gate_ptr, stride_output_row, stride_output_head, + stride_output_gate_row, + stride_output_gate_head, num_rows, HEAD_DIM: tl.constexpr, NUM_QUERY_HEADS: tl.constexpr, @@ -219,6 +237,16 @@ def _qsa_merge_splitk_kernel( ) merged = tl.sum(partial_output * weights[:, None], axis=0) merged = tl.where(denominator > 0, merged / denominator, 0.0) + # Preserve the unfused path's BF16 attention-output rounding before + # applying the gate in FP32. + merged = merged.to(output_ptr.dtype.element_ty) + output_gate = tl.load( + output_gate_ptr + + row * stride_output_gate_row + + head * stride_output_gate_head + + dim_offsets + ).to(tl.float32) + merged = merged.to(tl.float32) * tl.sigmoid(output_gate) tl.store( output_ptr + row * stride_output_row + head * stride_output_head + dim_offsets, merged, @@ -455,6 +483,8 @@ def qsa_sparse_paged_attention( token_to_req: torch.Tensor, use_prefill_config: bool, out: torch.Tensor | None = None, + *, + output_gate: torch.Tensor, ) -> torch.Tensor: """Run sparse GQA directly over paged BF16 K/V caches. @@ -496,6 +526,8 @@ def qsa_sparse_paged_attention( raise ValueError("QSA sparse output must match its query") assert out.dtype == q.dtype and out.device == q.device assert out.stride(2) == 1 + assert output_gate.is_contiguous() + output_gate_view = output_gate.view_as(q) if not q.shape[0]: return out @@ -533,6 +565,7 @@ def qsa_sparse_paged_attention( partial_output, partial_lse, out, + output_gate_view, q.stride(0), q.stride(1), k_cache.stride(0), @@ -545,6 +578,8 @@ def qsa_sparse_paged_attention( block_table.stride(0), out.stride(0), out.stride(1), + output_gate_view.stride(0), + output_gate_view.stride(1), q.shape[0], k_cache.shape[0], block_table.shape[0], @@ -568,8 +603,11 @@ def qsa_sparse_paged_attention( partial_output, partial_lse, out, + output_gate_view, out.stride(0), out.stride(1), + output_gate_view.stride(0), + output_gate_view.stride(1), q.shape[0], HEAD_DIM=q.shape[2], NUM_QUERY_HEADS=q.shape[1], @@ -632,6 +670,10 @@ def warmup_qsa_sparse_paged_attention( output_ptr = TritonWarmupTensor( torch.bfloat16, shape=(num_rows, num_query_heads, head_dim) ) + # The output gate is mandatory at runtime; warm the gated specialization. + output_gate_ptr = TritonWarmupTensor( + torch.bfloat16, shape=(num_rows, num_query_heads, head_dim) + ) head_stride = head_dim row_stride = num_query_heads * head_dim num_cache_blocks = triton_scalar_specialization_rep(kv_cache.shape[0]) @@ -659,6 +701,7 @@ def warmup_qsa_sparse_paged_attention( partial_output_ptr, partial_lse_ptr, output_ptr, + output_gate_ptr, row_stride, head_stride, key_cache.stride(0), @@ -671,6 +714,8 @@ def warmup_qsa_sparse_paged_attention( block_table.stride(0), row_stride, head_stride, + row_stride, + head_stride, num_rows, num_cache_blocks, num_requests, @@ -693,6 +738,9 @@ def warmup_qsa_sparse_paged_attention( partial_output_ptr, partial_lse_ptr, output_ptr, + output_gate_ptr, + row_stride, + head_stride, row_stride, head_stride, num_rows, diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 26622357ca..2ec8790e8d 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -195,6 +195,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): self, inputs: torch.Tensor, residual: torch.Tensor, + outer_residual: torch.Tensor, metadata: PleShortConvAttentionMetadata, conv_state: torch.Tensor, conv_weights: torch.Tensor, @@ -209,6 +210,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): has_non_spec = has_prefill or has_decode inputs = inputs[: metadata.num_actual_tokens] residual = residual[: metadata.num_actual_tokens] + outer_residual = outer_residual[: metadata.num_actual_tokens] spec_token_indices = None non_spec_token_indices = None @@ -235,6 +237,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): conv_state=conv_state, conv_weights=conv_weights, state_indices=spec_state_indices, + outer_residual=outer_residual, mode="spec", dilation=self.short_conv_dilation, query_start_loc=query_start_loc, @@ -259,11 +262,17 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): residual_d, residual_p = torch.split( residual, [num_decode_tokens, num_prefill_tokens], dim=0 ) + outer_residual_d, outer_residual_p = torch.split( + outer_residual, + [num_decode_tokens, num_prefill_tokens], + dim=0, + ) token_indices_d = None token_indices_p = None else: inputs_d = inputs_p = inputs residual_d = residual_p = residual + outer_residual_d = outer_residual_p = outer_residual token_indices_d, token_indices_p = torch.split( non_spec_token_indices, [num_decode_tokens, num_prefill_tokens], @@ -277,6 +286,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): conv_state=conv_state, conv_weights=conv_weights, state_indices=state_indices_d, + outer_residual=outer_residual_d, mode="decode", dilation=self.short_conv_dilation, has_initial_states=metadata.has_initial_states_d, @@ -298,6 +308,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): conv_state=conv_state, conv_weights=conv_weights, state_indices=state_indices_p, + outer_residual=outer_residual_p, mode="prefill", dilation=self.short_conv_dilation, query_start_loc=query_start_loc, @@ -316,6 +327,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): conv_state=conv_state, conv_weights=conv_weights, state_indices=state_indices[:num_decode_rows], + outer_residual=outer_residual, mode="decode", dilation=self.short_conv_dilation, has_initial_states=metadata.has_initial_states_d, @@ -324,12 +336,18 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): # State routing consumes the current request metadata on every replay. @eager_break_during_capture - def _short_conv(self, inputs: torch.Tensor, residual: torch.Tensor) -> None: + def _short_conv( + self, + inputs: torch.Tensor, + residual: torch.Tensor, + outer_residual: torch.Tensor, + ) -> None: forward_context = get_forward_context() attn_metadata = forward_context.attn_metadata - # Profiling omits all metadata or this Mamba entry. The residual - # already contains the gated output, so short convolution is a no-op. + # Profiling omits all metadata or this Mamba entry. Short convolution + # is a no-op there, but preserve the outer residual addition. if attn_metadata is None: + residual.add_(outer_residual) return if not isinstance(attn_metadata, dict): @@ -340,6 +358,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): layer_attn_metadata = attn_metadata.get(self.prefix) if layer_attn_metadata is None: + residual.add_(outer_residual) return if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata): raise TypeError( @@ -367,6 +386,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): self._short_conv_dilated_dispatch( inputs=inputs, residual=residual, + outer_residual=outer_residual, metadata=layer_attn_metadata, conv_state=conv_state, conv_weights=conv_weights.to(dtype=inputs.dtype), @@ -404,7 +424,9 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): self.norm_conv.weight, self.norm_key.eps, ) - self._short_conv(conv_input, gated_output) + # State routing depends on runtime request metadata and remains outside + # the graph; short convolution accumulates into gated_output. + self._short_conv(conv_input, gated_output, hidden_states) return gated_output diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 90449d5706..1c2808a6c4 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -121,6 +121,7 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): output: torch.Tensor, token_to_req: torch.Tensor, use_prefill_config: bool, + output_gate: torch.Tensor, output_scale: torch.Tensor | None = None, output_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: @@ -157,6 +158,7 @@ class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): token_to_req, use_prefill_config, output[:num_tokens], + output_gate=output_gate[:num_tokens], ) return output @@ -353,6 +355,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): key: torch.Tensor, value: torch.Tensor, output: torch.Tensor, + output_gate: torch.Tensor, ) -> None: metadata = get_forward_context().attn_metadata if isinstance(metadata, list): @@ -396,6 +399,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): output, token_to_req=side_metadata.token_to_req, use_prefill_config=main_metadata.max_query_len > self._max_decode_query_len, + output_gate=output_gate, ) def forward( @@ -405,6 +409,7 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v, gate = self._project_qkv_gate(qkv, positions) + assert gate is not None num_tokens = hidden_states.shape[0] query = q.view(num_tokens, self.num_heads, self.head_dim) key = k.view(num_tokens, self.num_kv_heads, self.head_dim) @@ -419,10 +424,9 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): key, value, attn_output, + gate, ) flat_output = attn_output.view(num_tokens, -1) - if gate is not None: - flat_output = flat_output * torch.sigmoid(gate) output, _ = self.o_proj(flat_output) return output