Skip to content

rocm/gfx1151: grad_clip_norm lane + flash_attn-backward perf ratchet baseline - #355

Merged
gstoner merged 2 commits into
mainfrom
rocm-gfx1151-grad-clip-perf-ladders
Jul 10, 2026
Merged

gstoner merged 2 commits into
mainfrom
rocm-gfx1151-grad-clip-perf-ladders

Conversation

@gstoner

@gstoner gstoner commented Jul 10, 2026

Copy link
Copy Markdown
Owner

The two follow-ups called out in PR #354's inventory §9.

grad_clip_norm (§5.5)

Global gradient-norm clipping g * min(1, max_norm/||g||) as a composed lane on both x86 and rocm: the L2 norm's global sum-of-squares runs on the device reduce kernel (the FLOP-heavy O(n) part, reusing _device_reduce_sum_rows), host does sqrt + the clip scale; norm_type=inf uses max|g|. Single-tensor form — the atomic unit of the pytree clip_grad_norm.

  • New rocm_grad_clip_compiled / x86_grad_clip_compiled lanes + execution-matrix rows + manifest entries + fixtures.
  • Matches optim.clip_grad_norm within f32 tolerance (the reference accumulates in f64).
  • 7/7 gfx1151, 4/4 x86.

Perf ladder — flash_attn backward ratchet

Extends the gfx1151 hot-path ratchet baseline (rocm_gfx1151_hot_paths.json) to cover the rocm_flash_attn_bwd_compiled lane. Measured on the box: ~44 ms / 87 ms for 1x8x512x64 / 1x16x1024x128≈15× the forward, honestly a correctness-first regression floor, not an MFU claim (Decision #26). The per-row ratchet (test_rocm_perf_ratchet) now re-times the backward rows too.

This is the honest reading of "perf ladders / MFU sign-off": the compiled lanes now carry measured baselines, but stay correctness-first — a real MFU ladder (register/LDS blocking) is where ROCm's lead-performance-target status (Decision #28) gets earned, and is future work.

Verification

  • 163-test regression sweep green (grad_clip x86+rocm, perf ratchet, matrix, inventory guard, normcompose, coverage).
  • Generated-doc drift gate in sync; ruff + mypy ratchet clean.

🤖 Generated with Claude Code

…baseline

The two follow-ups from PR #354's inventory §9.

grad_clip_norm (§5.5) — global gradient-norm clipping g*min(1, max_norm/||g||)
as a composed lane on BOTH x86 and rocm: the L2 norm's global sum-of-squares
runs on the device reduce kernel (the FLOP-heavy O(n) part, reusing
_device_reduce_sum_rows), host does sqrt + the clip scale; norm_type=inf uses
max|g|. Single-tensor form (the atomic unit of the pytree clip_grad_norm).
New rocm_grad_clip_compiled / x86_grad_clip_compiled lanes + execution-matrix
rows + manifest entries + fixtures. Matches optim.clip_grad_norm within f32
tolerance (7/7 gfx1151, 4/4 x86).

Perf ladder — extend the gfx1151 hot-path ratchet baseline to cover the
flash_attn BACKWARD lane (rocm_flash_attn_bwd_compiled). Measured on the box:
~44 ms / 87 ms for 1x8x512x64 / 1x16x1024x128, ~15x the forward — honestly a
correctness-first regression floor, NOT an MFU claim (Decision #26). The
per-row ratchet (test_rocm_perf_ratchet) now re-times the backward rows too.

Verified: 163-test regression sweep green; drift gate in sync; ruff + mypy clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

@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: 4eb62cc410

ℹ️ 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 python/tessera/runtime.py
Comment thread benchmarks/rocm/record_hot_path_baseline.py
…h_attn_bwd

Two PR #355 review fixes (both P2):

- grad_clip_norm: only L2 (norm_type==2) and inf are implemented — the norms
  optim.clip_grad_norm actually distinguishes (it computes tree_l2_norm for
  every finite norm_type). Reject other finite p-norms with a clear diagnostic
  instead of silently clipping by the L2 norm the caller did not ask for
  (Decision #21). + reject tests.

- perf ratchet: the flash_attn_bwd baseline rows were recorded but never
  re-timed live (the live checks only covered {"wmma"}/{"flash_attn"}), so a
  backward regression could not fail CI. Add test_live_flash_attn_bwd_within_
  ratchet (same flash lane gate) that re-times the backward rows against the
  committed baseline.

Verified on gfx1151: 25 grad_clip+ratchet tests + the slow live bwd ratchet
green; ruff + mypy clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@gstoner
gstoner merged commit b176e94 into main Jul 10, 2026
17 checks passed
@gstoner
gstoner deleted the rocm-gfx1151-grad-clip-perf-ladders branch July 10, 2026 20:14
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