Bulk AG Overlap for bf16 on gfx950 - #713
Conversation
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| _assert_numerics_passed(result) | ||
| fused = _reported_names(result.stdout, "UB FUSED NAMES: ") | ||
| assert fused is not None, f"harness printed no fused name set\n{result.stdout.decode()}" | ||
| assert "qkv_dgrad" in fused, fused |
There was a problem hiding this comment.
This assertion can't fail if the fused bulk dgrad path stops being taken.
_ub_fused_names is populated in add_ub() purely from the method config (method == "fused" and the region wasn't disabled). It says nothing about whether fused_bulk_ag_eligible() returned True for the call, or whether fused_overlap_bulk_ag ever ran. So if a shape-math regression in _fused_gemm_dims / _fused_gemm_shape_ok silently downgraded ub_bulk_dgrad to False, this test would still pass: numerics are correct on the fallback path, and qkv_dgrad is still in fused.
The gap is visible in the pair of tests: test_fused_layer_declines_ineligible_k (line 608) asserts the identical condition — "qkv_dgrad" in fused and not in disabled — while expecting the opposite outcome. The two tests can't distinguish "kernel ran" from "kernel declined", which is the property this PR is adding.
Consider having run_layer_with_overlap.py report the post-gate decision (e.g. print the effective ub_bulk_dgrad per layer after fused_bulk_ag_eligible) and assert on that instead, so the positive test fails when the path is not taken.
| """A Userbuffers region the fused backend cannot serve declines at setup." | ||
| """ |
There was a problem hiding this comment.
Stray " at the end of the first line — the docstring text ends with ... declines at setup." and then closes on the next line. It parses fine, but the quote looks like a typo.
| """A Userbuffers region the fused backend cannot serve declines at setup." | |
| """ | |
| """A Userbuffers region the fused backend cannot serve declines at setup.""" |
| dist_print("UB FUSED NAMES: " + " ".join(sorted(te.module.base._ub_fused_names))) | ||
| dist_print("UB DISABLED NAMES: " + " ".join(sorted(te.module.base._ub_disabled_names))) |
There was a problem hiding this comment.
These two dist_print calls have no src= argument, so every rank prints its own copy of the name lists. The consumer (_reported_names in test_comm_gemm_overlap.py) just takes the first substring match, so the duplicates are harmless but noisy — and if ranks ever disagreed, the test would silently key off whichever line landed first. Since the registered UB names are identical on all ranks, src=0 gives the same information with one line.
| if not IS_HIP_EXTENSION: | ||
| return True | ||
| if not _ub_is_fused(name): | ||
| return False |
There was a problem hiding this comment.
This early-out has the opposite polarity from its near-identical sibling fused_ag_gemm_eligible, which returns True for a non-fused name ("not our backend, don't veto") while this one returns False. The difference is deliberate and load-bearing — the caller uses this to enable the bulk path, so a non-fused region must decline — but the two functions sit next to each other with the same shape, and a future reader "fixing the inconsistency" would silently turn bulk dgrad on for every ROCm UB region. Worth a one-line comment stating that this predicate is opt-in whereas the other is opt-out.
| "fused bulk AG reached with a non-bf16 operand"); | ||
| NVTE_CHECK(ubuf.numel() != 0, "fused bulk AG reached without a gather destination"); | ||
|
|
||
| const size_t m = A.size(1); |
There was a problem hiding this comment.
m here comes from the GEMM operands, but the gather geometry is derived from chunk (the Userbuffers region), and nothing checks that chunk.size(0) * _tp_size equals the GEMM's row count. In fused_overlap_ag this can't drift because the ubuf is the A operand; in the bulk variant the gathered tensor and the dgrad GEMM are completely unrelated tensors that only happen to share a row count today (both are the layer input's sequence dim). The Python gate enforces the alignment rules but never compares the two.
If they ever disagree — a UB region registered with a different shape than the layer actually runs at, which the Python-side eligibility check doesn't catch — the gather writes a differently-sized region than the kernel grid was sized for, silently. An NVTE_CHECK that the ubuf rows match the GEMM rows would both document the coupling and turn a silent corruption into an error.
|
Reviewed the full diff of this PR against the base tip (16 files, +443/-63): the hipKittens bulk AG+GEMM kernel path ( Overall the approach looks sound — the Left 7 inline comments. The substantive ones are the new tests not actually distinguishing "fused kernel ran" from "fused kernel declined" (the reported name list reflects setup config, not the runtime gate), an unchecked shape coupling between the Userbuffers region and the dgrad GEMM in Copyright headers: OK. |
1 participant
Enables fused wgrad AG + dgrad GEMM for BF16 + gfx950.
Adds additional test coverage and layer tests, and makes fused/bulk_fused the default overlap scheme for gfx950, with fallback to no overlap at the moment for unaligned shapes, dtypes, and reduce-scatter.
1.282x Geomean speedup for wgrad AG + dgrad GEMM step