wehub-resource-sync
59a0a3844c
PR Test AMD / cancel-on-close (push) Has been skipped
PR Test NVIDIA ARM / scan (push) Has been skipped
PR Test NVIDIA / cancel-on-close (push) Has been skipped
PR Test AMD / scan (push) Has been skipped
PR Test NVIDIA ARM / cancel-on-close (push) Has been skipped
PR Test NVIDIA / scan (push) Has been skipped
Release Docker Images / build (cu129-torch-2.11.0) (push) Has been skipped
Release Docker Images / build (cu130-torch-2.11.0) (push) Has been skipped
Release PyPI / publish (push) Has been skipped
Scheduler Python Test / test (push) Successful in 27m19s
Docs / build (push) Successful in 28m8s
Scheduler C++ Test / test (push) Successful in 28m19s
Scheduler C++ Test / test-flat (push) Successful in 28m18s
Docs / deploy (push) Has been cancelled
PR Test AMD / finish (push) Has been cancelled
PR Test NVIDIA / finish (push) Has been cancelled
PR Test NVIDIA ARM / finish (push) Has been cancelled
PR Test NVIDIA ARM / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test AMD / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test NVIDIA / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
740 lines
26 KiB
Python
740 lines
26 KiB
Python
# Copyright (c) 2026 LightSeek Foundation
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
# copies of the Software, and to permit persons to whom the Software is
|
|
# furnished to do so, subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in
|
|
# all copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
# SOFTWARE.
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
import torch
|
|
from tokenspeed_kernel.ops.sampling.triton import (
|
|
_QRITA_PERCENTILE_TO_STD_TABLE,
|
|
gumbel_sample_from_pools,
|
|
gumbel_sample_from_pools_compact,
|
|
gumbel_sample_from_pools_generic,
|
|
gumbel_sample_min_p_from_pools,
|
|
gumbel_sample_min_p_from_pools_parallel,
|
|
gumbel_sample_top_k_top_p_from_pools,
|
|
gumbel_sample_top_k_top_p_qrita_from_pools,
|
|
gumbel_sample_top_p_parallel_from_pools,
|
|
)
|
|
|
|
# Sentinel matching tokenspeed.runtime.sampling.sampling_params._TOP_K_DISABLED.
|
|
_TOP_K_DISABLED = 1 << 30
|
|
|
|
|
|
def _gumbel_scratch(rows: int, vocab_size: int, device: str):
|
|
num_blocks = (vocab_size + 1023) // 1024
|
|
local_ids = torch.empty((rows, num_blocks), dtype=torch.int32, device=device)
|
|
local_scores = torch.empty((rows, num_blocks), dtype=torch.float32, device=device)
|
|
out = torch.empty((rows,), dtype=torch.int32, device=device)
|
|
return local_ids, local_scores, out
|
|
|
|
|
|
def _top_k_top_p_gumbel_scratch(rows: int, vocab_size: int, device: str):
|
|
num_blocks = (vocab_size + 2047) // 2048
|
|
num_candidates = num_blocks * 128
|
|
candidate_ids = torch.empty(
|
|
(rows, num_candidates), dtype=torch.int32, device=device
|
|
)
|
|
candidate_logits = torch.empty(
|
|
(rows, num_candidates), dtype=torch.float32, device=device
|
|
)
|
|
out = torch.empty((rows,), dtype=torch.int32, device=device)
|
|
return candidate_ids, candidate_logits, out
|
|
|
|
|
|
def _top_p_parallel_scratch(
|
|
rows: int, vocab_size: int, block_size: int, attempts: int, device: str
|
|
):
|
|
num_blocks = (vocab_size + block_size - 1) // block_size
|
|
local_max = torch.empty((rows, num_blocks), dtype=torch.float32, device=device)
|
|
local_sum = torch.empty_like(local_max)
|
|
local_argmax = torch.empty((rows, num_blocks), dtype=torch.int32, device=device)
|
|
local_scores = torch.empty(
|
|
(rows, num_blocks, attempts), dtype=torch.float32, device=device
|
|
)
|
|
local_logits = torch.empty_like(local_scores)
|
|
local_ids = torch.empty(
|
|
(rows, num_blocks, attempts), dtype=torch.int32, device=device
|
|
)
|
|
row_max = torch.empty((rows,), dtype=torch.float32, device=device)
|
|
row_total = torch.empty_like(row_max)
|
|
row_argmax = torch.empty((rows,), dtype=torch.int32, device=device)
|
|
row_candidate_logits = torch.empty(
|
|
(rows, attempts), dtype=torch.float32, device=device
|
|
)
|
|
row_candidate_ids = torch.empty((rows, attempts), dtype=torch.int32, device=device)
|
|
accepted = torch.empty((rows,), dtype=torch.int32, device=device)
|
|
out = torch.empty((rows,), dtype=torch.int32, device=device)
|
|
return (
|
|
local_max,
|
|
local_sum,
|
|
local_argmax,
|
|
local_scores,
|
|
local_logits,
|
|
local_ids,
|
|
row_max,
|
|
row_total,
|
|
row_argmax,
|
|
row_candidate_logits,
|
|
row_candidate_ids,
|
|
accepted,
|
|
out,
|
|
)
|
|
|
|
|
|
def _qrita_gumbel_scratch(rows: int, vocab_size: int, device: str):
|
|
buffer = torch.empty((rows, vocab_size), dtype=torch.float32, device=device)
|
|
table = torch.tensor(
|
|
_QRITA_PERCENTILE_TO_STD_TABLE, dtype=torch.float32, device=device
|
|
)
|
|
out = torch.empty((rows,), dtype=torch.int32, device=device)
|
|
return buffer, table, out
|
|
|
|
|
|
def _top_k_topp_allowed_ids(
|
|
logits: torch.Tensor, temperature: float, top_k: int, top_p: float
|
|
) -> set[int]:
|
|
top_k = min(top_k, logits.numel())
|
|
top_vals, top_ids = torch.topk(logits.float(), top_k, sorted=True)
|
|
probs = torch.softmax(top_vals / temperature, dim=-1)
|
|
cumulative_before = torch.cumsum(probs, dim=-1) - probs
|
|
keep = cumulative_before < top_p
|
|
return set(top_ids[keep].int().cpu().tolist())
|
|
|
|
|
|
def _top_k_topp_minp_allowed_ids(
|
|
logits: torch.Tensor, temperature: float, top_k: int, top_p: float, min_p: float
|
|
) -> set[int]:
|
|
top_k = min(top_k, logits.numel())
|
|
top_vals, top_ids = torch.topk(logits.float(), top_k, sorted=True)
|
|
scaled = top_vals / temperature
|
|
probs = torch.softmax(scaled, dim=-1)
|
|
cumulative_before = torch.cumsum(probs, dim=-1) - probs
|
|
keep = (cumulative_before < top_p) & (
|
|
scaled >= scaled.max() + torch.log(torch.tensor(min_p, device=logits.device))
|
|
)
|
|
return set(top_ids[keep].int().cpu().tolist())
|
|
|
|
|
|
def test_gumbel_sample_from_pools_takes_tail_token(device: str) -> None:
|
|
rows, vocab_size = 3, 1025
|
|
logits = torch.full((rows, vocab_size), -10.0, dtype=torch.float32, device=device)
|
|
logits[0, vocab_size - 1] = 1.0e6
|
|
logits[1, vocab_size - 17] = 1.0e6
|
|
logits[2, vocab_size - 33] = 1.0e6
|
|
req_pool_indices = torch.tensor([4, 2, 7], dtype=torch.int32, device=device)
|
|
pool_rows = 9
|
|
temperature_pool = torch.ones((pool_rows,), dtype=torch.float32, device=device)
|
|
seed_pool = torch.arange(123, 123 + pool_rows, dtype=torch.int64, device=device)
|
|
offsets_pool = torch.arange(17, 17 + pool_rows, dtype=torch.int64, device=device)
|
|
local_ids, local_scores, out = _gumbel_scratch(rows, vocab_size, device)
|
|
|
|
sampled = gumbel_sample_from_pools(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
local_ids,
|
|
local_scores,
|
|
out,
|
|
)
|
|
|
|
torch.testing.assert_close(
|
|
sampled.cpu(),
|
|
torch.tensor(
|
|
[vocab_size - 1, vocab_size - 17, vocab_size - 33], dtype=torch.int32
|
|
),
|
|
)
|
|
|
|
compact_out = torch.empty((rows,), dtype=torch.int32, device=device)
|
|
compact = gumbel_sample_from_pools_compact(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
compact_out,
|
|
block_size=1024,
|
|
).clone()
|
|
|
|
torch.testing.assert_close(compact.cpu(), sampled.cpu())
|
|
|
|
|
|
def test_gumbel_no_filter_verify_idx_mapping_matches_expanded_rows(
|
|
device: str,
|
|
) -> None:
|
|
torch.manual_seed(2027)
|
|
bs, n = 2, 3
|
|
rows = bs * n
|
|
vocab_size = 1025
|
|
logits = torch.randn((rows, vocab_size), dtype=torch.float32, device=device) * 2.0
|
|
req_pool_indices = torch.tensor([1, 3], dtype=torch.int32, device=device)
|
|
pool_rows = 5
|
|
temperature_pool = torch.linspace(
|
|
0.8, 1.2, pool_rows, dtype=torch.float32, device=device
|
|
)
|
|
seed_pool = torch.arange(701, 701 + pool_rows, dtype=torch.int64, device=device)
|
|
offsets_pool = torch.arange(19, 19 + pool_rows, dtype=torch.int64, device=device)
|
|
|
|
mapped_out = torch.empty((rows,), dtype=torch.int32, device=device)
|
|
mapped = gumbel_sample_from_pools_compact(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
mapped_out,
|
|
block_size=1024,
|
|
num_tokens_per_req=n,
|
|
).clone()
|
|
|
|
expanded_req = torch.arange(rows, dtype=torch.int32, device=device)
|
|
expanded_temperature = torch.empty((rows,), dtype=torch.float32, device=device)
|
|
expanded_seed = torch.empty((rows,), dtype=torch.int64, device=device)
|
|
expanded_offsets = torch.empty((rows,), dtype=torch.int64, device=device)
|
|
for row in range(rows):
|
|
pool_idx = int(req_pool_indices[row // n].item())
|
|
expanded_temperature[row] = temperature_pool[pool_idx]
|
|
expanded_seed[row] = seed_pool[pool_idx]
|
|
expanded_offsets[row] = offsets_pool[pool_idx] + row % n
|
|
|
|
expanded_out = torch.empty((rows,), dtype=torch.int32, device=device)
|
|
expanded = gumbel_sample_from_pools_compact(
|
|
logits,
|
|
expanded_req,
|
|
expanded_temperature,
|
|
expanded_seed,
|
|
expanded_offsets,
|
|
expanded_out,
|
|
block_size=1024,
|
|
).clone()
|
|
|
|
torch.testing.assert_close(mapped, expanded)
|
|
|
|
|
|
def test_gumbel_generic_mixed_batch_samples_allowed_set(device: str) -> None:
|
|
torch.manual_seed(321)
|
|
rows, vocab_size = 4, 257
|
|
logits = torch.randn((rows, vocab_size), dtype=torch.float32, device=device) * 2.5
|
|
req_pool_indices = torch.tensor([1, 2, 3, 4], dtype=torch.int32, device=device)
|
|
pool_rows = 5
|
|
temperature_pool = torch.linspace(
|
|
0.8, 1.2, pool_rows, dtype=torch.float32, device=device
|
|
)
|
|
top_k_pool = torch.tensor(
|
|
[1, _TOP_K_DISABLED, 16, _TOP_K_DISABLED, 32],
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
top_p_pool = torch.tensor(
|
|
[1.0, 1.0, 0.9, 0.75, 0.85],
|
|
dtype=torch.float32,
|
|
device=device,
|
|
)
|
|
seed_pool = torch.arange(99, 99 + pool_rows, dtype=torch.int64, device=device)
|
|
offsets_pool = torch.arange(7, 7 + pool_rows, dtype=torch.int64, device=device)
|
|
out = torch.empty((rows,), dtype=torch.int32, device=device)
|
|
|
|
first = gumbel_sample_from_pools_generic(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
top_k_pool,
|
|
top_p_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
out,
|
|
).clone()
|
|
second = gumbel_sample_from_pools_generic(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
top_k_pool,
|
|
top_p_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
out,
|
|
).clone()
|
|
|
|
torch.testing.assert_close(first, second)
|
|
for row, token_id in enumerate(first.cpu().tolist()):
|
|
pool_idx = int(req_pool_indices[row].item())
|
|
top_k = int(top_k_pool[pool_idx].item())
|
|
if top_k == _TOP_K_DISABLED:
|
|
top_k = logits.shape[1]
|
|
allowed = _top_k_topp_allowed_ids(
|
|
logits[row],
|
|
float(temperature_pool[pool_idx].item()),
|
|
top_k,
|
|
float(top_p_pool[pool_idx].item()),
|
|
)
|
|
assert token_id in allowed
|
|
|
|
|
|
@pytest.mark.parametrize("top_k,top_p", [(1, 1.0), (8, 1.0), (64, 0.9)])
|
|
def test_gumbel_top_k_top_p_samples_allowed_set(
|
|
device: str, top_k: int, top_p: float
|
|
) -> None:
|
|
torch.manual_seed(42 + top_k)
|
|
rows, vocab_size = 3, 513
|
|
logits = torch.randn((rows, vocab_size), dtype=torch.float32, device=device) * 3.0
|
|
# Mask a very tempting token to prove -inf logits stay out of the sample set.
|
|
logits[:, vocab_size - 1] = float("-inf")
|
|
req_pool_indices = torch.tensor([4, 2, 7], dtype=torch.int32, device=device)
|
|
pool_rows = 9
|
|
temperature_pool = torch.linspace(
|
|
0.7, 1.3, pool_rows, dtype=torch.float32, device=device
|
|
)
|
|
top_k_pool = torch.full((pool_rows,), top_k, dtype=torch.int32, device=device)
|
|
top_p_pool = torch.full((pool_rows,), top_p, dtype=torch.float32, device=device)
|
|
seed_pool = torch.arange(123, 123 + pool_rows, dtype=torch.int64, device=device)
|
|
offsets_pool = torch.arange(17, 17 + pool_rows, dtype=torch.int64, device=device)
|
|
candidate_ids, candidate_logits, out = _top_k_top_p_gumbel_scratch(
|
|
rows, vocab_size, device
|
|
)
|
|
|
|
sampled = gumbel_sample_top_k_top_p_from_pools(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
top_k_pool,
|
|
top_p_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
candidate_ids,
|
|
candidate_logits,
|
|
out,
|
|
)
|
|
|
|
if top_k == 1:
|
|
torch.testing.assert_close(
|
|
sampled, torch.argmax(logits, dim=-1).to(torch.int32)
|
|
)
|
|
return
|
|
|
|
for row, token_id in enumerate(sampled.cpu().tolist()):
|
|
pool_idx = int(req_pool_indices[row].item())
|
|
allowed = _top_k_topp_allowed_ids(
|
|
logits[row],
|
|
float(temperature_pool[pool_idx].item()),
|
|
top_k,
|
|
top_p,
|
|
)
|
|
assert token_id in allowed
|
|
|
|
|
|
def test_gumbel_top_k_top_p_min_p_samples_allowed_set(device: str) -> None:
|
|
torch.manual_seed(43)
|
|
rows, vocab_size = 3, 513
|
|
top_k, top_p, min_p = 64, 0.9, 0.15
|
|
logits = torch.randn((rows, vocab_size), dtype=torch.float32, device=device) * 3.0
|
|
req_pool_indices = torch.tensor([1, 2, 3], dtype=torch.int32, device=device)
|
|
pool_rows = 4
|
|
temperature_pool = torch.linspace(
|
|
0.7, 1.1, pool_rows, dtype=torch.float32, device=device
|
|
)
|
|
top_k_pool = torch.full((pool_rows,), top_k, dtype=torch.int32, device=device)
|
|
top_p_pool = torch.full((pool_rows,), top_p, dtype=torch.float32, device=device)
|
|
min_p_pool = torch.full((pool_rows,), min_p, dtype=torch.float32, device=device)
|
|
seed_pool = torch.arange(321, 321 + pool_rows, dtype=torch.int64, device=device)
|
|
offsets_pool = torch.arange(13, 13 + pool_rows, dtype=torch.int64, device=device)
|
|
candidate_ids, candidate_logits, out = _top_k_top_p_gumbel_scratch(
|
|
rows, vocab_size, device
|
|
)
|
|
|
|
sampled = gumbel_sample_top_k_top_p_from_pools(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
top_k_pool,
|
|
top_p_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
candidate_ids,
|
|
candidate_logits,
|
|
out,
|
|
min_p_pool=min_p_pool,
|
|
)
|
|
|
|
for row, token_id in enumerate(sampled.cpu().tolist()):
|
|
pool_idx = int(req_pool_indices[row].item())
|
|
allowed = _top_k_topp_minp_allowed_ids(
|
|
logits[row],
|
|
float(temperature_pool[pool_idx].item()),
|
|
top_k,
|
|
top_p,
|
|
min_p,
|
|
)
|
|
assert token_id in allowed
|
|
|
|
|
|
def test_gumbel_top_k_top_p_large_vocab_samples_allowed_set(device: str) -> None:
|
|
torch.manual_seed(2029)
|
|
rows, vocab_size = 2, 32768
|
|
logits = torch.randn((rows, vocab_size), dtype=torch.float32, device=device) * 2.0
|
|
logits[:, -1] = float("-inf")
|
|
req_pool_indices = torch.tensor([1, 2], dtype=torch.int32, device=device)
|
|
pool_rows = 3
|
|
temperature_pool = torch.ones((pool_rows,), dtype=torch.float32, device=device)
|
|
top_k_pool = torch.full((pool_rows,), 127, dtype=torch.int32, device=device)
|
|
top_p_pool = torch.full((pool_rows,), 0.9, dtype=torch.float32, device=device)
|
|
seed_pool = torch.arange(31, 31 + pool_rows, dtype=torch.int64, device=device)
|
|
offsets_pool = torch.arange(9, 9 + pool_rows, dtype=torch.int64, device=device)
|
|
candidate_ids, candidate_logits, out = _top_k_top_p_gumbel_scratch(
|
|
rows, vocab_size, device
|
|
)
|
|
|
|
sampled = gumbel_sample_top_k_top_p_from_pools(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
top_k_pool,
|
|
top_p_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
candidate_ids,
|
|
candidate_logits,
|
|
out,
|
|
)
|
|
|
|
for row, token_id in enumerate(sampled.cpu().tolist()):
|
|
pool_idx = int(req_pool_indices[row].item())
|
|
allowed = _top_k_topp_allowed_ids(
|
|
logits[row],
|
|
float(temperature_pool[pool_idx].item()),
|
|
int(top_k_pool[pool_idx].item()),
|
|
float(top_p_pool[pool_idx].item()),
|
|
)
|
|
assert token_id in allowed
|
|
|
|
|
|
def test_gumbel_top_k_top_p_masks_short_tail_block_candidates(device: str) -> None:
|
|
rows, vocab_size = 1, 2050
|
|
logits = torch.arange(vocab_size, dtype=torch.float32, device=device).unsqueeze(0)
|
|
req_pool_indices = torch.tensor([1], dtype=torch.int32, device=device)
|
|
pool_rows = 2
|
|
temperature_pool = torch.ones((pool_rows,), dtype=torch.float32, device=device)
|
|
top_k_pool = torch.full((pool_rows,), 100, dtype=torch.int32, device=device)
|
|
top_p_pool = torch.ones((pool_rows,), dtype=torch.float32, device=device)
|
|
seed_pool = torch.arange(71, 71 + pool_rows, dtype=torch.int64, device=device)
|
|
offsets_pool = torch.arange(5, 5 + pool_rows, dtype=torch.int64, device=device)
|
|
candidate_ids, candidate_logits, out = _top_k_top_p_gumbel_scratch(
|
|
rows, vocab_size, device
|
|
)
|
|
|
|
sampled = gumbel_sample_top_k_top_p_from_pools(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
top_k_pool,
|
|
top_p_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
candidate_ids,
|
|
candidate_logits,
|
|
out,
|
|
)
|
|
|
|
assert not torch.isnan(candidate_logits).any()
|
|
allowed = _top_k_topp_allowed_ids(logits[0], 1.0, 100, 1.0)
|
|
assert int(sampled.item()) in allowed
|
|
|
|
|
|
def test_gumbel_top_k_top_p_qrita_verify_idx_mapping_matches_expanded_rows(
|
|
device: str,
|
|
) -> None:
|
|
torch.manual_seed(5030)
|
|
bs, n = 2, 3
|
|
rows, vocab_size = bs * n, 1025
|
|
logits = torch.randn((rows, vocab_size), dtype=torch.float32, device=device) * 2.0
|
|
req_pool_indices = torch.tensor([1, 3], dtype=torch.int32, device=device)
|
|
pool_rows = 5
|
|
temperature_pool = torch.linspace(
|
|
0.8, 1.2, pool_rows, dtype=torch.float32, device=device
|
|
)
|
|
top_k_pool = torch.full((pool_rows,), 64, dtype=torch.int32, device=device)
|
|
top_p_pool = torch.full((pool_rows,), 0.85, dtype=torch.float32, device=device)
|
|
seed_pool = torch.arange(701, 701 + pool_rows, dtype=torch.int64, device=device)
|
|
offsets_pool = torch.arange(19, 19 + pool_rows, dtype=torch.int64, device=device)
|
|
|
|
mapped_buffer, mapped_table, mapped_out = _qrita_gumbel_scratch(
|
|
rows, vocab_size, device
|
|
)
|
|
mapped = gumbel_sample_top_k_top_p_qrita_from_pools(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
top_k_pool,
|
|
top_p_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
mapped_buffer,
|
|
mapped_table,
|
|
mapped_out,
|
|
num_tokens_per_req=n,
|
|
num_programs=rows,
|
|
).clone()
|
|
|
|
for row, token_id in enumerate(mapped.cpu().tolist()):
|
|
pool_idx = int(req_pool_indices[row // n].item())
|
|
allowed = _top_k_topp_allowed_ids(
|
|
logits[row],
|
|
float(temperature_pool[pool_idx].item()),
|
|
int(top_k_pool[pool_idx].item()),
|
|
float(top_p_pool[pool_idx].item()),
|
|
)
|
|
assert token_id in allowed
|
|
|
|
expanded_req = torch.arange(rows, dtype=torch.int32, device=device)
|
|
expanded_temperature = torch.empty((rows,), dtype=torch.float32, device=device)
|
|
expanded_top_k = torch.empty((rows,), dtype=torch.int32, device=device)
|
|
expanded_top_p = torch.empty((rows,), dtype=torch.float32, device=device)
|
|
expanded_seed = torch.empty((rows,), dtype=torch.int64, device=device)
|
|
expanded_offsets = torch.empty((rows,), dtype=torch.int64, device=device)
|
|
for row in range(rows):
|
|
pool_idx = int(req_pool_indices[row // n].item())
|
|
expanded_temperature[row] = temperature_pool[pool_idx]
|
|
expanded_top_k[row] = top_k_pool[pool_idx]
|
|
expanded_top_p[row] = top_p_pool[pool_idx]
|
|
expanded_seed[row] = seed_pool[pool_idx]
|
|
expanded_offsets[row] = offsets_pool[pool_idx] + row % n
|
|
|
|
expanded_buffer, expanded_table, expanded_out = _qrita_gumbel_scratch(
|
|
rows, vocab_size, device
|
|
)
|
|
expanded = gumbel_sample_top_k_top_p_qrita_from_pools(
|
|
logits,
|
|
expanded_req,
|
|
expanded_temperature,
|
|
expanded_top_k,
|
|
expanded_top_p,
|
|
expanded_seed,
|
|
expanded_offsets,
|
|
expanded_buffer,
|
|
expanded_table,
|
|
expanded_out,
|
|
num_programs=rows,
|
|
).clone()
|
|
|
|
torch.testing.assert_close(mapped, expanded)
|
|
|
|
|
|
def test_gumbel_top_p_parallel_samples_allowed_set(device: str) -> None:
|
|
torch.manual_seed(456)
|
|
rows, vocab_size = 6, 4097
|
|
block_size, attempts = 512, 1
|
|
logits = torch.randn((rows, vocab_size), dtype=torch.float32, device=device) * 2.0
|
|
logits[:, -1] = float("-inf")
|
|
req_pool_indices = torch.tensor(
|
|
[1, 2, 3, 4, 5, 6], dtype=torch.int32, device=device
|
|
)
|
|
pool_rows = 7
|
|
temperature_pool = torch.linspace(
|
|
0.7, 1.2, pool_rows, dtype=torch.float32, device=device
|
|
)
|
|
top_p_pool = torch.full((pool_rows,), 0.25, dtype=torch.float32, device=device)
|
|
seed_pool = torch.arange(333, 333 + pool_rows, dtype=torch.int64, device=device)
|
|
offsets_pool = torch.arange(11, 11 + pool_rows, dtype=torch.int64, device=device)
|
|
first_scratch = _top_p_parallel_scratch(
|
|
rows, vocab_size, block_size, attempts, device
|
|
)
|
|
second_scratch = _top_p_parallel_scratch(
|
|
rows, vocab_size, block_size, attempts, device
|
|
)
|
|
|
|
first = gumbel_sample_top_p_parallel_from_pools(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
top_p_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
*first_scratch,
|
|
block_size=block_size,
|
|
num_attempts=attempts,
|
|
)
|
|
second = gumbel_sample_top_p_parallel_from_pools(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
top_p_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
*second_scratch,
|
|
block_size=block_size,
|
|
num_attempts=attempts,
|
|
)
|
|
|
|
torch.testing.assert_close(first, second)
|
|
|
|
for row, token_id in enumerate(first.cpu().tolist()):
|
|
pool_idx = int(req_pool_indices[row].item())
|
|
allowed = _top_k_topp_allowed_ids(
|
|
logits[row],
|
|
float(temperature_pool[pool_idx].item()),
|
|
vocab_size,
|
|
float(top_p_pool[pool_idx].item()),
|
|
)
|
|
assert token_id in allowed
|
|
|
|
|
|
def test_gumbel_top_p_parallel_verify_idx_mapping_matches_expanded_rows(
|
|
device: str,
|
|
) -> None:
|
|
torch.manual_seed(987)
|
|
bs, n, vocab_size = 2, 3, 521
|
|
rows = bs * n
|
|
block_size, attempts = 128, 3
|
|
logits = torch.randn((rows, vocab_size), dtype=torch.float32, device=device) * 2.0
|
|
req_pool_indices = torch.tensor([2, 4], dtype=torch.int32, device=device)
|
|
pool_rows = 6
|
|
temperature_pool = torch.linspace(
|
|
0.75, 1.25, pool_rows, dtype=torch.float32, device=device
|
|
)
|
|
top_p_pool = torch.full((pool_rows,), 0.8, dtype=torch.float32, device=device)
|
|
seed_pool = torch.arange(700, 700 + pool_rows, dtype=torch.int64, device=device)
|
|
offsets_pool = torch.arange(10, 10 + pool_rows, dtype=torch.int64, device=device)
|
|
mapped_scratch = _top_p_parallel_scratch(
|
|
rows, vocab_size, block_size, attempts, device
|
|
)
|
|
mapped = gumbel_sample_top_p_parallel_from_pools(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
top_p_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
*mapped_scratch,
|
|
block_size=block_size,
|
|
num_attempts=attempts,
|
|
num_tokens_per_req=n,
|
|
).clone()
|
|
|
|
expanded_req_pool_indices = torch.arange(
|
|
1, rows + 1, dtype=torch.int32, device=device
|
|
)
|
|
expanded_temperature = torch.empty((rows + 1,), dtype=torch.float32, device=device)
|
|
expanded_top_p = torch.empty((rows + 1,), dtype=torch.float32, device=device)
|
|
expanded_seed = torch.empty((rows + 1,), dtype=torch.int64, device=device)
|
|
expanded_offsets = torch.empty((rows + 1,), dtype=torch.int64, device=device)
|
|
for row in range(rows):
|
|
req_row = row // n
|
|
spec_pos = row - req_row * n
|
|
src_pool = int(req_pool_indices[req_row].item())
|
|
dst_pool = row + 1
|
|
expanded_temperature[dst_pool] = temperature_pool[src_pool]
|
|
expanded_top_p[dst_pool] = top_p_pool[src_pool]
|
|
expanded_seed[dst_pool] = seed_pool[src_pool]
|
|
expanded_offsets[dst_pool] = offsets_pool[src_pool] + spec_pos
|
|
expanded_scratch = _top_p_parallel_scratch(
|
|
rows, vocab_size, block_size, attempts, device
|
|
)
|
|
expanded = gumbel_sample_top_p_parallel_from_pools(
|
|
logits,
|
|
expanded_req_pool_indices,
|
|
expanded_temperature,
|
|
expanded_top_p,
|
|
expanded_seed,
|
|
expanded_offsets,
|
|
*expanded_scratch,
|
|
block_size=block_size,
|
|
num_attempts=attempts,
|
|
).clone()
|
|
|
|
torch.testing.assert_close(mapped, expanded)
|
|
|
|
|
|
@pytest.mark.parametrize("parallel,vocab_size", [(False, 1025), (True, 4097)])
|
|
def test_gumbel_min_p_samples_allowed_set(
|
|
device: str, parallel: bool, vocab_size: int
|
|
) -> None:
|
|
torch.manual_seed(779)
|
|
rows = 4
|
|
logits = torch.randn((rows, vocab_size), dtype=torch.float32, device=device) * 2.0
|
|
req_pool_indices = torch.tensor([1, 2, 3, 4], dtype=torch.int32, device=device)
|
|
pool_rows = 5
|
|
temperature_pool = torch.linspace(
|
|
0.7, 1.2, pool_rows, dtype=torch.float32, device=device
|
|
)
|
|
min_p_pool = torch.full((pool_rows,), 0.2, dtype=torch.float32, device=device)
|
|
seed_pool = torch.arange(51, 51 + pool_rows, dtype=torch.int64, device=device)
|
|
offsets_pool = torch.arange(7, 7 + pool_rows, dtype=torch.int64, device=device)
|
|
out = torch.empty((rows,), dtype=torch.int32, device=device)
|
|
|
|
if parallel:
|
|
num_blocks = (vocab_size + 1023) // 1024
|
|
local_ids = torch.empty((rows, num_blocks), dtype=torch.int32, device=device)
|
|
local_scores = torch.empty(
|
|
(rows, num_blocks), dtype=torch.float32, device=device
|
|
)
|
|
row_max = torch.empty((rows,), dtype=torch.float32, device=device)
|
|
first = gumbel_sample_min_p_from_pools_parallel(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
min_p_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
local_ids,
|
|
local_scores,
|
|
row_max,
|
|
out,
|
|
).clone()
|
|
second = gumbel_sample_min_p_from_pools_parallel(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
min_p_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
local_ids,
|
|
local_scores,
|
|
row_max,
|
|
out,
|
|
).clone()
|
|
else:
|
|
first = gumbel_sample_min_p_from_pools(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
min_p_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
out,
|
|
).clone()
|
|
second = gumbel_sample_min_p_from_pools(
|
|
logits,
|
|
req_pool_indices,
|
|
temperature_pool,
|
|
min_p_pool,
|
|
seed_pool,
|
|
offsets_pool,
|
|
out,
|
|
).clone()
|
|
|
|
torch.testing.assert_close(first, second)
|
|
for row, token_id in enumerate(first.cpu().tolist()):
|
|
pool_idx = int(req_pool_indices[row].item())
|
|
scaled = logits[row].float() / float(temperature_pool[pool_idx].item())
|
|
probs = torch.softmax(scaled, dim=-1)
|
|
threshold = float(min_p_pool[pool_idx].item()) * probs.max()
|
|
assert probs[token_id] >= threshold
|