Skip to content

Commit 57ef38c

Browse files
physicsrobclaude
andcommitted
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

0 commit comments

Comments
 (0)