[Perf] Fuse DSV4.1 input metadata preparation with Triton (#56562)
Woosuk Kwon committed
Sep 12, 2026 at 07:20 UTC
13e221f8308416564a446909c2fd4805c78fbb8a
4 files changed
+274
-63
tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+105
@@ -32,6 +32,111 @@ from vllm.v1.worker.block_table import get_block_table_width
32
from vllm.v1.worker.utils import select_common_block_size
33
34
35
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
36
+@pytest.mark.parametrize("query_lens", [[1], [6] * 16, [0, 2, 6, 0, 4], [0, 0, 0]])
37
+@pytest.mark.parametrize("padding", [0, 17])
38
+def test_fused_indexer_decode_metadata(query_lens, padding):
39
+ """Flatten device query boundaries and clear graph padding on every replay."""
40
+ from vllm.v1.attention.ops.metadata import _indexer_decode_metadata_kernel
41
+
42
+ device = "cuda"
43
+ lengths = torch.tensor(query_lens, device=device, dtype=torch.int32)
44
+ qsl = torch.cat(
45
+ [torch.zeros(1, device=device, dtype=torch.int32), lengths.cumsum(0).int()]
46
+ )
47
+ reqs = len(query_lens)
48
+ n = sum(query_lens)
49
+ tokens = n + padding
50
+ capacity = max(tokens, reqs) + 31
51
+ seq = lengths + 4096
52
+ # Exercise block-table row padding and masked columns.
53
+ bt = torch.arange(reqs * 74, device=device, dtype=torch.int32).view(reqs, 74)[
54
+ :, :65
55
+ ]
56
+ out_bt = torch.full((capacity, 69), -99, device=device, dtype=torch.int32)
57
+ outputs = [
58
+ torch.full((capacity,), -99, device=device, dtype=torch.int32) for _ in range(4)
59
+ ]
60
+ out_seq, out_lens, indices, per_req = outputs
61
+ grid = max(reqs, tokens + (capacity - tokens + 255) // 256)
62
+ _indexer_decode_metadata_kernel[(grid,)](
63
+ qsl,
64
+ seq,
65
+ bt,
66
+ out_seq,
67
+ out_bt,
68
+ out_lens,
69
+ indices,
70
+ per_req,
71
+ reqs,
72
+ tokens,
73
+ capacity,
74
+ bt.stride(0),
75
+ out_bt.stride(0),
76
+ BLOCK_COLS=bt.shape[1],
77
+ num_warps=4,
78
+ )
79
+ expected_req = torch.repeat_interleave(torch.arange(reqs, device=device), lengths)
80
+ expected_seq = torch.cat(
81
+ [torch.arange(4097, 4097 + q, device=device) for q in query_lens]
82
+ )
83
+ torch.testing.assert_close(out_seq[:n], expected_seq.int(), rtol=0, atol=0)
84
+ assert torch.count_nonzero(out_seq[n:]) == 0
85
+ torch.testing.assert_close(out_bt[:n, :65], bt[expected_req], rtol=0, atol=0)
86
+ assert torch.count_nonzero(out_bt[n:tokens, :65]) == 0
87
+ torch.testing.assert_close(indices[:n], expected_req.int(), rtol=0, atol=0)
88
+ torch.testing.assert_close(
89
+ indices[n:tokens],
90
+ torch.arange(reqs, reqs + padding, device=device, dtype=torch.int32),
91
+ )
92
+ assert torch.all(out_lens[:tokens] == 1)
93
+ torch.testing.assert_close(per_req[:reqs], lengths)
94
+
95
+
96
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
97
+@pytest.mark.parametrize("query_lens", [[1], [0, 257, 1, 0, 3], [6] * 64, [2, 6, 4, 0]])
98
+def test_device_token_request_mapping(query_lens):
99
+ """Graph replay follows device boundaries even when CPU lengths are stale."""
100
+ lengths = torch.tensor(query_lens, device="cuda", dtype=torch.int32)
101
+ qsl = torch.cat(
102
+ [torch.zeros(1, device="cuda", dtype=torch.int32), lengths.cumsum(0).int()]
103
+ )
104
+ n = sum(query_lens)
105
+ output = torch.full((n + 7,), -99, device="cuda", dtype=torch.int32)
106
+ common = CommonAttentionMetadata(
107
+ query_start_loc=qsl,
108
+ query_start_loc_cpu=qsl.cpu(),
109
+ seq_lens=lengths,
110
+ num_reqs=len(query_lens),
111
+ num_actual_tokens=output.numel(),
112
+ max_query_len=max(query_lens),
113
+ max_seq_len=max(query_lens),
114
+ block_table_tensor=torch.empty(
115
+ (len(query_lens), 1), device="cuda", dtype=torch.int32
116
+ ),
117
+ slot_mapping=torch.full((n + 7,), -1, device="cuda", dtype=torch.int64),
118
+ )
119
+ result = common.token_to_req_indices(output)
120
+ assert result.data_ptr() == output.data_ptr()
121
+ expected = torch.repeat_interleave(
122
+ torch.arange(len(query_lens), device="cuda", dtype=torch.int32), lengths
123
+ )
124
+ torch.testing.assert_close(output[:n], expected)
125
+ assert torch.count_nonzero(output[n:]) == 0
126
+ graph = torch.cuda.CUDAGraph()
127
+ common._token_to_req_indices_cache = None
128
+ with torch.cuda.graph(graph):
129
+ common.token_to_req_indices(output)
130
+ reversed_lens = lengths.flip(0)
131
+ qsl[1:].copy_(reversed_lens.cumsum(0))
132
+ graph.replay()
133
+ expected = torch.repeat_interleave(
134
+ torch.arange(len(query_lens), device="cuda", dtype=torch.int32), reversed_lens
135
+ )
136
+ torch.testing.assert_close(output[:n], expected)
137
+ assert torch.count_nonzero(output[n:]) == 0
138
+
139
+
140
def test_indexer_shares_uncompressed_block_size_with_deepseek_v4_mla():
141
"""Packed MLA/indexer groups must retain 64 compressed rows per page."""
142
kernel_block_size = select_common_block_size(
vllm/v1/attention/backend.py
+13
-11
@@ -484,19 +484,21 @@ class CommonAttentionMetadata:
484
# Built from the device query_start_loc: adaptive verification decides the
485
# per-request draft split on device, so the CPU copy carries the right total
486
# but not the right per-request boundaries. Padding requests have a query
487
- # length of zero and drop out of the repeat.
487
+ # length of zero and are skipped by the device boundary search.
488
num_mapped_tokens = int(self.query_start_loc_cpu[-1])
489
- query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1]
490
- assert buffer.shape[0] >= max(num_mapped_tokens, num_tokens)
491
- token_to_req_indices = torch.repeat_interleave(
492
- torch.arange(query_lens.shape[0], dtype=torch.int32, device=buffer.device),
493
- query_lens,
494
- output_size=num_mapped_tokens,
489
+ from vllm.v1.attention.ops.metadata import _token_request_mapping_kernel
490
+
491
+ num_output_tokens = max(num_mapped_tokens, num_tokens)
492
+ assert buffer.shape[0] >= num_output_tokens
493
+ _token_request_mapping_kernel[((num_output_tokens + 255) // 256,)](
494
+ self.query_start_loc,
495
+ buffer,
496
+ self.query_start_loc.shape[0] - 1,
497
+ num_mapped_tokens,
498
+ num_output_tokens,
499
+ num_warps=4,
500
)
496
- buffer[:num_mapped_tokens].copy_(token_to_req_indices)
497
- if num_mapped_tokens < num_tokens:
498
- buffer[num_mapped_tokens:num_tokens].zero_()
499
- self._token_to_req_indices_cache = buffer[: max(num_mapped_tokens, num_tokens)]
501
+ self._token_to_req_indices_cache = buffer[:num_output_tokens]
502
return self._token_to_req_indices_cache[:num_tokens]
503
504
# TODO(lucas): remove once we have FULL-CG spec-decode support
vllm/v1/attention/backends/mla/indexer.py
+69
-52
@@ -960,28 +960,6 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
960
self.global_decode_seq_lens_buffer[actual_expanded:num_decode_tokens] = 0
961
return self.global_decode_seq_lens_buffer[:num_decode_tokens]
962
963
- def _build_varlen_decode_indices(
964
- self,
965
- decode_lens: torch.Tensor,
966
- decode_lens_cpu: torch.Tensor,
967
- num_decode_tokens: int,
968
- ) -> torch.Tensor:
969
- """Build request ids for flattened SM100 varlen rows."""
970
- indices = self.decode_indices_buffer[:num_decode_tokens]
971
- actual_expanded = int(decode_lens_cpu.sum().item())
972
- num_decodes = decode_lens.shape[0]
973
- indices[:actual_expanded] = torch.repeat_interleave(
974
- self.arange_buffer[:num_decodes],
975
- decode_lens,
976
- output_size=actual_expanded,
977
- )
978
- if actual_expanded < num_decode_tokens:
979
- pad = num_decode_tokens - actual_expanded
980
- indices[actual_expanded:num_decode_tokens] = (
981
- num_decodes + self.arange_buffer[:pad]
982
- )
983
- return indices
984
-
963
@staticmethod
964
def _split_indexer_prefill_chunks(
965
compressed_seq_lens_cpu: torch.Tensor,
@@ -1156,12 +1134,15 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1134
1135
decode_metadata = None
1136
if num_decodes > 0:
1159
- torch.diff(
1160
- common_attn_metadata.query_start_loc[: num_decodes + 1],
1161
- out=self.decode_lens_buffer[:num_decodes],
1162
- )
1137
+ if not self.supports_varlen:
1138
+ torch.diff(
1139
+ common_attn_metadata.query_start_loc[: num_decodes + 1],
1140
+ out=self.decode_lens_buffer[:num_decodes],
1141
+ )
1142
+ self.per_req_decode_lens_buffer[:num_decodes].copy_(
1143
+ self.decode_lens_buffer[:num_decodes]
1144
+ )
1145
decode_lens = self.decode_lens_buffer[:num_decodes]
1164
- self.per_req_decode_lens_buffer[:num_decodes].copy_(decode_lens)
1146
decode_lens_cpu = torch.diff(
1147
common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]
1148
)
@@ -1196,38 +1177,74 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1177
and step_next_n_ok
1178
)
1179
1199
- global_seq_lens_for_decode = self._prepare_global_decode_seq_lens(
1200
- global_seq_lens=global_seq_lens_for_decode,
1201
- decode_lens=decode_lens,
1202
- decode_lens_cpu=decode_lens_cpu,
1203
- query_start_loc=common_attn_metadata.query_start_loc[:num_decodes],
1204
- num_decode_tokens=num_decode_tokens,
1205
- use_native=use_native,
1206
- max_decode_len=max_decode_len,
1207
- )
1208
-
1209
- decode_indices = None
1210
- if self.supports_varlen:
1211
- decode_indices = self._build_varlen_decode_indices(
1212
- decode_lens=decode_lens,
1213
- decode_lens_cpu=decode_lens_cpu,
1214
- num_decode_tokens=num_decode_tokens,
1215
- )
1216
-
1217
- seq_lens, block_table, decode_lens, batch_size, requires_padding = (
1218
- self._prepare_decode_tensors(
1219
- seq_lens=seq_lens,
1220
- block_table=block_table,
1180
+ if not self.supports_varlen:
1181
+ global_seq_lens_for_decode = self._prepare_global_decode_seq_lens(
1182
+ global_seq_lens=global_seq_lens_for_decode,
1183
decode_lens=decode_lens,
1184
decode_lens_cpu=decode_lens_cpu,
1185
query_start_loc=common_attn_metadata.query_start_loc[:num_decodes],
1224
- num_decodes=num_decodes,
1186
num_decode_tokens=num_decode_tokens,
1187
use_native=use_native,
1227
- next_n=next_n,
1188
max_decode_len=max_decode_len,
1189
)
1230
- )
1190
+
1191
+ decode_indices = None
1192
+ if self.supports_varlen:
1193
+ from vllm.v1.attention.ops.metadata import (
1194
+ _indexer_decode_metadata_kernel,
1195
+ )
1196
+
1197
+ capacity = self.decode_seq_lens_buffer.numel()
1198
+ grid = max(
1199
+ num_decodes,
1200
+ num_decode_tokens + triton.cdiv(capacity - num_decode_tokens, 256),
1201
+ )
1202
+ _indexer_decode_metadata_kernel[(grid,)](
1203
+ query_start_loc,
1204
+ seq_lens,
1205
+ block_table,
1206
+ self.decode_seq_lens_buffer,
1207
+ self.expanded_block_table_buffer,
1208
+ self.decode_lens_buffer,
1209
+ self.decode_indices_buffer,
1210
+ self.per_req_decode_lens_buffer,
1211
+ num_decodes,
1212
+ num_decode_tokens,
1213
+ capacity,
1214
+ block_table.stride(0),
1215
+ self.expanded_block_table_buffer.stride(0),
1216
+ BLOCK_COLS=block_table.shape[1],
1217
+ num_warps=4,
1218
+ )
1219
+ seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
1220
+ block_table = self.expanded_block_table_buffer[:num_decode_tokens]
1221
+ decode_lens = self.decode_lens_buffer[:num_decode_tokens]
1222
+ decode_indices = self.decode_indices_buffer[:num_decode_tokens]
1223
+ requires_padding = False
1224
+ if global_seq_lens_for_decode is not None and max_decode_len > 1:
1225
+ self.global_decode_seq_lens_buffer[:num_decode_tokens].copy_(
1226
+ seq_lens
1227
+ )
1228
+ global_seq_lens_for_decode = self.global_decode_seq_lens_buffer[
1229
+ :num_decode_tokens
1230
+ ]
1231
+ else:
1232
+ seq_lens, block_table, decode_lens, batch_size, requires_padding = (
1233
+ self._prepare_decode_tensors(
1234
+ seq_lens=seq_lens,
1235
+ block_table=block_table,
1236
+ decode_lens=decode_lens,
1237
+ decode_lens_cpu=decode_lens_cpu,
1238
+ query_start_loc=common_attn_metadata.query_start_loc[
1239
+ :num_decodes
1240
+ ],
1241
+ num_decodes=num_decodes,
1242
+ num_decode_tokens=num_decode_tokens,
1243
+ use_native=use_native,
1244
+ next_n=next_n,
1245
+ max_decode_len=max_decode_len,
1246
+ )
1247
+ )
1248
1249
if self.compress_ratio > 1:
1250
kernel_block_size = self.kernel_block_size
vllm/v1/attention/ops/metadata.py
new
+87
@@ -0,0 +1,87 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Device-side request mapping and sparse indexer metadata."""
4
+
5
+from vllm.triton_utils import tl, triton
6
+
7
+
8
+@triton.jit
9
+def _token_request(query_start_loc, token, num_reqs):
10
+ lo = tl.full(token.shape, 0, tl.int32)
11
+ hi = tl.full(token.shape, num_reqs, tl.int32)
12
+ while tl.sum((lo < hi).to(tl.int32), 0) > 0:
13
+ mid = (lo + hi) // 2
14
+ end = tl.load(query_start_loc + mid + 1, mid < num_reqs, other=0x7FFFFFFF)
15
+ right = (lo < hi) & (end <= token)
16
+ lo = tl.where(right, mid + 1, lo)
17
+ hi = tl.where((lo < hi) & ~right, mid, hi)
18
+ return tl.minimum(lo, num_reqs - 1)
19
+
20
+
21
+@triton.jit(do_not_specialize=["num_reqs", "num_mapped", "num_tokens"])
22
+def _token_request_mapping_kernel(
23
+ query_start_loc, output, num_reqs, num_mapped, num_tokens
24
+):
25
+ token = tl.program_id(0) * 256 + tl.arange(0, 256)
26
+ req = _token_request(query_start_loc, token, num_reqs)
27
+ tl.store(output + token, tl.where(token < num_mapped, req, 0), token < num_tokens)
28
+
29
+
30
+@triton.jit(
31
+ do_not_specialize=[
32
+ "num_reqs",
33
+ "num_tokens",
34
+ "capacity",
35
+ "input_stride",
36
+ "output_stride",
37
+ ]
38
+)
39
+def _indexer_decode_metadata_kernel(
40
+ query_start_loc,
41
+ seq_lens,
42
+ block_table,
43
+ output_seq_lens,
44
+ output_block_table,
45
+ decode_lens,
46
+ indices,
47
+ per_req_lens,
48
+ num_reqs,
49
+ num_tokens,
50
+ capacity,
51
+ input_stride,
52
+ output_stride,
53
+ BLOCK_COLS: tl.constexpr,
54
+):
55
+ token = tl.program_id(0)
56
+ cols = tl.arange(0, triton.next_power_of_2(BLOCK_COLS))
57
+ actual_tokens = tl.load(query_start_loc + num_reqs)
58
+ if token < num_tokens:
59
+ req = _token_request(query_start_loc, tl.full((1,), token, tl.int32), num_reqs)
60
+ req = tl.sum(req, 0)
61
+ valid = token < actual_tokens
62
+ start = tl.load(query_start_loc + req)
63
+ end = tl.load(query_start_loc + req + 1)
64
+ seq = tl.load(seq_lens + req)
65
+ length = tl.where(valid, seq - (end - start) + token - start + 1, 0)
66
+ tl.store(output_seq_lens + token, length)
67
+ tl.store(decode_lens + token, 1)
68
+ tl.store(
69
+ indices + token, tl.where(valid, req, num_reqs + token - actual_tokens)
70
+ )
71
+ blocks = tl.load(
72
+ block_table + req.to(tl.int64) * input_stride + cols,
73
+ valid & (cols < BLOCK_COLS),
74
+ other=0,
75
+ )
76
+ tl.store(
77
+ output_block_table + token.to(tl.int64) * output_stride + cols,
78
+ blocks,
79
+ cols < BLOCK_COLS,
80
+ )
81
+ else:
82
+ tail = num_tokens + (token - num_tokens) * 256 + tl.arange(0, 256)
83
+ tl.store(output_seq_lens + tail, 0, tail < capacity)
84
+ if token < num_reqs:
85
+ start = tl.load(query_start_loc + token)
86
+ end = tl.load(query_start_loc + token + 1)
87
+ tl.store(per_req_lens + token, end - start)