From 13e221f8308416564a446909c2fd4805c78fbb8a Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sat, 12 Sep 2026 07:20:04 -0700 Subject: [PATCH] [Perf] Fuse DSV4.1 input metadata preparation with Triton (#56562) --- .../test_indexer_deepseek_v4_slot_mapping.py | 105 +++++++++++++++ vllm/v1/attention/backend.py | 24 ++-- vllm/v1/attention/backends/mla/indexer.py | 121 ++++++++++-------- vllm/v1/attention/ops/metadata.py | 87 +++++++++++++ 4 files changed, 274 insertions(+), 63 deletions(-) create mode 100644 vllm/v1/attention/ops/metadata.py diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index 800b0557de..aabe88eeee 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -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( diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index f4031a7ee0..addd68ec6a 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -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 diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 59fb109dac..7bc0b442c3 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -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 diff --git a/vllm/v1/attention/ops/metadata.py b/vllm/v1/attention/ops/metadata.py new file mode 100644 index 0000000000..f5baa51a92 --- /dev/null +++ b/vllm/v1/attention/ops/metadata.py @@ -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)