[GDN] Key the fused GDN-2 WY/dqkg backward autotuner on K and V - #1120
Merged
Conversation
chunk_gdn2_bwd_kernel_wy_dqkg_fused sweeps BK/BV tile configs, whose relative performance depends on the head dimensions K and V, but its autotune key was ['BT', 'STATE_V_FIRST']. The first (K, V) geometry tuned in a process therefore pinned its winning config for every other geometry sharing BT, silently running later shapes on a config tuned for a different problem size. Key on K and V as well, matching the wy_fast.py kernels in this op, and add a regression test asserting that two geometries produce two autotuner cache entries. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
Author
|
The same pattern — the config space sweeps tile sizes whose best choice depends on the head dimensions, while the autotune key omits them — appears elsewhere (found by scanning
|
5 tasks
zhiyuan1i
requested changes
Aug 27, 2026
zhiyuan1i
left a comment
zhiyuan1i
left a comment
Collaborator
There was a problem hiding this comment.
Confirmed this is the one remaining autotune-key gap left by #1171 — thanks for catching it. Two things before merge:
- Rebase needed: your new test calls
_rand_inputswith 6 args intests/ops/test_gdn2.py, but after #1165/#1166 the signature takes 7 (HV was added), so it will TypeError on current main. - Test cost: the new test pays a real cold-autotune in CI (~4m measured). Please follow the
_unwrap_autotunerprecedent intest_gdn.pyand assert the full key directly — that also pins K and V in the key explicitly.
zhiyuan1i
approved these changes
Sep 3, 2026
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
2 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.
Summary
chunk_gdn2_bwd_kernel_wy_dqkg_fusedautotunes overBK/BVtile configs(plus
num_warps/num_stages), but its autotune key was['BT', 'STATE_V_FIRST']—KandV, the head dimensions that bound thetiles and determine each config's occupancy, were not part of the key.
Consequently the first
(K, V)geometry that reaches the kernel in a processtunes it for every other geometry sharing
BT: a model withK=64, V=128heads running after a
K=128, V=256warmup silently inherits the largergeometry's winning config instead of tuning for its own. We hit this while
benchmarking GDN-2 at multiple head geometries — back-to-back shapes in one
process reproduced with the first shape's config, skewing every subsequent
measurement until each geometry was quarantined into its own process with a
fresh Triton cache.
This PR adds
KandVto the key, matching the sibling kernels infla/ops/gdn2/wy_fast.py(which key on['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN']), and adds a regressiontest. Only this kernel in
fla/ops/gdn2/sweepsBK/BVwhile omittingK/Vfrom its key.Cost: one extra autotune sweep per distinct
(K, V)geometry, on firstencounter only; results are cached as before. Single-geometry runs (the
common case) see no change.
Test plan
New in
tests/ops/test_gdn2.py:test_chunk_bwd_autotune_key_covers_head_dims— backwards through two(K, V)geometries at the same chunk size must leave two distinctautotuner cache entries. Fails on
main(the second geometry cache-hits onthe first's winner), passes on this branch.
Verified on H200: 1 failed on the base commit, 1 passed on this branch, same environment, with the key line as the only tree difference.
Dependent tests (
python scripts/find_dependent_tests.py fla/ops/gdn2/chunk_bwd.py):tests/ops/test_gdn2.py— 29 passed (full suite, incl. the new test)tests/layers/test_attn_varlen_pack_layout.py— 13 passed, 2 skippedLocal test environment:
Benchmark / NCU
Neutral for single-geometry processes: the key change does not alter which
config wins for a geometry tuned on its own — it only ensures each geometry
is tuned at all.
For multi-geometry processes it removes a measurable penalty. On H200
(bf16, T=32768 packed varlen in 32 documents, chunk 64, medians of 30 iters),
the per-geometry winners genuinely differ (
BK64/BV32/warps4/stages4atK=64/V=128 vs
BK64/BV64/warps4/stages4at K=128/V=256), and runningK=64/V=128/H=32 on the config leaked from a K=128/V=256 warmup (what
maindoes today) costs (the leaked run's tuner cache confirms it selected K128/V256's winner,
BK64/BV64/warps4/stages4, instead of the geometry's ownBK64/BV32):Breaking changes
None.
Checklist