Repository navigation
[ops] fix: stop NPU GDN Triton kernels recompiling per batch length - #1264
Merged
Coach257 merged 1 commit intoOct 17, 2026
Merged
Conversation
chunk_gated_delta_rule_fwd_kernel_h, prepare_wy_repr_bwd_kernel and recompute_w_u_fwd_kernel took NT (total chunk count of the packed batch) as a tl.constexpr, and chunk_fwd_kernel_o took N (sequence count) as one. Both change with every packed batch, so each new total length compiled a fresh set of kernels: ~10 s per GDN forward+backward on 910B2C. Warm caches hide it; a training run on varlen data sees new lengths every step. Pass them as runtime ints and add them, with the kkt kernel's NT/B/ TOTAL_TASKS, to do_not_specialize. Outputs and gradients are bit-identical to before over eight layouts; with a fresh cache, seven of eight new lengths then run without compiling. A 154-step GRPO run on DocQA-RL goes from 57-65 s to 48-54 s per step. Same bug class and same technique as fla-org/flash-linear-attention#1207, which fixed it in fla's own triton_ascend backend; these kernels came from MindSpeed-MM 39775d87 before that and carry a different set of offending constexprs, so the fix is reproduced here rather than ported.
|
|
|
Navigate logical layers of code changes, visualize relationships, and explore their blast radius. 🧰 Additional context used📚 Code guidelines (3)No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configuration
📒 Files selected for processing (3)
Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 4 remain after this review. 📝 WalkthroughWalkthroughSeveral Ascend gated delta rule Triton kernels update which arguments are excluded from specialization. Selected ChangesAscend kernel argument specialization
Priority: ➖ Normal Estimated code review effort: 2 (Simple) | ~8 minutes Change: Bug fix Merge Risk: ⚪ Minimal · up to The changed kernels retain aligned runtime arguments and compile-time tile sizes. No actionable merge-blocking regression is evident. Architecture SummaryArchitecture risk: 🔵 Low · up to The change affects 1 system. Changed systems: Architecture concerns Review detailsSystems and components
Before / after behavior
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Estrellajer
force-pushed
the
fix/npu-gdn-triton-recompile
branch
from
October 5, 2026 09:05
0fbf546 to
64167fe
Compare
8 tasks done
Coach257
approved these changes
Oct 16, 2026
sc1915
pushed a commit
to sc1915/VeOmni
that referenced
this pull request
Oct 23, 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
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 does this PR do?
Four of the vendored NPU GatedDeltaNet Triton kernels take a length-derived value as
tl.constexpr:NT, the chunk count of the whole packed batch, inchunk_gated_delta_rule_fwd_kernel_hand bothwy_fastkernels, andN, the sequence count, inchunk_fwd_kernel_o. Both change with everymicro-batch, so varlen training recompiles the whole set each time it meets a total length it has not
seen — about 9–11 s per GDN forward+backward on a 910B2C.
A warm cache hides this completely, which is why it is easy to read as "the first step is slow". It is
not: with
cu_seqlenspacking, every step is a new total length. This PR passes the two values asruntime arguments and adds them to
do_not_specialize, which changes Triton's recompilationgranularity and nothing else.
Checklist Before Starting
[ops, docs, ci] fix: support Ascend 910B4-1 GDR training #966, [ops, ci] fix: mask padded GDR dot lanes on Ascend #970, [dist] feat: add lossless headwise GDN context parallelism #1087 (same vendored
_ascend/tritonkernels).[{modules}] {type}: {description}:[ops] fix: stop NPU GDN Triton kernels recompiling per batch lengthdo_not_specializeand runtime arguments do not change what a kernelcomputes.
Test
Ascend 910B2C, CANN 9.1.1, torch 2.10.0, torch_npu 2.10.0.post2, triton-ascend 3.2.2. Each variant run
with its own empty
TRITON_CACHE_DIRover eight packed layouts with different total lengths, forwardand backward:
Outputs and all gradients are bit-identical to the unpatched kernels across the eight layouts
(max abs diff 0.0), which is the check that matters: the change is meant to be a no-op numerically.
End to end, a GRPO run on Qwen3.5-4B over 8×910B2C with variable-length DocQA data goes from 57–65 s
per step to 48–54 s. Any variable-length run gains from this: it is overhead removed from every step,
not a win specific to one workload.
API and Usage Example
No API change. Nothing to configure; the kernels are selected as before.
Design & Code Changes
veomni/ops/kernels/gated_delta_rule/_ascend/triton/:chunk_delta_h.py,wy_fast.py(two kernels),chunk_o.py: drop: tl.constexprfrom theNT/Nparameter and add the name to the kernel'sdo_not_specializelist. These parameters are onlyloop bounds and index arithmetic inside the kernel bodies; nothing specialises on their value.
chunk_scaled_dot_kkt.py:NT/B/TOTAL_TASKSwere already runtime arguments but were notlisted in
do_not_specialize, so Triton still re-specialised them by divisibility. Listed now, forthe same reason.
This bug class is known: fla-org/flash-linear-attention#1207 (2026-09-03) fixed the same thing in
fla's own
triton_ascendbackend by the same technique, reporting Triton cache misses dropping from2647 to 203 over 20 steps. That patch does not transfer file-by-file — that backend is a rewritten
lineage with a different set of offending symbols — so the fix is applied here directly.
Checklist Before Submitting
ruff check,ruff format)tasks/training scripts moved or renamedtests cover correctness and the outputs are bit-identical.
Summary by CodeRabbit