[FA4] Harden block_logit: scaled logit, pack_gqa guard, docs, tests

[ForFishes]

Summary

Hardens the HySparse block-score (block_logit) fusion in the SM100/Blackwell FA4 forward kernel. Commits on this branch:

1. 4a7d31a Fuse HySparse block-score (block_logit) into the SM100 forward epilogue.
2. 6f7e843 Harden the fused block-score: guards, docs, and an accuracy/contract test suite.
3. ef194c8 Unify block_logit to the scaled logit + reject pack_gqa.

What changed

• Scaled logit semantics: block_logit now stores the SCALED attention logit softmax_scale * q·kᵀ (+ score_mod bias) — the exact value fed into softmax — instead of the raw q·kᵀ. This puts every head on one head-independent scale, so downstream block_logit - LSE gives log(max attention weight in the block), comparable across heads for cross-head Top-K block selection (HySparse eq(3)).
  • The score_mod path already folds softmax_scale into the logit (apply_score_mod_inner); the no-score_mod path carries raw q·kᵀ, so we multiply by softmax_scale (= scale_log2 · ln2) there only. -inf masked columns stay -inf.
• pack_gqa guard: assert not pack_gqa when block_logit is requested. block_logit is indexed by the query head (head_idx) and query row (m_block*m_block_size+tidx); under pack_gqa=True the query heads are packed into the M/row dim and head_idx becomes the KV head, so writes would land in wrong locations. Since pack_gqa defaults to (qhead_per_kvhead > 1), GQA callers must pass pack_gqa=False explicitly.
• Normalization stays downstream: the kernel keeps returning block_logit + LSE; no in-kernel normalization.

Downstream: computing the unified-scale Top-K block index

The kernel returns two tensors that together let you recover a head-independent per-block score:

• block_logit : [B, H, S, num_blocks] fp32 — per-(query, key-block) max of the scaled logit
  m̃_block = max_{j∈block} ( softmax_scale · qᵀk_j (+ bias) ).
• lse : [B, H, S] — the softmax normalizer over the full row (natural log):
  LSE = m_i + ln(ℓ_i), where m_i = max_j s_ij and ℓ_i = Σ_j exp(s_ij − m_i).

Why raw block_logit is NOT cross-head comparable. The raw per-block max logit has a per-row/per-head offset (each head has its own m_i, ℓ_i), so a larger logit in head A does not mean a larger attention weight than head B. You cannot max raw logits across heads.

The unifying step — subtract LSE (log-domain normalization). Because LSE is the exact softmax denominator of the same scaled logits,

log_weight = block_logit − LSE[..., None]
           = m̃_block − (m_i + ln ℓ_i)
           = ln( max_{j∈block} exp(s_ij − m_i) / ℓ_i )
           = ln( max_{j∈block} softmax_j )      # HySparse eq(3), in log space

i.e. log_weight is the log of the largest softmax attention weight inside the block. Softmax weights are probabilities in (0, 1], so log_weight ∈ (−∞, 0] lives on one common scale for every head → it can be maxed/compared across heads. exp(log_weight) recovers the actual peak weight. Masked / unvisited blocks are −inf and are never selected.

Top-K index (cross-head). log is monotonic, so Top-K on log_weight equals Top-K on the true weight — no need to exp first.

out, lse = flash_attn_fwd(..., block_logit=block_logit, block_size=BS, pack_gqa=False)
# block_logit: [B, H, S, num_blocks],  lse: [B, H, S]

log_weight = block_logit - lse.unsqueeze(-1)      # [B,H,S,nb] = log(max softmax weight/block)
block_score = log_weight.max(axis=1)              # [B,S,nb]  max across heads (unified scale)
topk_idx = block_score.topk(k, axis=-1)[1]        # [B,S,k]   selected key-block indices

Equivalent in the weight domain (same result): weight = paddle.exp(log_weight), then max-over-heads and topk.

Notes:

• lse is natural-log and taken over the SCALED logits, matching block_logit's units exactly.
• Caller must pre-fill block_logit with -inf for blocks the causal loop never visits (documented contract; see test_block_logit_unvisited_future_blocks_untouched_causal).

Tests

tests/test_block_score_fusion.py:

• References now scale by softmax_scale; tolerances scale accordingly.
• New negative test test_block_logit_rejects_pack_gqa (explicit pack_gqa=True and defaulted None) verifies the assert fires.
• Accuracy / contract tests: multi-dim, block-size variants, corner cases (ragged seqlens, partial blocks, single token, seqlen_q≠seqlen_k), does-not-change-output, GQA, large multi-tile non-causal, and unvisited-future-block causal contract. All pass (pack_gqa=False).

Runs only on SM 10.x (Blackwell).