From b24ac2b5bbf91bd73342d275f39aa9b6bb3f338a Mon Sep 17 00:00:00 2001 From: halt95 Date: Mon, 28 Sep 2026 20:46:22 +0800 Subject: [PATCH 065/103] ple: reduce PLE metadata construction overhead (port of upstream #58114) Port of the upstream vLLM change #58114 (Apache-2.0). Applied cleanly, including the NVIDIA and AMD PLE layers; the pull transport does not read the changed metadata fields. --- tests/models/qwen4_exp/test_ple.py | 125 ++++++++++++++++++ vllm/models/qwen4_exp/amd/ple_layer.py | 9 +- vllm/models/qwen4_exp/nvidia/ple_layer.py | 3 +- vllm/v1/attention/backends/mamba_attn.py | 20 +-- vllm/v1/attention/backends/short_conv_attn.py | 82 +++--------- 5 files changed, 161 insertions(+), 78 deletions(-) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index e7d0e72..281275e 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -14,6 +14,9 @@ from torch.nn import functional as F import vllm.model_executor.layers.vocab_parallel_embedding as embedding_module import vllm.model_executor.parameter as parameter_module import vllm.models.qwen4_exp.nvidia.ngram_embedding as ngram_embedding_module +from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata +from vllm.config import SpeculativeConfig, VllmConfig +from vllm.config.compilation import CUDAGraphMode from vllm.model_executor.layers.quantization.fp8 import Fp8Config from vllm.model_executor.layers.quantization.modelopt import ( ModelOptMixedPrecisionConfig, @@ -35,8 +38,125 @@ from vllm.models.qwen4_exp.nvidia.ngram_embedding import ( from vllm.models.qwen4_exp.nvidia.ple_layer import Qwen4ExpPLELayer from vllm.v1.attention.backends.short_conv_attn import ( PleShortConvAttentionMetadata, + PleShortConvAttentionMetadataBuilder, ) from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from vllm.v1.kv_cache_interface import MambaSpec + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("full_graph", [False, True]) +@pytest.mark.parametrize( + "query_lens,drafts,spec_rows,non_spec_rows,token_order", + [ + pytest.param([4, 2, 0], [3, 1, -1], [0, 1], [], None, id="spec-padding"), + pytest.param( + [3, 1, 2, 5, 0], + [2, -1, 1, -1, -1], + [0, 2], + [1, 3], + [0, 1, 2, 4, 5, 3, 6, 7, 8, 9, 10], + id="mixed-padding", + ), + pytest.param([1, 5], [-1, -1], [], [0, 1], None, id="non-spec"), + ], +) +def test_ple_metadata_preserves_request_and_token_order( + device, full_graph, query_lens, drafts, spec_rows, non_spec_rows, token_order +): + """Preserve state routing, acceptance and prefill offsets without GPU sync.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + device = torch.device(device) + config = VllmConfig() + config.speculative_config = SpeculativeConfig( + method="ngram", num_speculative_tokens=3 + ) + config.compilation_config.cudagraph_mode = ( + CUDAGraphMode.FULL if full_graph else CUDAGraphMode.NONE + ) + config.compilation_config.max_cudagraph_capture_size = 32 + builder = PleShortConvAttentionMetadataBuilder( + MambaSpec(block_size=16, shapes=((16, 64),), dtypes=(torch.float32,)), + ["ple"], + config, + device, + ) + common = create_common_attn_metadata( + BatchSpec(seq_lens=[q + 16 for q in query_lens], query_lens=query_lens), + 16, + device, + arange_block_indices=True, + ) + common.is_prefilling = torch.tensor( + [q > 1 and d < 0 for q, d in zip(query_lens, drafts)] + ) + # Token padding is independent of the sum of per-request query lengths. + common.num_actual_tokens += 3 + accepted = torch.arange(1, len(query_lens) + 1, dtype=torch.int32, device=device) + drafts_cpu = torch.tensor(drafts, dtype=torch.int32) + sync_mode = torch.cuda.get_sync_debug_mode() if device.type == "cuda" else None + if sync_mode is not None: + torch.cuda.set_sync_debug_mode("error") + try: + metadata = builder.build( + 0, + common, + num_accepted_tokens=accepted, + num_decode_draft_tokens_cpu=drafts_cpu, + ) + finally: + if sync_mode is not None: + torch.cuda.set_sync_debug_mode(sync_mode) + + def check(actual, expected): + torch.testing.assert_close( + actual, torch.tensor(expected, dtype=actual.dtype, device=device) + ) + + assert metadata.num_spec_decodes == len(spec_rows) + assert metadata.num_decodes == sum(query_lens[i] == 1 for i in non_spec_rows) + assert metadata.num_prefills == sum(query_lens[i] > 1 for i in non_spec_rows) + if spec_rows: + torch.testing.assert_close( + metadata.spec_state_indices_tensor[: len(spec_rows)], + common.block_table_tensor[spec_rows, 0], + ) + torch.testing.assert_close( + metadata.num_accepted_tokens[: len(spec_rows)], accepted[spec_rows] + ) + check( + metadata.spec_query_start_loc[: len(spec_rows) + 1], + [0, *accumulate(query_lens[i] for i in spec_rows)], + ) + if full_graph and not non_spec_rows: + check(metadata.spec_state_indices_tensor[len(spec_rows) :], [NULL_BLOCK_ID]) + check(metadata.num_accepted_tokens[len(spec_rows) :], [1]) + check( + metadata.spec_query_start_loc[len(spec_rows) + 1 :], [sum(query_lens)] + ) + assert ( + metadata.spec_state_indices_tensor.data_ptr() + == builder.spec_state_indices_tensor.data_ptr() + ) + if non_spec_rows: + torch.testing.assert_close( + metadata.state_indices_tensor, common.block_table_tensor[non_spec_rows, 0] + ) + prefill_lens = [query_lens[i] for i in non_spec_rows if query_lens[i] > 1] + check(metadata.query_start_loc_p, [0, *accumulate(prefill_lens)]) + check(metadata.has_initial_states_p, [True] * len(prefill_lens)) + if token_order is not None: + check( + torch.cat((metadata.spec_token_indx, metadata.non_spec_token_indx)), + token_order, + ) + else: + assert metadata.spec_token_indx is None + assert metadata.non_spec_token_indx is None + assert metadata.nums_dict is None + assert metadata.batch_ptr is None + assert metadata.token_chunk_offset_ptr is None def _mock_etp_group( @@ -1501,6 +1621,11 @@ def _make_conv_metadata( else None ), non_spec_query_start_loc=(non_spec_query_start_loc if has_non_spec else None), + query_start_loc_p=( + non_spec_query_start_loc[case.num_decodes :] - case.num_decodes + if num_prefills + else None + ), has_initial_states_p=( torch.arange(num_prefills, device=device) % 2 == 0 if num_prefills else None ), diff --git a/vllm/models/qwen4_exp/amd/ple_layer.py b/vllm/models/qwen4_exp/amd/ple_layer.py index fc494be..c686e34 100644 --- a/vllm/models/qwen4_exp/amd/ple_layer.py +++ b/vllm/models/qwen4_exp/amd/ple_layer.py @@ -616,14 +616,9 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): num_decode_tokens: int, num_prefill_tokens: int, ) -> torch.Tensor: - # ``non_spec_query_start_loc`` covers the non-spec (decode + prefill) - # requests and equals ``query_start_loc`` when spec-decode is inactive. - non_spec_query_start_loc = metadata.non_spec_query_start_loc - if non_spec_query_start_loc is None: + query_start_loc_p = metadata.query_start_loc_p + if query_start_loc_p is None: raise ValueError("query_start_loc is required for prefill short-conv") - query_start_loc_p = ( - non_spec_query_start_loc[-num_prefills - 1 :] - num_decode_tokens - ) # The metadata builder guarantees that the prefill query offsets start # at 0 and end at num_prefill_tokens. Avoid reading those values here, # since doing so would force a device-to-host synchronization. diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 94d908e..4925e03 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -320,10 +320,9 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase): token_indices=token_indices_d, ) - query_start_loc = metadata.non_spec_query_start_loc + query_start_loc = metadata.query_start_loc_p if query_start_loc is None: raise ValueError("query_start_loc is required for prefill short-conv") - query_start_loc = query_start_loc[-num_prefills - 1 :] - num_decode_tokens has_initial_states = metadata.has_initial_states_p if has_initial_states is None: raise ValueError( diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index e0b5fab..817773a 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -94,6 +94,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): # Will be disabled if speculative decoding is used supports_update_block_table: bool = True + needs_causal_conv1d_metadata: bool = True def __init__( self, @@ -602,10 +603,6 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): if num_computed_tokens is None: num_computed_tokens = common_attn_metadata.compute_num_computed_tokens() - query_start_loc_p_cpu = ( - common_attn_metadata.query_start_loc_cpu[-num_prefills - 1 :] - - num_decode_tokens - ) query_start_loc_p = ( common_attn_metadata.query_start_loc[-num_prefills - 1 :] - num_decode_tokens @@ -614,12 +611,17 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): num_computed_tokens[num_reqs - num_prefills : num_reqs] > 0 ) - nums_dict, batch_ptr, token_chunk_offset_ptr = ( - compute_causal_conv1d_metadata( - query_start_loc_p_cpu, - device=common_attn_metadata.query_start_loc.device, + if self.needs_causal_conv1d_metadata: + query_start_loc_p_cpu = ( + common_attn_metadata.query_start_loc_cpu[-num_prefills - 1 :] + - num_decode_tokens + ) + nums_dict, batch_ptr, token_chunk_offset_ptr = ( + compute_causal_conv1d_metadata( + query_start_loc_p_cpu, + device=common_attn_metadata.query_start_loc.device, + ) ) - ) if self.vllm_config.cache_config.mamba_cache_mode == "all": assert num_computed_tokens is not None diff --git a/vllm/v1/attention/backends/short_conv_attn.py b/vllm/v1/attention/backends/short_conv_attn.py index 5a0c6f7..f2cc3da 100644 --- a/vllm/v1/attention/backends/short_conv_attn.py +++ b/vllm/v1/attention/backends/short_conv_attn.py @@ -18,7 +18,6 @@ from vllm.v1.attention.backends.mamba_attn import ( ) from vllm.v1.attention.backends.utils import ( NULL_BLOCK_ID, - compute_causal_conv1d_metadata, mamba_get_block_table_tensor, ) from vllm.v1.kv_cache_interface import MambaSpec @@ -108,6 +107,7 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH reorder_batch_threshold: int = 1 supports_update_block_table = False + needs_causal_conv1d_metadata = False def __init__( self, @@ -281,18 +281,10 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): spec_sequence_masks_cpu, device=query_start_loc.device ) - # For causal_conv1d (non-spec prefill Triton kernel metadata). - nums_dict = None - batch_ptr = None - token_chunk_offset_ptr = None has_initial_states_p = None has_initial_states_d = None num_computed_tokens_p = None - # Original request indices of the non-spec requests, ordered - # [decodes, prefills]. Used to gather per-request data consistently. - non_spec_req_idx_cpu: torch.Tensor | None = None - query_lens = torch.diff(query_start_loc) # Per-request classification by mask, NOT by position. With # spec-decode, the front decode group can contain both spec-decode # requests and plain non-spec single-token decodes. Spec requests are @@ -317,37 +309,32 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): else 0 ) - # Original request indices grouped as - # [spec | non-spec decode | non-spec prefill]; each group keeps the - # original (already reordered) relative order via a stable nonzero. - spec_req_idx_cpu = spec_sequence_masks_cpu.nonzero(as_tuple=True)[0] - decode_req_idx_cpu = decode_mask_cpu.nonzero(as_tuple=True)[0] - prefill_req_idx_cpu = prefill_mask_cpu.nonzero(as_tuple=True)[0] - non_spec_req_idx_cpu = torch.cat((decode_req_idx_cpu, prefill_req_idx_cpu)) - spec_req_idx = async_tensor_h2d(spec_req_idx_cpu, device=query_start_loc.device) + assert num_accepted_tokens is not None non_spec_req_idx: torch.Tensor | None = None if num_decodes == 0 and num_prefills == 0: - # Pure speculative-decode batch: all real tokens are spec tokens. - spec_token_indx = torch.arange( - num_spec_decode_tokens, - dtype=torch.int32, - device=query_start_loc.device, - ) - non_spec_token_indx = torch.empty( - 0, dtype=torch.int32, device=query_start_loc.device - ) - spec_state_indices_tensor = block_table_tensor[spec_req_idx, 0] + # All real requests are speculative; zero-length padding is trailing. + spec_token_indx = None + non_spec_token_indx = None + spec_state_indices_tensor = block_table_tensor[:num_spec_decodes, 0] + num_accepted_tokens = num_accepted_tokens[:num_spec_decodes] non_spec_state_indices_tensor = None spec_query_start_loc = query_start_loc[: num_spec_decodes + 1] non_spec_query_start_loc = None - non_spec_query_start_loc_cpu = None else: # Mixed batch: build a per-token group key consistent with the # request grouping above (spec=0 | decode=1 | prefill=2) and a # stable sort, so tokens of each request stay contiguous and in # request order. This yields spec tokens first, then the non-spec # [decode, prefill] tokens. + query_lens = torch.diff(query_start_loc) + spec_req_idx_cpu = spec_sequence_masks_cpu.nonzero(as_tuple=True)[0] + decode_req_idx_cpu = decode_mask_cpu.nonzero(as_tuple=True)[0] + prefill_req_idx_cpu = prefill_mask_cpu.nonzero(as_tuple=True)[0] + non_spec_req_idx_cpu = torch.cat((decode_req_idx_cpu, prefill_req_idx_cpu)) + spec_req_idx = async_tensor_h2d( + spec_req_idx_cpu, device=query_start_loc.device + ) non_spec_req_idx = async_tensor_h2d( non_spec_req_idx_cpu, device=query_start_loc.device ) @@ -360,14 +347,17 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): dtype=torch.int64, device=query_start_loc.device, ) - req_group[spec_req_idx] = 0 - req_group[decode_req_idx] = 1 - token_group = torch.repeat_interleave(req_group, query_lens) + req_group.index_fill_(0, spec_req_idx, 0) + req_group.index_fill_(0, decode_req_idx, 1) + token_group = torch.repeat_interleave( + req_group, query_lens, output_size=int(query_start_loc_cpu[-1]) + ) token_perm = torch.argsort(token_group, stable=True) spec_token_indx = token_perm[:num_spec_decode_tokens] non_spec_token_indx = token_perm[num_spec_decode_tokens:] spec_state_indices_tensor = block_table_tensor[spec_req_idx, 0] + num_accepted_tokens = num_accepted_tokens[spec_req_idx] non_spec_state_indices_tensor = block_table_tensor[non_spec_req_idx, 0] spec_query_start_loc = torch.zeros( num_spec_decodes + 1, @@ -385,22 +375,8 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): dim=0, out=non_spec_query_start_loc[1:], ) - non_spec_query_start_loc_cpu = torch.zeros( - num_decodes + num_prefills + 1, dtype=torch.int32 - ) - torch.cumsum( - query_lens_cpu[non_spec_req_idx_cpu], - dim=0, - out=non_spec_query_start_loc_cpu[1:], - ) - assert num_accepted_tokens is not None - # Accepted-token counts must follow the same request order as the - # speculative state indices. - num_accepted_tokens = num_accepted_tokens[spec_req_idx] - - # Compute the conv-state slots for the non-spec decode/prefill split, - # plus the initial-state masks and Triton causal_conv1d metadata. + # Compute the conv-state slots for the non-spec decode/prefill split. if non_spec_state_indices_tensor is None: state_indices_tensor = block_table_tensor[:0, 0] else: @@ -431,20 +407,9 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): ] has_initial_states_p = num_computed_tokens_p > 0 assert non_spec_query_start_loc is not None - assert non_spec_query_start_loc_cpu is not None query_start_loc_p = ( non_spec_query_start_loc[num_decodes:] - num_decode_tokens ) - query_start_loc_p_cpu = ( - non_spec_query_start_loc_cpu[num_decodes:] - num_decode_tokens - ) - if query_start_loc.device.type != "cpu": - nums_dict, batch_ptr, token_chunk_offset_ptr = ( - compute_causal_conv1d_metadata( - query_start_loc_p_cpu, - device=query_start_loc.device, - ) - ) # Prepare persistent tensors for CUDA graph capture and replay. # ``m.num_actual_tokens`` is already padded by the model runner. @@ -509,9 +474,6 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): non_spec_token_indx=non_spec_token_indx, num_accepted_tokens=num_accepted_tokens, num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, - nums_dict=nums_dict, - batch_ptr=batch_ptr, - token_chunk_offset_ptr=token_chunk_offset_ptr, query_start_loc_p=query_start_loc_p, query_start_loc_d=query_start_loc_d, state_indices_tensor_p=state_indices_tensor_p, -- 2.43.0