Repository navigation
ci: validate exact commit from #6890 - #7924
Closed
balasaajay wants to merge 6 commits into
Closed
balasaajay wants to merge 6 commits into
balasaajay wants to merge 6 commits into
Conversation
All five per-head bounds masks in the fused MLA YaRN RoPE kernels are built
from the intra-block head index while the base pointer has already been
advanced to the program's head block, so the comparison reduces to
`local_head < head_num` and is true for every lane of every block.
The launch grid is `cdiv(nheads, BLOCK_H)` and `BLOCK_H` is autotuned over
{1,2,...,128}, so whenever the chosen block size does not divide the head
count the final program runs unmasked. It then rotates
`cdiv(head_num, BLOCK_H) * BLOCK_H - head_num` head rows that do not exist:
because `stride_x_seq == head_num * stride_x_nheads` those rows are the next
token's leading heads, and for the last token they are past the end of the
tensor. The result therefore depends on the autotuned block size, which is
meant to be a pure tiling choice.
Bound the absolute head index instead, which is the idiom already used by
the kernels added to this file by NVIDIA#6343. The fifth site is inside
`if pid_head == 0:` and iterates `for i in tl.static_range(...)`, so it is
offset by `i` rather than by `pid_head`; going stride-free also fixes the
reuse of one mask across dKV/dK/dV, whose head strides differ.
Signed-off-by: Rui Zhu <rui.zhu.rz399@yale.edu>
Two changes from review. 1. `other=0.0` on the four masked `tl.load` calls in the `dEMB` reduction of `_mla_rope_bwd_kv_split_kernel`. Without it a masked-out lane is undefined, and these values are added unconditionally into the accumulators and then reduced with `tl.sum`. Every other masked load in this file feeds a masked store, which discards the invalid lanes; a reduction cannot. 2. `TestFusedMLARopePartialHeadBlock`, which forces a partial final head block instead of hoping for one: `head_num=12` with `BLOCK_H` pinned to 8 by swapping each kernel's autotuner config list, so the tuner cannot pick a divisor. `BLOCK_H=4` is the control. Over both `REMOVE_INTERLEAVING` branches it covers the forward output and every gradient against the unfused reference in `sbhd` and `thd` with `inverse` both ways, including the KV-split `dEMB` gradient; a bitwise `BLOCK_H=8` vs `BLOCK_H=4` equality check that needs no reference implementation; and guard-buffer integrity, where the rotated tensor is a view of the head of a larger buffer whose poisoned tail must come back untouched. All 16 cases fail with this kernel file restored to main's version and pass with the fix. Before the fix the guard buffer counts 255 elements written past the end of the tensor against the 256 the mechanism predicts, one bf16 store having landed on the poison value it replaced. No performance claim: nothing here was timed, and no number is compared to any number upstream reports. Signed-off-by: Rui Zhu <rui.zhu.rz399@yale.edu>
The new kernel-determinism gate asks for this file when `fused_mla_yarn_rope_apply.py` changes, and it is the right place: the defect this PR fixes is a determinism defect. The per-head bounds masks were built from the intra-block head index, so whenever the autotuned `BLOCK_H` did not divide the head count the tail program ran unmasked -- which makes the result a function of a tiling choice made on timing. Why the existing cases could not see it: both pinned `heads = 32`, and 32 is divisible by every `BLOCK_H` in the autotune list that does not collapse the grid to a single program, so the partial final block was never built here at all. Two additions: * `heads` is parametrized over 32 and 12 on both replay tests. At 12 the final program is half empty. * `test_fused_mla_rope_is_independent_of_block_h` pins `BLOCK_H` to 8 and then 4 on one fixed input and requires the same bits, over the query path and the KV-split path. A replay cannot check this: Triton caches the config the first call chose, so every replay reuses one tiling. Measured on one NVIDIA L4, torch 2.12.0+cu130, triton 3.7.0: all 14 cases pass here, and 4 of them fail with this kernel restored to main's version -- both `is_independent_of_block_h` cases and the out-of-place replay at 12 heads, whose fresh output buffer makes the overflow differ between replays. Bit-exactness assertions only; nothing here was timed and no number of ours is compared to any number upstream reports. Signed-off-by: Rui Zhu <rui.zhu.rz399@yale.edu>
Review feedback on the two test files, no change to the kernel. * The shared config constants are gone. Both files declare what each test needs inside that test, per "Tests Too DRY Make Bugs WET"; the comments that say why 12 and why two tilings moved with them, since that is the part a reader needs. * `block_heads_partial` and `block_heads_exact` instead of the abbreviated names, and `num_heads` for the head count. * The module aliases are imported at file top level, not inside the contextmanager that uses them. * One guarded import block per file instead of two, and it catches `ImportError` rather than `Exception` -- the form the determinism tests already use for the same module. The merged block binds every name its failure path needs. Verified on one NVIDIA L4, torch 2.12.0+cu130, triton 3.7.0: unchanged in both directions. 16 partial-head-block cases and 14 MLA determinism cases pass, and with the kernel restored to 5e85e2b's version the same 16 and the same 4 fail as before this commit. No assertion, tolerance or parametrize value was touched. Nothing here was timed. Signed-off-by: Rui Zhu <rui.zhu.rz399@yale.edu>
Resolves the conflict with NVIDIA#5963, which landed in the same two in-place kernels. The kernel merged cleanly and both changes are intact: the absolute-index bounds masks and the zero-filled dEMB loads from this branch, and the load-before-store barriers from NVIDIA#5963. The conflict was in tests/unit_tests/fusions/test_mla_yarn_rope_apply.py, where both sides added code at the same places. Kept both: NVIDIA#5963's race test and this branch's partial-head-block tests. The guarded import binds the union of both sides' names. One change git did not flag: NVIDIA#5963 moved autotuning off `_mla_rope_{fwd,bwd}_inplace_kernel` onto new `_autotuned_*` wrappers, so the BLOCK_H pinning in both this file and the determinism test now targets the wrappers. Without that the pin would read `.configs` off a plain JITFunction and fail. Signed-off-by: Rui Zhu <rui.zhu.rz399@yale.edu>
Contributor
Author
|
/ok to test f703e6f |
This branch was previously deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
This draft is a CI validation copy of #6890 at its exact head commit,
f703e6f77338554891fb133a343e2c69a629d0a8.Do not merge this PR. The source change, discussion, review, and merge remain in #6890.
Issue tracking
Related to #6890
Validation