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

[Perf] Fuse DSV4.1 input metadata preparation with Triton (#56562)

Browse Source
This commit is contained in:
Woosuk Kwon
2026-09-12 07:20:04 -07:00
committed by GitHub
parent 2f59050eda
commit 13e221f830
4 changed files with 274 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+105
View File
@@ -32,6 +32,111 @@ from vllm.v1.worker.block_table import get_block_table_width
from vllm.v1.worker.utils import select_common_block_size
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
@pytest.mark.parametrize("query_lens", [[1], [6] * 16, [0, 2, 6, 0, 4], [0, 0, 0]])
@pytest.mark.parametrize("padding", [0, 17])
def test_fused_indexer_decode_metadata(query_lens, padding):
"""Flatten device query boundaries and clear graph padding on every replay."""
from vllm.v1.attention.ops.metadata import _indexer_decode_metadata_kernel
device = "cuda"
lengths = torch.tensor(query_lens, device=device, dtype=torch.int32)
qsl = torch.cat(
[torch.zeros(1, device=device, dtype=torch.int32), lengths.cumsum(0).int()]
)
reqs = len(query_lens)
n = sum(query_lens)
tokens = n + padding
capacity = max(tokens, reqs) + 31
seq = lengths + 4096
# Exercise block-table row padding and masked columns.
bt = torch.arange(reqs * 74, device=device, dtype=torch.int32).view(reqs, 74)[
:, :65
]
out_bt = torch.full((capacity, 69), -99, device=device, dtype=torch.int32)
outputs = [
torch.full((capacity,), -99, device=device, dtype=torch.int32) for _ in range(4)
]
out_seq, out_lens, indices, per_req = outputs
grid = max(reqs, tokens + (capacity - tokens + 255) // 256)
_indexer_decode_metadata_kernel[(grid,)](
qsl,
seq,
bt,
out_seq,
out_bt,
out_lens,
indices,
per_req,
reqs,
tokens,
capacity,
bt.stride(0),
out_bt.stride(0),
BLOCK_COLS=bt.shape[1],
num_warps=4,
)
expected_req = torch.repeat_interleave(torch.arange(reqs, device=device), lengths)
expected_seq = torch.cat(
[torch.arange(4097, 4097 + q, device=device) for q in query_lens]
)
torch.testing.assert_close(out_seq[:n], expected_seq.int(), rtol=0, atol=0)
assert torch.count_nonzero(out_seq[n:]) == 0
torch.testing.assert_close(out_bt[:n, :65], bt[expected_req], rtol=0, atol=0)
assert torch.count_nonzero(out_bt[n:tokens, :65]) == 0
torch.testing.assert_close(indices[:n], expected_req.int(), rtol=0, atol=0)
torch.testing.assert_close(
indices[n:tokens],
torch.arange(reqs, reqs + padding, device=device, dtype=torch.int32),
)
assert torch.all(out_lens[:tokens] == 1)
torch.testing.assert_close(per_req[:reqs], lengths)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
@pytest.mark.parametrize("query_lens", [[1], [0, 257, 1, 0, 3], [6] * 64, [2, 6, 4, 0]])
def test_device_token_request_mapping(query_lens):
"""Graph replay follows device boundaries even when CPU lengths are stale."""
lengths = torch.tensor(query_lens, device="cuda", dtype=torch.int32)
qsl = torch.cat(
[torch.zeros(1, device="cuda", dtype=torch.int32), lengths.cumsum(0).int()]
)
n = sum(query_lens)
output = torch.full((n + 7,), -99, device="cuda", dtype=torch.int32)
common = CommonAttentionMetadata(
query_start_loc=qsl,
query_start_loc_cpu=qsl.cpu(),
seq_lens=lengths,
num_reqs=len(query_lens),
num_actual_tokens=output.numel(),
max_query_len=max(query_lens),
max_seq_len=max(query_lens),
block_table_tensor=torch.empty(
(len(query_lens), 1), device="cuda", dtype=torch.int32
),
slot_mapping=torch.full((n + 7,), -1, device="cuda", dtype=torch.int64),
)
result = common.token_to_req_indices(output)
assert result.data_ptr() == output.data_ptr()
expected = torch.repeat_interleave(
torch.arange(len(query_lens), device="cuda", dtype=torch.int32), lengths
)
torch.testing.assert_close(output[:n], expected)
assert torch.count_nonzero(output[n:]) == 0
graph = torch.cuda.CUDAGraph()
common._token_to_req_indices_cache = None
with torch.cuda.graph(graph):
common.token_to_req_indices(output)
reversed_lens = lengths.flip(0)
qsl[1:].copy_(reversed_lens.cumsum(0))
graph.replay()
expected = torch.repeat_interleave(
torch.arange(len(query_lens), device="cuda", dtype=torch.int32), reversed_lens
)
torch.testing.assert_close(output[:n], expected)
assert torch.count_nonzero(output[n:]) == 0
def test_indexer_shares_uncompressed_block_size_with_deepseek_v4_mla():
"""Packed MLA/indexer groups must retain 64 compressed rows per page."""
kernel_block_size = select_common_block_size(
vllm/v1/attention/backend.py
+13 -11
View File
@@ -484,19 +484,21 @@ class CommonAttentionMetadata:
# Built from the device query_start_loc: adaptive verification decides the
# per-request draft split on device, so the CPU copy carries the right total
# but not the right per-request boundaries. Padding requests have a query
# length of zero and drop out of the repeat.
# length of zero and are skipped by the device boundary search.
num_mapped_tokens = int(self.query_start_loc_cpu[-1])
query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1]
assert buffer.shape[0] >= max(num_mapped_tokens, num_tokens)
token_to_req_indices = torch.repeat_interleave(
torch.arange(query_lens.shape[0], dtype=torch.int32, device=buffer.device),
query_lens,
output_size=num_mapped_tokens,
from vllm.v1.attention.ops.metadata import _token_request_mapping_kernel
num_output_tokens = max(num_mapped_tokens, num_tokens)
assert buffer.shape[0] >= num_output_tokens
_token_request_mapping_kernel[((num_output_tokens + 255) // 256,)](
self.query_start_loc,
buffer,
self.query_start_loc.shape[0] - 1,
num_mapped_tokens,
num_output_tokens,
num_warps=4,
)
buffer[:num_mapped_tokens].copy_(token_to_req_indices)
if num_mapped_tokens < num_tokens:
buffer[num_mapped_tokens:num_tokens].zero_()
self._token_to_req_indices_cache = buffer[: max(num_mapped_tokens, num_tokens)]
self._token_to_req_indices_cache = buffer[:num_output_tokens]
return self._token_to_req_indices_cache[:num_tokens]
# TODO(lucas): remove once we have FULL-CG spec-decode support
vllm/v1/attention/backends/mla/indexer.py
+69 -52
View File
@@ -960,28 +960,6 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
self.global_decode_seq_lens_buffer[actual_expanded:num_decode_tokens] = 0
return self.global_decode_seq_lens_buffer[:num_decode_tokens]
def _build_varlen_decode_indices(
self,
decode_lens: torch.Tensor,
decode_lens_cpu: torch.Tensor,
num_decode_tokens: int,
) -> torch.Tensor:
"""Build request ids for flattened SM100 varlen rows."""
indices = self.decode_indices_buffer[:num_decode_tokens]
actual_expanded = int(decode_lens_cpu.sum().item())
num_decodes = decode_lens.shape[0]
indices[:actual_expanded] = torch.repeat_interleave(
self.arange_buffer[:num_decodes],
decode_lens,
output_size=actual_expanded,
)
if actual_expanded < num_decode_tokens:
pad = num_decode_tokens - actual_expanded
indices[actual_expanded:num_decode_tokens] = (
num_decodes + self.arange_buffer[:pad]
)
return indices
@staticmethod
def _split_indexer_prefill_chunks(
compressed_seq_lens_cpu: torch.Tensor,
@@ -1156,12 +1134,15 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
decode_metadata = None
if num_decodes > 0:
torch.diff(
common_attn_metadata.query_start_loc[: num_decodes + 1],
out=self.decode_lens_buffer[:num_decodes],
)
if not self.supports_varlen:
torch.diff(
common_attn_metadata.query_start_loc[: num_decodes + 1],
out=self.decode_lens_buffer[:num_decodes],
)
self.per_req_decode_lens_buffer[:num_decodes].copy_(
self.decode_lens_buffer[:num_decodes]
)
decode_lens = self.decode_lens_buffer[:num_decodes]
self.per_req_decode_lens_buffer[:num_decodes].copy_(decode_lens)
decode_lens_cpu = torch.diff(
common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]
)
@@ -1196,38 +1177,74 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
and step_next_n_ok
)
global_seq_lens_for_decode = self._prepare_global_decode_seq_lens(
global_seq_lens=global_seq_lens_for_decode,
decode_lens=decode_lens,
decode_lens_cpu=decode_lens_cpu,
query_start_loc=common_attn_metadata.query_start_loc[:num_decodes],
num_decode_tokens=num_decode_tokens,
use_native=use_native,
max_decode_len=max_decode_len,
)
decode_indices = None
if self.supports_varlen:
decode_indices = self._build_varlen_decode_indices(
decode_lens=decode_lens,
decode_lens_cpu=decode_lens_cpu,
num_decode_tokens=num_decode_tokens,
)
seq_lens, block_table, decode_lens, batch_size, requires_padding = (
self._prepare_decode_tensors(
seq_lens=seq_lens,
block_table=block_table,
if not self.supports_varlen:
global_seq_lens_for_decode = self._prepare_global_decode_seq_lens(
global_seq_lens=global_seq_lens_for_decode,
decode_lens=decode_lens,
decode_lens_cpu=decode_lens_cpu,
query_start_loc=common_attn_metadata.query_start_loc[:num_decodes],
num_decodes=num_decodes,
num_decode_tokens=num_decode_tokens,
use_native=use_native,
next_n=next_n,
max_decode_len=max_decode_len,
)
)
decode_indices = None
if self.supports_varlen:
from vllm.v1.attention.ops.metadata import (
_indexer_decode_metadata_kernel,
)
capacity = self.decode_seq_lens_buffer.numel()
grid = max(
num_decodes,
num_decode_tokens + triton.cdiv(capacity - num_decode_tokens, 256),
)
_indexer_decode_metadata_kernel[(grid,)](
query_start_loc,
seq_lens,
block_table,
self.decode_seq_lens_buffer,
self.expanded_block_table_buffer,
self.decode_lens_buffer,
self.decode_indices_buffer,
self.per_req_decode_lens_buffer,
num_decodes,
num_decode_tokens,
capacity,
block_table.stride(0),
self.expanded_block_table_buffer.stride(0),
BLOCK_COLS=block_table.shape[1],
num_warps=4,
)
seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
block_table = self.expanded_block_table_buffer[:num_decode_tokens]
decode_lens = self.decode_lens_buffer[:num_decode_tokens]
decode_indices = self.decode_indices_buffer[:num_decode_tokens]
requires_padding = False
if global_seq_lens_for_decode is not None and max_decode_len > 1:
self.global_decode_seq_lens_buffer[:num_decode_tokens].copy_(
seq_lens
)
global_seq_lens_for_decode = self.global_decode_seq_lens_buffer[
:num_decode_tokens
]
else:
seq_lens, block_table, decode_lens, batch_size, requires_padding = (
self._prepare_decode_tensors(
seq_lens=seq_lens,
block_table=block_table,
decode_lens=decode_lens,
decode_lens_cpu=decode_lens_cpu,
query_start_loc=common_attn_metadata.query_start_loc[
:num_decodes
],
num_decodes=num_decodes,
num_decode_tokens=num_decode_tokens,
use_native=use_native,
next_n=next_n,
max_decode_len=max_decode_len,
)
)
if self.compress_ratio > 1:
kernel_block_size = self.kernel_block_size
vllm/v1/attention/ops/metadata.py
+87
View File
@@ -0,0 +1,87 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Device-side request mapping and sparse indexer metadata."""
from vllm.triton_utils import tl, triton
@triton.jit
def _token_request(query_start_loc, token, num_reqs):
lo = tl.full(token.shape, 0, tl.int32)
hi = tl.full(token.shape, num_reqs, tl.int32)
while tl.sum((lo < hi).to(tl.int32), 0) > 0:
mid = (lo + hi) // 2
end = tl.load(query_start_loc + mid + 1, mid < num_reqs, other=0x7FFFFFFF)
right = (lo < hi) & (end <= token)
lo = tl.where(right, mid + 1, lo)
hi = tl.where((lo < hi) & ~right, mid, hi)
return tl.minimum(lo, num_reqs - 1)
@triton.jit(do_not_specialize=["num_reqs", "num_mapped", "num_tokens"])
def _token_request_mapping_kernel(
query_start_loc, output, num_reqs, num_mapped, num_tokens
):
token = tl.program_id(0) * 256 + tl.arange(0, 256)
req = _token_request(query_start_loc, token, num_reqs)
tl.store(output + token, tl.where(token < num_mapped, req, 0), token < num_tokens)
@triton.jit(
do_not_specialize=[
"num_reqs",
"num_tokens",
"capacity",
"input_stride",
"output_stride",
]
)
def _indexer_decode_metadata_kernel(
query_start_loc,
seq_lens,
block_table,
output_seq_lens,
output_block_table,
decode_lens,
indices,
per_req_lens,
num_reqs,
num_tokens,
capacity,
input_stride,
output_stride,
BLOCK_COLS: tl.constexpr,
):
token = tl.program_id(0)
cols = tl.arange(0, triton.next_power_of_2(BLOCK_COLS))
actual_tokens = tl.load(query_start_loc + num_reqs)
if token < num_tokens:
req = _token_request(query_start_loc, tl.full((1,), token, tl.int32), num_reqs)
req = tl.sum(req, 0)
valid = token < actual_tokens
start = tl.load(query_start_loc + req)
end = tl.load(query_start_loc + req + 1)
seq = tl.load(seq_lens + req)
length = tl.where(valid, seq - (end - start) + token - start + 1, 0)
tl.store(output_seq_lens + token, length)
tl.store(decode_lens + token, 1)
tl.store(
indices + token, tl.where(valid, req, num_reqs + token - actual_tokens)
)
blocks = tl.load(
block_table + req.to(tl.int64) * input_stride + cols,
valid & (cols < BLOCK_COLS),
other=0,
)
tl.store(
output_block_table + token.to(tl.int64) * output_stride + cols,
blocks,
cols < BLOCK_COLS,
)
else:
tail = num_tokens + (token - num_tokens) * 256 + tl.arange(0, 256)
tl.store(output_seq_lens + tail, 0, tail < capacity)
if token < num_reqs:
start = tl.load(query_start_loc + token)
end = tl.load(query_start_loc + token + 1)
tl.store(per_req_lens + token, end - start)