@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d154d90d6c

[Model Runner V2] batch-sharded sample (#50465)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>

Giancarlo Delfin committed Aug 24, 2026 at 15:49 UTC d154d90d6c4bcf26a0c78ac4f3e43621c14333ba
19 files changed +1781 -22
.buildkite/test_areas/model_runner_v2.yaml
+6 -1
@@ -63,7 +63,7 @@ steps:
63
64 - label: ":nvidia: (L4) Model Runner V2 Distributed"
65 key: model-runner-v2-distributed-2-gpus
66 - timeout_in_minutes: 30
66 + timeout_in_minutes: 40
67 working_dir: "/vllm-workspace/tests"
68 num_devices: 2
69 source_file_dependencies:
@@ -72,6 +72,8 @@ steps:
72 - tests/basic_correctness/test_basic_correctness.py
73 - tests/v1/distributed/test_async_llm_dp.py
74 - tests/v1/distributed/test_eagle_dp.py
75 + - tests/v1/e2e/general/test_sharded_sampling.py
76 + - tests/v1/e2e/spec_decode/test_sharded_sampling.py
77 commands:
78 - set -x
79 - export VLLM_USE_V2_MODEL_RUNNER=1
@@ -81,6 +83,9 @@ steps:
83 - export NCCL_CUMEM_HOST_ENABLE=0
84 - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
85 - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
86 + - pytest -v -s v1/e2e/general/test_sharded_sampling.py
87 + - pytest -v -s v1/e2e/spec_decode/test_sharded_sampling.py
88 +
89 mirror:
90 amd:
91 label: ":amd: (MI300) Model Runner V2 Distributed"
tests/v1/e2e/general/test_sharded_sampling.py new
+162
@@ -0,0 +1,162 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Batch-sharded sampling must be a drop-in for replicated sampling.
4 +
5 +The comparison is tolerance-based, not bitwise: engine boots land in
6 +slightly different kernel/collective states (measured: same-state boots are
7 +bitwise identical across the two modes; cross-state boots differ by up to
8 +~0.25 logprob and can flip a near-tie token). A real sharding bug — rows
9 +routed to the wrong request, misaligned gathers — produces many-nats
10 +divergence and scrambled top-k sets, far outside these bounds."""
11 +
12 +import pytest
13 +
14 +from tests.utils import multi_gpu_test
15 +from vllm import LLM, SamplingParams
16 +from vllm.distributed import cleanup_dist_env_and_memory
17 +
18 +MODEL = "meta-llama/Llama-3.2-1B-Instruct"
19 +
20 +# Boot-state numeric noise bound (measured <= 0.25 on identical contexts).
21 +LOGPROB_TOL = 0.5
22 +# Near-tie greedy/seeded flips from boot-state noise; measured <= 1 per
23 +# engine boot across ~16 completions, plus one of headroom.
24 +MAX_DIVERGENT_PROMPTS = 3
25 +
26 +PROMPTS = [
27 + "The capital of France is",
28 + "The three primary colors are",
29 + "Write a haiku about the ocean.",
30 + "1, 1, 2, 3, 5, 8,",
31 + "The chemical symbol for gold is",
32 + "Once upon a time, in a distant kingdom,",
33 + "Python's GIL stands for",
34 + "The fastest land animal is",
35 +]
36 +
37 +# Exercise greedy, seeded random, top-k/top-p, penalties, and logprobs.
38 +SAMPLING_PARAMS = [
39 + SamplingParams(temperature=0.0, max_tokens=32, logprobs=5, prompt_logprobs=1),
40 + SamplingParams(temperature=0.8, seed=42, top_p=0.9, max_tokens=32, logprobs=3),
41 + SamplingParams(temperature=0.7, seed=7, top_k=50, min_p=0.05, max_tokens=32),
42 + SamplingParams(
43 + temperature=0.9,
44 + seed=123,
45 + frequency_penalty=0.5,
46 + repetition_penalty=1.2,
47 + max_tokens=32,
48 + logprobs=2,
49 + ),
50 +]
51 +
52 +
53 +def _is_sharded_sampling_active(worker) -> bool:
54 + return worker.model_runner.batch_sharder is not None
55 +
56 +
57 +def _generate(monkeypatch: pytest.MonkeyPatch, disable_sharding: bool):
58 + llm = LLM(
59 + model=MODEL,
60 + tensor_parallel_size=2,
61 + enable_batch_sharded_sampling=not disable_sharding,
62 + max_model_len=1024,
63 + # More requests than ranks, so round-robin slot ownership splits the
64 + # batch across both ranks.
65 + max_num_seqs=len(PROMPTS),
66 + enforce_eager=True,
67 + enable_prefix_caching=False,
68 + gpu_memory_utilization=0.7,
69 + # Timing-based kernel selection is the dominant source of cross-boot
70 + # numeric noise; disable it so the comparison bounds stay tight.
71 + kernel_config={"enable_flashinfer_autotune": False},
72 + )
73 + try:
74 + # Guard against a vacuous comparison: verify every worker is in the
75 + # intended sampling mode.
76 + modes = llm.llm_engine.collective_rpc(_is_sharded_sampling_active)
77 + assert all(mode == (not disable_sharding) for mode in modes), modes
78 + params = [
79 + SAMPLING_PARAMS[i % len(SAMPLING_PARAMS)] for i in range(len(PROMPTS))
80 + ]
81 + # Two waves: the second reuses the request slots freed by the first
82 + # (in whatever order requests finished), covering the cross-rank
83 + # slot-recycling path that request ownership derives from.
84 + wave1 = llm.generate(PROMPTS, params)
85 + wave2 = llm.generate(PROMPTS[::-1], params[::-1])
86 + return wave1 + wave2
87 + finally:
88 + del llm
89 + cleanup_dist_env_and_memory()
90 +
91 +
92 +def _assert_logprob_dicts_close(ref_lps, out_lps, what: str) -> None:
93 + # Allow one boundary entry of the top-k set to swap at a near-tie.
94 + common = set(ref_lps) & set(out_lps)
95 + assert len(common) >= max(len(ref_lps), len(out_lps)) - 1, (
96 + f"{what}: top-k token sets diverge: {sorted(ref_lps)} vs {sorted(out_lps)}"
97 + )
98 + for token_id in common:
99 + diff = abs(ref_lps[token_id].logprob - out_lps[token_id].logprob)
100 + assert diff <= LOGPROB_TOL, (
101 + f"{what}[{token_id}]: logprob diff {diff} "
102 + f"({ref_lps[token_id].logprob} vs {out_lps[token_id].logprob})"
103 + )
104 +
105 +
106 +@multi_gpu_test(num_gpus=2)
107 +def test_sharded_sampling_outputs_match(monkeypatch: pytest.MonkeyPatch):
108 + """Generation, logprobs, and prompt logprobs match between batch-sharded
109 + sampling and the replicated fallback, up to measured boot-state noise."""
110 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
111 + # Required for the collective_rpc mode check below.
112 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
113 +
114 + ref_outputs = _generate(monkeypatch, disable_sharding=True)
115 + shard_outputs = _generate(monkeypatch, disable_sharding=False)
116 +
117 + assert len(ref_outputs) == len(shard_outputs) == 2 * len(PROMPTS)
118 + num_divergent = 0
119 + for i, (ref, out) in enumerate(zip(ref_outputs, shard_outputs)):
120 + ref_completion = ref.outputs[0]
121 + out_completion = out.outputs[0]
122 + ref_ids = list(ref_completion.token_ids)
123 + out_ids = list(out_completion.token_ids)
124 +
125 + # Logits are only comparable at positions whose context (prompt +
126 + # previously sampled tokens) is identical: every position up to and
127 + # including the first divergence.
128 + prefix = 0
129 + while prefix < min(len(ref_ids), len(out_ids)):
130 + if ref_ids[prefix] != out_ids[prefix]:
131 + break
132 + prefix += 1
133 + if ref_ids != out_ids:
134 + num_divergent += 1
135 +
136 + if ref_completion.logprobs is not None:
137 + assert out_completion.logprobs is not None
138 + comparable = min(prefix + 1, len(ref_ids), len(out_ids))
139 + for pos in range(comparable):
140 + _assert_logprob_dicts_close(
141 + ref_completion.logprobs[pos],
142 + out_completion.logprobs[pos],
143 + f"prompt {i} logprobs[{pos}]",
144 + )
145 +
146 + # The prompt is fixed, so every prompt-logprob position is comparable.
147 + if ref.prompt_logprobs is not None:
148 + assert out.prompt_logprobs is not None
149 + for pos, (ref_lps, out_lps) in enumerate(
150 + zip(ref.prompt_logprobs, out.prompt_logprobs)
151 + ):
152 + if ref_lps is None or out_lps is None:
153 + assert ref_lps is None and out_lps is None
154 + continue
155 + _assert_logprob_dicts_close(
156 + ref_lps, out_lps, f"prompt {i} prompt_logprobs[{pos}]"
157 + )
158 +
159 + assert num_divergent <= MAX_DIVERGENT_PROMPTS, (
160 + f"{num_divergent}/{2 * len(PROMPTS)} prompts diverged: beyond near-tie "
161 + "boot-state noise; sharded sampling is likely misrouting requests"
162 + )
tests/v1/e2e/spec_decode/test_sharded_sampling.py new
+102
@@ -0,0 +1,102 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Batch-sharded sampling must be a drop-in for replicated sampling under
4 +speculative decoding.
5 +
6 +The non-spec arm of this claim lives in tests/v1/e2e/general/
7 +test_sharded_sampling.py. Spec decode is the harder case: draft tokens give
8 +each request a different number of logits rows, so the shard plan's
9 +per-request arithmetic and the rejection sampler are both exercised.
10 +"""
11 +
12 +import pytest
13 +import torch
14 +
15 +from tests.utils import wait_for_memory_to_settle
16 +from vllm import LLM, SamplingParams
17 +from vllm.distributed import cleanup_dist_env_and_memory
18 +
19 +from .utils import _skip_if_insufficient_gpus_for_tp, get_test_prompts, greedy_sampling
20 +
21 +
22 +def _is_sharded_sampling_active(worker) -> bool:
23 + return worker.model_runner.batch_sharder is not None
24 +
25 +
26 +def test_mtp_sharded_sampling_equivalence(monkeypatch: pytest.MonkeyPatch):
27 + """Batch-sharded sampling must be a bit-exact drop-in for replicated
28 + sampling under MTP spec decoding: the collectives move the same logits
29 + bytes, Gumbel keys derive from (request slot, position, seed), and slot
30 + assignment is rank-deterministic. Both runs here are spec decode with
31 + identical math, so outputs must match exactly."""
32 + tp_size = 2
33 + _skip_if_insufficient_gpus_for_tp(tp_size)
34 + model_name = "Qwen/Qwen3.5-0.8B-Base"
35 + test_prompts = get_test_prompts(mm_enabled=False, num_prompts=20)
36 +
37 + with monkeypatch.context() as m:
38 + # Batch-sharded sampling only exists in the V2 model runner.
39 + m.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
40 + # Required for the collective_rpc mode check below.
41 + m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
42 +
43 + def run(disable_sharding: bool):
44 + llm = LLM(
45 + model=model_name,
46 + tensor_parallel_size=tp_size,
47 + enable_batch_sharded_sampling=not disable_sharding,
48 + speculative_config={
49 + "method": "mtp",
50 + "num_speculative_tokens": 1,
51 + "max_model_len": 2048,
52 + },
53 + max_model_len=2048,
54 + limit_mm_per_prompt={"image": 0, "video": 0},
55 + # More requests than ranks, so round-robin slot ownership
56 + # splits the batch across both ranks.
57 + max_num_seqs=len(test_prompts),
58 + # Timing-based kernel selection is the dominant source of
59 + # cross-boot numeric noise; disable for a tight comparison.
60 + kernel_config={"enable_flashinfer_autotune": False},
61 + )
62 + try:
63 + # Guard against a vacuous comparison: verify every worker is
64 + # in the intended sampling mode.
65 + modes = llm.llm_engine.collective_rpc(_is_sharded_sampling_active)
66 + assert all(mode == (not disable_sharding) for mode in modes)
67 + greedy = llm.chat(test_prompts, greedy_sampling())
68 + seeded = llm.chat(
69 + test_prompts,
70 + SamplingParams(temperature=1.0, seed=33, max_tokens=10),
71 + )
72 + return greedy, seeded
73 + finally:
74 + del llm
75 + torch.accelerator.empty_cache()
76 + cleanup_dist_env_and_memory()
77 + wait_for_memory_to_settle()
78 +
79 + ref_greedy, ref_seeded = run(disable_sharding=True)
80 + shard_greedy, shard_seeded = run(disable_sharding=False)
81 +
82 + # Engine boots are not bitwise deterministic (kernel/collective
83 + # selection shifts logits by ~1 ulp, flipping near-tie tokens), so
84 + # allow a small number of divergent prompts. A sharding bug produces
85 + # wholesale divergence, not isolated near-tie flips.
86 + for name, ref_outputs, shard_outputs in (
87 + ("greedy", ref_greedy, shard_greedy),
88 + ("seeded", ref_seeded, shard_seeded),
89 + ):
90 + num_divergent = 0
91 + for i, (ref, out) in enumerate(zip(ref_outputs, shard_outputs)):
92 + if list(ref.outputs[0].token_ids) != list(out.outputs[0].token_ids):
93 + num_divergent += 1
94 + print(
95 + f"{name} prompt {i} diverged:\n"
96 + f" replicated: {ref.outputs[0].text!r}\n"
97 + f" sharded: {out.outputs[0].text!r}"
98 + )
99 + assert num_divergent <= 2, (
100 + f"{name}: {num_divergent}/{len(ref_outputs)} prompts diverged, "
101 + "beyond near-tie boot noise"
102 + )
tests/v1/worker/test_gpu_batch_shard.py new
+613
@@ -0,0 +1,613 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Routing-math tests for batch-sharded sampling.
4 +
5 +These simulate every TP rank in-process (no torch.distributed): the local
6 +sub-batch, the all-to-all send/recv layout, and the padded result gather must
7 +compose to a bit-exact round trip of the replicated computation.
8 +"""
9 +
10 +from types import SimpleNamespace
11 +from unittest import mock
12 +
13 +import numpy as np
14 +import pytest
15 +import torch
16 +
17 +from vllm.v1.core.sched.output import GrammarOutput
18 +from vllm.v1.worker.gpu.input_batch import InputBatch
19 +from vllm.v1.worker.gpu.sample import batch_shard
20 +from vllm.v1.worker.gpu.sample.batch_shard import (
21 + BatchSharder,
22 + _shard_grammar_output,
23 +)
24 +
25 +DEVICE = "cuda"
26 +
27 +# The owner-sorted layout is expanded by a Triton kernel, so the routing math
28 +# can only be exercised on device.
29 +requires_cuda = pytest.mark.skipif(
30 + not torch.cuda.is_available(), reason="requires CUDA"
31 +)
32 +
33 +
34 +def _np(t: torch.Tensor) -> np.ndarray:
35 + return t.cpu().numpy()
36 +
37 +
38 +def _make_batch(
39 + rng: np.random.Generator,
40 + num_reqs: int,
41 + max_num_reqs: int,
42 + max_spec: int,
43 + slot_pool: np.ndarray | None = None,
44 +) -> InputBatch:
45 + """Consistent batch with random unique request slots, 1 + k_i logits
46 + rows per request at the tail of each query segment."""
47 + if slot_pool is None:
48 + slot_pool = np.arange(max_num_reqs)
49 + slots = rng.choice(slot_pool, size=num_reqs, replace=False)
50 + slots = slots.astype(np.int32)
51 + k = rng.integers(0, max_spec + 1, size=num_reqs)
52 + num_logits_per_req = (1 + k).astype(np.int32)
53 + cu = np.zeros(num_reqs + 1, dtype=np.int32)
54 + np.cumsum(num_logits_per_req, out=cu[1:])
55 +
56 + query_lens = num_logits_per_req + rng.integers(0, 4, size=num_reqs).astype(np.int32)
57 + qsl = np.zeros(num_reqs + 1, dtype=np.int32)
58 + np.cumsum(query_lens, out=qsl[1:])
59 + num_tokens = int(qsl[-1])
60 +
61 + logits_indices_np = np.concatenate(
62 + [
63 + np.arange(qsl[i + 1] - num_logits_per_req[i], qsl[i + 1])
64 + for i in range(num_reqs)
65 + ]
66 + ).astype(np.int64)
67 + expanded_idx_np = np.repeat(slots, num_logits_per_req)
68 + expanded_pos_np = np.concatenate(
69 + [np.arange(n, dtype=np.int32) for n in num_logits_per_req]
70 + )
71 + seq_lens_np = (query_lens + rng.integers(0, 100, size=num_reqs)).astype(np.int32)
72 +
73 + def dev(a: np.ndarray) -> torch.Tensor:
74 + return torch.from_numpy(a).to(DEVICE)
75 +
76 + return InputBatch(
77 + req_ids=[f"req-{i}" for i in range(num_reqs)],
78 + num_reqs=num_reqs,
79 + num_reqs_after_padding=num_reqs,
80 + idx_mapping=dev(slots),
81 + idx_mapping_np=slots,
82 + expanded_idx_mapping=dev(expanded_idx_np),
83 + expanded_local_pos=dev(expanded_pos_np),
84 + num_scheduled_tokens=query_lens,
85 + num_tokens=num_tokens,
86 + num_tokens_after_padding=num_tokens,
87 + num_draft_tokens=int(k.sum()),
88 + num_draft_tokens_per_req=k.astype(np.int32),
89 + query_start_loc=dev(qsl),
90 + query_start_loc_np=qsl,
91 + seq_lens=dev(seq_lens_np),
92 + seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_np.copy()),
93 + dcp_local_seq_lens=None,
94 + num_computed_tokens_np=np.zeros(num_reqs, dtype=np.int32),
95 + prefill_len_np=np.zeros(num_reqs, dtype=np.int32),
96 + num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32),
97 + is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_),
98 + has_prefill=False,
99 + max_seq_len_np=None,
100 + input_ids=torch.zeros(num_tokens, dtype=torch.int32, device=DEVICE),
101 + positions=torch.arange(num_tokens, dtype=torch.int64, device=DEVICE),
102 + is_padding=torch.zeros(num_tokens, dtype=torch.bool, device=DEVICE),
103 + logits_indices=dev(logits_indices_np),
104 + cu_num_logits=dev(cu),
105 + cu_num_logits_np=cu,
106 + has_structured_output_reqs=False,
107 + prompt_lens=None,
108 + )
109 +
110 +
111 +def _shard_all_ranks(batch: InputBatch, max_num_reqs: int, tp_size: int):
112 + """Simulate every TP rank by faking get_tp_group() per BatchSharder."""
113 + results = []
114 + for rank in range(tp_size):
115 + fake_group = SimpleNamespace(rank_in_group=rank, world_size=tp_size)
116 + with mock.patch.object(batch_shard, "get_tp_group", return_value=fake_group):
117 + sharder = BatchSharder(
118 + max_num_reqs=max_num_reqs,
119 + max_num_logits_per_req=8,
120 + device=torch.device(DEVICE),
121 + )
122 + local_batch, sorted_logits_indices, _, metadata = sharder.shard_sampler_inputs(
123 + batch, None
124 + )
125 + results.append((local_batch, sorted_logits_indices, metadata))
126 + return results
127 +
128 +
129 +def _owned_batch_indices(local_batch: InputBatch) -> np.ndarray:
130 + return np.array([int(r.split("-")[1]) for r in local_batch.req_ids], dtype=np.int64)
131 +
132 +
133 +def _owned_rows(cu: np.ndarray, owned: np.ndarray) -> np.ndarray:
134 + if len(owned) == 0:
135 + return np.array([], dtype=np.int64)
136 + return np.concatenate([np.arange(cu[r], cu[r + 1]) for r in owned])
137 +
138 +
139 +@requires_cuda
140 +@pytest.mark.parametrize("tp_size", [2, 4, 8])
141 +@pytest.mark.parametrize("seed", [0, 1])
142 +def test_local_batch_partition(tp_size: int, seed: int):
143 + """Local sub-batches partition the batch by owner, preserve batch order
144 + within an owner, and carry exactly the owned requests' entries of every
145 + per-request and per-row array."""
146 + rng = np.random.default_rng(seed)
147 + max_num_reqs = 64
148 + batch = _make_batch(rng, num_reqs=23, max_num_reqs=max_num_reqs, max_spec=3)
149 + num_logits = int(batch.cu_num_logits_np[-1])
150 + results = _shard_all_ranks(batch, max_num_reqs, tp_size)
151 + cu = batch.cu_num_logits_np
152 +
153 + metadata0 = results[0][2]
154 + assert sum(metadata0.num_logits_per_rank) == num_logits
155 + # Gather width = max logits rows of any request (bounds num_sampled).
156 + assert metadata0.max_num_logits_per_req == int(np.diff(cu).max())
157 +
158 + all_owned: list[int] = []
159 + for rank, (local, _, metadata) in enumerate(results):
160 + owned = _owned_batch_indices(local)
161 + assert local.num_reqs == len(owned) == metadata.num_local_reqs
162 + assert (batch.idx_mapping_np[owned] % tp_size == rank).all()
163 + # Stable sort: batch order is preserved within an owner.
164 + assert (np.diff(owned) > 0).all()
165 + all_owned.extend(owned.tolist())
166 +
167 + rows = _owned_rows(cu, owned)
168 + assert (
169 + metadata.num_local_logits
170 + == len(rows)
171 + == metadata0.num_logits_per_rank[rank]
172 + )
173 +
174 + expected_idx = batch.idx_mapping_np[owned]
175 + assert (local.idx_mapping_np == expected_idx).all()
176 + assert _np(local.idx_mapping).tolist() == expected_idx.tolist()
177 + assert (np.diff(local.cu_num_logits_np) == np.diff(cu)[owned]).all()
178 + assert _np(local.cu_num_logits).tolist() == local.cu_num_logits_np.tolist()
179 + assert (
180 + _np(local.expanded_idx_mapping) == _np(batch.expanded_idx_mapping)[rows]
181 + ).all()
182 + assert (
183 + _np(local.expanded_local_pos) == _np(batch.expanded_local_pos)[rows]
184 + ).all()
185 + assert (_np(local.logits_indices) == _np(batch.logits_indices)[rows]).all()
186 + assert (_np(local.seq_lens) == _np(batch.seq_lens)[owned]).all()
187 + assert batch.num_draft_tokens_per_req is not None
188 + assert local.num_draft_tokens_per_req is not None
189 + assert (
190 + local.num_draft_tokens_per_req == batch.num_draft_tokens_per_req[owned]
191 + ).all()
192 + assert local.num_draft_tokens == int(
193 + batch.num_draft_tokens_per_req[owned].sum()
194 + )
195 + # Triton specializes kernels on pointer alignment: every packed-buffer
196 + # view fed to kernels must stay 16-byte aligned across batch shapes,
197 + # or serving JITs new kernel variants mid-collective.
198 + for t in (
199 + local.idx_mapping,
200 + local.expanded_idx_mapping,
201 + local.expanded_local_pos,
202 + local.logits_indices,
203 + local.cu_num_logits,
204 + metadata.gathered_src_indices,
205 + ):
206 + assert t.data_ptr() % 16 == 0
207 + assert sorted(all_owned) == list(range(batch.num_reqs))
208 +
209 + # sorted_logits_indices is identical on all ranks: the owner-sorted logits row
210 + # positions, i.e. logits_indices permuted rank-major.
211 + sorted_logits_indices = _np(results[0][1])
212 + for _, gi, _ in results:
213 + assert (_np(gi) == sorted_logits_indices).all()
214 + rows_all = np.concatenate(
215 + [_owned_rows(cu, _owned_batch_indices(local)) for local, _, _ in results]
216 + )
217 + assert (
218 + sorted_logits_indices.tolist() == _np(batch.logits_indices)[rows_all].tolist()
219 + )
220 +
221 +
222 +@requires_cuda
223 +@pytest.mark.parametrize("tp_size", [2, 4])
224 +def test_all_to_all_reassembly_simulated(tp_size: int):
225 + """Vocab shards computed from owner-sorted hidden states and routed
226 + through the simulated all-to-all reassemble into exactly the full-vocab
227 + logits of each rank's owned rows."""
228 + rng = np.random.default_rng(7)
229 + max_num_reqs = 32
230 + batch = _make_batch(rng, num_reqs=11, max_num_reqs=max_num_reqs, max_spec=4)
231 + vocab = 64 * tp_size
232 + shard_w = vocab // tp_size
233 + token_logits = torch.randn(batch.num_tokens, vocab, device=DEVICE)
234 +
235 + results = _shard_all_ranks(batch, max_num_reqs, tp_size)
236 + sorted_logits_indices = results[0][1]
237 + num_logits_per_rank = results[0][2].num_logits_per_rank
238 + bounds = np.zeros(tp_size + 1, dtype=np.int64)
239 + np.cumsum(num_logits_per_rank, out=bounds[1:])
240 +
241 + # send_r = rank r's vocab shard of the owner-sorted rows, with no further
242 + # permute; chunk d of send_r goes to rank d.
243 + sends = [
244 + token_logits[sorted_logits_indices][:, r * shard_w : (r + 1) * shard_w]
245 + for r in range(tp_size)
246 + ]
247 + for rank, (local, _, metadata) in enumerate(results):
248 + recv = torch.cat(
249 + [sends[src][bounds[rank] : bounds[rank + 1]] for src in range(tp_size)]
250 + )
251 + # Reassemble source-rank-major chunks into full-vocab rows, mirroring
252 + # all_to_all_logits (source-rank order is vocab-shard order).
253 + reassembled = (
254 + recv.view(tp_size, metadata.num_local_logits, shard_w)
255 + .permute(1, 0, 2)
256 + .reshape(metadata.num_local_logits, vocab)
257 + )
258 + expected = token_logits[local.logits_indices]
259 + assert torch.equal(reassembled, expected)
260 +
261 +
262 +@requires_cuda
263 +def test_shard_no_spec():
264 + """Batches without draft tokens (one logits row per request) degenerate
265 + cleanly: rows mirror requests."""
266 + rng = np.random.default_rng(11)
267 + tp_size = 4
268 + max_num_reqs = 32
269 + batch = _make_batch(rng, num_reqs=10, max_num_reqs=max_num_reqs, max_spec=0)
270 + results = _shard_all_ranks(batch, max_num_reqs, tp_size)
271 +
272 + assert results[0][2].max_num_logits_per_req == 1
273 + for local, _, metadata in results:
274 + assert metadata.num_local_logits == metadata.num_local_reqs == local.num_reqs
275 + owned = _owned_batch_indices(local)
276 + assert _np(local.logits_indices).tolist() == (
277 + _np(batch.logits_indices)[owned].tolist()
278 + )
279 +
280 +
281 +@pytest.mark.parametrize("tp_size", [2, 4])
282 +def test_shard_grammar_output(tp_size: int):
283 + """The grammar filter keeps exactly the owned requests' bitmask rows, in
284 + grammar order, and the per-rank pieces partition the global bitmask."""
285 + rng = np.random.default_rng(5)
286 + max_num_reqs = 32
287 + num_reqs = 12
288 + slots = rng.choice(np.arange(max_num_reqs), size=num_reqs, replace=False)
289 + k = rng.integers(0, 4, size=num_reqs)
290 + cu = np.zeros(num_reqs + 1, dtype=np.int32)
291 + np.cumsum(1 + k, out=cu[1:])
292 + owner = slots % tp_size
293 +
294 + # Grammar requests: every other request, in a scrambled order.
295 + req_ids = [f"req-{i}" for i in range(num_reqs)]
296 + grammar_idx = list(rng.permutation(np.arange(0, num_reqs, 2)))
297 + grammar_req_ids = [req_ids[i] for i in grammar_idx]
298 + num_grammar_rows = sum(int(cu[i + 1] - cu[i]) for i in grammar_idx)
299 + bitmask = rng.integers(0, 2**31, size=(num_grammar_rows, 4), dtype=np.int32)
300 + input_batch = SimpleNamespace(req_ids=req_ids, cu_num_logits_np=cu)
301 +
302 + # Global row offset of each grammar request within the bitmask.
303 + row_offsets = {}
304 + cursor = 0
305 + for i in grammar_idx:
306 + row_offsets[i] = cursor
307 + cursor += int(cu[i + 1] - cu[i])
308 +
309 + grammar_output = GrammarOutput(
310 + structured_output_request_ids=grammar_req_ids,
311 + grammar_bitmask=bitmask,
312 + )
313 +
314 + kept_total = 0
315 + for rank in range(tp_size):
316 + local_req_ids = [req_ids[i] for i in range(num_reqs) if owner[i] == rank]
317 + local = _shard_grammar_output(grammar_output, input_batch, local_req_ids)
318 + expected_idx = [i for i in grammar_idx if owner[i] == rank]
319 + if not expected_idx:
320 + assert local is None
321 + continue
322 + assert local is not None
323 + assert local.structured_output_request_ids == [req_ids[i] for i in expected_idx]
324 + expected_rows = np.concatenate(
325 + [
326 + np.arange(row_offsets[i], row_offsets[i] + cu[i + 1] - cu[i])
327 + for i in expected_idx
328 + ]
329 + )
330 + assert (local.grammar_bitmask == bitmask[expected_rows.astype(int)]).all()
331 + kept_total += local.grammar_bitmask.shape[0]
332 + assert kept_total == num_grammar_rows
333 +
334 +
335 +@requires_cuda
336 +@pytest.mark.parametrize("tp_size", [2, 4])
337 +def test_gather_src_indices_round_trip(tp_size: int):
338 + """Padded per-rank result blocks indexed by gathered_src_indices restore
339 + batch order, including when some rank owns zero requests."""
340 + rng = np.random.default_rng(3)
341 + max_num_reqs = 40
342 + all_slots = np.arange(max_num_reqs)
343 + batch = _make_batch(
344 + rng,
345 + num_reqs=17,
346 + max_num_reqs=max_num_reqs,
347 + max_spec=2,
348 + slot_pool=all_slots[all_slots % tp_size != 0], # rank 0 owns nothing
349 + )
350 + results = _shard_all_ranks(batch, max_num_reqs, tp_size)
351 + metadata0 = results[0][2]
352 + assert metadata0.num_local_reqs == 0
353 +
354 + # Value of request j is j; each rank packs its owned requests' values into
355 + # a zero-padded (max_num_reqs_per_rank, 1) block, gathered rank-major.
356 + pad = metadata0.max_num_reqs_per_rank
357 + blocks = []
358 + for local, _, _ in results:
359 + owned = _owned_batch_indices(local)
360 + block = torch.zeros(pad, 1, dtype=torch.int64, device=DEVICE)
361 + block[: len(owned), 0] = torch.from_numpy(owned).to(DEVICE)
362 + blocks.append(block)
363 + gathered = torch.cat(blocks)
364 +
365 + restored = gathered[metadata0.gathered_src_indices][:, 0]
366 + assert restored.tolist() == list(range(batch.num_reqs))
367 +
368 +
369 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
370 +@pytest.mark.parametrize("tp_size", [2, 4])
371 +def test_gather_sampler_output_kernels(tp_size: int):
372 + """The fused pack/unpack kernels round-trip per-rank sampler results
373 + into batch order: token ids zero-filled past the local width, counts
374 + cast back to int32, identical on every rank."""
375 + from dataclasses import replace as dc_replace
376 +
377 + from vllm.v1.worker.gpu.sample.output import SamplerOutput
378 +
379 + rng = np.random.default_rng(9)
380 + max_num_reqs = 32
381 + batch = _make_batch(rng, num_reqs=13, max_num_reqs=max_num_reqs, max_spec=3)
382 + results = _shard_all_ranks(batch, max_num_reqs, tp_size)
383 + device = torch.device("cuda")
384 + num_reqs = batch.num_reqs
385 + width = results[0][2].max_num_logits_per_req
386 +
387 + # Distinctive per-rank local outputs. Even ranks mimic the regular
388 + # sampler (one column, exercising zero-fill); odd ranks mimic the
389 + # rejection sampler's over-allocation (exercising truncation).
390 + local_outputs: list[SamplerOutput | None] = []
391 + expected_ids = torch.zeros(num_reqs, width, dtype=torch.int64)
392 + expected_sampled = torch.zeros(num_reqs, dtype=torch.int32)
393 + expected_rejected = torch.zeros(num_reqs, dtype=torch.int32)
394 + for rank, (local, _, metadata) in enumerate(results):
395 + if metadata.num_local_reqs == 0:
396 + local_outputs.append(None)
397 + continue
398 + owned = torch.from_numpy(_owned_batch_indices(local))
399 + src_width = 1 if rank % 2 == 0 else width + 3
400 + ids = owned[:, None] * 1000 + torch.arange(src_width)[None, :]
401 + num_sampled = (owned % width + 1).to(torch.int32)
402 + num_rejected = (owned % 3).to(torch.int32)
403 + local_outputs.append(
404 + SamplerOutput(
405 + sampled_token_ids=ids.to(device),
406 + logprobs_tensors=None,
407 + num_nans=None,
408 + num_sampled=num_sampled.to(device),
409 + num_rejected=num_rejected.to(device),
410 + )
411 + )
412 + copy_width = min(src_width, width)
413 + expected_ids[owned, :copy_width] = ids[:, :copy_width]
414 + expected_sampled[owned] = num_sampled
415 + expected_rejected[owned] = num_rejected
416 +
417 + # Two passes over a mocked all-gather: capture each rank's packed send
418 + # block, then replay with the concatenated blocks as the gathered result.
419 + packed_blocks: dict[int, torch.Tensor] = {}
420 + gathered_full: list[torch.Tensor] = []
421 +
422 + def run(rank: int) -> SamplerOutput:
423 + metadata = dc_replace(
424 + results[rank][2],
425 + gathered_src_indices=results[rank][2].gathered_src_indices.to(device),
426 + )
427 +
428 + def fake_all_gather(x: torch.Tensor, dim: int = 0) -> torch.Tensor:
429 + packed_blocks[rank] = x.clone()
430 + if gathered_full:
431 + return gathered_full[0]
432 + return torch.zeros(
433 + tp_size * x.shape[0], x.shape[1], dtype=x.dtype, device=x.device
434 + )
435 +
436 + with mock.patch.object(
437 + batch_shard, "tensor_model_parallel_all_gather", fake_all_gather
438 + ):
439 + return batch_shard.gather_sampler_output(
440 + local_outputs[rank],
441 + metadata,
442 + device,
443 + global_batch=batch,
444 + local_batch=results[rank][0],
445 + )
446 +
447 + for rank in range(tp_size):
448 + run(rank)
449 + # Zero the uninitialized padding rows so the comparison is deterministic;
450 + # the real path never reads them (gathered_src_indices skips padding).
451 + for rank, (local, _, metadata) in enumerate(results):
452 + packed_blocks[rank][metadata.num_local_reqs :] = 0
453 + gathered_full.append(torch.cat([packed_blocks[r] for r in range(tp_size)]))
454 +
455 + for rank in range(tp_size):
456 + out = run(rank)
457 + assert out.sampled_token_ids.dtype == torch.int64
458 + assert out.num_sampled is not None and out.num_sampled.dtype == torch.int32
459 + assert out.num_rejected is not None
460 + assert torch.equal(out.sampled_token_ids.cpu(), expected_ids)
461 + assert torch.equal(out.num_sampled.cpu(), expected_sampled)
462 + assert torch.equal(out.num_rejected.cpu(), expected_rejected.to(torch.int32))
463 +
464 +
465 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
466 +def test_gather_sampler_output_logprobs_and_nans():
467 + """num_nans is reduced per request and gathered; LogprobsTensors are
468 + packed into fixed per-request blocks, width-aligned across ranks (pad
469 + and truncate paths), and restored to global row order."""
470 + from dataclasses import replace as dc_replace
471 +
472 + from vllm.v1.outputs import LogprobsTensors
473 + from vllm.v1.worker.gpu.sample.output import SamplerOutput
474 +
475 + tp_size = 2
476 + rng = np.random.default_rng(21)
477 + max_num_reqs = 32
478 + batch = _make_batch(rng, num_reqs=6, max_num_reqs=max_num_reqs, max_spec=3)
479 + results = _shard_all_ranks(batch, max_num_reqs, tp_size)
480 + device = torch.device("cuda")
481 + num_reqs = batch.num_reqs
482 + cu = batch.cu_num_logits_np
483 + num_logits = int(cu[-1])
484 + width = results[0][2].max_num_logits_per_req
485 + num_cols = 3
486 +
487 + def cuda_fields(b: InputBatch) -> InputBatch:
488 + return dc_replace(
489 + b,
490 + cu_num_logits=b.cu_num_logits.to(device),
491 + expanded_local_pos=b.expanded_local_pos.to(device),
492 + )
493 +
494 + global_cuda = cuda_fields(batch)
495 +
496 + local_outputs: list[SamplerOutput | None] = []
497 + local_batches: list[InputBatch] = []
498 + expected_ids = torch.zeros(num_logits, num_cols, dtype=torch.int64)
499 + expected_logprobs = torch.full(
500 + (num_logits, num_cols), float("-inf"), dtype=torch.float32
501 + )
502 + expected_ranks = torch.zeros(num_logits, dtype=torch.int64)
503 + expected_nans = torch.zeros(num_reqs, dtype=torch.int32)
504 + for rank, (local, _, metadata) in enumerate(results):
505 + local_batches.append(cuda_fields(local))
506 + if metadata.num_local_reqs == 0:
507 + local_outputs.append(None)
508 + continue
509 + owned = _owned_batch_indices(local)
510 + rows = torch.from_numpy(_owned_rows(cu, owned))
511 + # Rank 0 sends fewer columns (pad path), rank 1 more (truncate path).
512 + local_cols = num_cols - 1 if rank == 0 else num_cols + 2
513 + ids = rows[:, None] * 10 + torch.arange(local_cols)[None, :]
514 + logprobs = ids.float() * 0.5
515 + ranks_t = rows + 7
516 + nans = (rows % 4).to(torch.int32)
517 + local_outputs.append(
518 + SamplerOutput(
519 + sampled_token_ids=torch.zeros(
520 + metadata.num_local_reqs, width, dtype=torch.int64, device=device
521 + ),
522 + logprobs_tensors=LogprobsTensors(
523 + logprob_token_ids=ids.to(device),
524 + logprobs=logprobs.to(device),
525 + selected_token_ranks=ranks_t.to(device),
526 + cu_num_generated_tokens=None,
527 + ),
528 + num_nans=nans.to(device),
529 + num_sampled=torch.ones(
530 + metadata.num_local_reqs, dtype=torch.int32, device=device
531 + ),
532 + num_rejected=torch.zeros(
533 + metadata.num_local_reqs, dtype=torch.int32, device=device
534 + ),
535 + )
536 + )
537 + copy_cols = min(local_cols, num_cols)
538 + expected_ids[rows, :copy_cols] = ids[:, :copy_cols]
539 + expected_logprobs[rows, :copy_cols] = logprobs[:, :copy_cols]
540 + expected_ranks[rows] = ranks_t
541 + for j in owned:
542 + in_req = (rows >= int(cu[j])) & (rows < int(cu[j + 1]))
543 + expected_nans[j] = int(nans[in_req].sum())
544 +
545 + recorded: dict[int, list[torch.Tensor]] = {}
546 + gathered_full: list[torch.Tensor] = []
547 +
548 + def run(rank: int) -> SamplerOutput:
549 + metadata = dc_replace(
550 + results[rank][2],
551 + gathered_src_indices=results[rank][2].gathered_src_indices.to(device),
552 + )
553 + calls = {"i": 0}
554 +
555 + def fake_all_gather(x: torch.Tensor, dim: int = 0) -> torch.Tensor:
556 + i = calls["i"]
557 + calls["i"] += 1
558 + if gathered_full:
559 + return gathered_full[i]
560 + recorded.setdefault(rank, []).append(x.clone())
561 + return torch.zeros(
562 + tp_size * x.shape[0], *x.shape[1:], dtype=x.dtype, device=x.device
563 + )
564 +
565 + with mock.patch.object(
566 + batch_shard, "tensor_model_parallel_all_gather", fake_all_gather
567 + ):
568 + return batch_shard.gather_sampler_output(
569 + local_outputs[rank],
570 + metadata,
571 + device,
572 + global_batch=global_cuda,
573 + local_batch=local_batches[rank],
574 + gather_num_nans=True,
575 + logprobs_dims=(num_cols - 1, 0),
576 + )
577 +
578 + for rank in range(tp_size):
579 + run(rank)
580 + gathered_full.extend(
581 + torch.cat([recorded[r][i] for r in range(tp_size)]) for i in range(3)
582 + )
583 +
584 + for rank in range(tp_size):
585 + out = run(rank)
586 + assert out.num_nans is not None
587 + assert torch.equal(out.num_nans.cpu(), expected_nans)
588 + lp = out.logprobs_tensors
589 + assert lp is not None
590 + assert torch.equal(lp.logprob_token_ids.cpu(), expected_ids)
591 + assert torch.equal(lp.logprobs.cpu(), expected_logprobs)
592 + assert torch.equal(lp.selected_token_ranks.cpu(), expected_ranks)
593 + if num_logits != num_reqs:
594 + assert lp.cu_num_generated_tokens == cu.tolist()
595 + else:
596 + assert lp.cu_num_generated_tokens is None
597 +
598 +
599 +def test_finish_requests_frees_slots_in_sorted_order():
600 + """Request slots must be freed in the same order on every TP rank:
601 + ownership derives from slot indices, and `finished_req_ids` is a set
602 + whose iteration order is per-process hash-randomized."""
603 + from vllm.v1.worker.gpu.model_runner import GPUModelRunner
604 +
605 + removed: list[str] = []
606 + fake_runner = SimpleNamespace(_remove_request=removed.append, pooling_runner=None)
607 + scheduler_output = SimpleNamespace(
608 + finished_req_ids={"req-b", "req-c", "req-a"},
609 + preempted_req_ids={"req-e", "req-d"},
610 + )
611 + GPUModelRunner.finish_requests(fake_runner, scheduler_output)
612 + assert removed == sorted(removed)
613 + assert set(removed) == {"req-a", "req-b", "req-c", "req-d", "req-e"}
tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+1
@@ -60,6 +60,7 @@ def _make_runner(**overrides: Any) -> Any:
60 runner.parallel_config = SimpleNamespace(
61 enable_eplb=True,
62 enable_elastic_ep=False,
63 + enable_batch_sharded_sampling=False,
64 eplb_config=SimpleNamespace(log_balancedness=True),
65 )
66 runner.vllm_config = SimpleNamespace(
vllm/config/parallel.py
+7
@@ -164,6 +164,13 @@ class ParallelConfig:
164 """Whether the deployed model is MoE (if known)."""
165 enable_expert_parallel: bool = False
166 """Use expert parallelism instead of tensor parallelism for MoE layers."""
167 + enable_batch_sharded_sampling: bool | None = None
168 + """Use sharded sampling across tensor parallel ranks. Each rank samples
169 + a slice of the batch instead of every rank sampling all of it. Currently
170 + defaults to False if not set. Enabling it explicitly raises when the config
171 + cannot support it (`tensor_parallel_size` must be > 1, `max_num_seqs` at
172 + least `tensor_parallel_size`, and `max_logprobs` non-negative). Models opt in
173 + by implementing `compute_logits_local`."""
174 enable_ep_weight_filter: bool = False
175 """Skip non-local expert weights during model loading when expert
176 parallelism is active. Each rank only reads its own expert shard from
vllm/config/vllm.py
+53
@@ -1675,6 +1675,8 @@ class VllmConfig:
1675 "Remove VLLM_USE_V2_MODEL_RUNNER=0."
1676 )
1677
1678 + self._validate_batch_sharded_sampling()
1679 +
1680 # Re-compute compile ranges after platform-specific config updates
1681 # (e.g., XPU may lower max_num_batched_tokens when MLA is enabled)
1682 self._set_compile_ranges()
@@ -2552,6 +2554,57 @@ class VllmConfig:
2554
2555 return unsupported
2556
2557 + def _validate_batch_sharded_sampling(self) -> None:
2558 + """Validate `enable_batch_sharded_sampling` against the rest of the config."""
2559 + if not self.parallel_config.enable_batch_sharded_sampling:
2560 + # Default to False if not set.
2561 + self.parallel_config.enable_batch_sharded_sampling = False
2562 + return
2563 +
2564 + blockers: list[str] = []
2565 + tp_size = self.parallel_config.tensor_parallel_size
2566 +
2567 + if not self.use_v2_model_runner:
2568 + blockers.append("it is only implemented by Model Runner V2")
2569 +
2570 + if tp_size <= 1:
2571 + blockers.append("tensor_parallel_size is 1, so there is nothing to shard")
2572 + elif self.scheduler_config.max_num_seqs < tp_size:
2573 + # Requests are assigned to ranks whole, so fewer slots than ranks
2574 + # leaves some ranks without work in every step.
2575 + blockers.append(
2576 + f"max_num_seqs ({self.scheduler_config.max_num_seqs}) is below "
2577 + f"tensor_parallel_size ({tp_size})"
2578 + )
2579 +
2580 + if self.model_config is not None and self.model_config.max_logprobs < 0:
2581 + # max_logprobs == -1 allows vocab-size logprob requests, which the
2582 + # fixed-width logprobs gather cannot reasonably size for.
2583 + blockers.append("max_logprobs is -1, allowing vocab-size logprob requests")
2584 +
2585 + if (
2586 + self.speculative_config is not None
2587 + and self.speculative_config.enable_adaptive_verification
2588 + ):
2589 + # Adaptive verification picks the per-request draft split on the GPU,
2590 + # so cu_num_logits_np is only an upper bound, while the shard plan is
2591 + # built from that CPU array. The two disagree once the budget binds.
2592 + # TODO(TheEpicDolphin): Support adaptive verification with batch-sharded
2593 + # sampling.
2594 + blockers.append(
2595 + "it does not yet work with adaptive verification, which decides "
2596 + "the per-request logits counts on the GPU, where the CPU-side "
2597 + "shard plan cannot see them"
2598 + )
2599 +
2600 + if blockers:
2601 + raise ValueError(
2602 + "Batch-sharded sampling was explicitly enabled via "
2603 + "the --enable-batch-sharded-sampling flag, but is not supported "
2604 + "in this configuration for the following reason(s): "
2605 + f"{'; '.join(blockers)}."
2606 + )
2607 +
2608 def _validate_v2_model_runner(self) -> None:
2609 """Check for features not yet supported by the V2 model runner."""
2610 if not HAS_TRITON:
vllm/engine/arg_utils.py
+8
@@ -496,6 +496,9 @@ class EngineArgs:
496 data_parallel_multi_port_external_lb: bool = False
497 data_parallel_backend: DataParallelBackend = ParallelConfig.data_parallel_backend
498 enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel
499 + enable_batch_sharded_sampling: bool | None = (
500 + ParallelConfig.enable_batch_sharded_sampling
501 + )
502 enable_ep_weight_filter: bool = ParallelConfig.enable_ep_weight_filter
503 moe_backend: MoEBackend = KernelConfig.moe_backend
504 linear_backend: LinearBackend = KernelConfig.linear_backend
@@ -1156,6 +1159,10 @@ class EngineArgs:
1159 "-ep",
1160 **parallel_kwargs["enable_expert_parallel"],
1161 )
1162 + parallel_group.add_argument(
1163 + "--enable-batch-sharded-sampling",
1164 + **parallel_kwargs["enable_batch_sharded_sampling"],
1165 + )
1166 parallel_group.add_argument(
1167 "--enable-ep-weight-filter",
1168 **parallel_kwargs["enable_ep_weight_filter"],
@@ -2279,6 +2286,7 @@ class EngineArgs:
2286 data_parallel_hybrid_lb=self.data_parallel_hybrid_lb,
2287 is_moe_model=model_config.is_moe,
2288 enable_expert_parallel=self.enable_expert_parallel,
2289 + enable_batch_sharded_sampling=self.enable_batch_sharded_sampling,
2290 enable_ep_weight_filter=self.enable_ep_weight_filter,
2291 all2all_backend=self.all2all_backend,
2292 enable_elastic_ep=self.enable_elastic_ep,
vllm/lora/layers/logits_processor.py
+7
@@ -155,7 +155,14 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
155 hidden_states: torch.Tensor,
156 lm_head: VocabParallelEmbedding,
157 embedding_bias: torch.Tensor | None = None,
158 + skip_gather: bool = False,
159 ) -> torch.Tensor | None:
160 + # The LoRA delta is accumulated into the full gathered logits, so the
161 + # TP gather cannot be skipped here.
162 + if skip_gather:
163 + raise NotImplementedError(
164 + "Skipping the logits TP gather is not supported with an lm_head LoRA."
165 + )
166 # Get the logits for the next tokens.
167 if hasattr(lm_head, "base_layer"):
168 actual_lm_head = lm_head.base_layer
vllm/model_executor/layers/logits_processor.py
+7 -1
@@ -99,12 +99,15 @@ class LogitsProcessor(PluggableLayer):
99 lm_head: VocabParallelEmbedding,
100 hidden_states: torch.Tensor,
101 embedding_bias: torch.Tensor | None = None,
102 + skip_gather: bool = False,
103 ) -> torch.Tensor | None:
104 if self.logits_as_input:
105 logits = hidden_states
106 else:
107 # Get the logits for the next tokens.
107 - logits = self._get_logits(hidden_states, lm_head, embedding_bias)
108 + logits = self._get_logits(
109 + hidden_states, lm_head, embedding_bias, skip_gather
110 + )
111 if logits is not None:
112 if self.soft_cap is not None:
113 logits = logits / self.soft_cap
@@ -173,9 +176,12 @@ class LogitsProcessor(PluggableLayer):
176 hidden_states: torch.Tensor,
177 lm_head: VocabParallelEmbedding,
178 embedding_bias: torch.Tensor | None,
179 + skip_gather: bool = False,
180 ) -> torch.Tensor | None:
181 # Get the logits for the next tokens.
182 logits = self._apply_head(lm_head, hidden_states, embedding_bias)
183 + if skip_gather:
184 + return logits
185
186 # Gather logits for TP
187 if lm_head.tp_size > 1:
vllm/model_executor/models/llama.py
+6
@@ -532,6 +532,12 @@ class LlamaForCausalLM(
532 logits = self.logits_processor(self.lm_head, hidden_states)
533 return logits
534
535 + def compute_logits_local(
536 + self,
537 + hidden_states: torch.Tensor,
538 + ) -> torch.Tensor:
539 + return self.logits_processor(self.lm_head, hidden_states, skip_gather=True)
540 +
541 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
542 loader = AutoWeightsLoader(self)
543 return loader.load_weights(weights)
vllm/model_executor/models/qwen3_5.py
+9
@@ -429,6 +429,12 @@ class Qwen3_5ForCausalLMBase(
429 ) -> torch.Tensor | None:
430 return self.logits_processor(self.lm_head, hidden_states)
431
432 + def compute_logits_local(
433 + self,
434 + hidden_states: torch.Tensor,
435 + ) -> torch.Tensor:
436 + return self.logits_processor(self.lm_head, hidden_states, skip_gather=True)
437 +
438 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
439 loader = AutoWeightsLoader(self)
440 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
@@ -591,6 +597,9 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid)
597
598 return hidden_states
599
600 + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
601 + return self.language_model.compute_logits_local(hidden_states)
602 +
603 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
604 loader = AutoWeightsLoader(self)
605 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/models/deepseek_v4/nvidia/model.py
+6
@@ -1765,6 +1765,12 @@ class DeepseekV4ForCausalLM(
1765 logits = self.logits_processor(self.lm_head, hidden_states)
1766 return logits
1767
1768 + def compute_logits_local(
1769 + self,
1770 + hidden_states: torch.Tensor,
1771 + ) -> torch.Tensor:
1772 + return self.logits_processor(self.lm_head, hidden_states, skip_gather=True)
1773 +
1774 def forward(
1775 self,
1776 input_ids: torch.Tensor,
vllm/models/inkling/amd/logits_processor.py
+11 -4
@@ -63,6 +63,7 @@ class InklingLogitsProcessor(LogitsProcessor):
63 lm_head: VocabParallelEmbedding,
64 hidden_states: torch.Tensor,
65 embedding_bias: torch.Tensor | None = None,
66 + skip_gather: bool = False,
67 ) -> torch.Tensor | None:
68 # ``base_layer`` exists only on the LogitsProcessorWithLoRA wrapper,
69 # which calls this forward with ``self`` bound to the wrapper. The
@@ -72,15 +73,16 @@ class InklingLogitsProcessor(LogitsProcessor):
73 # lives on this class).
74 if hasattr(self, "base_layer"):
75 return type(self.base_layer)._lora_forward(
75 - self, lm_head, hidden_states, embedding_bias
76 + self, lm_head, hidden_states, embedding_bias, skip_gather
77 )
77 - return self._base_forward(lm_head, hidden_states, embedding_bias)
78 + return self._base_forward(lm_head, hidden_states, embedding_bias, skip_gather)
79
80 def _lora_forward(
81 self,
82 lm_head: VocabParallelEmbedding,
83 hidden_states: torch.Tensor,
84 embedding_bias: torch.Tensor | None = None,
85 + skip_gather: bool = False,
86 ) -> torch.Tensor | None:
87 # ``self`` is the LogitsProcessorWithLoRA wrapper here: ``_get_logits``
88 # returns the base logits plus the lm_head LoRA delta. Apply the muP
@@ -90,7 +92,9 @@ class InklingLogitsProcessor(LogitsProcessor):
92 if self.logits_as_input:
93 logits = hidden_states
94 else:
93 - logits = self._get_logits(hidden_states, lm_head, embedding_bias)
95 + logits = self._get_logits(
96 + hidden_states, lm_head, embedding_bias, skip_gather
97 + )
98 # TODO: fuse this multiplication
99 if logits is not None and mup:
100 assert self.base_layer.soft_cap is None
@@ -103,10 +107,11 @@ class InklingLogitsProcessor(LogitsProcessor):
107 lm_head: VocabParallelEmbedding,
108 hidden_states: torch.Tensor,
109 embedding_bias: torch.Tensor | None = None,
110 + skip_gather: bool = False,
111 ) -> torch.Tensor | None:
112 mup = self.logits_mup_width_multiplier
113 if not mup:
109 - return super().forward(lm_head, hidden_states, embedding_bias)
114 + return super().forward(lm_head, hidden_states, embedding_bias, skip_gather)
115 # Fold the muP width divisor into the lm_head GEMM alpha (fp32 epilogue):
116 # no separate elementwise kernel, no bf16 rounding of scaled logits, and
117 # no weight mutation. Overfit to the served checkpoint: bf16 lm_head, no
@@ -123,6 +128,8 @@ class InklingLogitsProcessor(LogitsProcessor):
128 beta=0.0,
129 alpha=1.0 / mup,
130 )
131 + if skip_gather:
132 + return logits
133 logits = self._gather_logits(logits)
134 if logits is not None:
135 logits = logits[..., : self.org_vocab_size]
vllm/models/inkling/nvidia/logits_processor.py
+11 -4
@@ -63,6 +63,7 @@ class InklingLogitsProcessor(LogitsProcessor):
63 lm_head: VocabParallelEmbedding,
64 hidden_states: torch.Tensor,
65 embedding_bias: torch.Tensor | None = None,
66 + skip_gather: bool = False,
67 ) -> torch.Tensor | None:
68 # ``base_layer`` exists only on the LogitsProcessorWithLoRA wrapper,
69 # which calls this forward with ``self`` bound to the wrapper. The
@@ -72,15 +73,16 @@ class InklingLogitsProcessor(LogitsProcessor):
73 # lives on this class).
74 if hasattr(self, "base_layer"):
75 return type(self.base_layer)._lora_forward(
75 - self, lm_head, hidden_states, embedding_bias
76 + self, lm_head, hidden_states, embedding_bias, skip_gather
77 )
77 - return self._base_forward(lm_head, hidden_states, embedding_bias)
78 + return self._base_forward(lm_head, hidden_states, embedding_bias, skip_gather)
79
80 def _lora_forward(
81 self,
82 lm_head: VocabParallelEmbedding,
83 hidden_states: torch.Tensor,
84 embedding_bias: torch.Tensor | None = None,
85 + skip_gather: bool = False,
86 ) -> torch.Tensor | None:
87 # ``self`` is the LogitsProcessorWithLoRA wrapper here: ``_get_logits``
88 # returns the base logits plus the lm_head LoRA delta. Apply the muP
@@ -90,7 +92,9 @@ class InklingLogitsProcessor(LogitsProcessor):
92 if self.logits_as_input:
93 logits = hidden_states
94 else:
93 - logits = self._get_logits(hidden_states, lm_head, embedding_bias)
95 + logits = self._get_logits(
96 + hidden_states, lm_head, embedding_bias, skip_gather
97 + )
98 # TODO: fuse this multiplication
99 if logits is not None and mup:
100 assert self.base_layer.soft_cap is None
@@ -103,10 +107,11 @@ class InklingLogitsProcessor(LogitsProcessor):
107 lm_head: VocabParallelEmbedding,
108 hidden_states: torch.Tensor,
109 embedding_bias: torch.Tensor | None = None,
110 + skip_gather: bool = False,
111 ) -> torch.Tensor | None:
112 mup = self.logits_mup_width_multiplier
113 if not mup:
109 - return super().forward(lm_head, hidden_states, embedding_bias)
114 + return super().forward(lm_head, hidden_states, embedding_bias, skip_gather)
115 # Fold the muP width divisor into the lm_head GEMM alpha (fp32 epilogue):
116 # no separate elementwise kernel, no bf16 rounding of scaled logits, and
117 # no weight mutation. Overfit to the served checkpoint: bf16 lm_head, no
@@ -123,6 +128,8 @@ class InklingLogitsProcessor(LogitsProcessor):
128 beta=0.0,
129 alpha=1.0 / mup,
130 )
131 + if skip_gather:
132 + return logits
133 logits = self._gather_logits(logits)
134 if logits is not None:
135 logits = logits[..., : self.org_vocab_size]
vllm/models/minimax_m3/nvidia/model.py
+6
@@ -1059,6 +1059,9 @@ class MiniMaxM3SparseForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
1059 def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
1060 return self.logits_processor(self.lm_head, hidden_states)
1061
1062 + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
1063 + return self.logits_processor(self.lm_head, hidden_states, skip_gather=True)
1064 +
1065 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1066 return self.model.get_expert_mapping()
1067
@@ -1272,6 +1275,9 @@ class MiniMaxM3SparseForConditionalGeneration(
1275 def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
1276 return self.language_model.compute_logits(hidden_states)
1277
1278 + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
1279 + return self.language_model.compute_logits_local(hidden_states)
1280 +
1281 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1282 return self.language_model.get_expert_mapping()
1283
vllm/v1/worker/gpu/model_runner.py
+72 -4
@@ -127,6 +127,11 @@ from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras
127 from vllm.v1.worker.gpu.model_states import init_model_state
128 from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner
129 from vllm.v1.worker.gpu.pp_utils import PPHandler
130 +from vllm.v1.worker.gpu.sample.batch_shard import (
131 + BatchSharder,
132 + all_to_all_logits,
133 + gather_sampler_output,
134 +)
135 from vllm.v1.worker.gpu.sample.output import SamplerOutput
136 from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
137 from vllm.v1.worker.gpu.sample.sampler import Sampler
@@ -310,6 +315,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
315 # model_state exists (num_new_sampled_tokens_per_step from ModelState).
316 self.sampler: Sampler | None = None
317 self.rejection_sampler: RejectionSampler | None = None
318 + self.batch_sharder: BatchSharder | None = None
319 self.prompt_logprobs_worker: PromptLogprobsWorker | None = None
320 self.structured_outputs_worker: StructuredOutputsWorker | None = None
321 self.cudagraph_manager: ModelCudaGraphManager | None = None
@@ -404,6 +410,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):
410 + self.model_state.num_new_sampled_tokens_per_step
411 )
412
413 + if self.parallel_config.enable_batch_sharded_sampling:
414 + if hasattr(self.model, "compute_logits_local"):
415 + self.batch_sharder = BatchSharder(
416 + max_num_reqs=self.max_num_reqs,
417 + max_num_logits_per_req=self.decode_query_len,
418 + device=self.device,
419 + )
420 + logger.info("Batch-sharded sampling enabled.")
421 + else:
422 + logger.warning_once(
423 + "Disabling batch-sharded sampling: %s does not implement "
424 + "compute_logits_local",
425 + type(self.model).__name__,
426 + )
427 +
428 # Initialize samplers. Model states may override via custom_sampler().
429 if self.is_last_pp_rank and not self.is_pooling_model:
430 self.sampler = Sampler(
@@ -935,7 +956,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
956 preempted_req_ids = scheduler_output.preempted_req_ids
957 if preempted_req_ids:
958 finished_req_ids = finished_req_ids.union(preempted_req_ids)
938 - for req_id in finished_req_ids:
959 + # Sorted so every TP rank frees request slots in the same order.
960 + # Features like batch-sharded sampling derive rank request ownership
961 + # from the slot index.
962 + for req_id in sorted(finished_req_ids):
963 self._remove_request(req_id)
964
965 def free_states(self, scheduler_output: SchedulerOutput) -> None:
@@ -1335,8 +1359,25 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1359 input_batch: InputBatch,
1360 grammar_output: GrammarOutput | None,
1361 ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]:
1338 - sample_hidden_states = hidden_states[input_batch.logits_indices]
1339 - logits = self.model.compute_logits(sample_hidden_states)
1362 + shard_metadata = None
1363 + global_input_batch = input_batch
1364 + if self.batch_sharder is not None:
1365 + # Shard the inputs along the batch dimension to sample in parallel
1366 + # across TP ranks.
1367 + input_batch, sorted_logits_indices, grammar_output, shard_metadata = (
1368 + self.batch_sharder.shard_sampler_inputs(input_batch, grammar_output)
1369 + )
1370 + # The hidden states must be gathered in rank-owner-sorted order
1371 + # before computing the partial-vocab logits, so that the all-to-all
1372 + # produces full-vocab logits for just the locally-owned requests.
1373 + sample_hidden_states = hidden_states[sorted_logits_indices]
1374 + local_logits = self.model.compute_logits_local(sample_hidden_states)
1375 + logits = all_to_all_logits(local_logits, shard_metadata)
1376 + logits = logits[:, : self.vocab_size]
1377 + else:
1378 + sample_hidden_states = hidden_states[input_batch.logits_indices]
1379 + logits = self.model.compute_logits(sample_hidden_states)
1380 +
1381 if grammar_output is not None:
1382 # Apply grammar bitmask to the logits in-place.
1383 assert self.structured_outputs_worker is not None
@@ -1347,7 +1388,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1388 grammar_output.grammar_bitmask,
1389 )
1390
1350 - if input_batch.num_draft_tokens == 0 or self.rejection_sampler is None:
1391 + sampler_output: SamplerOutput | None
1392 + if input_batch.num_reqs == 0:
1393 + # This rank owns no requests this step. It contributes an
1394 + # all-padding block to the gather below.
1395 + sampler_output = None
1396 + elif input_batch.num_draft_tokens == 0 or self.rejection_sampler is None:
1397 assert self.sampler is not None
1398 sampler_output = self.sampler(logits, input_batch)
1399 else:
@@ -1361,6 +1407,28 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1407 self.speculator.draft_logits,
1408 )
1409
1410 + if shard_metadata is not None:
1411 + # Gather the sharded sampler outputs from the TP ranks into a single
1412 + # sampler output.
1413 + assert self.sampler is not None
1414 + sampler_output = gather_sampler_output(
1415 + sampler_output,
1416 + shard_metadata,
1417 + device=self.device,
1418 + global_batch=global_input_batch,
1419 + local_batch=input_batch,
1420 + gather_num_nans=self.sampler.compute_nans,
1421 + logprobs_dims=self.sampler.get_logprobs_dims(
1422 + global_input_batch.idx_mapping_np,
1423 + # Rejection sampler does not return logprob token ids.
1424 + include_token_ids=(
1425 + global_input_batch.num_draft_tokens == 0
1426 + or self.rejection_sampler is None
1427 + ),
1428 + ),
1429 + )
1430 +
1431 + assert sampler_output is not None
1432 return sampler_output, sampler_output.num_sampled, sampler_output.num_rejected
1433
1434 def postprocess_sampled(
vllm/v1/worker/gpu/sample/batch_shard.py new
+674
@@ -0,0 +1,674 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from dataclasses import dataclass, replace
5 +
6 +import numpy as np
7 +import torch
8 +import torch.distributed
9 +
10 +from vllm.distributed.communication_op import tensor_model_parallel_all_gather
11 +from vllm.distributed.parallel_state import get_tp_group
12 +from vllm.triton_utils import tl, triton
13 +from vllm.v1.core.sched.output import GrammarOutput
14 +from vllm.v1.outputs import LogprobsTensors
15 +from vllm.v1.worker.gpu.input_batch import InputBatch
16 +from vllm.v1.worker.gpu.sample.output import SamplerOutput
17 +
18 +
19 +@dataclass
20 +class BatchShardMetadata:
21 + """Collective layout for one sampling step.
22 +
23 + Requests (and their logits) are owner-sorted: all requests owned by
24 + rank 0 first, then rank 1, etc. Within an owner, batch order is
25 + preserved (stable sort). Every field is a pure function of the
26 + replicated idx_mapping and cu_num_logits, so all ranks build identical
27 + plans without communication — which relies on request slots
28 + (idx_mapping) being assigned identically on every rank (see
29 + finish_requests in the model runner).
30 + """
31 +
32 + tp_size: int
33 + # Per-owner-rank logits counts (all-to-all send splits).
34 + num_logits_per_rank: list[int]
35 + num_local_logits: int
36 + num_local_reqs: int
37 + # The number of request entries contributed by each rank to the
38 + # gathered result (shorter shards are padded up to this value).
39 + max_num_reqs_per_rank: int
40 + # [num_reqs] For batch request i, its source index in the rank-major
41 + # gathered results. Indexing the gathered tensors with this restores
42 + # the original batch order.
43 + gathered_src_indices: torch.Tensor
44 + # The number of entries occupied by each request in the gathered
45 + # tensor. Derived from the global cu_num_logits, so each rank agrees
46 + # ont he gather shapes, including empty shards.
47 + max_num_logits_per_req: int
48 +
49 +
50 +@triton.jit(
51 + do_not_specialize=["num_reqs", "local_logits_start", "max_num_reqs_per_rank"]
52 +)
53 +def _build_shard_plan_kernel(
54 + idx_mapping_ptr,
55 + cu_num_logits_ptr,
56 + query_start_loc_ptr,
57 + seq_lens_ptr,
58 + sorted_logits_indices_ptr,
59 + gathered_src_indices_ptr,
60 + local_idx_mapping_ptr,
61 + local_cu_num_logits_ptr,
62 + local_logits_indices_ptr,
63 + local_expanded_idx_mapping_ptr,
64 + local_expanded_local_pos_ptr,
65 + local_seq_lens_ptr,
66 + num_reqs,
67 + local_logits_start,
68 + max_num_reqs_per_rank,
69 + TP_SIZE: tl.constexpr,
70 + TP_RANK: tl.constexpr,
71 + PADDED_NUM_REQS: tl.constexpr,
72 + PADDED_NUM_LOGITS_PER_REQ: tl.constexpr,
73 +):
74 + req_idx = tl.program_id(0)
75 + req_state_idx = tl.load(idx_mapping_ptr + req_idx)
76 + owner = req_state_idx % TP_SIZE
77 + num_logits = tl.load(cu_num_logits_ptr + req_idx + 1) - tl.load(
78 + cu_num_logits_ptr + req_idx
79 + )
80 +
81 + req_block = tl.arange(0, PADDED_NUM_REQS)
82 + req_mask = req_block < num_reqs
83 + owners = tl.load(idx_mapping_ptr + req_block, mask=req_mask, other=0) % TP_SIZE
84 + num_logits_per_req = tl.load(
85 + cu_num_logits_ptr + req_block + 1, mask=req_mask, other=0
86 + ) - tl.load(cu_num_logits_ptr + req_block, mask=req_mask, other=0)
87 +
88 + # Owner-sorted layout is rank 0's requests in batch order, then rank 1's,
89 + # etc, so a request's starting offset can be derived from counting.
90 + lower = (owners < owner) & req_mask
91 + earlier = (owners == owner) & (req_block < req_idx) & req_mask
92 +
93 + # Compute the logits row offset for the current request in the sorted
94 + # batch.
95 + req_logits_start = tl.sum(tl.where(lower, num_logits_per_req, 0)) + tl.sum(
96 + tl.where(earlier, num_logits_per_req, 0)
97 + )
98 +
99 + # Copy the current request logit indices into the sorted batch.
100 + query_start = tl.load(query_start_loc_ptr + req_idx + 1) - num_logits
101 + logit_block = tl.arange(0, PADDED_NUM_LOGITS_PER_REQ)
102 + logit_mask = logit_block < num_logits
103 + logits_indices = (query_start + logit_block).to(tl.int64)
104 + tl.store(
105 + sorted_logits_indices_ptr + req_logits_start + logit_block,
106 + logits_indices,
107 + mask=logit_mask,
108 + )
109 +
110 + # Write the mapping of output tensor request indices => gathered tensor
111 + # request indices.
112 + # Batch request i's result lives in the padded, gathered results at
113 + # rank_of(i) * max_num_reqs_per_rank + offset_within_rank(i).
114 + local_req_idx = tl.sum(earlier.to(tl.int32))
115 + tl.store(
116 + gathered_src_indices_ptr + req_idx,
117 + (owner * max_num_reqs_per_rank + local_req_idx).to(tl.int64),
118 + )
119 +
120 + if req_idx == 0:
121 + tl.store(local_cu_num_logits_ptr, 0)
122 +
123 + if owner == TP_RANK:
124 + local_start = req_logits_start - local_logits_start
125 + tl.store(local_idx_mapping_ptr + local_req_idx, req_state_idx)
126 + tl.store(local_cu_num_logits_ptr + local_req_idx + 1, local_start + num_logits)
127 + tl.store(
128 + local_logits_indices_ptr + local_start + logit_block,
129 + logits_indices,
130 + mask=logit_mask,
131 + )
132 + tl.store(
133 + local_expanded_idx_mapping_ptr + local_start + logit_block,
134 + req_state_idx,
135 + mask=logit_mask,
136 + )
137 + tl.store(
138 + local_expanded_local_pos_ptr + local_start + logit_block,
139 + logit_block,
140 + mask=logit_mask,
141 + )
142 + tl.store(local_seq_lens_ptr + local_req_idx, tl.load(seq_lens_ptr + req_idx))
143 +
144 +
145 +class BatchSharder:
146 + """Shards the sampler inputs across TP ranks along the batch dimension."""
147 +
148 + def __init__(
149 + self,
150 + max_num_reqs: int,
151 + max_num_logits_per_req: int,
152 + device: torch.device,
153 + ):
154 + tp_group = get_tp_group()
155 + self.tp_rank = tp_group.rank_in_group
156 + self.tp_size = tp_group.world_size
157 + self.device = device
158 + self._padded_num_reqs = triton.next_power_of_2(max_num_reqs)
159 + self._padded_num_logits_per_req = triton.next_power_of_2(max_num_logits_per_req)
160 + self._num_warps = max(1, min(8, self._padded_num_reqs // 128))
161 +
162 + def shard_sampler_inputs(
163 + self,
164 + input_batch: InputBatch,
165 + grammar_output: GrammarOutput | None,
166 + ) -> tuple[InputBatch, torch.Tensor, GrammarOutput | None, BatchShardMetadata]:
167 + """Owner-sort the batch and build this rank's local sampler inputs.
168 +
169 + Returns the local sub-batch, the owner-sorted logits_indices (gather
170 + the sampling hidden states with these so `compute_logits_local` emits
171 + logits in all-to-all send order), the local grammar output (None if
172 + this rank owns none of the structured-output requests), and the shard
173 + metadata (used for all-gathering the sampler outputs).
174 + """
175 + tp_rank = self.tp_rank
176 + tp_size = self.tp_size
177 + num_reqs = input_batch.idx_mapping_np.shape[0]
178 + num_logits = int(input_batch.cu_num_logits_np[-1])
179 +
180 + # Deterministically assign requests to ranks, round-robin over slot
181 + # indices so ownership stays balanced when the slot allocator fills
182 + # low slots first (partial occupancy).
183 + # NOTE: This mirrors the assignment used in _build_shard_plan_kernel.
184 + req_owner_np = input_batch.idx_mapping_np % tp_size
185 + local_req_indices_np = np.flatnonzero(req_owner_np == tp_rank)
186 + local_idx_mapping_np = input_batch.idx_mapping_np[local_req_indices_np]
187 + num_local_reqs = local_req_indices_np.shape[0]
188 + num_reqs_per_rank_np = np.bincount(req_owner_np, minlength=tp_size)
189 + max_num_reqs_per_rank = int(num_reqs_per_rank_np.max()) if num_reqs else 1
190 + # Derive the number of logits owned by each rank, as well as the
191 + # local rank.
192 + num_logits_per_req_np = np.diff(input_batch.cu_num_logits_np)
193 + num_logits_per_rank_np = np.bincount(
194 + req_owner_np, weights=num_logits_per_req_np, minlength=tp_size
195 + ).astype(np.int64)
196 + num_local_logits = int(num_logits_per_rank_np[tp_rank])
197 + local_logits_start = int(num_logits_per_rank_np[:tp_rank].sum())
198 + local_cu_num_logits_np = np.zeros(num_local_reqs + 1, dtype=np.int32)
199 + np.cumsum(
200 + num_logits_per_req_np[local_req_indices_np], out=local_cu_num_logits_np[1:]
201 + )
202 + max_num_logits_per_req = int(num_logits_per_req_np.max()) if num_reqs else 1
203 +
204 + # Shard the input batch GPU tensors.
205 + sorted_logits_indices = torch.empty(
206 + num_logits, dtype=torch.int64, device=self.device
207 + )
208 + gathered_src_indices = torch.empty(
209 + num_reqs, dtype=torch.int64, device=self.device
210 + )
211 + local_logits_indices = torch.empty(
212 + num_local_logits, dtype=torch.int64, device=self.device
213 + )
214 + local_idx_mapping = torch.empty(
215 + num_local_reqs, dtype=torch.int32, device=self.device
216 + )
217 + local_cu_num_logits = torch.empty(
218 + num_local_reqs + 1, dtype=torch.int32, device=self.device
219 + )
220 + local_expanded_idx_mapping = torch.empty(
221 + num_local_logits, dtype=torch.int32, device=self.device
222 + )
223 + local_expanded_local_pos = torch.empty(
224 + num_local_logits, dtype=torch.int32, device=self.device
225 + )
226 + local_seq_lens = torch.empty(
227 + num_local_reqs, dtype=torch.int32, device=self.device
228 + )
229 + if num_reqs > 0:
230 + _build_shard_plan_kernel[(num_reqs,)](
231 + input_batch.idx_mapping,
232 + input_batch.cu_num_logits,
233 + input_batch.query_start_loc,
234 + input_batch.seq_lens,
235 + sorted_logits_indices,
236 + gathered_src_indices,
237 + local_idx_mapping,
238 + local_cu_num_logits,
239 + local_logits_indices,
240 + local_expanded_idx_mapping,
241 + local_expanded_local_pos,
242 + local_seq_lens,
243 + num_reqs,
244 + local_logits_start,
245 + max_num_reqs_per_rank,
246 + TP_SIZE=tp_size,
247 + TP_RANK=tp_rank,
248 + PADDED_NUM_REQS=self._padded_num_reqs,
249 + PADDED_NUM_LOGITS_PER_REQ=self._padded_num_logits_per_req,
250 + num_warps=self._num_warps,
251 + )
252 +
253 + # Compute the local number of draft tokens.
254 + num_draft_tokens_per_req = None
255 + num_draft_tokens = 0
256 + if input_batch.num_draft_tokens_per_req is not None:
257 + num_draft_tokens_per_req = input_batch.num_draft_tokens_per_req[
258 + local_req_indices_np
259 + ]
260 + num_draft_tokens = int(num_draft_tokens_per_req.sum())
261 +
262 + local_req_ids = [input_batch.req_ids[i] for i in local_req_indices_np.tolist()]
263 + local_batch = replace(
264 + input_batch,
265 + req_ids=local_req_ids,
266 + num_reqs=num_local_reqs,
267 + idx_mapping=local_idx_mapping,
268 + idx_mapping_np=local_idx_mapping_np,
269 + expanded_idx_mapping=local_expanded_idx_mapping,
270 + expanded_local_pos=local_expanded_local_pos,
271 + seq_lens=local_seq_lens,
272 + logits_indices=local_logits_indices,
273 + cu_num_logits=local_cu_num_logits,
274 + cu_num_logits_np=local_cu_num_logits_np,
275 + num_draft_tokens=num_draft_tokens,
276 + num_draft_tokens_per_req=num_draft_tokens_per_req,
277 + )
278 + local_grammar_output = None
279 + if grammar_output is not None:
280 + local_grammar_output = _shard_grammar_output(
281 + grammar_output, input_batch, local_req_ids
282 + )
283 + metadata = BatchShardMetadata(
284 + tp_size=tp_size,
285 + num_logits_per_rank=num_logits_per_rank_np.tolist(),
286 + num_local_logits=num_local_logits,
287 + num_local_reqs=num_local_reqs,
288 + max_num_reqs_per_rank=max_num_reqs_per_rank,
289 + gathered_src_indices=gathered_src_indices,
290 + max_num_logits_per_req=max_num_logits_per_req,
291 + )
292 + return local_batch, sorted_logits_indices, local_grammar_output, metadata
293 +
294 +
295 +def all_to_all_logits(
296 + logits_shard: torch.Tensor,
297 + metadata: BatchShardMetadata,
298 +) -> torch.Tensor:
299 + recv = logits_shard.new_empty(
300 + metadata.tp_size * metadata.num_local_logits, logits_shard.shape[-1]
301 + )
302 + torch.distributed.all_to_all_single(
303 + recv,
304 + logits_shard.contiguous(),
305 + output_split_sizes=[metadata.num_local_logits] * metadata.tp_size,
306 + input_split_sizes=metadata.num_logits_per_rank,
307 + group=get_tp_group().device_group,
308 + )
309 + shard_width = recv.shape[-1]
310 + return (
311 + recv.view(metadata.tp_size, metadata.num_local_logits, shard_width)
312 + .permute(1, 0, 2)
313 + .reshape(metadata.num_local_logits, metadata.tp_size * shard_width)
314 + )
315 +
316 +
317 +def _shard_grammar_output(
318 + grammar_output: GrammarOutput,
319 + input_batch: InputBatch,
320 + local_req_ids: list[str],
321 +) -> GrammarOutput | None:
322 + owned = set(local_req_ids)
323 + req_id_to_idx = {req_id: i for i, req_id in enumerate(input_batch.req_ids)}
324 + cu_num_logits_np = input_batch.cu_num_logits_np
325 +
326 + local_ids: list[str] = []
327 + keep_indices: list[int] = []
328 + cursor = 0
329 + for req_id in grammar_output.structured_output_request_ids:
330 + req_idx = req_id_to_idx[req_id]
331 + num_req_logits = int(cu_num_logits_np[req_idx + 1] - cu_num_logits_np[req_idx])
332 + if req_id in owned:
333 + local_ids.append(req_id)
334 + keep_indices.extend(range(cursor, cursor + num_req_logits))
335 + cursor += num_req_logits
336 + if not local_ids:
337 + return None
338 + return GrammarOutput(
339 + structured_output_request_ids=local_ids,
340 + grammar_bitmask=grammar_output.grammar_bitmask[keep_indices],
341 + )
342 +
343 +
344 +@triton.jit(
345 + do_not_specialize=["max_num_logits_per_req", "num_src_cols", "packed_stride"]
346 +)
347 +def _pack_sampler_output_kernel(
348 + packed_ptr,
349 + packed_stride,
350 + sampled_token_ids_ptr,
351 + sampled_token_ids_stride,
352 + num_sampled_ptr,
353 + num_rejected_ptr,
354 + num_nans_ptr,
355 + local_cu_num_logits_ptr,
356 + max_num_logits_per_req,
357 + num_src_cols,
358 + BLOCK_SIZE: tl.constexpr,
359 +):
360 + req_idx = tl.program_id(0)
361 + cols = tl.arange(0, BLOCK_SIZE)
362 + row_ptr = packed_ptr + req_idx * packed_stride
363 +
364 + token_ids = tl.load(
365 + sampled_token_ids_ptr + req_idx * sampled_token_ids_stride + cols,
366 + mask=cols < num_src_cols,
367 + other=0,
368 + )
369 + tl.store(row_ptr + cols, token_ids, mask=cols < max_num_logits_per_req)
370 + num_sampled = tl.load(num_sampled_ptr + req_idx)
371 + tl.store(row_ptr + max_num_logits_per_req, num_sampled.to(tl.int64))
372 + num_rejected = tl.load(num_rejected_ptr + req_idx)
373 + tl.store(row_ptr + max_num_logits_per_req + 1, num_rejected.to(tl.int64))
374 + if num_nans_ptr is not None:
375 + # num_nans is per logits row; reduce it per request. The replicated
376 + # samplers report per-row counts that downstream zips per request,
377 + # so a per-request sum is the well-defined equivalent.
378 + start = tl.load(local_cu_num_logits_ptr + req_idx)
379 + num_req_logits = tl.load(local_cu_num_logits_ptr + req_idx + 1) - start
380 + nans = tl.load(num_nans_ptr + start + cols, mask=cols < num_req_logits, other=0)
381 + tl.store(
382 + row_ptr + max_num_logits_per_req + 2, tl.sum(nans.to(tl.int64), axis=0)
383 + )
384 +
385 +
386 +@triton.jit(do_not_specialize=["max_num_logits_per_req", "gathered_stride"])
387 +def _unpack_gathered_output_kernel(
388 + gathered_ptr,
389 + gathered_stride,
390 + gathered_src_indices_ptr,
391 + sampled_token_ids_ptr,
392 + num_sampled_ptr,
393 + num_rejected_ptr,
394 + num_nans_ptr,
395 + max_num_logits_per_req,
396 + BLOCK_SIZE: tl.constexpr,
397 +):
398 + req_idx = tl.program_id(0)
399 + src = tl.load(gathered_src_indices_ptr + req_idx)
400 + row_ptr = gathered_ptr + src * gathered_stride
401 + cols = tl.arange(0, BLOCK_SIZE)
402 + mask = cols < max_num_logits_per_req
403 +
404 + token_ids = tl.load(row_ptr + cols, mask=mask)
405 + tl.store(
406 + sampled_token_ids_ptr + req_idx * max_num_logits_per_req + cols,
407 + token_ids,
408 + mask=mask,
409 + )
410 + num_sampled = tl.load(row_ptr + max_num_logits_per_req)
411 + tl.store(num_sampled_ptr + req_idx, num_sampled.to(tl.int32))
412 + num_rejected = tl.load(row_ptr + max_num_logits_per_req + 1)
413 + tl.store(num_rejected_ptr + req_idx, num_rejected.to(tl.int32))
414 + if num_nans_ptr is not None:
415 + num_nans = tl.load(row_ptr + max_num_logits_per_req + 2)
416 + tl.store(num_nans_ptr + req_idx, num_nans.to(tl.int32))
417 +
418 +
419 +@triton.jit(
420 + do_not_specialize=[
421 + "max_num_logits_per_req",
422 + "num_logprob_cols",
423 + "num_src_cols",
424 + "send_ids_stride",
425 + "send_logprobs_stride",
426 + "src_ids_stride",
427 + "src_logprobs_stride",
428 + ]
429 +)
430 +def _pack_logprobs_kernel(
431 + send_ids_ptr,
432 + send_ids_stride,
433 + send_logprobs_ptr,
434 + send_logprobs_stride,
435 + src_ids_ptr,
436 + src_ids_stride,
437 + src_logprobs_ptr,
438 + src_logprobs_stride,
439 + selected_token_ranks_ptr,
440 + local_cu_num_logits_ptr,
441 + max_num_logits_per_req,
442 + num_logprob_cols,
443 + num_src_cols,
444 + BLOCK_SIZE: tl.constexpr,
445 +):
446 + req_idx = tl.program_id(0)
447 + pos = tl.program_id(1)
448 + start = tl.load(local_cu_num_logits_ptr + req_idx)
449 + if pos < tl.load(local_cu_num_logits_ptr + req_idx + 1) - start:
450 + src = start + pos
451 + dst = req_idx * max_num_logits_per_req + pos
452 + cols = tl.arange(0, BLOCK_SIZE)
453 + # A rank whose sub-batch has no draft tokens runs the regular sampler
454 + # while others run the rejection sampler, and their column counts can
455 + # differ (e.g. logprob_token_ids widening); the two masks pad or
456 + # truncate to the agreed width.
457 + src_mask = cols < num_src_cols
458 + dst_mask = cols < num_logprob_cols
459 + ids = tl.load(src_ids_ptr + src * src_ids_stride + cols, mask=src_mask, other=0)
460 + logprobs = tl.load(
461 + src_logprobs_ptr + src * src_logprobs_stride + cols,
462 + mask=src_mask,
463 + other=float("-inf"),
464 + )
465 + ids_row = send_ids_ptr + dst * send_ids_stride
466 + tl.store(ids_row + cols, ids.to(tl.int64), mask=dst_mask)
467 + rank = tl.load(selected_token_ranks_ptr + src)
468 + tl.store(ids_row + num_logprob_cols, rank.to(tl.int64))
469 + tl.store(
470 + send_logprobs_ptr + dst * send_logprobs_stride + cols,
471 + logprobs.to(tl.float32),
472 + mask=dst_mask,
473 + )
474 +
475 +
476 +@triton.jit(
477 + do_not_specialize=[
478 + "max_num_logits_per_req",
479 + "num_logprob_cols",
480 + "gathered_ids_stride",
481 + "gathered_logprobs_stride",
482 + ]
483 +)
484 +def _unpack_logprobs_kernel(
485 + gathered_ids_ptr,
486 + gathered_ids_stride,
487 + gathered_logprobs_ptr,
488 + gathered_logprobs_stride,
489 + gathered_src_indices_ptr,
490 + cu_num_logits_ptr,
491 + logprob_token_ids_ptr,
492 + logprobs_ptr,
493 + selected_token_ranks_ptr,
494 + max_num_logits_per_req,
495 + num_logprob_cols,
496 + BLOCK_SIZE: tl.constexpr,
497 +):
498 + req_idx = tl.program_id(0)
499 + pos = tl.program_id(1)
500 + start = tl.load(cu_num_logits_ptr + req_idx)
501 + if pos < tl.load(cu_num_logits_ptr + req_idx + 1) - start:
502 + src = tl.load(gathered_src_indices_ptr + req_idx) * max_num_logits_per_req + pos
503 + dst = start + pos
504 + cols = tl.arange(0, BLOCK_SIZE)
505 + mask = cols < num_logprob_cols
506 + src_ids_row = gathered_ids_ptr + src * gathered_ids_stride
507 + ids = tl.load(src_ids_row + cols, mask=mask)
508 + tl.store(logprob_token_ids_ptr + dst * num_logprob_cols + cols, ids, mask=mask)
509 + rank = tl.load(src_ids_row + num_logprob_cols)
510 + tl.store(selected_token_ranks_ptr + dst, rank)
511 + logprobs = tl.load(
512 + gathered_logprobs_ptr + src * gathered_logprobs_stride + cols, mask=mask
513 + )
514 + tl.store(logprobs_ptr + dst * num_logprob_cols + cols, logprobs, mask=mask)
515 +
516 +
517 +def _gather_logprobs_tensors(
518 + local_output: SamplerOutput | None,
519 + metadata: BatchShardMetadata,
520 + global_batch: InputBatch,
521 + local_batch: InputBatch,
522 + logprobs_dims: tuple[int, int],
523 + device: torch.device,
524 +) -> LogprobsTensors:
525 + num_logprobs, max_token_ids = logprobs_dims
526 + num_logprob_cols = 1 + max(num_logprobs, max_token_ids)
527 + max_num_logits_per_req = metadata.max_num_logits_per_req
528 + num_send_rows = metadata.max_num_reqs_per_rank * max_num_logits_per_req
529 + block_size = triton.next_power_of_2(num_logprob_cols)
530 +
531 + # Padded slots are never read back: gathered_src_indices only addresses
532 + # rows that an owner rank wrote.
533 + send_ids_ranks = torch.empty(
534 + num_send_rows, num_logprob_cols + 1, dtype=torch.int64, device=device
535 + )
536 + send_logprobs = torch.empty(
537 + num_send_rows, num_logprob_cols, dtype=torch.float32, device=device
538 + )
539 + if local_output is not None and local_output.logprobs_tensors is not None:
540 + lp = local_output.logprobs_tensors
541 + _pack_logprobs_kernel[(metadata.num_local_reqs, max_num_logits_per_req)](
542 + send_ids_ranks,
543 + send_ids_ranks.stride(0),
544 + send_logprobs,
545 + send_logprobs.stride(0),
546 + lp.logprob_token_ids,
547 + lp.logprob_token_ids.stride(0),
548 + lp.logprobs,
549 + lp.logprobs.stride(0),
550 + lp.selected_token_ranks,
551 + local_batch.cu_num_logits,
552 + max_num_logits_per_req,
553 + num_logprob_cols,
554 + lp.logprob_token_ids.shape[1],
555 + BLOCK_SIZE=block_size,
556 + )
557 +
558 + gathered_ids_ranks = tensor_model_parallel_all_gather(send_ids_ranks, dim=0)
559 + gathered_logprobs = tensor_model_parallel_all_gather(send_logprobs, dim=0)
560 +
561 + num_reqs = global_batch.num_reqs
562 + num_logits = int(global_batch.cu_num_logits_np[-1])
563 + logprob_token_ids = torch.empty(
564 + num_logits, num_logprob_cols, dtype=torch.int64, device=device
565 + )
566 + logprobs = torch.empty(
567 + num_logits, num_logprob_cols, dtype=torch.float32, device=device
568 + )
569 + selected_token_ranks = torch.empty(num_logits, dtype=torch.int64, device=device)
570 + _unpack_logprobs_kernel[(num_reqs, max_num_logits_per_req)](
571 + gathered_ids_ranks,
572 + gathered_ids_ranks.stride(0),
573 + gathered_logprobs,
574 + gathered_logprobs.stride(0),
575 + metadata.gathered_src_indices,
576 + global_batch.cu_num_logits,
577 + logprob_token_ids,
578 + logprobs,
579 + selected_token_ranks,
580 + max_num_logits_per_req,
581 + num_logprob_cols,
582 + BLOCK_SIZE=block_size,
583 + )
584 + return LogprobsTensors(
585 + logprob_token_ids=logprob_token_ids,
586 + logprobs=logprobs,
587 + selected_token_ranks=selected_token_ranks,
588 + cu_num_generated_tokens=(
589 + global_batch.cu_num_logits_np.tolist() if num_logits != num_reqs else None
590 + ),
591 + )
592 +
593 +
594 +def gather_sampler_output(
595 + local_output: SamplerOutput | None,
596 + metadata: BatchShardMetadata,
597 + device: torch.device,
598 + global_batch: InputBatch,
599 + local_batch: InputBatch,
600 + gather_num_nans: bool = False,
601 + logprobs_dims: tuple[int, int] | None = None,
602 +) -> SamplerOutput:
603 + max_num_logits_per_req = metadata.max_num_logits_per_req
604 + num_packed_cols = max_num_logits_per_req + 2 + (1 if gather_num_nans else 0)
605 + block_size = triton.next_power_of_2(max_num_logits_per_req)
606 +
607 + # Pack the sampler output tensors (excluding logprobs) into a single
608 + # tensor for the all-gather.
609 + packed = torch.empty(
610 + metadata.max_num_reqs_per_rank,
611 + num_packed_cols,
612 + dtype=torch.int64,
613 + device=device,
614 + )
615 + if local_output is not None:
616 + assert local_output.num_sampled is not None
617 + assert local_output.num_rejected is not None
618 + assert not gather_num_nans or local_output.num_nans is not None
619 + num_src_cols = min(
620 + local_output.sampled_token_ids.shape[1], max_num_logits_per_req
621 + )
622 + _pack_sampler_output_kernel[(metadata.num_local_reqs,)](
623 + packed,
624 + packed.stride(0),
625 + local_output.sampled_token_ids,
626 + local_output.sampled_token_ids.stride(0),
627 + local_output.num_sampled,
628 + local_output.num_rejected,
629 + local_output.num_nans if gather_num_nans else None,
630 + local_batch.cu_num_logits if gather_num_nans else None,
631 + max_num_logits_per_req,
632 + num_src_cols,
633 + BLOCK_SIZE=block_size,
634 + )
635 +
636 + # All-gather the sampler output tensors (excluding logprobs).
637 + gathered = tensor_model_parallel_all_gather(packed, dim=0)
638 + # All-gather the logprobs tensors.
639 + logprobs_tensors = None
640 + if logprobs_dims is not None:
641 + logprobs_tensors = _gather_logprobs_tensors(
642 + local_output, metadata, global_batch, local_batch, logprobs_dims, device
643 + )
644 +
645 + # Unpack the gathered tensor into a sampler output object.
646 + num_reqs = metadata.gathered_src_indices.shape[0]
647 + sampled_token_ids = torch.empty(
648 + num_reqs, max_num_logits_per_req, dtype=torch.int64, device=device
649 + )
650 + num_sampled = torch.empty(num_reqs, dtype=torch.int32, device=device)
651 + num_rejected = torch.empty(num_reqs, dtype=torch.int32, device=device)
652 + num_nans = (
653 + torch.empty(num_reqs, dtype=torch.int32, device=device)
654 + if gather_num_nans
655 + else None
656 + )
657 + _unpack_gathered_output_kernel[(num_reqs,)](
658 + gathered,
659 + gathered.stride(0),
660 + metadata.gathered_src_indices,
661 + sampled_token_ids,
662 + num_sampled,
663 + num_rejected,
664 + num_nans,
665 + max_num_logits_per_req,
666 + BLOCK_SIZE=block_size,
667 + )
668 + return SamplerOutput(
669 + sampled_token_ids=sampled_token_ids,
670 + logprobs_tensors=logprobs_tensors,
671 + num_nans=num_nans,
672 + num_sampled=num_sampled,
673 + num_rejected=num_rejected,
674 + )
vllm/v1/worker/gpu/sample/sampler.py
+20 -8
@@ -103,6 +103,22 @@ class Sampler:
103 if self.trace_replay_state is not None:
104 self.trace_replay_state.apply_staged_writes()
105
106 + def get_logprobs_dims(
107 + self, idx_mapping_np: np.ndarray, include_token_ids: bool = True
108 + ) -> tuple[int, int] | None:
109 + """(num_logprobs, max_per_req_token_ids) for the given requests, or
110 + None when none of them want logprobs."""
111 + max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)
112 + max_token_ids = (
113 + self.logprob_token_ids_state.max_num_token_ids(idx_mapping_np)
114 + if include_token_ids
115 + else 0
116 + )
117 + if max_num_logprobs == NO_LOGPROBS and max_token_ids == 0:
118 + return None
119 + num_logprobs = max_num_logprobs if max_num_logprobs != NO_LOGPROBS else 0
120 + return num_logprobs, max_token_ids
121 +
122 def __call__(
123 self,
124 logits: torch.Tensor,
@@ -120,11 +136,7 @@ class Sampler:
136 # that num_nans is computed before applying penalties and temperature.
137 num_nans = get_num_nans(logits) if self.compute_nans else None
138
123 - max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)
124 - max_per_req_token_ids = self.logprob_token_ids_state.max_num_token_ids(
125 - idx_mapping_np
126 - )
127 - return_logprobs = max_num_logprobs != NO_LOGPROBS or max_per_req_token_ids > 0
139 + logprobs_dims = self.get_logprobs_dims(idx_mapping_np)
140
141 sampled, processed_logits = self.sample(
142 logits,
@@ -134,7 +146,7 @@ class Sampler:
146 pos,
147 input_ids,
148 expanded_local_pos,
137 - return_logprobs=return_logprobs,
149 + return_logprobs=logprobs_dims is not None,
150 )
151
152 if self.trace_replay_state is not None:
@@ -142,12 +154,12 @@ class Sampler:
154 # computed logprobs reflect the real distribution of the forced token.
155 self.trace_replay_state.apply_trace(sampled, idx_mapping)
156
145 - if return_logprobs:
157 + if logprobs_dims is not None:
158 + num_logprobs, max_per_req_token_ids = logprobs_dims
159 if self.logprobs_mode in PROCESSED_LOGPROBS_MODES:
160 logits = processed_logits
161 expanded_logits = logits.shape[0] != idx_mapping_np.shape[0]
162 cu_num_logits = cu_num_logits_np.tolist() if expanded_logits else None
150 - num_logprobs = max_num_logprobs if max_num_logprobs != NO_LOGPROBS else 0
163 logprobs_tensors = compute_topk_scores(
164 logits,
165 num_logprobs,