Fix SM90 WY input buffer reuse before WGMMA completion

[zheyang0825]

The dq, dw, and dk inter-chunk loops release their current input pipeline stages while the newest WGMMA group may still be reading shared memory. The first iteration skips the wait entirely; subsequent iterations use wait_group(1), which permits the current group to remain outstanding. A producer can then overwrite an input stage before its consumer finishes.

Wait for all outstanding WGMMA groups with wait_group(0) before these three releases. Preserve the existing h-buffer lifetime across dq/dw, and update the now-stale overlap comment. This conservative fix sacrifices some overlap between V iterations to make stage reuse safe; it does not change the output epilogue handshake or backward intra kernel.

Validation: Python syntax, Ruff lint/format, and git diff --check passed. The equivalent fix was previously exercised on H200 with CuTeDSL 4.6.2–4.8.0 as part of a broader synchronization patch; GPU tests have not been rerun on this public-tree port. No claim is made that this alone resolves every previously observed dq discrepancy.

Summary by CodeRabbit

• Bug Fixes
  • Improved synchronization during SM90 backward computations to ensure shared-memory stages remain available until processing is complete.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: inclusionAI/cuLA/.coderabbit.yaml
• Review profile: CHILL
• Plan: Advanced
• Run ID: f17b8aea-fda4-4118-bf0e-7edc3874737c

📥 Commits

Reviewing files that changed from the base of the PR and between b1fa803 and b710ec8.

📒 Files selected for processing: 1

• cula/ops/kda/sm90/bwd_wy_dqkg.py

Included review availability: This review used your included allowance. Your plan provides up to 10 included reviews per hour; 9 remain after this review.

📝 Walkthrough

The SM90 backward kernel’s dq, dw, and dk V-loops now call warpgroup.wait_group(0) before releasing their current shared-memory stages to DMA. This replaces conditional waits for wait_group(1).

Changes: SM90 backward kernel

Layer / File(s)	Summary
Wait before V-loop stage release cula/ops/kda/sm90/bwd_wy_dqkg.py	The dq, dw, and dk V-loops wait for all outstanding WGMMA work before releasing their current shared-memory stages to DMA. The h stage remains held for reuse by the dw loop.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix

Merge Risk: ⚪ Minimal · up to b710e

The changed loops wait for WGMMA reads to finish before reusing shared-memory stages, while preserving the h-stage handoff. No actionable merge-blocking risk is established; GPU runtime tests have not been rerun for this port.

Security architecture risk: 🔵 Low · up to b710e

The change strengthens synchronization before input buffers can be reused and does not expand public access or authority. No introduced security defect was identified, but hardware behavior has not been validated on this revision.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

• inferred — The changed exposure is bounded to synchronization timing for existing SM90 backward invocations. The inspected comparison does not widen caller reachability, tensor access, or kernel authority; no cross-tenant exposure was established.

Trust Boundaries and Controls

• observed — Inputs continue through the same wrapper and compiled kernel call chain. The changed control strengthens the existing producer-consumer ownership boundary by postponing stage release until WGMMA completion rather than bypassing a validation or access control.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 1 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title is concise, specific, and accurately describes the main change: preventing SM90 WY input-buffer reuse before WGMMA completion.
Description check		The description clearly explains the bug, fix, scope, synchronization tradeoff, and validation status. It omits the template headings for Related Issues, the pre-commit checklist, Performance, and Rev…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches: 💡 1

📝 Generate docstrings: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

[icavan]

LGTM

[fkuner]

LGTM