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
224 lines
7.1 KiB
Python
224 lines
7.1 KiB
Python
"""Numerical equivalence between TorchNativeBackend and CudaBackend.
|
|
|
|
Covers training forward, inference prefill, inference decode (mixed
|
|
seq_lens with padding mask), and end-to-end scheduler.run_batch.
|
|
"""
|
|
|
|
import torch
|
|
|
|
from astrai.extension import ATTN_BACKEND, attn_backend
|
|
from astrai.inference.cache import PagePool, TaskCacheManager
|
|
from astrai.inference.workspace import InferenceWorkspace
|
|
from tests.extension.conftest import D, skip_no_kernel
|
|
|
|
|
|
def _mk_task_cache(pool: PagePool) -> TaskCacheManager:
|
|
return TaskCacheManager(pool)
|
|
|
|
|
|
def _ws(pool: PagePool) -> InferenceWorkspace:
|
|
return InferenceWorkspace(
|
|
pool.max_batch_size,
|
|
pool.max_seq_len,
|
|
max_q_heads=2,
|
|
head_dim=64,
|
|
device=pool.device,
|
|
dtype=pool.dtype,
|
|
)
|
|
|
|
|
|
@skip_no_kernel
|
|
def test_training_forward_matches_torch(cuda_model):
|
|
"""Training forward (kv_cache=None) uses torch-native SDPA.
|
|
|
|
CudaBackend does not support training (requires kv_cache).
|
|
Torch-native backend must match default (which falls back to torch).
|
|
"""
|
|
|
|
model, _ = cuda_model
|
|
input_ids = torch.randint(0, 1000, (2, 16), device="cuda")
|
|
|
|
with torch.no_grad():
|
|
out_default = model(input_ids)
|
|
|
|
with attn_backend(ATTN_BACKEND.TORCH_NATIVE):
|
|
with torch.no_grad():
|
|
out_torch = model(input_ids)
|
|
|
|
torch.testing.assert_close(
|
|
out_torch["logits"], out_default["logits"], atol=1e-6, rtol=1e-6
|
|
)
|
|
assert out_default["logits"].shape[0] == 2
|
|
|
|
|
|
@skip_no_kernel
|
|
def test_prefill_with_kv_cache_matches_torch(cuda_model):
|
|
"""Inference prefill with KV cache should match torch backend."""
|
|
model, _ = cuda_model
|
|
prompt_ids = [[1, 2, 3, 4, 5, 6, 7, 8], [10, 11, 12, 13, 14, 15]]
|
|
max_len = max(len(p) for p in prompt_ids)
|
|
batch = len(prompt_ids)
|
|
|
|
device = "cuda"
|
|
input_ids = torch.zeros(batch, max_len, dtype=torch.long, device=device)
|
|
input_mask = torch.zeros(batch, max_len, dtype=torch.bool, device=device)
|
|
position_ids = torch.zeros(batch, max_len, dtype=torch.long, device=device)
|
|
for i, p in enumerate(prompt_ids):
|
|
input_ids[i, : len(p)] = torch.tensor(p, device=device)
|
|
input_mask[i, : len(p)] = True
|
|
position_ids[i, : len(p)] = torch.arange(len(p), device=device)
|
|
|
|
cache = PagePool(
|
|
n_layers=2,
|
|
n_kv_heads=1,
|
|
head_dim=D,
|
|
max_batch_size=4,
|
|
max_seq_len=64,
|
|
device=device,
|
|
dtype=torch.bfloat16,
|
|
)
|
|
|
|
task_cache = _mk_task_cache(cache)
|
|
ws = _ws(cache)
|
|
task_cache.task_alloc("t1", prompt_ids[0])
|
|
task_cache.task_alloc("t2", prompt_ids[1])
|
|
kv1 = task_cache.bind(["t1", "t2"], ws, start_pos=0)
|
|
with torch.inference_mode():
|
|
out_torch = model(
|
|
input_ids, input_mask=input_mask, kv_cache=kv1, position_ids=position_ids
|
|
)
|
|
|
|
task_cache.task_free("t1")
|
|
task_cache.task_free("t2")
|
|
task_cache.task_alloc("t1", prompt_ids[0])
|
|
task_cache.task_alloc("t2", prompt_ids[1])
|
|
kv2 = task_cache.bind(["t1", "t2"], ws, start_pos=0)
|
|
with attn_backend(ATTN_BACKEND.CUDA):
|
|
with torch.inference_mode():
|
|
out_cuda = model(
|
|
input_ids,
|
|
input_mask=input_mask,
|
|
kv_cache=kv2,
|
|
position_ids=position_ids,
|
|
)
|
|
|
|
for i, p in enumerate(prompt_ids):
|
|
d = (
|
|
(
|
|
out_torch["logits"][i, : len(p)].float()
|
|
- out_cuda["logits"][i, : len(p)].float()
|
|
)
|
|
.abs()
|
|
.max()
|
|
.item()
|
|
)
|
|
assert d == 0.0, f"Prefill diff for sample {i}: {d}"
|
|
|
|
|
|
@skip_no_kernel
|
|
def test_decode_mixed_seq_lens_matches_torch(cuda_model):
|
|
"""Decode with mixed seq_lens in batch — padding mask must produce correct output."""
|
|
model, _ = cuda_model
|
|
device = "cuda"
|
|
|
|
prompt_ids = [[1, 2, 3, 4, 5, 6, 7, 8], [10, 11, 12, 13, 14, 15]]
|
|
cache = PagePool(
|
|
n_layers=2,
|
|
n_kv_heads=1,
|
|
head_dim=D,
|
|
max_batch_size=4,
|
|
max_seq_len=64,
|
|
device=device,
|
|
dtype=torch.bfloat16,
|
|
)
|
|
|
|
# Prefill to populate cache
|
|
max_len = max(len(p) for p in prompt_ids)
|
|
batch = len(prompt_ids)
|
|
input_ids = torch.zeros(batch, max_len, dtype=torch.long, device=device)
|
|
input_mask = torch.zeros(batch, max_len, dtype=torch.bool, device=device)
|
|
position_ids = torch.zeros(batch, max_len, dtype=torch.long, device=device)
|
|
for i, p in enumerate(prompt_ids):
|
|
input_ids[i, : len(p)] = torch.tensor(p, device=device)
|
|
input_mask[i, : len(p)] = True
|
|
position_ids[i, : len(p)] = torch.arange(len(p), device=device)
|
|
|
|
task_cache = _mk_task_cache(cache)
|
|
ws = _ws(cache)
|
|
task_cache.task_alloc("t1", prompt_ids[0])
|
|
task_cache.task_alloc("t2", prompt_ids[1])
|
|
kv = task_cache.bind(["t1", "t2"], ws, start_pos=0)
|
|
with torch.inference_mode():
|
|
model(input_ids, input_mask=input_mask, kv_cache=kv, position_ids=position_ids)
|
|
|
|
# Decode step — seq_lens are 9 and 7 (after extending)
|
|
dec_ids = torch.tensor([[99], [98]], dtype=torch.long, device=device)
|
|
dec_pos = torch.tensor([[8], [6]], dtype=torch.long, device=device)
|
|
total_len = 9
|
|
dec_mask = dec_pos[:, None, None] >= torch.arange(total_len, device=device)
|
|
|
|
task_cache.task_extend("t1", 8)
|
|
task_cache.task_extend("t2", 6)
|
|
kv_t = task_cache.bind(["t1", "t2"], ws)
|
|
with torch.inference_mode():
|
|
out_torch = model(
|
|
dec_ids, input_mask=dec_mask, kv_cache=kv_t, position_ids=dec_pos
|
|
)
|
|
|
|
kv_c = task_cache.bind(["t1", "t2"], ws)
|
|
with attn_backend(ATTN_BACKEND.CUDA):
|
|
with torch.inference_mode():
|
|
out_cuda = model(
|
|
dec_ids, input_mask=dec_mask, kv_cache=kv_c, position_ids=dec_pos
|
|
)
|
|
|
|
diff = (out_torch["logits"].float() - out_cuda["logits"].float()).abs().max().item()
|
|
assert diff < 0.05, f"Decode diff (mixed seq_lens): {diff}"
|
|
|
|
|
|
@skip_no_kernel
|
|
def test_run_batch_cuda_matches_torch_greedy(cuda_model):
|
|
"""Greedy decode (temperature=0) should produce identical tokens."""
|
|
from astrai.inference.scheduler import InferenceScheduler
|
|
from tests.helpers import FakeTokenizer
|
|
|
|
model, _ = cuda_model
|
|
tokenizer = FakeTokenizer()
|
|
|
|
prompts = [[1, 2, 3, 4, 5], [10, 11, 12, 13, 14, 15, 16]]
|
|
|
|
sched = InferenceScheduler(
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
max_batch_size=4,
|
|
max_seq_len=64,
|
|
device="cuda",
|
|
dtype=torch.bfloat16,
|
|
)
|
|
out_torch = sched.run_batch(prompts, max_tokens=5, temperature=0.0)
|
|
sched.stop()
|
|
|
|
cache_cuda = PagePool(
|
|
n_layers=2,
|
|
n_kv_heads=1,
|
|
head_dim=D,
|
|
max_batch_size=4,
|
|
max_seq_len=64,
|
|
device="cuda",
|
|
dtype=torch.bfloat16,
|
|
)
|
|
sched2 = InferenceScheduler(
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
max_batch_size=4,
|
|
max_seq_len=64,
|
|
device="cuda",
|
|
dtype=torch.bfloat16,
|
|
cache=cache_cuda,
|
|
)
|
|
with attn_backend(ATTN_BACKEND.CUDA):
|
|
out_cuda = sched2.run_batch(prompts, max_tokens=5, temperature=0.0)
|
|
sched2.stop()
|
|
|
|
assert out_torch == out_cuda, f"Torch={out_torch} != CUDA={out_cuda}"
|