CSA compressor: review-response fixups for the ratio=128 kernels (follow-up to #427)

[zkyue]

What

Two review-response fixups from the CodeRabbit round on #427 that were ready on the
review branch but missed the merge window (#427 merged at head d77e48e). Both were
built on exactly that head; the cherry-pick onto develop (b950af1) is clean and the
CSA files are byte-identical to the reviewed pre-merge state, except a one-line
stale-comment fix in the test file (flagged in review; no functional change).

• 678fca3f — bench(csa): use the DSL's documented __ptx__ accessor in the r128
  register probe. fn.artifacts is nvidia-cutlass-dsl's internal
  JitFunctionArtifacts dataclass; the documented accessor on the compiled handle is
  the __ptx__ property (JitCompiledFunction), which is None when PTX retention
  (CUTE_DSL_KEEP=ptx) is off instead of poisoning ptxas input. Kernels that expose
  no PTX are now skipped with a clear message and a nonzero exit (a kernel that
  wasn't probed wasn't verified). Verified against nvidia-cutlass-dsl 4.6.1;
  identical 16-kernel table, ALL 0 spill / 0 stack.

• 037a53ba — CSA compressor: adopt per-bucket grad_out register reuse in the r128
  backward. The backward loaded each row's grad_out vec twice (phases 2 and 4).
  Loading it once and holding the register vec across the phase barriers is now a
  per-bucket schedule field (goreuse), decided by measured time and adopted in
  5 of 6 backward buckets; the one measured loser (c1d128 default) keeps the
  reload and compiles to byte-identical PTX vs the pre-change kernel. dKV/dScore are
  bitwise-identical with and without the field on every schedule; no schedule spills
  either way (0 spill / 0 stack, sm_100a).

Performance (goreuse commit)

Methodology: nsys pure-kernel per-launch durations (kernel time only — no wrapper,
launch, or sync overhead), median of 100 after 30 warmup; both variants interleaved
in one process on the same clocks and the same input tensors; reproduced
order-swapped. One B200. Measured on the review branch at e70f6140, whose CSA files
are byte-identical to this branch apart from the one-line test-comment fix, and the
adopted buckets compile to byte-identical PTX vs the measured kernels — byte-identical
PTX means the measured kernels are the shipped kernels (absolute timings vary within
the session-noise envelope below).

bwd bucket	shape	base → reuse (µs)	delta (run1/run2)	decision
c1d128 small	1x8192	8.45 → 6.27	+25.8% / +25.8%	adopt
c1d128 small	3x8192	16.10 → 11.68	+27.4% / +27.6%	adopt
c1d128 default	1x131072	63.87 → 64.70	−1.3% / −1.6%	keep reload
c2d128 small	1x8192	11.30 → 10.66	+5.7% / +7.1%	adopt
c2d128 default	1x65536	61.25 → 60.32	+1.5% / +1.5%	adopt
c1d512 default	1x8192	17.65 → 17.28	+2.1% / +2.0% (65k tie)	adopt
c2d512 default	1x8192	28.27 → 27.84	+1.5% / +1.9% (65k tie)	adopt

Docs tables updated for the two configurations (both tables) that moved outside
session noise (c1d128 8192: bwd 8.4 → 6.3 µs; c2d128 8192: 11.2 → 10.5 µs
kernel-time); long-context rows measured within the ±3% session-noise envelope of
their published cells and stand.

Responds to CodeRabbit feedback on #427

• Comment 1 (fn.artifacts.PTX is a non-public surface) → fixed in 678fca3f
  as described above.
• Comment 2 (bwd grad_out loaded twice) → measured per schedule bucket and
  adopted where it wins, 037a53ba. The register table was a poor predictor: the
  bucket with the largest growth (c1d128 small, 48 → 64 regs) wins the most because
  its small grid underfills the machine, while the one loser has unchanged registers
  (80 → 80).
• Comment 3 (eager reference duplicated between the gate and the test suite) →
  no code change; the duplication is deliberate and staked in the gate's docstring.
  The gate is the PR's numerics authority, so its reference should be auditable in
  place rather than imported across benchmark/ ↔ test/python (neither is a
  package; conftest import-order constraints). Drift is checked, not trusted: every
  gate run first cross-checks the copied reference against the shipped production
  ratio=4 backward (bitwise dKV/dScore, validate_reference()) and fails before
  gating anything if the copy has drifted.

Validation (re-run on this branch, develop + these 2 commits, B200, 2026-07-30)

• ratio=128 contract gate: PASS 21/21 (3-run determinism, tolerance vs fp32
  eager, fp64-oracle parity, overflow NaN-pattern case; schedule coverage fwd 10/10,
  bwd 6/6 asserted through the dispatch tables, including the goreuse fields).
• pytest -m "L0 or L1" fe_api/csa/test_CSA_compressor.py: 98 passed, 1 skipped.
• register probe (via the fixed __ptx__ path): 16 kernels probed (sm_100a);
  ALL 0 spill / 0 stack.

Summary by CodeRabbit

• New Features

  • Added an optional backward-kernel scheduling setting that can improve register reuse and performance.
  • Updated ratio-128 scheduling to select and validate the new setting across supported configurations.

• Documentation

  • Clarified backward scheduling behavior, launch-time row selection, and register reuse effects.
  • Updated ratio-128 performance measurements.

• Bug Fixes

  • Improved benchmark reporting when kernel PTX is unavailable, ensuring incomplete results are reported as failures.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 12fa5e0e-2e8c-4bd7-9e28-ca3d88f3b6b0

📥 Commits

Reviewing files that changed from the base of the PR and between b950af1 and 037a53b.

📒 Files selected for processing (4)

• benchmark/csa/reg_probe_csa_compressor_r128.py
• docs/fe-oss-apis/csa.md
• python/cudnn/csa/compressor/compressor_sm100_r128.py
• test/python/fe_api/csa/test_CSA_compressor.py

---

📝 Walkthrough

Walkthrough

The ratio-128 CSA backward kernel adds a goreuse scheduling option that controls grad_out register reuse. Schedule tables, dispatch tests, and documentation are updated, while the PTX probe now treats missing retained PTX as a failed outcome.

Changes

Ratio-128 CSA updates

Layer / File(s)	Summary
Goreuse backward-kernel implementation python/cudnn/csa/compressor/compressor_sm100_r128.py	The backward launch and kernel accept goreuse; enabled schedules reuse a registered grad_out, while disabled schedules reload it.
Backward schedule contract and validation python/cudnn/csa/compressor/compressor_sm100_r128.py, test/python/fe_api/csa/test_CSA_compressor.py, docs/fe-oss-apis/csa.md	Schedule selection returns (vec, tchunks, threads_x, fastexp, goreuse), bucket expectations include the new field, and ratio-128 behavior and timing tables are updated.
PTX retention probe verdicts benchmark/csa/reg_probe_csa_compressor_r128.py	Kernels without valid retained PTX are skipped and counted, and the probe exits successfully only when all kernels are clean and none were skipped.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Schedule as _bwd_schedule_r128
  participant Launch as _compressor_bwd_r128_launch
  participant Kernel as _compressor_bwd_r128_kernel
  participant GradOut as mGO
  Schedule->>Launch: select goreuse
  Launch->>Kernel: pass compile-time goreuse
  Kernel->>GradOut: load grad_out
  Kernel->>Kernel: reuse fr_go when goreuse is enabled
  Kernel->>GradOut: reload grad_out when goreuse is disabled

Loading

Suggested labels: mod-cutedsl

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly names the main area and nature of the change and is concise.
Description check	✅ Passed	The description covers the change, why, related review context, validation, and compatibility notes, though it doesn't follow the template headings exactly.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[Anerudhan]

Thanks @zkyue for the request.

Launching the CI for the PR.

[Anerudhan]

@cudnn-ci-bot run

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-452-037a53b
Pipeline: 60919841