Skip to content

engine: hoist prefill sliding-window mask build out of the per-layer attention call #92

Description

@tlkahn

Context

Stage D0 phase 2 (#88, PR #91) wires sliding-window attention into fwd_attention. On prefill, the window mask (fwd_sliding_window_mask) is built inside the per-layer attention call, so an identical [1,1,q_len,kv_len] mask is rebuilt for every local layer in the model on every prefill chunk.

Deferred in both review rounds of PR #91: the clean fix needs per-step scratch state owned by model_forward, which does not exist yet.

Proposal

Build the mask once per forward step (keyed on q_len / kv_len / window, which are identical across local layers of a uniform- or pattern-windowed model) and pass it down to the per-layer calls. Natural home is a model_forward scratch struct, which llama3 freqs precomputation (D1) will also want.

Related

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions