Revert "[TRITON][GLUON] Prefill MQA Logits kernel tuning for GLM 5.x … - #5077
Draft
zhuyuhua-v wants to merge 5 commits into
Draft
zhuyuhua-v wants to merge 5 commits into
zhuyuhua-v wants to merge 5 commits into
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags: |
zhuyuhua-v
force-pushed
the
dev/glm5.2
branch
3 times, most recently
from
August 31, 2026 11:22
150ebf3 to
34d0afc
Compare
zhuyuhua-v
force-pushed
the
dev/glm5.2
branch
from
September 1, 2026 06:30
34d0afc to
142afca
Compare
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Comment on lines
226
to
227
| block_m = 1 | ||
| other = {"LOOP_VARIANT": loop_variant} |
Contributor
* [FlyDSL] Add DCP decode TopK merge kernel
Under Decode Context Parallel each rank holds 1/W of the KV and takes a
local top-k, so the ranks must agree on the global top-k before attention.
The existing path reconstructs that global set on every rank and then drops
the ~(W-1)/W of it the rank does not own -- W times the work, most of it
discarded.
This op emits only the rank's OWNED physical KV slots, already localized and
compacted, plus the kv_indptr attention needs. Two kernels: a per-row radix
select for the global threshold, then a per-row pack that also computes the
cross-row prefix sum. Two properties make the fusion possible:
* ownership is positional -- candidate column c came from rank c // k_loc --
so gids never have to travel and the caller can exchange scores alone;
* the block-table slot formula collapses to
block_table[j // page] * page + j % page, independent of W and of the
interleave size, so the localize folds into the emit loop.
The prefix sum lives in pack rather than in its own kernel: a grid=(1,) scan
walking the rows on a single thread cost 37% of the op at rows=128 (15.7 of
42.3 us) to sum 128 integers. Each pack block instead recomputes the prefix it
needs in parallel -- rows blocks each scanning rows counts is redundant, but a
block scan over a few hundred values is a handful of barriers, and it removes
both the serial loop and a launch. Chunked, so rows > 256 stays correct.
That is also why select and pack cannot merge further: pack needs
kv_indptr[row], which depends on every OTHER row's count. Closing that gap
needs either a cross-block barrier (unsafe -- block scheduling order is not
guaranteed) or grid=(1,), which would serialise select by rows.
Allocates no device scratch (the caller owns `staging`), so it is safe inside
a captured CUDAGraph.
Measured against the equivalent six-op sequence at W=8, k_loc=topk=2048:
4.82-5.14x across bs in [1,4,16,32,64,128], and half the exchange payload.
The curve is flat -- 33.6-34.3 us end to end regardless of bs -- where the
six-op path and the earlier three-kernel form both grew with rows.
Co-Authored-By: Claude <noreply@anthropic.com>
* [HIP] Add optional score output to top_k_per_row_decode
top_k_per_row_decode now takes an optional `values` tensor and writes each
selected index's logit alongside it, mirroring top_k_per_row_prefill. The
one-block kernel already supported this -- WRITE_TOPK_VALUES is a template
parameter of dispatch_topk_oneblock -- but the decode entry point hardcoded
it to false with /*out=*/nullptr. Wiring it through the header, the pybind
arg list and the Python wrapper is the whole change; the default path
compiles to the same instantiation as before.
This lets a DCP decode indexer skip the separate gather whose only job was
to read these same scores back out of the logits.
Also fixes the padding value for short rows: where the index is padded with
-1, the score was padded with 0.0. A score that pads to 0.0 does not sort
below real ones -- logits are routinely negative, so the padding outranks
them, and a consumer that ranks these scores (a DCP merge across ranks) lets
padding steal real candidates' slots. Pad with -inf instead.
Both short-row pads are changed, not just decode's. The one-block and the
multi-block kernels each have one, and which of the two runs is a perf
heuristic (should_use_mulblocks) -- leaving them disagreeing would give the
same call two different paddings across a batch-size boundary.
The line is shared with prefill. aiter's own prefill caller
(sampling.py:_select_topk) does pass a real `values` tensor, so
WRITE_TOPK_VALUES is true there; the short-row branch is nonetheless
unreachable for it, because that fast path requires vocab >= 65536 and
k <= 256, so row_len <= k never holds. An earlier version of this message
claimed all prefill callers pass values=None -- that was wrong, and would have
cleared the change for the wrong reason.
Co-Authored-By: Claude <noreply@anthropic.com>
---------
Co-authored-by: Claude <noreply@anthropic.com>
* [FlyDSL] Add variable-length decode TopK Add a compile-time-specialized multi-CTA radix-select path with optional deterministic output for the existing per-row decode interface. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] Optimize and extend decode TopK Reuse radix histograms for stable selection, add wide-count support for long contexts, and align launcher caching and validation with other FlyDSL operators. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] Speed up stable decode TopK writes Vectorize stable writes with a compile-time width choice and remove redundant scan barriers to improve long-context latency without regressing short rows. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] Fix short-row decode TopK padding Initialize empty-row state and overwrite unused output slots so CUDA Graph replays cannot retain stale indices. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] Optimize decode TopK radix selection Reduce long-row scan and dispatch overhead with adaptive three-pass radix, vectorized histogram reduction, fused stable prefixing, cached workspace, and short-row direct fill. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] Tune decode TopK launch geometry Use a 1024-thread, 16-chunk data geometry with a 256-thread reduce to improve three-pass radix performance across stable and unstable batches. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] Unify decode TopK on three-pass radix Remove four-pass and shape-specific routing so every decode path uses the tuned 11-bit radix geometry. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] Reduce decode TopK dispatch overhead Add single-launch stable paths and reuse dynamic-N multi-kernel binaries so model workloads avoid repeated launches and per-width compilation. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] Avoid persistent decode TopK deadlocks Route long single-row calls through the safe multi-launch path because concurrent graph replays can deadlock non-cooperative barriers, while retaining the barrier-free short path. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] Add gfx950 decode TopK dispatch Route measured gfx950 shapes through dynamic FlyDSL kernels with HIP fallback, and remove unused persistent paths to keep graph-safe execution maintainable. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] Use one-workgroup TopK for short unordered rows Reuse the faster deterministic short-row kernel for unordered calls, whose output remains valid for the weaker ordering contract. Co-authored-by: Cursor <cursoragent@cursor.com> * [CI] Fix Black formatting for decode TopK Apply the formatter output required by the code-style check. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] Streamline decode TopK tests Keep focused dispatch, end-to-end accuracy, and CUDA Graph coverage while removing redundant performance and implementation-detail cases. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] Add conservative gfx942 decode TopK gates Enable only the measured short and long gfx942 ranges while retaining HIP fallback for unstable intermediate shapes. Co-authored-by: Cursor <cursoragent@cursor.com> * Tighten gfx942 stable FlyDSL decode TopK row gate to 64. Reduce the short-context stable dispatch limit on gfx942 so FlyDSL is only selected for smaller batch sizes where it reliably wins. Co-authored-by: Cursor <cursoragent@cursor.com> * Align gfx942 decode TopK gate tests with 64-row limit. Update gate and e2e cases so short-context stable dispatch boundaries match the tightened gfx942 FlyDSL row gate. Co-authored-by: Cursor <cursoragent@cursor.com> * Reuse dpp_utils.update_dpp_i32 in decode TopK warp prefix scan. Replace inline fly_rocdl.update_dpp calls with the shared DPP helper. Co-authored-by: Cursor <cursoragent@cursor.com> * Fix import formatting in decode TopK kernel for Ruff/Black CI. Add the blank line Ruff I001 expects between third-party and local imports. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] Add optional values output to decode TopK Specialize kernels at compile time to write selected values inline while preserving the no-values path and signed-zero ordering. Co-authored-by: Cursor <cursoragent@cursor.com> * [FlyDSL] Harden decode TopK dispatch and edge coverage Fall back to HIP for unsupported FlyDSL signatures and align edge-case ordering while covering ties, NaNs, and MTP row lengths. Co-authored-by: Cursor <cursoragent@cursor.com> * Fix Black formatting in decode TopK tests Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: lirui927 <ruili@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Zhu Yuhua <yuhzhu@amd.com>
Keep row count out of the FlyDSL builder cache key so varying decode batches reuse compiled kernels. Co-authored-by: Cursor <cursoragent@cursor.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
3 participants
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.
…(#4563)"
This reverts commit aded0f8.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist