ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: standardize packed 3d inference

Browse Source 
- keep training attention on dense 4d tensors
- use packed 3d tensors with KV cache for inference
- extend CUDA rotary embedding to packed 3d inputs
- adapt torch, CUDA and FlashAttention backend dispatch
This commit is contained in:
ViperEkura
2026-08-16 13:24:02 +08:00
parent 0dd9a417b7
commit 3406157431
15 changed files with 301 additions and 248 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/extension/test_backend_equivalence.py
+22 -42
View File
@@ -59,17 +59,9 @@ 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)
input_ids = torch.tensor(sum(prompt_ids, []), dtype=torch.long, device=device)
position_ids = torch.cat([torch.arange(len(p), device=device) for p in prompt_ids])
cache = PagePool(
n_layers=2,
@@ -88,7 +80,7 @@ def test_prefill_with_kv_cache_matches_torch(cuda_model):
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
input_ids, kv_cache=kv1, position_ids=position_ids, fwd="prefill"
)
task_cache.task_free("t1")
@@ -100,22 +92,24 @@ def test_prefill_with_kv_cache_matches_torch(cuda_model):
with torch.inference_mode():
out_cuda = model(
input_ids,
input_mask=input_mask,
kv_cache=kv2,
position_ids=position_ids,
fwd="prefill",
)
offset = 0
for i, p in enumerate(prompt_ids):
d = (
(
out_torch["logits"][i, : len(p)].float()
- out_cuda["logits"][i, : len(p)].float()
out_torch["logits"][offset : offset + len(p)].float()
- out_cuda["logits"][offset : offset + len(p)].float()
)
.abs()
.max()
.item()
)
assert d == 0.0, f"Prefill diff for sample {i}: {d}"
offset += len(p)
@skip_no_kernel
@@ -136,15 +130,8 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
)
# 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)
input_ids = torch.tensor(sum(prompt_ids, []), dtype=torch.long, device=device)
position_ids = torch.cat([torch.arange(len(p), device=device) for p in prompt_ids])
task_cache = _mk_task_cache(cache)
ws = _ws(cache)
@@ -152,28 +139,22 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
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)
model(input_ids, kv_cache=kv, position_ids=position_ids, fwd="prefill")
# 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)
dec_ids = torch.tensor([99, 98], dtype=torch.long, device=device)
dec_pos = torch.tensor([8, 6], dtype=torch.long, 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
)
out_torch = model(dec_ids, kv_cache=kv_t, position_ids=dec_pos, fwd="decode")
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
)
out_cuda = model(dec_ids, kv_cache=kv_c, position_ids=dec_pos, fwd="decode")
diff = (out_torch["logits"].float() - out_cuda["logits"].float()).abs().max().item()
assert diff < 0.05, f"Decode diff (mixed seq_lens): {diff}"
@@ -198,16 +179,15 @@ def test_decode_cuda_graph_replay_is_exact(cuda_model):
ws = _ws(cache)
task_cache.task_alloc("t1", prompt_ids)
input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=device)
position_ids = torch.arange(len(prompt_ids), device=device).unsqueeze(0)
input_mask = torch.ones(1, len(prompt_ids), dtype=torch.bool, device=device)
input_ids = torch.tensor(prompt_ids, dtype=torch.long, device=device)
position_ids = torch.arange(len(prompt_ids), device=device)
with attn_backend(ATTN_BACKEND.CUDA), torch.inference_mode():
model(
input_ids,
input_mask=input_mask,
position_ids=position_ids,
kv_cache=task_cache.bind(["t1"], ws, start_pos=0),
fwd="prefill",
)
task_cache.task_extend("t1", len(prompt_ids))
@@ -217,16 +197,16 @@ def test_decode_cuda_graph_replay_is_exact(cuda_model):
assert kv_cache.out_cache_loc.dtype == torch.int32
decode_args = {
"input_ids": torch.tensor([[9]], dtype=torch.long, device=device),
"input_mask": torch.ones(1, 1, 64, dtype=torch.bool, device=device),
"position_ids": torch.tensor([[len(prompt_ids)]], device=device),
"input_ids": torch.tensor([9], dtype=torch.long, device=device),
"position_ids": torch.tensor([len(prompt_ids)], device=device),
"kv_cache": kv_cache,
"fwd": "decode",
}
graph = CudaGraphContext(enabled=True)
graph.forward(model, key=(1,), **decode_args)
graph.forward(model, key=(1,), **decode_args)
first = graph.forward(model, key=(1,), **decode_args)["logits"].clone()
slot = kv_cache.out_cache_loc[0, 0]
slot = kv_cache.out_cache_loc[0]
first_k = kv_cache.k_buffer[:, slot].clone()
first_v = kv_cache.v_buffer[:, slot].clone()
tests/inference/test_cache.py
+2 -2
View File
@@ -279,7 +279,7 @@ def test_page_pool_contiguous_bind_tasks_prefill():
task_cache.task_alloc("t1", list(range(10)))
task_cache.task_alloc("t2", list(range(10)))
kv = task_cache.bind(["t1", "t2"], _ws(pool), start_pos=0)
assert kv.out_cache_loc.shape == (2, 10)
assert kv.out_cache_loc.shape == (20,)
assert kv.out_cache_loc.dtype == torch.int32
assert kv.seq_lens.tolist() == [10, 10]
assert kv.req_pool_indices.shape == (2,)
@@ -295,7 +295,7 @@ def test_page_pool_contiguous_bind_tasks_decode():
assert task_cache.task_extend("t1", 10)
assert task_cache.task_extend("t2", 8)
kv = task_cache.bind(["t1", "t2"], _ws(pool))
assert kv.out_cache_loc.shape == (2, 1)
assert kv.out_cache_loc.shape == (2,)
assert kv.seq_lens.tolist() == [11, 9]
tests/module/test_model_forward.py
+46
View File
@@ -39,6 +39,52 @@ def _make_model(config=None) -> AutoRegressiveLM:
return AutoRegressiveLM(config)
def test_model_forward_contract_uses_dense_training_and_packed_inference():
from astrai.inference.cache import PagePool, TaskCacheManager
from astrai.inference.workspace import InferenceWorkspace
config = AutoRegressiveLMConfig(**TINY_CONFIG)
model = AutoRegressiveLM(config).eval()
dense = model(torch.tensor([[1, 2, 3]]))
assert dense["logits"].shape == (1, 3, config.vocab_size)
pool = PagePool(
n_layers=config.num_hidden_layers,
n_kv_heads=config.num_key_value_heads,
head_dim=config.hidden_size // config.num_attention_heads,
max_batch_size=1,
max_seq_len=config.max_position_embeddings,
device="cpu",
dtype=torch.float32,
)
cache = TaskCacheManager(pool)
workspace = InferenceWorkspace(
1,
config.max_position_embeddings,
config.num_attention_heads,
config.hidden_size // config.num_attention_heads,
torch.device("cpu"),
torch.float32,
)
assert cache.task_alloc("t", [1, 2, 3])
packed = model(
torch.tensor([1, 2, 3]),
position_ids=torch.arange(3),
kv_cache=cache.bind(["t"], workspace, start_pos=0),
fwd="prefill",
)
assert packed["logits"].shape == (3, config.vocab_size)
with pytest.raises(ValueError, match="training input_ids"):
model(torch.tensor([1, 2, 3]))
with pytest.raises(ValueError, match="inference input_ids"):
model(
torch.tensor([[1, 2, 3]]),
kv_cache=cache.bind(["t"], workspace, start_pos=0),
fwd="prefill",
)
def _router_stats(probs, topk_indices):
return {"probs": probs, "topk_indices": topk_indices}