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/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index 026dd1ff2b..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, @@ -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, 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): 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 "