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
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 amodel_forwardscratch struct, which llama3 freqs precomputation (D1) will also want.Related