[Refactor][MRV2] Unify uniform decode token count helper (#51917)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Lucas Wilkinson committed
Aug 12, 2026 at 01:58 UTC
02ac17851dfb603c1ba788545200ae3e43a4f61b
3 files changed
+12
-24
tests/v1/spec_decode/test_dynamic_sd_cug.py
+8
-5
@@ -124,10 +124,11 @@ def test_dynamic_sd_full_cudagraph_covers_all_uniform_decode_shapes(monkeypatch)
124
# Uniform decode means every request contributes the same number of
125
# tokens, so the total token count is exactly num_reqs * query_len.
126
num_tokens = num_reqs * max_query_len
127
- uniform_tok_count = gpu_cudagraph_utils.get_uniform_token_count(
127
+ uniform_tok_count = get_uniform_decode_token_count(
128
num_reqs,
129
num_tokens,
130
max_query_len,
131
+ has_prefill=False,
132
)
133
134
# The scheduler should mark every one of these shapes as a uniform
@@ -233,8 +234,8 @@ def test_prompt_chunks_shaped_like_spec_decode_miss_the_full_graph(monkeypatch):
234
235
# The batch shape is indistinguishable from a full batch of spec decodes.
236
assert (
236
- gpu_cudagraph_utils.get_uniform_token_count(
237
- num_reqs, num_tokens, decode_query_len
237
+ get_uniform_decode_token_count(
238
+ num_reqs, num_tokens, decode_query_len, has_prefill=False
239
)
240
== decode_query_len
241
)
@@ -305,10 +306,11 @@ def test_basic_sd_does_not_capture_shorter_full_decode_shapes(monkeypatch):
306
# These are still uniform decode batches, but basic SD should only
307
# have FULL graphs for query_len == max_decode_query_len.
308
num_tokens = num_reqs * max_query_len
308
- uniform_tok_count = gpu_cudagraph_utils.get_uniform_token_count(
309
+ uniform_tok_count = get_uniform_decode_token_count(
310
num_reqs,
311
num_tokens,
312
max_query_len,
313
+ has_prefill=False,
314
)
315
assert uniform_tok_count == max_query_len
316
@@ -370,10 +372,11 @@ def test_dynamic_sd_only_captures_scheduled_query_lengths(monkeypatch):
372
for num_reqs in range(1, max_num_seqs + 1):
373
for max_query_len in range(1, max_decode_query_len + 1):
374
num_tokens = num_reqs * max_query_len
373
- uniform_tok_count = gpu_cudagraph_utils.get_uniform_token_count(
375
+ uniform_tok_count = get_uniform_decode_token_count(
376
num_reqs,
377
num_tokens,
378
max_query_len,
379
+ has_prefill=False,
380
)
381
assert uniform_tok_count == max_query_len
382
vllm/v1/worker/gpu/cudagraph_utils.py
+1
-17
@@ -35,7 +35,7 @@ from vllm.v1.worker.gpu.block_table import BlockTables
35
from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
36
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
37
from vllm.v1.worker.gpu.model_states.interface import ModelState
38
-from vllm.v1.worker.utils import AttentionGroup, is_uniform_query_len
38
+from vllm.v1.worker.utils import AttentionGroup
39
40
logger = init_logger(__name__)
41
@@ -89,22 +89,6 @@ def _is_compatible(
89
)
90
91
92
-def get_uniform_token_count(
93
- num_reqs: int, num_tokens: int, max_query_len: int
94
-) -> int | None:
95
- """
96
- Return the uniform token count if batch is uniform, else None.
97
- A batch is uniform if all requests have the same number of tokens.
98
-
99
- Shape test only, valid for batches known to be decoding by construction
100
- (e.g. dummy runs). Scheduled batches must use
101
- `get_uniform_decode_token_count`.
102
- """
103
- if is_uniform_query_len(num_reqs, num_tokens, max_query_len):
104
- return max_query_len
105
- return None
106
-
107
-
92
class CudaGraphManager:
93
def __init__(
94
self,
vllm/v1/worker/gpu/model_runner.py
+3
-2
@@ -88,7 +88,6 @@ from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
88
from vllm.v1.worker.gpu.cudagraph_utils import (
89
BatchExecutionDescriptor,
90
ModelCudaGraphManager,
91
- get_uniform_token_count,
91
)
92
from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
93
from vllm.v1.worker.gpu.ec_connector import get_ec_connector
@@ -991,7 +990,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
990
991
if dummy_run:
992
# Dummy batches are uniform by construction.
994
- return None, get_uniform_token_count(num_reqs, num_toks, max_query_len)
993
+ return None, get_uniform_decode_token_count(
994
+ num_reqs, num_toks, max_query_len, has_prefill=False
995
+ )
996
997
# batch_idx -> req_id
998
req_ids = sort_batch_req_ids(num_tokens_per_req, self.decode_query_len)