Conversation
| } | ||
| ptx::warpgroup_commit_batch(); | ||
| // Keep register A operands live until the asynchronous WGMMA completes. | ||
| ptx::warpgroup_wait<0>(); |
There was a problem hiding this comment.
🔵 suggestion: Placing ptx::warpgroup_wait<0>() after warpgroup_commit_batch() is the correct minimal fix: it keeps the single-buffered register operands a[...] live until the async WGMMA has consumed them. Consequence is that every K stage now fully drains its WGMMA before the next A load, so WGMMA and the next stage's A-load/reduction cannot overlap. This is not a regression versus the buggy code (which already waited each iteration) and correctness must come first, but it is worth a short follow-up note: if the kernel is compute bound, double-buffering a and using warpgroup_wait<1>() would restore overlap while still protecting the register source.
🤖 v4
|
|
||
| ptx::warpgroup_wait<0>(); | ||
| if (s > 0) | ||
| empty_barriers[(s - 1) % kNumStages]->arrive(); |
There was a problem hiding this comment.
🔵 suggestion: With the wait moved to the end of the loop body, empty_barriers[(s - 1) % kNumStages]->arrive() no longer has an explicit preceding warpgroup_wait<0>() in this iteration; it relies on the wait performed at the end of the previous iteration (i.e. the stage s-1 WGMMA is already complete). That is correct today, but the dependency is now implicit and a future refactor could silently reintroduce the race. Please add a one-line comment next to the arrive stating that stage s-1's WGMMA has already been waited on, so it is safe to release smem_a/smem_b of that stage.
🤖 v4
| # repeatability across repeated launches (a correct deterministic kernel | ||
| # must produce identical results every time). | ||
| previous_pdl = deep_gemm.get_pdl() | ||
| previous_tf32 = torch.backends.cuda.matmul.allow_tf32 |
There was a problem hiding this comment.
🟡 warning: The test saves/restores torch.backends.cuda.matmul.allow_tf32, but test_hc_prenorm_gemm() (line 18) sets torch.backends.cudnn.allow_tf32 = True and never restores it. As a result the pre-existing test still leaks global cudnn TF32 state to the rest of the suite, and the new test does not fully isolate the backend flags either. Since the new test intentionally disables TF32 for the reference, please save/restore both torch.backends.cuda.matmul.allow_tf32 and torch.backends.cudnn.allow_tf32 (ideally in test_hc_prenorm_gemm() as well) so the two tests cannot affect each other.
🤖 v4
| deep_gemm.set_pdl(enable_pdl) | ||
| stream = torch.cuda.Stream() if use_side_stream else torch.cuda.current_stream() | ||
| with torch.cuda.stream(stream): | ||
| generator = torch.Generator(device='cuda').manual_seed(123) |
There was a problem hiding this comment.
🔵 suggestion: generator is recreated with manual_seed(123) inside all four nested loops, so each (pdl, side_stream) configuration runs on the exact same a/b values. That is fine for the repeatability assertion, but it means the sweep never exercises different data across configurations. Since the race is timing/data-order sensitive, deriving the seed from the loop variables (or using a different seed per configuration) would broaden coverage without hurting determinism.
🤖 v4
| torch.testing.assert_close(sqr_sum, ref_s, rtol=1e-5, atol=2e-3) | ||
| if previous_d is not None: | ||
| # A correct deterministic kernel is bit-exact across launches. | ||
| torch.testing.assert_close(d, previous_d, rtol=0, atol=0) |
There was a problem hiding this comment.
🔵 suggestion: assert_close(d, previous_d, rtol=0, atol=0) asserts bit-exact determinism across launches. The SM90 kernel is currently deterministic (outputs are partitioned per split with no atomics), so this holds, but the assertion is stronger than typical GEMM tests and would become a false failure if split-K reduction is ever changed to a non-deterministic scheme. Please add a brief rationale comment or a named helper (assert_bit_exact) so the intent and fragility are explicit.
🤖 v4
🤖 ds-review-bot Code Reviewv6v5v4The MR fixes a real SM90 register-source WGMMA race in sm90_tf32_hc_prenorm_gemm. Previously the kernel loaded the FP32 A fragment into single-buffered Files reviewed: 2 |
2 participants
Fix a WGMMA register reuse race in SM90 tf32_hc_prenorm_gemm by waiting for completion before reusing A registers.
Add a regression test covering split-K, PDL, and side streams, with reference and repeatability checks.
Validation: tests/test_hyperconnection.py passed on an SM90 GPU.