- add logits_positions to AutoRegressiveLM.forward, gathering rows before the final norm so the lm_head GEMM covers only the positions prefill samples from
- execute_prefill builds last_token_indices up front and passes them in, dropping the post-forward gather of a [tokens, vocab] tensor
- prefill graph warmup passes a single index; decode stays untouched (every row is sampled) and prefill itself runs eager, so graph capture is unaffected
- update the ragged-prefill fake to slice by the received index and add a packed-row exact-equality test
Benchmark: NVIDIA L20 (idle), CUDA 12.8, torch 2.11.0+cu128, 1.2B bf16 checkpoint, 512-token prompts, greedy; prefill B=32: 368.1 -> 323.5 ms (44.5k -> 50.6k tok/s, +13.8%), B=8: 89.3 -> 78.9 ms (+13.2%), B=1: 12.3 -> 11.4 ms (+7.9%); decode step unchanged; full suite: 897 passed
- Convert protocol and MoE test classes to plain functions
- Add real server/engine integration and generate_async tests
- Isolate test model per test and use pytest tmp_path
- Reset FastAPI engine state after inference tests
- Fix generate_async StopIteration handling on Python 3.12
- Fix HF adapter MoE dense/shared and Gemma qk_norm mapping
- Correct dev dependency httpx2 to httpx
- extract preprocessing config factories into tests/data/factories.py
- keep conftest.py fixtures-only; stop importing builders from it
- promote temp_dir fixture to root conftest for cross-directory reuse
- unify duplicate BPE tokenizer builders into build_test_tokenizer
- merge grpo/dpo online e2e tests into one parametrized integration test
- extract engine mock factory and shared model batch builders
- drop local tempfile usage in favor of shared fixtures
No behavior change: 519 tests pass.
- 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
ViperEkura