Expose multi-view rollout layout options

[aidanfnv]

Summary

Follow-up to #617 that exposes two layout controls for optimized multi-view inference:

• Forward configurable block-mask geometry through ChunkRollout and run_rollout, including asymmetric (Q, KV) block sizes.
• Allow FixedSlotKVCache to use caller-owned K/V buffers, validating all layers before mutation and preserving aliased prefill data while clearing unused capacity.

Testing

• Multi-view cache and rollout tests: 25 passed.
• Ruff checks, formatting, and ty passed.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[aidanfnv]

/ok to test efb9950

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; the previously reported cache initialization issues are fully addressed and no new actionable failures remain.

Summary

This PR exposes configurable block-mask geometry throughout multi-view rollout and adds validated caller-owned buffers to the fixed-slot KV cache.

• Forwards scalar or asymmetric (Q, KV) block sizes through run_rollout and ChunkRollout.
• Validates all caller-owned cache layers before mutation.
• Copies prefill data before clearing unused capacity, preserving prefill views into the same buffer’s suffix.
• Adds tests for external storage retention, atomic validation, alias-safe initialization, and asymmetric block geometry.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    Caller[Rollout caller] -->|mask_block_size| Run[run_rollout]
    Run --> Chunk[ChunkRollout]
    Chunk --> Metadata[Attention metadata]
    Metadata -->|block_size| Mask[BlockMask]

    Prefill[Prefilled K/V] --> Cache[FixedSlotKVCache]
    Storage[Caller-owned K/V buffers] -->|validate all layers| Cache
    Cache -->|copy prefixes first| Prefix[Initialized cache prefixes]
    Cache -->|clear unused suffixes| Suffix[Zeroed spare capacity]

Loading

_{Reviews (3) · Last reviewed commit: "fix: preserve aliased cache prefill data"}

[aidanfnv]

/ok to test f5d0d1d