Compile CP softmax LSE corrections with dynamic shapes - #695
Open
JessicaJiang-123 wants to merge 1 commit into
Open
Compile CP softmax LSE corrections with dynamic shapes#695JessicaJiang-123 wants to merge 1 commit into
JessicaJiang-123 wants to merge 1 commit into
Conversation
JessicaJiang-123
requested review from
ipanfilo,
wangye805 and
wenchenvincent
as code owners
August 6, 2026 19:54
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.
Description
Fixes the flaky ROCm nightly test
test_qwen3_30B_A3B/test_baseline.pyin miles, where rollout-0log_probsandref_log_probsintermittently fail theabs_tol=1e-8check.The two CP softmax LSE correction functions use
@jit_fuser, which wrapstorch.compile:flash_attn_fwd_softmax_lse_correctionflash_attn_fwd_second_half_softmax_lse_correctionDynamo may compile multiple variants of these functions after seeing different input shapes. On gfx950, those variants can select different launch configurations, and
libdevice.log1pis not bit-exact across them.In failing runs, the reference and actor forwards receive bit-identical operands but use different
log1ppaths. The resulting approximately 1-ULP difference is repeatedly introduced across the model's attention layers and produces a final1e-5to2e-4log-probability mismatch.This PR compiles only these two functions with
dynamic=Truefrom their first invocation. This avoids the initial static specialization and keeps both forwards on the same numerical path.CP=1is unaffected because the LSE merge is only used whenCP > 1.Type of change
Changes
jit_fuser_dynamictojit.py, mirroringjit_fuserbut usingtorch.compile(..., dynamic=True).jit_fuser_dynamiconly to the two CP softmax LSE correction functions.Verification
Tested
test_qwen3_30B_A3B/test_baseline.pyon 4× MI355X with:CP=2/PP=2/TP=1/EP=2Comparing rollout-0 reference and actor log-probabilities:
Checklist: