Commit 57ef38c
feat(ops): add attend_argmin_above_in_bucket selection primitive
A reusable content-based attention head that, per query position, selects
the smallest-score earlier row that is valid, in a runtime-requested bucket,
and strictly above a runtime-requested threshold, and reads that row's value.
Bucket = one-hot equality rendezvous; threshold = a "run of 1s up to the
score" table read out by a one-hot column picker.
- Decoupled identity V/O so a wide value never enlarges d_qk
(= 2 + n_buckets + n_thresholds); the compiler splits value over
ceil(d_v/d_head) physical heads.
- Op-local bonuses (_VALIDITY_BONUS / _BUCKET_BONUS / _ABOVE_MATCH_BONUS = 256)
sized just above this op's score swing (_QUERY_GAIN*S = 96); distinct names
so they cannot rebind the inherited 1000-unit _ABOVE_BONUS that
attend_argmin_above_integer reads.
- Documents the supported score-diameter contract (< 32 at the shipped
constants) in the docstring; past it a filter-failing low-score row can
silently outscore a match.
- Corrects two stale "bf16" comments in attention_ops.py that contradicted the
live fp32 (SDPA MATH backend, TF32-off) precision policy.
Tests: 22 exact-math/structure tests in tests/ops/test_attention_ops.py
(filters, boundaries, no-match/degenerate, near-one-hot tolerance, single-column
widths, worst-case score-gap bonus dominance, all-zero selectors, width/V-O and
bonus-magnitude+sizing-invariant regressions, unsafe-scalar vs safe-one-hot
presence recompute) plus 4 compiled fp32 round-trip tests in
tests/compile/forward/test_bucketed_argmin.py (adjacent-score recovery,
wide-value head split with second-head winner pass-through, all-invalid finite,
self-row-only finite). Design rationale in bucketed_argmin.md.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>1 parent e0ca246 commit 57ef38c
4 files changed
Lines changed: 1545 additions & 6 deletions
File tree
- tests
- compile/forward
- ops
- torchwright/ops
0 commit comments