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

[Spec Decode][Qwen3.8] Enable fused draft metadata updates for the QSA side cache

Browse Source 
QSAMetadataBuilder did not set supports_draft_decode_metadata_update, so
MTP with num_speculative_tokens > 1 fell back to rebuilding attention
metadata between draft steps for the whole model.

Opt the builder in. update_draft_decode_metadata re-launches the builder's
own _build_qsa_metadata_kernel in place on the metadata's persistent
tensors, so draft-step token_to_req, logical_positions, visible_blocks,
slot_mapping and k_work_metadata come from the same code path as build().
The kernel launch is factored into _launch_qsa_metadata_kernel, and
QSAForwardMetadata carries common_slot_mapping and num_mapped_tokens so
the update does not read query_start_loc_cpu. The PyTorch fallback builder
keeps the rebuild path.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: bilikas <valdas@vbtech.eu>
This commit is contained in:
bilikas
2026-09-23 20:58:12 +03:00
co-authored by Claude Opus 5.5
parent 15ed1262e7
commit 31d78a21b9
2 changed files with 135 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/qwen4_exp/test_qsa_reference.py
+61
View File
@@ -456,6 +456,67 @@ def test_qsa_compressed_metadata_keeps_dummy_slots_inert() -> None:
@requires_qsa_kernels
@requires_qsa_kernels
@pytest.mark.parametrize("layout", ["plain", "compressed", "circular"])
def test_qsa_draft_decode_metadata_update_matches_rebuild(layout: str) -> None:
"""A fused draft step's in-place update equals a fresh build at seq_lens + 1."""
device = torch.device("cuda")
compress_ratio = 4 if layout != "plain" else 1
def make_builder() -> QSAMetadataBuilder:
builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder)
builder.compress_ratio = compress_ratio
builder.reorder_batch_threshold = 1
builder.is_circular_buffer = layout == "circular"
builder.kv_cache_spec = SimpleNamespace(block_size=4)
builder.storage_block_size = 16 if layout == "compressed" else 4
builder.token_to_req_buffer = torch.empty(8, dtype=torch.int32, device=device)
builder.slot_mapping_buffer = torch.empty(8, dtype=torch.int64, device=device)
builder.logical_positions_buffer = torch.empty(
8, dtype=torch.int64, device=device
)
builder.visible_blocks_buffer = torch.empty(8, dtype=torch.int32, device=device)
builder.request_capacity = 8
builder.k_work_metadata_buffer = torch.empty(
8 if layout == "compressed" else 0, 2, dtype=torch.int32, device=device
)
return builder
query_start_loc = torch.tensor([0, 1, 2, 3], dtype=torch.int32, device=device)
token_to_req = torch.tensor([0, 1, 2], dtype=torch.int32, device=device)
def make_common(seq_lens: torch.Tensor) -> SimpleNamespace:
return SimpleNamespace(
num_actual_tokens=3,
num_reqs=3,
max_query_len=1,
max_seq_len=int(seq_lens.max()),
query_start_loc=query_start_loc,
query_start_loc_cpu=query_start_loc.cpu(),
seq_lens=seq_lens,
slot_mapping=torch.tensor([5, 9, 20], dtype=torch.int64, device=device),
block_table_tensor=torch.tensor(
[[3, 1], [0, 4], [2, 5]], dtype=torch.int32, device=device
),
token_to_req_indices=lambda buffer: buffer.copy_(token_to_req),
)
# 7 -> 8 crosses a compression boundary, so visible_blocks must change.
seq_lens = torch.tensor([7, 11, 16], dtype=torch.int32, device=device)
builder = make_builder()
metadata = builder.build(0, make_common(seq_lens))
positions_before = metadata.logical_positions.clone()
seq_lens.add_(1)
builder.update_draft_decode_metadata(metadata)
expected = make_builder().build(0, make_common(seq_lens.clone()))
assert not torch.equal(metadata.logical_positions, positions_before)
for name in ("token_to_req", "logical_positions", "visible_blocks", "slot_mapping"):
torch.testing.assert_close(getattr(metadata, name), getattr(expected, name))
torch.testing.assert_close(metadata.k_work_metadata, expected.k_work_metadata)
@pytest.mark.usefixtures("default_vllm_config")
def test_qsa_unfused_cache_update_ignores_padded_qk() -> None:
"""Padded projected Q/K rows must not affect either side cache."""
vllm/models/qwen4_exp/common/qsa_cache.py
+74 -12
View File
@@ -398,12 +398,53 @@ def build_qsa_metadata_triton(
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""Build QSA side-cache and optional pre-indexer work metadata."""
num_tokens = common_attn_metadata.num_actual_tokens
num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1])
token_to_req = token_to_req_buffer[:num_tokens]
logical_positions = logical_positions_buffer[:num_tokens]
visible_blocks = visible_blocks_buffer[:num_tokens]
slot_mapping = slot_mapping_buffer[:num_tokens]
num_reqs = common_attn_metadata.query_start_loc.shape[0] - 1
if num_tokens == 0 and k_work_metadata_buffer is None:
return token_to_req, logical_positions, visible_blocks, slot_mapping
_launch_qsa_metadata_kernel(
common_attn_metadata.query_start_loc,
common_attn_metadata.seq_lens,
common_attn_metadata.slot_mapping,
common_attn_metadata.block_table_tensor,
token_to_req,
logical_positions,
visible_blocks,
slot_mapping,
k_work_metadata_buffer,
num_mapped_tokens=int(common_attn_metadata.query_start_loc_cpu[-1]),
storage_block_size=storage_block_size,
compress_ratio=compress_ratio,
circular_buffer_size=circular_buffer_size,
request_capacity=request_capacity,
)
if circular_buffer_size == 0 and compress_ratio == 1:
slot_mapping = common_attn_metadata.slot_mapping[:num_tokens]
return token_to_req, logical_positions, visible_blocks, slot_mapping
def _launch_qsa_metadata_kernel(
query_start_loc: torch.Tensor,
seq_lens: torch.Tensor,
common_slot_mapping: torch.Tensor,
block_table: torch.Tensor,
token_to_req: torch.Tensor,
logical_positions: torch.Tensor,
visible_blocks: torch.Tensor,
slot_mapping: torch.Tensor,
k_work_metadata_buffer: torch.Tensor | None,
*,
num_mapped_tokens: int,
storage_block_size: int,
compress_ratio: int,
circular_buffer_size: int,
request_capacity: int | None,
) -> None:
"""Fill QSA metadata in place; capture-safe given fixed shapes and scalars."""
num_tokens = token_to_req.shape[0]
num_reqs = query_start_loc.shape[0] - 1
assert num_reqs > 0
if k_work_metadata_buffer is not None:
@@ -417,10 +458,6 @@ def build_qsa_metadata_triton(
request_scan_size = 1
max_num_work = 0
if num_tokens == 0 and k_work_metadata_buffer is None:
return token_to_req, logical_positions, visible_blocks, slot_mapping
block_table = common_attn_metadata.block_table_tensor
num_search_steps = int(math.ceil(math.log2(num_reqs)))
work_search_steps = int(math.ceil(math.log2(num_reqs)))
# The same grid covers token tiles and, for the compressed cache, work tiles.
@@ -429,9 +466,9 @@ def build_qsa_metadata_triton(
cdiv(max_num_work, 256) if k_work_metadata_buffer is not None else 0
)
_build_qsa_metadata_kernel[(max(num_token_blocks, num_work_blocks, 1),)](
common_attn_metadata.query_start_loc,
common_attn_metadata.seq_lens,
common_attn_metadata.slot_mapping,
query_start_loc,
seq_lens,
common_slot_mapping,
block_table,
token_to_req,
logical_positions,
@@ -456,9 +493,6 @@ def build_qsa_metadata_triton(
WORK_BLOCK_SIZE=256,
num_warps=4,
)
if circular_buffer_size == 0 and compress_ratio == 1:
slot_mapping = common_attn_metadata.slot_mapping[:num_tokens]
return token_to_req, logical_positions, visible_blocks, slot_mapping
def _build_qsa_metadata_torch(
@@ -581,6 +615,8 @@ class QSAForwardMetadata(AttentionMetadata):
logical_positions: torch.Tensor
visible_blocks: torch.Tensor
k_work_metadata: torch.Tensor
common_slot_mapping: torch.Tensor
num_mapped_tokens: int
num_actual_tokens: int
num_decodes: int
num_decode_tokens: int
@@ -597,6 +633,7 @@ class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):
"""Build QSA metadata from vLLM's cache-group-specific common metadata."""
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
supports_draft_decode_metadata_update = HAS_TRITON
def __init__(
self,
@@ -708,6 +745,8 @@ class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):
logical_positions=logical_positions,
visible_blocks=visible_blocks,
k_work_metadata=k_work_metadata,
common_slot_mapping=common_attn_metadata.slot_mapping,
num_mapped_tokens=int(common_attn_metadata.query_start_loc_cpu[-1]),
num_actual_tokens=num_tokens,
num_decodes=num_decodes,
num_decode_tokens=num_decode_tokens,
@@ -720,6 +759,29 @@ class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):
compress_ratio=self.compress_ratio,
)
def update_draft_decode_metadata(self, metadata: QSAForwardMetadata) -> None:
if metadata.num_actual_tokens == 0:
return
build_k_work = not self.is_circular_buffer and self.compress_ratio != 1
_launch_qsa_metadata_kernel(
metadata.query_start_loc,
metadata.seq_lens,
metadata.common_slot_mapping,
metadata.block_table,
metadata.token_to_req,
metadata.logical_positions,
metadata.visible_blocks,
metadata.slot_mapping,
metadata.k_work_metadata if build_k_work else None,
num_mapped_tokens=metadata.num_mapped_tokens,
storage_block_size=self.storage_block_size,
compress_ratio=self.compress_ratio,
circular_buffer_size=(
self.kv_cache_spec.block_size if self.is_circular_buffer else 0
),
request_capacity=self.request_capacity if build_k_work else None,
)
class QSAStateBackend(AttentionBackend):
"""Key-only dummy backend for out-of-band QSA side-cache operations."""