[feature] one-token decode fast path: skip FlashAttention's prefill-tuned tiling for q_len=1

[joaovictor91123]

Benchmark phase

attention-foundation-v1-rtx5070ti

Workload affected

• prefill
• decode
• both

Current limitation

attention.foundation.attention_forward dispatches every workload — including
the three one-token decode workloads (q_len=1, kv_len ∈ {1024, 4096, 8192}, heads=16, dim=128, fp16, causal=False) — straight to
F.scaled_dot_product_attention, whose fused backends (FlashAttention /
mem-efficient) are tuned for prefill: they tile the q_len × kv_len score
matrix and run an online (streaming) softmax to avoid materializing it.

For q_len = 1 that machinery is pure overhead. The scores are a single
(1 × kv_len) row, so there is nothing to tile and no memory pressure to stream
around. Decode is entirely HBM-bandwidth bound — its cost is reading K and
V once (2 · heads · kv_len · dim elements). The flash tiling, per-tile
rescaling, and kernel-launch path add fixed latency that a plain
(1×d)·(d×kv) → softmax → (1×kv)·(kv×d) sequence (two cuBLAS GEMV-shaped
matmuls that read K/V exactly once) does not pay. This is the same reason
dedicated "flash-decoding" / decode kernels exist alongside prefill flash.

Proposed implementation

A single, self-contained q_len == 1 fast path inside attention_forward:
when the query length is 1, compute attention directly as
softmax(q @ Kᵀ / √d, dim=-1) @ V (fp32 softmax accumulation for numerical
safety, cast back to the input dtype), instead of calling SDPA. All other
shapes — every prefill workload and both guards — fall through unchanged to
the existing F.scaled_dot_product_attention call, so prefill cannot move.

One independently measurable feature: a decode-only dispatch in
attention/foundation.py (plus focused tests). No new dependencies, no changes
to eval/, the manifest, or any protected path.

Correctness risks

• Shape: guarded strictly on q.shape[-2] == 1; q_len ≠ kv_len is expected
  (decode passes the full visible KV cache). Output shape (…, 1, dim) matches
  SDPA.
• Causality: decode is causal=False (a one-token query attends to the whole
  visible cache); the fast path only triggers when causal is False, and falls
  through to SDPA otherwise so a q_len==1, causal=True caller is never mis-handled.
• dtype: fp16 in/out; softmax accumulated in fp32 then cast back — keeps
  relative-Frobenius error well under the manifest's 0.005 and abs under 0.05.
• Finite outputs: standard softmax (max-subtracted) over a non-empty KV row;
  no -inf masking, so no all-masked NaN rows.
• Layouts / ragged: the two guard workloads are prefill (q_len > 1) and
  never enter the fast path; contiguous and strided/non-contiguous inputs are
  handled by torch.matmul/softmax exactly as SDPA handles them.

Expected measurement

• Improve (≥ the 5% minimum, per-workload): decode-b1-h16-kv1024-d128,
  decode-b1-h16-kv4096-d128, decode-b1-h16-kv8192-d128.
• Must not regress (protected): all four causal prefill workloads and both
  guards — they keep the identical SDPA call, so their latency/VRAM are unchanged
  within calibrated noise (≤5% latency, ≤2% VRAM).
• Correctness: finite_output_required, rel-Frobenius ≤ 0.005, abs ≤ 0.05 vs the
  fp32 oracle and the SDPA reference, verified on the protected RTX 5070 Ti run.

Scope

Expected implementation paths: attention/ (a decode branch in
attention/foundation.py) and focused tests/. Do not implement until a
maintainer confirms that the benchmark phase is ACTIVE and applies
status:phase-approved to this issue.