ViperEkura
0c1b7664c1
- Eliminate core/ directory into cache/, runtime/, network/ subpackages plus flat modules
- Split cache.py (647 lines) into cache/{buffer,strategy,pool}.py by layer
- Add explicit ContiguousStrategy, make AllocationStrategy a real ABC
- Move TaskCacheState to cache/strategy.py, drop string forward references
- Rename api/ to network/, server.py to app.py
- Move sample.py into runtime/ alongside executor and graph
- Simplify TaskCacheManager.__init__ to single pool param
- Expose pool.strategy and pool.req_pool as public properties
- Fix KVCache import in attention_backend.py (TYPE_CHECKING guard)
- Fix steady-state decode reading uninitialized position_ids on first step
285 lines
8.9 KiB
Python
285 lines
8.9 KiB
Python
"""Unit tests for inference sampling strategies."""
|
|
|
|
import torch
|
|
|
|
from astrai.inference.runtime.sample import (
|
|
FrequencyPenaltyStrategy,
|
|
SamplingPipeline,
|
|
TemperatureStrategy,
|
|
TopKStrategy,
|
|
TopPStrategy,
|
|
sample,
|
|
)
|
|
|
|
|
|
def test_temperature_scalar():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0]])
|
|
s = TemperatureStrategy(0.5)
|
|
result = s.apply(logits.clone())
|
|
assert torch.allclose(result, logits / 0.5)
|
|
|
|
|
|
def test_temperature_skip_when_one():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0]])
|
|
s = TemperatureStrategy(1.0)
|
|
result = s.apply(logits.clone())
|
|
assert torch.equal(result, logits)
|
|
|
|
|
|
def test_temperature_per_sample_tensor():
|
|
logits = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
|
|
s = TemperatureStrategy(torch.tensor([0.5, 0.5]))
|
|
result = s.apply(logits.clone())
|
|
assert torch.allclose(result, logits / 0.5)
|
|
|
|
|
|
def test_top_k_keeps_top():
|
|
logits = torch.tensor([[0.1, 0.5, 0.3, 0.9, 0.2]])
|
|
s = TopKStrategy(top_k=2)
|
|
result = s.apply(logits.clone(), filter_value=-1e9)
|
|
kept = (result > -1e9).sum().item()
|
|
assert kept == 2
|
|
|
|
|
|
def test_top_k_skip_when_zero():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0]])
|
|
s = TopKStrategy(top_k=0)
|
|
result = s.apply(logits.clone())
|
|
assert torch.equal(result, logits)
|
|
|
|
|
|
def test_top_k_batch_tensor():
|
|
"""Each row respects its own top_k."""
|
|
logits = torch.tensor([[0.1, 0.5, 0.3], [0.9, 0.2, 0.1]])
|
|
s = TopKStrategy(top_k=torch.tensor([2, 1]))
|
|
result = s.apply(logits.clone(), filter_value=-1e9)
|
|
assert (result[0] > -1e9).sum() == 2
|
|
assert (result[1] > -1e9).sum() == 1
|
|
|
|
|
|
def test_top_p_nucleus_filtering():
|
|
logits = torch.tensor([[10.0, 1.0, 1.0, 1.0, 1.0]])
|
|
s = TopPStrategy(top_p=0.5)
|
|
result = s.apply(logits.clone(), filter_value=-1e9)
|
|
kept = (result > -1e9).sum().item()
|
|
assert kept >= 1
|
|
|
|
|
|
def test_top_p_skip_when_one():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0]])
|
|
s = TopPStrategy(top_p=1.0)
|
|
result = s.apply(logits.clone())
|
|
assert torch.equal(result, logits)
|
|
|
|
|
|
def test_top_p_filter_all_except_max_when_zero():
|
|
logits = torch.tensor([[0.1, 0.5, 0.3, 0.9, 0.2]])
|
|
s = TopPStrategy(top_p=0.0)
|
|
result = s.apply(logits.clone(), filter_value=-1e9)
|
|
kept = (result > -1e9).sum().item()
|
|
assert kept == 1
|
|
|
|
|
|
def test_sampling_pipeline_composes_strategies():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]])
|
|
pipeline = SamplingPipeline(
|
|
[
|
|
TemperatureStrategy(0.8),
|
|
TopKStrategy(3),
|
|
TopPStrategy(0.95),
|
|
]
|
|
)
|
|
result = pipeline.apply(logits.clone(), filter_value=-1e9)
|
|
kept = (result > -1e9).sum().item()
|
|
assert 1 <= kept <= 3
|
|
|
|
|
|
def test_sampling_pipeline_sample_returns_valid_token():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]])
|
|
pipeline = SamplingPipeline(
|
|
[
|
|
TemperatureStrategy(0.8),
|
|
TopKStrategy(3),
|
|
TopPStrategy(0.95),
|
|
]
|
|
)
|
|
tokens = pipeline.sample(logits)
|
|
assert tokens.shape == (1,)
|
|
assert 0 <= tokens[0] < logits.size(-1)
|
|
|
|
|
|
def test_module_sample_shortcut():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0]])
|
|
tokens = sample(logits, temperature=0.8, top_k=3, top_p=0.95)
|
|
assert tokens.shape == (1,)
|
|
assert 0 <= tokens[0] < logits.size(-1)
|
|
|
|
|
|
def test_module_sample_batch():
|
|
logits = torch.tensor(
|
|
[
|
|
[1.0, 2.0, 3.0, 4.0, 5.0],
|
|
[5.0, 4.0, 3.0, 2.0, 1.0],
|
|
]
|
|
)
|
|
tokens = sample(logits, temperature=0.8, top_k=3, top_p=0.95)
|
|
assert tokens.shape == (2,)
|
|
for t in tokens:
|
|
assert 0 <= t < logits.size(-1)
|
|
|
|
|
|
def test_frequency_penalty_noop_when_zero():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0]])
|
|
input_ids = torch.tensor([[0, 2]])
|
|
s = FrequencyPenaltyStrategy(penalty=0.0)
|
|
result = s.apply(logits.clone(), input_ids=input_ids)
|
|
assert torch.equal(result, logits)
|
|
|
|
|
|
def test_frequency_penalty_noop_when_no_input_ids():
|
|
logits = torch.tensor([[1.0, 2.0, 3.0]])
|
|
s = FrequencyPenaltyStrategy(penalty=0.5)
|
|
result = s.apply(logits.clone())
|
|
assert torch.equal(result, logits)
|
|
|
|
|
|
def test_frequency_penalty_single_occurrence():
|
|
logits = torch.tensor([[4.0, 1.0, 2.0]])
|
|
input_ids = torch.tensor([[0, 2]])
|
|
input_mask = torch.tensor([[True, True]])
|
|
s = FrequencyPenaltyStrategy(penalty=0.5)
|
|
result = s.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask)
|
|
assert result[0, 0] == 3.5
|
|
assert result[0, 1] == 1.0
|
|
assert result[0, 2] == 1.5
|
|
|
|
|
|
def test_frequency_penalty_multiple_occurrences():
|
|
logits = torch.tensor([[4.0, 1.0, 2.0]])
|
|
input_ids = torch.tensor([[0, 2, 0]])
|
|
input_mask = torch.tensor([[True, True, True]])
|
|
s = FrequencyPenaltyStrategy(penalty=0.5)
|
|
result = s.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask)
|
|
assert result[0, 0] == 3.0
|
|
assert result[0, 1] == 1.0
|
|
assert result[0, 2] == 1.5
|
|
|
|
|
|
def test_frequency_penalty_respects_padding_mask():
|
|
logits = torch.tensor([[4.0, 1.0, 2.0]])
|
|
input_ids = torch.tensor([[0, 2, 0]])
|
|
input_mask = torch.tensor([[True, True, False]])
|
|
s = FrequencyPenaltyStrategy(penalty=0.5)
|
|
result = s.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask)
|
|
assert result[0, 0] == 3.5
|
|
assert result[0, 1] == 1.0
|
|
assert result[0, 2] == 1.5
|
|
|
|
|
|
def test_frequency_penalty_batch_tensor():
|
|
logits = torch.tensor(
|
|
[
|
|
[4.0, 1.0, 2.0],
|
|
[3.0, 5.0, 1.0],
|
|
]
|
|
)
|
|
input_ids = torch.tensor([[0, 2, 0], [1, 1, 0]])
|
|
input_mask = torch.tensor([[True, True, True], [True, True, False]])
|
|
s = FrequencyPenaltyStrategy(penalty=torch.tensor([0.5, 1.0]))
|
|
result = s.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask)
|
|
assert result[0, 0] == 3.0
|
|
assert result[0, 2] == 1.5
|
|
assert result[1, 1] == 3.0
|
|
|
|
|
|
def test_frequency_penalty_negative_penalty_boosts_repeats():
|
|
logits = torch.tensor([[4.0, 1.0, 2.0]])
|
|
input_ids = torch.tensor([[0, 0]])
|
|
input_mask = torch.tensor([[True, True]])
|
|
s = FrequencyPenaltyStrategy(penalty=-0.5)
|
|
result = s.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask)
|
|
assert result[0, 0] == 5.0
|
|
|
|
|
|
def test_frequency_penalty_in_pipeline():
|
|
logits = torch.tensor([[5.0, 1.0, 2.0, 3.0]])
|
|
input_ids = torch.tensor([[0, 2, 0]])
|
|
input_mask = torch.tensor([[True, True, True]])
|
|
pipeline = SamplingPipeline(
|
|
[
|
|
TemperatureStrategy(1.0),
|
|
FrequencyPenaltyStrategy(0.5),
|
|
]
|
|
)
|
|
result = pipeline.apply(logits.clone(), input_ids=input_ids, input_mask=input_mask)
|
|
assert result[0, 0] == 4.0
|
|
assert result[0, 2] == 1.5
|
|
|
|
|
|
def test_sample_with_frequency_penalty():
|
|
logits = torch.tensor([[5.0, 1.0, 2.0, 3.0]])
|
|
input_ids = torch.tensor([[0, 2, 0]])
|
|
input_mask = torch.tensor([[True, True, True]])
|
|
tokens = sample(
|
|
logits,
|
|
temperature=1.0,
|
|
top_k=0,
|
|
top_p=1.0,
|
|
frequency_penalty=0.5,
|
|
input_ids=input_ids,
|
|
input_mask=input_mask,
|
|
)
|
|
assert tokens.shape == (1,)
|
|
assert 0 <= tokens[0] < logits.size(-1)
|
|
|
|
|
|
def test_sample_return_logprobs_shape():
|
|
"""``return_logprobs=True`` returns ``[batch]`` logprobs aligned to tokens."""
|
|
logits = torch.tensor([[1.0, 2.0, 3.0], [3.0, 2.0, 1.0]])
|
|
out = sample(logits, temperature=1.0, return_logprobs=True)
|
|
tokens, logprobs = out
|
|
assert tokens.shape == (2,)
|
|
assert logprobs.shape == (2,)
|
|
|
|
|
|
def test_sample_return_logprobs_nonpositive():
|
|
"""Probabilities never exceed 1, so logprobs are always ≤ 0."""
|
|
torch.manual_seed(0)
|
|
logits = torch.randn(4, 50)
|
|
_, logprobs = sample(
|
|
logits, temperature=0.8, top_k=20, top_p=0.9, return_logprobs=True
|
|
)
|
|
assert torch.all(logprobs <= 1e-5)
|
|
|
|
|
|
def test_sample_return_logprobs_greedy_path():
|
|
"""Greedy decode (temperature 0) also returns logprobs."""
|
|
logits = torch.tensor([[1.0, 5.0, 2.0]])
|
|
tokens, logprobs = sample(logits, temperature=0.0, return_logprobs=True)
|
|
assert tokens[0].item() == 1
|
|
# log p(token=1) should equal log_softmax(logits)[1]
|
|
expected = torch.log_softmax(logits.float(), dim=-1)[0, 1]
|
|
assert torch.allclose(logprobs[0], expected, atol=1e-5)
|
|
|
|
|
|
def test_sample_return_logprobs_matches_manual_computation():
|
|
"""Returned logprob equals log_softmax(transformed_logits)[token]."""
|
|
torch.manual_seed(1)
|
|
logits = torch.randn(2, 30)
|
|
tokens, logprobs = sample(logits, temperature=0.7, top_p=0.95, return_logprobs=True)
|
|
# Recompute with the same pipeline
|
|
from astrai.inference.runtime.sample import (
|
|
SamplingPipeline,
|
|
TemperatureStrategy,
|
|
TopPStrategy,
|
|
)
|
|
|
|
pipeline = SamplingPipeline([TemperatureStrategy(0.7), TopPStrategy(0.95)])
|
|
transformed = pipeline.apply(logits.clone())
|
|
expected = torch.gather(
|
|
torch.log_softmax(transformed.float(), dim=-1),
|
|
-1,
|
|
tokens.unsqueeze(-1),
|
|
).squeeze(-1)
|
|
assert torch.allclose(logprobs, expected, atol=1e-5)
|