[Perf][SM70][DeepSeek-V4] Keep the cuBLAS decode indexer above 131,072 tokens of context

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)

tests/v1/cudagraph/test_cudagraph_dispatch.py

@@ -15,6 +15,7 @@
     CompilationConfig,
     CompilationMode,
     CUDAGraphMode,
+    ObservabilityConfig,
     ParallelConfig,
     SchedulerConfig,
     VllmConfig,
@@ -26,6 +27,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
 
@@ -57,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
@@ -418,7 +422,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 +435,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 +491,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,

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):

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 "