iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Canlin Guo
2026-09-14 12:16:25 +00:00
committed by GitHub
co-authored by zjy0516 Kimi Code
parent c612e2bff0
commit 3f55ad2f07
7 changed files with 111 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/qwen4_exp/test_ple.py
+12 -3
View File
@@ -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:],
)
tests/models/qwen4_exp/test_qsa_reference.py
+3
View File
@@ -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)
vllm/models/qwen4_exp/nvidia/model.py
+1 -1
View File
@@ -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,
vllm/models/qwen4_exp/nvidia/ops/ple.py
+15 -2
View File
@@ -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,
vllm/models/qwen4_exp/nvidia/ops/qsa.py
+48
View File
@@ -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,
vllm/models/qwen4_exp/nvidia/ple_layer.py
+26 -4
View File
@@ -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
vllm/models/qwen4_exp/nvidia/qsa.py
+6 -2
View File
@@ -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