422 lines
20 KiB
Diff
422 lines
20 KiB
Diff
From b24ac2b5bbf91bd73342d275f39aa9b6bb3f338a Mon Sep 17 00:00:00 2001
|
|
From: halt95 <halt95@users.noreply.github.com>
|
|
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
|
|
|