Skip to content

docs: Block AttnRes model + ROCm-first plan, with executable contract tests - #552

Merged
gstoner merged 2 commits into
mainfrom
plan/block-attnres-rocm-first
Aug 12, 2026
Merged

gstoner merged 2 commits into
mainfrom
plan/block-attnres-rocm-first

Conversation

@gstoner

@gstoner gstoner commented Aug 12, 2026

Copy link
Copy Markdown
Owner

Summary

Lands the Block AttnRes (arXiv 2603.15031) work as a plan + executable mathematical contract, sequenced ROCm-first per fleet routing (Strix Halo owns the first hardware proof on gfx1151).

  • docs/audit/compiler/BLOCK_ATTNRES_ROCM_PLAN.md — formal model with the paper's gaps filled and marked ([GAP-1..5]): last-block convention, first-of-block partial exclusion (no zero placeholder), attention-based output aggregation, numeric policy in Decision-#15a terms, and the derived VJP (the paper ships no backward). Plus the algorithm set (naive forward, attn_with_stats/softmax_merge primitive pair, exact two-phase schedule, VJP, sharded prefill) and the lowering plan: Graph IR op-set (softmax_merge as the reassociable primitive, return_stats on FlashAttnOp promoting the existing FA-4 LSE checkpoint, thin depth_attn with its decomposition as declared oracle per Decision Apple GPU: Gumbel-max inference sampler (#18-safe, reproducible) #31), ROCm execution lane, then schedule-level transforms. Appendices: sequence-mixer facet mapping (recurrence axis as facet), public-implementation survey (no faithful Block AttnRes exists publicly).
  • tests/unit/test_block_attnres_model.py — 14 pure-numpy tests, the executable contract and future F4 oracles: VJP vs central finite differences (≤1e-7), merge-lemma exactness over random partition trees (atol 1e-12), properties P1–P4 (convexity, zero-init uniformity, key scale-invariance, gain absorption), two-phase == naive across block partitions, the Full-AttnRes S=1 limit, GAP-2 source-count semantics, and input-validation errors.

No registry, catalog, or generated-doc changes — those begin at plan Phase 1.

Test plan

  • python3 -m pytest tests/unit/test_block_attnres_model.py -q → 14 passed (numpy-only; no torch, runs on the Strix box venv and the Mac Homebrew env)
  • ruff check clean on the new test file
  • bash scripts/check_generated_docs.sh → in sync; pre-push spec/freshness gate passes

🤖 Generated with Claude Code

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f579d0f6cd

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread docs/audit/compiler/BLOCK_ATTNRES_ROCM_PLAN.md
Comment thread docs/audit/compiler/BLOCK_ATTNRES_ROCM_PLAN.md
Comment thread docs/audit/compiler/BLOCK_ATTNRES_ROCM_PLAN.md
Comment thread docs/audit/compiler/BLOCK_ATTNRES_ROCM_PLAN.md Outdated
Comment thread docs/audit/compiler/BLOCK_ATTNRES_ROCM_PLAN.md Outdated
gstoner and others added 2 commits August 12, 2026 10:34
… tests

Formal model of Block AttnRes (arXiv 2603.15031) with the paper's gaps
filled explicitly (first-of-block partial exclusion, attention-based output
aggregation, derived VJP eqs. 11-13, numeric policy per Decision #15a), the
algorithm set (naive forward, attn-with-stats + softmax_merge primitive
pair, exact two-phase schedule, VJP), and a lowering plan sequenced
ROCm-first per fleet routing: the mathematical contract lands host-free
(Phases 0-3), the first hardware proof on gfx1151 (Phase 4).

tests/unit/test_block_attnres_model.py is the executable contract: VJP vs
central finite differences (<=1e-7), merge-lemma exactness over random
partition trees (atol 1e-12), properties P1-P4, two-phase == naive across
block partitions, Full-AttnRes limit, and GAP-2 source-count semantics.
These numpy references double as the F4 oracles for the later ROCm/Apple
kernels. No registry or generated-doc changes in this PR (those start at
plan Phase 1).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@gstoner
gstoner force-pushed the plan/block-attnres-rocm-first branch from f579d0f to d34957d Compare August 12, 2026 15:45
@gstoner
gstoner merged commit 18b72c3 into main Aug 12, 2026
18 checks passed
@gstoner
gstoner deleted the plan/block-attnres-rocm-first branch August 12, 2026 22:13
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant