onnxruntime
9e9dfb2e - LightningIndexer: bound the cache and score passes by the live rows

Commit
59 days ago
LightningIndexer: bound the cache and score passes by the live rows The cache is sized by max_seq_len / ratio, fixed at export time, so a 256K-capable export serving a 1.5K context leaves 65,280 of its 65,664 rows never written by any step. Two passes paid for all of them every decode step. The cache update copied the whole capacity forward and widened it to float. Slots above last_slot hold nothing any query can reach, so the copy now stops at the live extent, read from the device and therefore exact under graph replay. present_cache above last_slot becomes unspecified, which the schema now says. The scoring GEMM takes its row count on the host, and under capture the host-side clamp cannot bind, so it scored all 65,664 rows into a 100 MB buffer that a reduce pass then read back. A short step is small enough to score and reduce in one kernel that takes its bound from past_lens on the device. A lane owns a head and carries its whole dot product, so a row costs one cross-lane reduction rather than one per head. Prefill keeps the GEMM: its host clamp is exact and its arithmetic is large. 1024 prompt, 512 generation, code corpus, tokens/step identical at 4.6455: target verify 17.401 -> 14.798 ms (-15.0%), speculative step 20.999 -> 18.375 ms (-12.5%), decode 221.2 -> 252.8 tps. Holds across context: verify -13.8% at 1K, -13.9% at 4K, -8.9% at 16K, where the clamp starts to stop binding. MMLU-Pro 800 532/800 both before and after with zero discordant items, and zero rank argmax disagreements in every run.
Author
Parents
Loading