The optimized path's cache: its self-attention K/V arrays per layer, so a pipelined loop can schedule their evaluation with the step outputs.
Readonly
Tokens already in the self-attention cache.
Beam reorder: keep rows indices (take along the batch axis).
indices
The optimized path's cache: its self-attention K/V arrays per layer, so a pipelined loop can schedule their evaluation with the step outputs.