Repository navigation
[Triton-Gluon-MLA-GFX950] return_lse: full decode + merged fp32 lse - #3542
Merged
Merged
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates the Gluon mla_decode_gluon implementation so return_lse=True runs the full decode (stage-1 + stage-2 reduce) and returns a merged fp32 LSE tensor [batch, nhead], avoiding bf16 precision loss during the stage-2 merge by keeping per-split LSE in a separate fp32 buffer.
Changes:
- Update
mla_decode_gluonto store per-split fp32 LSE inmid_lseand have stage-2 writefinal_lse(with a NUM_KV_SPLITS==1 fast path writingfinal_lsein stage-1). - Update
op_tests/test_mla.pyto validate the returned merged LSE againstlse_reffor the bh16bn{64,128} Gluon paths, and adjust comments/gating. - Update Gluon README to reflect the new
return_lsebehavior and performance numbers.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
op_tests/test_mla.py |
Removes the old stage-1-only host-merge path and adds LSE validation for the updated return_lse semantics. |
aiter/ops/triton/gluon/README.md |
Updates documentation to match full-decode + merged fp32 LSE behavior. |
aiter/ops/triton/gluon/mla_decode_gluon.py |
Implements fp32 mid_lse + merged final_lse, and updates return_lse behavior to always run full decode. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
…DCP) return_lse now runs the full decode (stage-1 + stage-2 reduce) and returns the merged fp32 lse [batch, nhead] (e_max + log(e_sum) over the whole sequence), matching the asm/torch references. Removes the stage-1-only DCP path. Supported for all regimes. - per-split lse kept in a separate fp32 mid_lse buffer (not packed into the bf16 O), fixing the bf16 precision loss in the stage-2 merge - stage-2 (_mla_softmax_reducev_kernel) writes final_lse; NUM_KV_SPLITS==1 fast path writes it directly in stage-1 - test_mla.py: validate returned lse vs lse_ref; `o` is now identical with and without -lse, so the gluon bh16bn64 gate is unified to just non-empty splits (ctx >= 256 // batch_size). README updated. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Dewei-Wang-sh
force-pushed
the
mla-gluon-lse
branch
from
June 4, 2026 15:33
34c9b8d to
e7a08af
Compare
Trim verbose comments along the call chain (drop repeated "separate fp32 tensor" / "matching logsumexp" prose, use "split" not "K"). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Dewei-Wang-sh
force-pushed
the
mla-gluon-lse
branch
from
June 4, 2026 15:47
e7a08af to
44e2f19
Compare
This was referenced Jun 16, 2026
valarLip
pushed a commit
that referenced
this pull request
Sep 11, 2026
…el for chunked prefill (#3606) * fix(mla): always write partial output + LSE for prefill PS tiles The PS scheduler in v1_2_host.cuh's generate_work_info tagged tiles that a single threadgroup could absorb from block 0 with the sentinel partial_o_loc = -1, causing generate_reduce_info to exclude them from reduce_indptr / reduce_final_map / reduce_partial_map. As a result mla_reduce_v1 never wrote final_lse[qo_start:qo_end, :] for those tiles, leaving stale workspace bytes that downstream consumers (vLLM's chunked-context MLA prefill on MI350) read as garbage LSE and fed into merge_attn_states, causing a ~2-3pp gsm8k accuracy regression on DeepSeek-V3. Always assign a real partial_o_loc so every tile produces a partial output + partial LSE that reduce processes (num_partials==1 for what used to be the fast-path tiles). The kernel already supports this path -- only the host scheduler needed changing; no assembly edits. The if(work.partial_o_loc == -1) continue; guard in generate_reduce_info is left in as defensive coding. Mirrors PR #3542 ("[Triton-Gluon-MLA-GFX950] return_lse: full decode + merged fp32 lse"), which fixed the analogous bug in the triton-gluon decode kernel's NUM_KV_SPLITS==1 fast path. * fix: dont skip reduction of tiles with num_split==1 Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: increase max_partial now that we support split tiles Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: compute the work and partial based on actual kv size Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: return final_lse instead of partial attn_lse from kernel wrapper Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * test: add proper final lse validation test Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * test: add non-causal chunk context to correctness test Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * pass max_kvlen to correctness test Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: simplify comments Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: remove redundant comment Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: remove redundant handling of partial_o_loc=-1 Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: simplify comment Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: indent Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: indentation Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: avoid running reduce when num_splits==0 Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * chore: add comment about fullgraphs Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: size without mxa kvlen Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * chore: fix docstring Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: rename lse_ref in test Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: add need_lse flag This avoid avoid unnecessary reduction when lse is not needed. Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * chore: simplify comments Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: comment in get_ps_metadata_info_v1 Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * chore: restructure to minimzie diff Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: use qo_len in test when available instead of ctx_len Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * chore: remove empty eol spaces Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: set max_partials to one per qo-tile + splits Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: set lower lse test tolerances Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: typo Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: set lse error atol to 3e-2 to be above the fp8 rounding onise Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * chore: fix black erro Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: fill final_lse buffer with -inf (#3) Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: remove redundant asserts (#6) Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: include cassert in v1_2_host.cuh (#7) Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: add return_lse param and check that ensures final_lse is populated completely if return_lse (#3 & #4) Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: align flydsl reduce kernel with the cu variant Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * chore: indent py bindings Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: convert asserts to AITER_CHECK Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: address comments about tests Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: add test cases to flydsl reduce to ensre parity Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: switch to int64 slot offsets Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> * fix: add test of aiter.mla.mla_prefill_ps_fwd python wrapper Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> --------- Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
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.
return_lse now runs the full decode (stage-1 + stage-2 reduce) and returns the merged fp32 lse [batch, nhead] (e_max + log(e_sum) over the whole sequence), supported for all regimes.