Make auto-CP planning stable across head shards - #32
Draft
Micuks wants to merge 2 commits into
Draft
Conversation
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 join this conversation on GitHub.
Already have an account?
Sign in to comment
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 changed
Auto-CP planning can now use a logical head count from
FLASH_QLA_AUTO_CP_PLANNING_HEADS. Kernel calls still use their local headcount. The override only affects the occupancy estimate and segment plan.
The default behavior is unchanged when the variable is unset. Invalid values
raise a clear error, and the logical head count must be a positive multiple of
the local head count.
This PR also adds a standalone regression and records the H800 reproduction.
The test calls FlashQLA directly and does not depend on a model repository.
Why
A logical H16 operation may be issued as four H4 calls after head sharding.
The current scheduler plans each call from its local head count, so the H16 and
H4 calls can choose different segment boundaries.
For the reproduced mixed-length input, H16 used 18 auto-CP segments while H4
used 34. Their BF16 outputs differed. Disabling auto-CP made the same inputs
bit-exact, which isolates the difference to the auto-CP path.
Setting
FLASH_QLA_AUTO_CP_PLANNING_HEADS=16gives both launch shapes the same18-segment plan and restores bit-exact output while keeping auto-CP enabled.
Validation
The standalone test ran on an NVIDIA H800 with PyTorch 2.8.0+cu128:
I also ran:
The regression covers one mixed-length BF16 forward case on SM90. It does not
claim parity for every shape, GPU architecture, backward path, or full
training run.