From 7a8f92949fb11990da3c1a5fb484dfbde8a10c5b Mon Sep 17 00:00:00 2001 From: Peuqui <43776522+Peuqui@users.noreply.github.com> Date: Fri, 16 Oct 2026 13:33:55 +0200 Subject: [PATCH 1/3] [Perf][SM70][DeepSeek-V4] Score the cuBLAS decode indexer in key tiles The cuBLAS decode route of the SM70 indexer gathers the keys once and scores them into one [rows * heads, keys] FP32 tile, admitted only when both fit the indexer workspace budget (indexer_prefill_tile_mb, default 192 MiB). Beyond that the decode falls back to the paged Triton kernel, which at 140k tokens of context took 1.27 ms per call on a V100 and 6.9 ms on a Quadro RTX 8000, against about 0.11 ms for the GEMM of the 131,072-token bucket. For six verifier rows of 64 heads the default budget admits about 112,000 compressed keys; the 64 MiB a PP5 deployment uses to keep its V100 prefill stages clear of OOM admits 32,768 keys (131,072 tokens). The gathered keys stay whole and the scores are computed in key tiles that fit next to them. A key bound whose scores fit keeps its single tile, so the route and its result are unchanged wherever it was admitted before. Signed-off-by: Peuqui <43776522+Peuqui@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GJxVXkjSHhfq2D71DoYSK7 --- .../kernels/test_deepseek_v4_sm70_indexer.py | 100 ++++++++++++++++++ vllm/models/deepseek_v4/sm70/indexer.py | 80 +++++++++----- 2 files changed, 151 insertions(+), 29 deletions(-) diff --git a/tests/kernels/test_deepseek_v4_sm70_indexer.py b/tests/kernels/test_deepseek_v4_sm70_indexer.py index 860d25cddd..1a46cba71f 100644 --- a/tests/kernels/test_deepseek_v4_sm70_indexer.py +++ b/tests/kernels/test_deepseek_v4_sm70_indexer.py @@ -411,3 +411,103 @@ def test_sm70_indexer_decode_valid_logits_do_not_depend_on_workspace_width() -> assert narrow.shape == (num_queries, 16) assert wide.shape == (num_queries, 64) torch.testing.assert_close(narrow[:, :num_keys], wide[:, :num_keys], rtol=0, atol=0) + + +class _RecordingWorkspace: + """Allocates every requested buffer and records the requests.""" + + def __init__(self): + self.requests = [] + + def get_simultaneous(self, *specs): + self.requests.append(specs) + return tuple( + torch.empty(shape, dtype=dtype, device="cuda") for shape, dtype in specs + ) + + +@requires_sm70 +def test_decode_cublas_score_tiles_match_one_tile(monkeypatch): + """Above the budget for whole FP32 scores the cuBLAS decode route scores + key tiles. Each row keeps its top-k keys and its logits within FP32 + rounding (cuBLAS may pick another GEMM for the narrower tiles), also when + the tiled route is captured and replayed as a CUDA graph.""" + torch.manual_seed(20261016) + generator = torch.Generator().manual_seed(20261016) + num_heads = 64 + rows = 6 + graph_width = 6144 + live_seq_len = 5000 + blocks = (live_seq_len + _BLOCK_SIZE - 1) // _BLOCK_SIZE + + bits, _ = _random_fp8_keys(blocks * _BLOCK_SIZE, generator) + value_bits = bits.reshape(blocks, _BLOCK_SIZE, _HEAD_DIM) + scales = torch.rand((blocks, _BLOCK_SIZE), generator=generator).cuda() + 0.5 + cache = _build_paged_index_cache(value_bits, scales) + table = torch.full( + (1, graph_width // _BLOCK_SIZE), -1, dtype=torch.int32, device="cuda" + ) + table[0, :blocks] = torch.randperm(blocks, generator=generator).to( + device="cuda", dtype=torch.int32 + ) + # The flattened speculative layout DeepSeek-V4 decodes with. + block_table = table.repeat(rows, 1).contiguous() + seq_lens = torch.arange( + live_seq_len - rows + 1, live_seq_len + 1, dtype=torch.int32, device="cuda" + ).view(rows, 1) + q, weights = _make_queries(rows, num_heads) + + monkeypatch.setattr(sparse_policy(), "indexer_decode_cublas_enabled", True) + workspace = _RecordingWorkspace() + monkeypatch.setattr(sm70_indexer, "current_workspace_manager", lambda: workspace) + + def decode(): + return sm70_indexer_decode_logits( + q, + cache, + weights, + seq_lens, + block_table, + graph_width, + rows, + indexer_decode_cublas=True, + ) + + single = decode() + assert workspace.requests[-1][2][0] == (rows * num_heads, graph_width) + + # Gathered keys (1.5 MiB) plus a 1.5 MiB score tile: six tiles of 1024 keys. + monkeypatch.setattr(sparse_policy(), "indexer_prefill_tile_mb", 3) + tiled = decode() + assert workspace.requests[-1][2][0] == (rows * num_heads, 1024) + + capture_stream = torch.cuda.Stream() + capture_stream.wait_stream(current_stream()) + with torch.cuda.stream(capture_stream): + for _ in range(3): + decode() + capture_stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + replayed = decode() + graph.replay() + torch.accelerator.synchronize() + + for candidate in (tiled, replayed): + for row, row_len in enumerate(seq_lens.reshape(-1).tolist()): + torch.testing.assert_close( + candidate[row, :row_len], single[row, :row_len], rtol=1e-4, atol=1e-5 + ) + expected = torch.topk(single[row, :row_len], 512).indices.sort().values + actual = torch.topk(candidate[row, :row_len], 512).indices.sort().values + assert torch.equal(actual, expected) + + # Not even one 256-key tile next to the gathered keys: the paged kernel. + monkeypatch.setattr(sparse_policy(), "indexer_prefill_tile_mb", 1) + requests_before = len(workspace.requests) + paged = decode() + assert len(workspace.requests) == requests_before + for row, row_len in enumerate(seq_lens.reshape(-1).tolist()): + expected = torch.topk(single[row, :row_len], 512).indices.sort().values + actual = torch.topk(paged[row, :row_len], 512).indices.sort().values + assert torch.equal(actual, expected) diff --git a/vllm/models/deepseek_v4/sm70/indexer.py b/vllm/models/deepseek_v4/sm70/indexer.py index 899770affd..43813ede60 100644 --- a/vllm/models/deepseek_v4/sm70/indexer.py +++ b/vllm/models/deepseek_v4/sm70/indexer.py @@ -695,6 +695,20 @@ def sm70_indexer_prefill_logits( return torch.mm(weighted_q, k_fp16.t(), out_dtype=torch.float32) +def _decode_score_tile_keys(max_seq_len: int, rows: int, num_heads: int) -> int: + """Keys per FP32 score tile of the cuBLAS decode route, 0 if none fits. + + The gathered keys stay whole; the [rows * heads, keys] scores, by far the + largest buffer, are tiled so the gather plus one tile fit the indexer + workspace budget. A key bound whose scores fit keeps a single tile. + """ + budget = sparse_policy().value("indexer_prefill_tile_mb") * 2**20 + tile = (budget - max_seq_len * _INDEX_HEAD_DIM * 2) // (rows * num_heads * 4) + if tile >= max_seq_len: + return max_seq_len + return max(0, tile // _EPILOGUE_BLOCK_K * _EPILOGUE_BLOCK_K) + + def _decode_logits_cublas( q: torch.Tensor, cache: torch.Tensor, @@ -708,16 +722,18 @@ def _decode_logits_cublas( `block_table` holds `table_rows_per_request` identical rows per request (the flattened speculative decode) or one (the native layout); `q`, - `weights` and `flat_lens` hold the request's rows back to back. + `weights` and `flat_lens` hold the request's rows back to back. The FP32 + scores are computed in key tiles of `_decode_score_tile_keys`. """ total_rows, num_heads, head_dim = q.shape num_requests = block_table.shape[0] // table_rows_per_request rows = total_rows // num_requests + tile_keys = _decode_score_tile_keys(max_seq_len, rows, num_heads) workspace = current_workspace_manager() gathered_k, k_scales, scores = workspace.get_simultaneous( ((max_seq_len, head_dim), torch.float16), ((max_seq_len,), torch.float32), - ((rows * num_heads, max_seq_len), torch.float32), + ((rows * num_heads, tile_keys), torch.float32), ) out = torch.empty((total_rows, max_seq_len), dtype=torch.float32, device=q.device) for request in range(num_requests): @@ -769,29 +785,35 @@ def _score_request_cublas( ) q2 = q.reshape(total_rows * num_heads, head_dim) - torch.mm( - q2, - gathered_k.t(), - out=scores, - out_dtype=torch.float32, - ) - _relu_weight_headsum_kernel[ - (total_rows, triton.cdiv(max_seq_len, _EPILOGUE_BLOCK_K)) - ]( - scores, - weights, - k_scales, - out, - scores.stride(0) * num_heads, - scores.stride(0), - weights.stride(0), - out.stride(0), - max_seq_len, - num_heads=num_heads, - BLOCK_K=_EPILOGUE_BLOCK_K, - BLOCK_H=_EPILOGUE_BLOCK_H, - num_warps=4, - ) + tile_keys = scores.shape[1] + for start in range(0, max_seq_len, tile_keys): + width = min(tile_keys, max_seq_len - start) + # A contiguous [rows * heads, width] view keeps the rows of the last, + # narrower tile back to back, as the epilogue's strides expect. + tile_scores = scores.view(-1)[: q2.shape[0] * width].view(q2.shape[0], width) + torch.mm( + q2, + gathered_k[start : start + width].t(), + out=tile_scores, + out_dtype=torch.float32, + ) + _relu_weight_headsum_kernel[ + (total_rows, triton.cdiv(width, _EPILOGUE_BLOCK_K)) + ]( + tile_scores, + weights, + k_scales[start : start + width], + out[:, start : start + width], + tile_scores.stride(0) * num_heads, + tile_scores.stride(0), + weights.stride(0), + out.stride(0), + width, + num_heads=num_heads, + BLOCK_K=_EPILOGUE_BLOCK_K, + BLOCK_H=_EPILOGUE_BLOCK_H, + num_warps=4, + ) def _decode_cublas_blocker( @@ -863,10 +885,10 @@ def rows_per_request() -> int: and weights.dtype == torch.float32, ), ( - "shared gather and FP32 scores within the indexer workspace budget", - lambda: max_seq_len - * (_INDEX_HEAD_DIM * 2 + rows_per_request() * q.shape[1] * 4) - <= policy.value("indexer_prefill_tile_mb") * 2**20, + "the shared key gather and one FP32 score tile within the indexer " + "workspace budget", + lambda: _decode_score_tile_keys(max_seq_len, rows_per_request(), q.shape[1]) + > 0, ), ( f"a uint8 cache [blocks, block_size, {_INDEX_CACHE_BYTES}] with " From d2451124f8ccc2433babe711d1249197b95f3630 Mon Sep 17 00:00:00 2001 From: Peuqui <43776522+Peuqui@users.noreply.github.com> Date: Fri, 16 Oct 2026 13:33:55 +0200 Subject: [PATCH 2/3] [Perf][SM70][DeepSeek-V4] End the compressed-index decode buckets at the model length The automatic buckets stopped at 64 times the short bound (131,072 tokens for DeepSeek-V4). Longer contexts replayed the full-length graph, where the indexer stays on its paged kernel by design. A last bucket at the maximum model length keeps those contexts on the cuBLAS decode route. Signed-off-by: Peuqui <43776522+Peuqui@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GJxVXkjSHhfq2D71DoYSK7 --- tests/v1/cudagraph/test_cudagraph_dispatch.py | 32 +++++++++++++++++-- vllm/model_executor/models/graph_contract.py | 6 +++- 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index 026dd1ff2b..d8b2b7119f 100644 --- a/tests/v1/cudagraph/test_cudagraph_dispatch.py +++ b/tests/v1/cudagraph/test_cudagraph_dispatch.py @@ -26,6 +26,7 @@ BatchDescriptor, set_forward_context, ) +from vllm.model_executor.models.graph_contract import GraphModelContract from vllm.platforms import current_platform from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher @@ -418,7 +419,7 @@ def test_dsv4_decode_context_bucket_is_derived_on_sm70(self, monkeypatch): uniform_decode_query_len=1, ) - assert dispatcher.sm70_dsv4_decode_context_buckets == (2048,) + assert dispatcher.sm70_dsv4_decode_context_buckets == (2048, 4096) bounded = BatchDescriptor( num_tokens=1, num_reqs=1, @@ -431,6 +432,33 @@ def test_dsv4_decode_context_bucket_is_derived_on_sm70(self, monkeypatch): disabled = CudagraphDispatcher(config) assert not disabled.sm70_dsv4_decode_context_buckets + @pytest.mark.parametrize( + ("max_model_len", "expected"), + [ + (4096, (2048, 4096)), + (131072, (2048, 4096, 16384, 65536, 131072)), + (307200, (2048, 4096, 16384, 65536, 131072, 307200)), + ], + ) + def test_dsv4_decode_context_buckets_end_at_the_model_length( + self, max_model_len, expected + ): + # Without a bucket at the maximum length, contexts above the largest + # multiple replay the full-length graph, where the indexer leaves its + # cuBLAS decode route. + contract = GraphModelContract( + speculative=True, + max_model_len=max_model_len, + num_heads=64, + num_kv_heads=1, + head_dim=512, + index_topk=512, + compress_ratios=(4, 128), + cache_dtype="fp8", + attention_backend=None, + ) + assert contract.compressed_context_buckets() == expected + @pytest.mark.parametrize("speculative_tokens", [3, 5, 7]) def test_dsv4_context_buckets_are_derived_for_mtp_on_sm70( self, monkeypatch, speculative_tokens @@ -460,7 +488,7 @@ def test_dsv4_context_buckets_are_derived_for_mtp_on_sm70( dispatcher = CudagraphDispatcher(config) # Speculative graphs bucket like single-row decode: the sparse indexer # needs a bounded key length to take its cuBLAS route under full graphs. - assert dispatcher.sm70_dsv4_decode_context_buckets == (2048,) + assert dispatcher.sm70_dsv4_decode_context_buckets == (2048, 4096) dispatcher.initialize_cudagraph_keys( cudagraph_mode=comp_config.cudagraph_mode, uniform_decode_query_len=query_rows, diff --git a/vllm/model_executor/models/graph_contract.py b/vllm/model_executor/models/graph_contract.py index cc918aebd2..8c45779104 100644 --- a/vllm/model_executor/models/graph_contract.py +++ b/vllm/model_executor/models/graph_contract.py @@ -33,9 +33,13 @@ def compressed_context_buckets(self): short = self.index_topk * min(self.compress_ratios) if short <= 0 or self.max_model_len <= short: return () - return tuple( + buckets = tuple( short * m for m in (1, 2, 8, 32, 64) if short * m < self.max_model_len ) + # Above the largest bucket a decode graph would be bound by the model's + # maximum length, which keeps the indexer off its cuBLAS decode route; + # a last bucket at the maximum length keeps it on. + return (*buckets, self.max_model_len) def model_graph_contract(cfg): From f755ff05abce205bea70c657f4340689ed251eab Mon Sep 17 00:00:00 2001 From: Peuqui <43776522+Peuqui@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:16:39 +0200 Subject: [PATCH 3/3] [Test] Give the CUDA-graph wrapper tests a real observability config The dispatcher tests build VllmConfig as a MagicMock and set the configs the code under test reads. CUDAGraphWrapper now also builds a decode event tracer from observability_config.runtime_trace; on the mock that policy was a MagicMock, so tracing ran and compared the elapsed time against a MagicMock threshold. Three tests failed with "'<' not supported between instances of 'float' and 'MagicMock'" (TestCUDAGraphWrapper:: test_capture_and_replay, test_capture_replay_bypass_logic, test_nested_wrappers). The mock now carries a default ObservabilityConfig, as other config tests do. Signed-off-by: Peuqui <43776522+Peuqui@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GJxVXkjSHhfq2D71DoYSK7 --- tests/v1/cudagraph/test_cudagraph_dispatch.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index d8b2b7119f..46058f72ff 100644 --- a/tests/v1/cudagraph/test_cudagraph_dispatch.py +++ b/tests/v1/cudagraph/test_cudagraph_dispatch.py @@ -15,6 +15,7 @@ CompilationConfig, CompilationMode, CUDAGraphMode, + ObservabilityConfig, ParallelConfig, SchedulerConfig, VllmConfig, @@ -58,6 +59,8 @@ def _create_vllm_config( max_num_seqs=max_num_seqs, ) mock_config.parallel_config = ParallelConfig() + # CUDAGraphWrapper builds a decode event tracer from the runtime trace policy. + mock_config.observability_config = ObservabilityConfig() mock_config.speculative_config = None # No speculative decoding if not lora_config: mock_config.lora_config = None