[Perf][Qwen3.8] Reduce PLE metadata construction overhead (#58114)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Jiangyun Zhu committed
Sep 28, 2026 at 20:46 UTC
20b52e9f5b5793d56586b91df9b93cbe30a24040
5 files changed
+161
-78
tests/models/qwen4_exp/test_ple.py
+125
@@ -14,6 +14,9 @@ from torch.nn import functional as F
14
import vllm.model_executor.layers.vocab_parallel_embedding as embedding_module
15
import vllm.model_executor.parameter as parameter_module
16
import vllm.models.qwen4_exp.common.ngram_embedding as ngram_embedding_module
17
+from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
18
+from vllm.config import SpeculativeConfig, VllmConfig
19
+from vllm.config.compilation import CUDAGraphMode
20
from vllm.config.quantization import QuantizationConfigArgs
21
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
22
from vllm.model_executor.layers.quantization.modelopt import (
@@ -43,8 +46,125 @@ from vllm.models.qwen4_exp.nvidia.ple_layer import Qwen4ExpPLELayer
46
from vllm.utils.torch_utils import weak_ref_tensor
47
from vllm.v1.attention.backends.short_conv_attn import (
48
PleShortConvAttentionMetadata,
49
+ PleShortConvAttentionMetadataBuilder,
50
)
51
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
52
+from vllm.v1.kv_cache_interface import MambaSpec
53
+
54
+
55
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
56
+@pytest.mark.parametrize("full_graph", [False, True])
57
+@pytest.mark.parametrize(
58
+ "query_lens,drafts,spec_rows,non_spec_rows,token_order",
59
+ [
60
+ pytest.param([4, 2, 0], [3, 1, -1], [0, 1], [], None, id="spec-padding"),
61
+ pytest.param(
62
+ [3, 1, 2, 5, 0],
63
+ [2, -1, 1, -1, -1],
64
+ [0, 2],
65
+ [1, 3],
66
+ [0, 1, 2, 4, 5, 3, 6, 7, 8, 9, 10],
67
+ id="mixed-padding",
68
+ ),
69
+ pytest.param([1, 5], [-1, -1], [], [0, 1], None, id="non-spec"),
70
+ ],
71
+)
72
+def test_ple_metadata_preserves_request_and_token_order(
73
+ device, full_graph, query_lens, drafts, spec_rows, non_spec_rows, token_order
74
+):
75
+ """Preserve state routing, acceptance and prefill offsets without GPU sync."""
76
+ if device == "cuda" and not torch.cuda.is_available():
77
+ pytest.skip("CUDA required")
78
+ device = torch.device(device)
79
+ config = VllmConfig()
80
+ config.speculative_config = SpeculativeConfig(
81
+ method="ngram", num_speculative_tokens=3
82
+ )
83
+ config.compilation_config.cudagraph_mode = (
84
+ CUDAGraphMode.FULL if full_graph else CUDAGraphMode.NONE
85
+ )
86
+ config.compilation_config.max_cudagraph_capture_size = 32
87
+ builder = PleShortConvAttentionMetadataBuilder(
88
+ MambaSpec(block_size=16, shapes=((16, 64),), dtypes=(torch.float32,)),
89
+ ["ple"],
90
+ config,
91
+ device,
92
+ )
93
+ common = create_common_attn_metadata(
94
+ BatchSpec(seq_lens=[q + 16 for q in query_lens], query_lens=query_lens),
95
+ 16,
96
+ device,
97
+ arange_block_indices=True,
98
+ )
99
+ common.is_prefilling = torch.tensor(
100
+ [q > 1 and d < 0 for q, d in zip(query_lens, drafts)]
101
+ )
102
+ # Token padding is independent of the sum of per-request query lengths.
103
+ common.num_actual_tokens += 3
104
+ accepted = torch.arange(1, len(query_lens) + 1, dtype=torch.int32, device=device)
105
+ drafts_cpu = torch.tensor(drafts, dtype=torch.int32)
106
+ sync_mode = torch.cuda.get_sync_debug_mode() if device.type == "cuda" else None
107
+ if sync_mode is not None:
108
+ torch.cuda.set_sync_debug_mode("error")
109
+ try:
110
+ metadata = builder.build(
111
+ 0,
112
+ common,
113
+ num_accepted_tokens=accepted,
114
+ num_decode_draft_tokens_cpu=drafts_cpu,
115
+ )
116
+ finally:
117
+ if sync_mode is not None:
118
+ torch.cuda.set_sync_debug_mode(sync_mode)
119
+
120
+ def check(actual, expected):
121
+ torch.testing.assert_close(
122
+ actual, torch.tensor(expected, dtype=actual.dtype, device=device)
123
+ )
124
+
125
+ assert metadata.num_spec_decodes == len(spec_rows)
126
+ assert metadata.num_decodes == sum(query_lens[i] == 1 for i in non_spec_rows)
127
+ assert metadata.num_prefills == sum(query_lens[i] > 1 for i in non_spec_rows)
128
+ if spec_rows:
129
+ torch.testing.assert_close(
130
+ metadata.spec_state_indices_tensor[: len(spec_rows)],
131
+ common.block_table_tensor[spec_rows, 0],
132
+ )
133
+ torch.testing.assert_close(
134
+ metadata.num_accepted_tokens[: len(spec_rows)], accepted[spec_rows]
135
+ )
136
+ check(
137
+ metadata.spec_query_start_loc[: len(spec_rows) + 1],
138
+ [0, *accumulate(query_lens[i] for i in spec_rows)],
139
+ )
140
+ if full_graph and not non_spec_rows:
141
+ check(metadata.spec_state_indices_tensor[len(spec_rows) :], [NULL_BLOCK_ID])
142
+ check(metadata.num_accepted_tokens[len(spec_rows) :], [1])
143
+ check(
144
+ metadata.spec_query_start_loc[len(spec_rows) + 1 :], [sum(query_lens)]
145
+ )
146
+ assert (
147
+ metadata.spec_state_indices_tensor.data_ptr()
148
+ == builder.spec_state_indices_tensor.data_ptr()
149
+ )
150
+ if non_spec_rows:
151
+ torch.testing.assert_close(
152
+ metadata.state_indices_tensor, common.block_table_tensor[non_spec_rows, 0]
153
+ )
154
+ prefill_lens = [query_lens[i] for i in non_spec_rows if query_lens[i] > 1]
155
+ check(metadata.query_start_loc_p, [0, *accumulate(prefill_lens)])
156
+ check(metadata.has_initial_states_p, [True] * len(prefill_lens))
157
+ if token_order is not None:
158
+ check(
159
+ torch.cat((metadata.spec_token_indx, metadata.non_spec_token_indx)),
160
+ token_order,
161
+ )
162
+ else:
163
+ assert metadata.spec_token_indx is None
164
+ assert metadata.non_spec_token_indx is None
165
+ assert metadata.nums_dict is None
166
+ assert metadata.batch_ptr is None
167
+ assert metadata.token_chunk_offset_ptr is None
168
169
170
def _mock_etp_group(
@@ -1541,6 +1661,11 @@ def _make_conv_metadata(
1661
else None
1662
),
1663
non_spec_query_start_loc=(non_spec_query_start_loc if has_non_spec else None),
1664
+ query_start_loc_p=(
1665
+ non_spec_query_start_loc[case.num_decodes :] - case.num_decodes
1666
+ if num_prefills
1667
+ else None
1668
+ ),
1669
has_initial_states_p=(
1670
torch.arange(num_prefills, device=device) % 2 == 0 if num_prefills else None
1671
),
vllm/models/qwen4_exp/amd/ple_layer.py
+2
-7
@@ -672,14 +672,9 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
672
num_decode_tokens: int,
673
num_prefill_tokens: int,
674
) -> torch.Tensor:
675
- # ``non_spec_query_start_loc`` covers the non-spec (decode + prefill)
676
- # requests and equals ``query_start_loc`` when spec-decode is inactive.
677
- non_spec_query_start_loc = metadata.non_spec_query_start_loc
678
- if non_spec_query_start_loc is None:
675
+ query_start_loc_p = metadata.query_start_loc_p
676
+ if query_start_loc_p is None:
677
raise ValueError("query_start_loc is required for prefill short-conv")
680
- query_start_loc_p = (
681
- non_spec_query_start_loc[-num_prefills - 1 :] - num_decode_tokens
682
- )
678
# The metadata builder guarantees that the prefill query offsets start
679
# at 0 and end at num_prefill_tokens. Avoid reading those values here,
680
# since doing so would force a device-to-host synchronization.
vllm/models/qwen4_exp/nvidia/ple_layer.py
+1
-2
@@ -292,10 +292,9 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
292
token_indices=token_indices_d,
293
)
294
295
- query_start_loc = metadata.non_spec_query_start_loc
295
+ query_start_loc = metadata.query_start_loc_p
296
if query_start_loc is None:
297
raise ValueError("query_start_loc is required for prefill short-conv")
298
- query_start_loc = query_start_loc[-num_prefills - 1 :] - num_decode_tokens
298
has_initial_states = metadata.has_initial_states_p
299
if has_initial_states is None:
300
raise ValueError(
vllm/v1/attention/backends/mamba_attn.py
+11
-9
@@ -97,6 +97,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
97
98
# Will be disabled if speculative decoding is used
99
supports_update_block_table: bool = True
100
+ needs_causal_conv1d_metadata: bool = True
101
102
def __init__(
103
self,
@@ -603,10 +604,6 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
604
if num_computed_tokens is None:
605
num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()
606
606
- query_start_loc_p_cpu = (
607
- common_attn_metadata.query_start_loc_cpu[-num_prefills - 1 :]
608
- - num_decode_tokens
609
- )
607
query_start_loc_p = (
608
common_attn_metadata.query_start_loc[-num_prefills - 1 :]
609
- num_decode_tokens
@@ -615,12 +612,17 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
612
num_computed_tokens[num_reqs - num_prefills : num_reqs] > 0
613
)
614
618
- nums_dict, batch_ptr, token_chunk_offset_ptr = (
619
- compute_causal_conv1d_metadata(
620
- query_start_loc_p_cpu,
621
- device=common_attn_metadata.query_start_loc.device,
615
+ if self.needs_causal_conv1d_metadata:
616
+ query_start_loc_p_cpu = (
617
+ common_attn_metadata.query_start_loc_cpu[-num_prefills - 1 :]
618
+ - num_decode_tokens
619
+ )
620
+ nums_dict, batch_ptr, token_chunk_offset_ptr = (
621
+ compute_causal_conv1d_metadata(
622
+ query_start_loc_p_cpu,
623
+ device=common_attn_metadata.query_start_loc.device,
624
+ )
625
)
623
- )
626
627
if self.vllm_config.cache_config.mamba_cache_mode == "all":
628
assert num_computed_tokens is not None
vllm/v1/attention/backends/short_conv_attn.py
+22
-60
@@ -18,7 +18,6 @@ from vllm.v1.attention.backends.mamba_attn import (
18
)
19
from vllm.v1.attention.backends.utils import (
20
NULL_BLOCK_ID,
21
- compute_causal_conv1d_metadata,
21
mamba_get_block_table_tensor,
22
)
23
from vllm.v1.kv_cache_interface import MambaSpec
@@ -108,6 +107,7 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
107
_cudagraph_support = AttentionCGSupport.UNIFORM_BATCH
108
reorder_batch_threshold: int = 1
109
supports_update_block_table = False
110
+ needs_causal_conv1d_metadata = False
111
112
def __init__(
113
self,
@@ -281,18 +281,10 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
281
spec_sequence_masks_cpu, device=query_start_loc.device
282
)
283
284
- # For causal_conv1d (non-spec prefill Triton kernel metadata).
285
- nums_dict = None
286
- batch_ptr = None
287
- token_chunk_offset_ptr = None
284
has_initial_states_p = None
285
has_initial_states_d = None
286
num_computed_tokens_p = None
291
- # Original request indices of the non-spec requests, ordered
292
- # [decodes, prefills]. Used to gather per-request data consistently.
293
- non_spec_req_idx_cpu: torch.Tensor | None = None
287
295
- query_lens = torch.diff(query_start_loc)
288
# Per-request classification by mask, NOT by position. With
289
# spec-decode, the front decode group can contain both spec-decode
290
# requests and plain non-spec single-token decodes. Spec requests are
@@ -317,37 +309,32 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
309
else 0
310
)
311
320
- # Original request indices grouped as
321
- # [spec | non-spec decode | non-spec prefill]; each group keeps the
322
- # original (already reordered) relative order via a stable nonzero.
323
- spec_req_idx_cpu = spec_sequence_masks_cpu.nonzero(as_tuple=True)[0]
324
- decode_req_idx_cpu = decode_mask_cpu.nonzero(as_tuple=True)[0]
325
- prefill_req_idx_cpu = prefill_mask_cpu.nonzero(as_tuple=True)[0]
326
- non_spec_req_idx_cpu = torch.cat((decode_req_idx_cpu, prefill_req_idx_cpu))
327
- spec_req_idx = async_tensor_h2d(spec_req_idx_cpu, device=query_start_loc.device)
312
+ assert num_accepted_tokens is not None
313
non_spec_req_idx: torch.Tensor | None = None
314
315
if num_decodes == 0 and num_prefills == 0:
331
- # Pure speculative-decode batch: all real tokens are spec tokens.
332
- spec_token_indx = torch.arange(
333
- num_spec_decode_tokens,
334
- dtype=torch.int32,
335
- device=query_start_loc.device,
336
- )
337
- non_spec_token_indx = torch.empty(
338
- 0, dtype=torch.int32, device=query_start_loc.device
339
- )
340
- spec_state_indices_tensor = block_table_tensor[spec_req_idx, 0]
316
+ # All real requests are speculative; zero-length padding is trailing.
317
+ spec_token_indx = None
318
+ non_spec_token_indx = None
319
+ spec_state_indices_tensor = block_table_tensor[:num_spec_decodes, 0]
320
+ num_accepted_tokens = num_accepted_tokens[:num_spec_decodes]
321
non_spec_state_indices_tensor = None
322
spec_query_start_loc = query_start_loc[: num_spec_decodes + 1]
323
non_spec_query_start_loc = None
344
- non_spec_query_start_loc_cpu = None
324
else:
325
# Mixed batch: build a per-token group key consistent with the
326
# request grouping above (spec=0 | decode=1 | prefill=2) and a
327
# stable sort, so tokens of each request stay contiguous and in
328
# request order. This yields spec tokens first, then the non-spec
329
# [decode, prefill] tokens.
330
+ query_lens = torch.diff(query_start_loc)
331
+ spec_req_idx_cpu = spec_sequence_masks_cpu.nonzero(as_tuple=True)[0]
332
+ decode_req_idx_cpu = decode_mask_cpu.nonzero(as_tuple=True)[0]
333
+ prefill_req_idx_cpu = prefill_mask_cpu.nonzero(as_tuple=True)[0]
334
+ non_spec_req_idx_cpu = torch.cat((decode_req_idx_cpu, prefill_req_idx_cpu))
335
+ spec_req_idx = async_tensor_h2d(
336
+ spec_req_idx_cpu, device=query_start_loc.device
337
+ )
338
non_spec_req_idx = async_tensor_h2d(
339
non_spec_req_idx_cpu, device=query_start_loc.device
340
)
@@ -360,14 +347,17 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
347
dtype=torch.int64,
348
device=query_start_loc.device,
349
)
363
- req_group[spec_req_idx] = 0
364
- req_group[decode_req_idx] = 1
365
- token_group = torch.repeat_interleave(req_group, query_lens)
350
+ req_group.index_fill_(0, spec_req_idx, 0)
351
+ req_group.index_fill_(0, decode_req_idx, 1)
352
+ token_group = torch.repeat_interleave(
353
+ req_group, query_lens, output_size=int(query_start_loc_cpu[-1])
354
+ )
355
token_perm = torch.argsort(token_group, stable=True)
356
spec_token_indx = token_perm[:num_spec_decode_tokens]
357
non_spec_token_indx = token_perm[num_spec_decode_tokens:]
358
359
spec_state_indices_tensor = block_table_tensor[spec_req_idx, 0]
360
+ num_accepted_tokens = num_accepted_tokens[spec_req_idx]
361
non_spec_state_indices_tensor = block_table_tensor[non_spec_req_idx, 0]
362
spec_query_start_loc = torch.zeros(
363
num_spec_decodes + 1,
@@ -385,22 +375,8 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
375
dim=0,
376
out=non_spec_query_start_loc[1:],
377
)
388
- non_spec_query_start_loc_cpu = torch.zeros(
389
- num_decodes + num_prefills + 1, dtype=torch.int32
390
- )
391
- torch.cumsum(
392
- query_lens_cpu[non_spec_req_idx_cpu],
393
- dim=0,
394
- out=non_spec_query_start_loc_cpu[1:],
395
- )
378
397
- assert num_accepted_tokens is not None
398
- # Accepted-token counts must follow the same request order as the
399
- # speculative state indices.
400
- num_accepted_tokens = num_accepted_tokens[spec_req_idx]
401
-
402
- # Compute the conv-state slots for the non-spec decode/prefill split,
403
- # plus the initial-state masks and Triton causal_conv1d metadata.
379
+ # Compute the conv-state slots for the non-spec decode/prefill split.
380
if non_spec_state_indices_tensor is None:
381
state_indices_tensor = block_table_tensor[:0, 0]
382
else:
@@ -431,20 +407,9 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
407
]
408
has_initial_states_p = num_computed_tokens_p > 0
409
assert non_spec_query_start_loc is not None
434
- assert non_spec_query_start_loc_cpu is not None
410
query_start_loc_p = (
411
non_spec_query_start_loc[num_decodes:] - num_decode_tokens
412
)
438
- query_start_loc_p_cpu = (
439
- non_spec_query_start_loc_cpu[num_decodes:] - num_decode_tokens
440
- )
441
- if query_start_loc.device.type != "cpu":
442
- nums_dict, batch_ptr, token_chunk_offset_ptr = (
443
- compute_causal_conv1d_metadata(
444
- query_start_loc_p_cpu,
445
- device=query_start_loc.device,
446
- )
447
- )
413
414
# Prepare persistent tensors for CUDA graph capture and replay.
415
# ``m.num_actual_tokens`` is already padded by the model runner.
@@ -509,9 +474,6 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
474
non_spec_token_indx=non_spec_token_indx,
475
num_accepted_tokens=num_accepted_tokens,
476
num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
512
- nums_dict=nums_dict,
513
- batch_ptr=batch_ptr,
514
- token_chunk_offset_ptr=token_chunk_offset_ptr,
477
query_start_loc_p=query_start_loc_p,
478
query_start_loc_d=query_start_loc_d,
479
state_indices_tensor_p=state_indices_tensor_p,