Skip to content

Remove DeepGEMM for indexer GEMM in piecewise NSA path - #26494

Merged
Fridge003 merged 1 commit into
mainfrom
brayden/remove-deepgemm-piecewise
May 28, 2026
Merged

Remove DeepGEMM for indexer GEMM in piecewise NSA path#26494
Fridge003 merged 1 commit into
mainfrom
brayden/remove-deepgemm-piecewise

Conversation

@b8zhong

@b8zhong b8zhong commented May 27, 2026

Copy link
Copy Markdown
Collaborator

Motivation

Remove it in piecewise CUDA graph path, which seems to be accidentally reintroduced in #23351 due to PRs merging at different times. Additionally, it will increase the warmup time

Modifications

Remove it.

Accuracy Tests

It's already been validated in #23856

Speed Tests and Profiling

Unit benchmark with CUPTI.
Results: (only N = 64 and N = 32 are relevant for indexer for DSV3.2 and GLM-5), the rest are Deepseek V4 shapes, which shows we could remove it from Deepseek V4 path later (currently not used on default path anyway)
image


CI States

Latest PR Test (Base): 🚫 Run #26527866613
Latest PR Test (Extra): ❌ Run #26527866300

@b8zhong

b8zhong commented May 27, 2026

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request replaces a custom deep_gemm_wrapper matrix multiplication with PyTorch's native torch.mm in logits_head_gate_pcg. However, standard torch.mm does not support the out_dtype keyword argument, which will cause a runtime TypeError. It is recommended to perform the multiplication and cast the result using .to(torch.float32) instead.

Comment thread python/sglang/srt/layers/attention/dsa/dsa_indexer.py
@Fridge003
Fridge003 merged commit e31ea50 into main May 28, 2026
200 of 229 checks passed
@Fridge003
Fridge003 deleted the brayden/remove-deepgemm-piecewise branch May 28, 2026 07:15
mqhc2020 pushed a commit to mqhc2020/sglang that referenced this pull request Jun 2, 2026
…6494)

Co-authored-by: b8zhong <b8zhong@users.noreply.github.com>
Chronostasys pushed a commit to MindLab-Research/sglang that referenced this pull request Aug 24, 2026
…6494)

Co-authored-by: b8zhong <b8zhong@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants