[Fix] Fix DSA indexer fusion for NeoX RoPE - #30111
Merged
Merged
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request refactors the DSA indexer by making the indexer fusion flag instance-specific (self.use_dsa_indexer_fusion) and disabling it when is_neox_style is active. It also increases the accuracy threshold in the DeepSeek-V32 index cache test from 0.93 to 0.935. There are no review comments, and we have no feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Fridge003
changed the title
Fridge003
force-pushed
the
codex/skip-neox-dsa-indexer-fusion
branch
from
July 4, 2026 09:46
982bd4a to
2d8cf0a
Compare
Collaborator
Author
|
/rerun-test test/registered/8-gpu-models/test_deepseek_v32_indexcache.py test/registered/models_e2e/test_dsa_glm52_tp_mtp.py test/registered/models_e2e/test_dsa_glm52_dp_mtp.py test/registered/models_e2e/test_dsa_glm52_nvfp4_tp_mtp.py test/registered/models_e2e/test_dsa_glm52_nvfp4_dp_mtp.py |
Fridge003
marked this pull request as ready for review
July 4, 2026 09:52
Fridge003
requested review from
1am9trash,
HaiShaw,
YAMY1234,
hubertlu-tw,
kkHuang-amd and
rainj-me
as code owners
July 4, 2026 09:52
Contributor
|
Results for 🚀 🚀 |
vincentzed
added a commit
to vincentzed/sglang
that referenced
this pull request
Jul 8, 2026
… for DeepSeek-V3.2 The DSA lightning-indexer fused Q/K kernels (PR sgl-project#27705) hardcoded the interleave/GPT-J RoPE pairing (2i, 2i+1), so is_neox_style=True models (DeepSeek-V3.2) were parked on the slow split path by sgl-project#30111 after the wrong rotation corrupted indexer top-k selection (~2.4pt GSM8K drop, 0.955 -> 0.931, sgl-project#30088). Add a compile-time kIsNeox branch to both kernels that pairs dim i with i + rope_dim/2 via a lane shuffle (partner pack lives at lane ^ 8) and loads the cos/sin halves layout with per-element frequency indexing; the cache layout itself is identical between the two conventions. Plumb is_neox through the JIT wrappers and drop the `and not is_neox_style` fusion gate clause. The interleave path is untouched (GLM-5.x byte-identical); Hadamard behavior (kHadamard) and the DeepSeek-V4 native indexer are unchanged. Tests: parametrize the fused-kernel unit tests over is_neox and add a fused-vs-eager NeoX top-k checksum guard (an interleave-only rotation collapses overlap to ~topk/3).
This was referenced Aug 20, 2026
Chronostasys
pushed a commit
to MindLab-Research/sglang
that referenced
this pull request
Aug 24, 2026
sumin-hong
added a commit
to moreh-dev/sglang
that referenced
this pull request
Aug 27, 2026
…to one aiter kernel The DSA indexer Q/K fusion has been CUDA-only. On ROCm the indexer still runs the unfused path: rope on q, LayerNorm + rope on k, an fp8 quant of each, the head-gate scale, and the index-K cache write, as separate launches per layer per step. All of it is launch-bound. aiter's indexer_qk_rope_quant_and_cache does the whole thing in one launch, so extend the existing fusion switch to ROCm rather than build a second one: the fused wk_weights_proj GEMM, the no-Hadamard invariant, the weight loader and the graph split-op plumbing are all platform-agnostic already. - Probe aiter for the kernel at import and fall back with a warning; aiter is pinned per image and an older one would only fail at the first forward. - Build k_norm in fp32 whenever fusion is on. The kernel requires fp32 norm params, as the CUDA kernels already do; bf16 stays for the unfused ROCm path, where matching x.dtype is what selects aiter's CK layernorm. - Read cos/sin off aiter's rope module, which keeps them apart as [max_position, 1, 1, rope_dim/2] instead of one cos_sin_cache. - Give _fused_k_prepare_and_store a ROCm branch. Its CUDA fallback is a JIT kernel, and the k-only decode fast path reaches it whenever fusion is on. - Gate the fused-store branch on _is_cuda so non-CUDA stops paying for a JIT compile attempt that can only fail. Two adjacent fixes: - lora_manager imported _use_dsa_indexer_fusion, which sgl-project#30111 deleted, so the indexer-LoRA guard raised ImportError instead of its intended error. Restored as dsa_indexer_fusion_supported(). - The ROCm branch of _store_index_k_cache read forward_batch.out_cache_loc, ignoring the sliced out_cache_loc its caller passes under the graph split-op contract. The two K-cache writers stay live at once -- the k-only decode path writes unfused -- and they are not byte-identical: the unfused path rounds to bf16 before quantizing while the fused kernel goes fp32 -> fp8 directly, so a value near an fp8 midpoint can land on either neighbour. Measured on 41x128: scales identical, 15/5248 elements differ, each by one fp8 code. Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Sumin Hong <sumin.hong@moreh.io>
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
1 participant
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
is_neox_style=Trueindexers while keeping the global fusion env enabled by default.self.use_dsa_indexer_fusionas the only fusion branch guard; the module-level_use_dsa_indexer_fusionhelper was removed.SGLANG_DISABLE_DSA_INDEXER_FUSIONtoEnvBool(False)on latestmain.0.935.Root Cause
DSA indexer fusion was controlled only by the global
SGLANG_DISABLE_DSA_INDEXER_FUSIONenv flag. NeoX-style RoPE models could still take the fused indexer path, but that path is not compatible withis_neox_style=Truerotation handling, causing the observed accuracy regression.CI States
Latest PR Test (Base): ❌ Run #28702478657
Latest PR Test (Extra): ❌ Run #28702478568