Conversation
…path works when SGLANG_ENABLE_JIT_DEEPGEMM=0.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a local import of deep_gemm inside the tf32_hc_prenorm_gemm function. The feedback notes that because this function is on the hot path, importing the module on every invocation adds unnecessary overhead. The reviewer suggests lazily importing and caching the module globally to optimize performance.
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.
| import deep_gemm | ||
|
|
||
| deep_gemm.tf32_hc_prenorm_gemm(x, fn, out, sqrsum, num_splits=num_splits) |
There was a problem hiding this comment.
In high-performance LLM serving, tf32_hc_prenorm_gemm is executed on the hot path (potentially once per layer/step). Running import deep_gemm on every single invocation introduces unnecessary overhead from sys.modules lookups and import lock acquisition.
By declaring global deep_gemm and checking if it is already present in globals(), we can lazily import it on the first call and reuse the cached global reference on all subsequent calls. This reduces the overhead to a simple, extremely fast dictionary lookup.
| import deep_gemm | |
| deep_gemm.tf32_hc_prenorm_gemm(x, fn, out, sqrsum, num_splits=num_splits) | |
| global deep_gemm | |
| if "deep_gemm" not in globals(): | |
| import deep_gemm | |
| deep_gemm.tf32_hc_prenorm_gemm(x, fn, out, sqrsum, num_splits=num_splits) |
There was a problem hiding this comment.
global cache call saves ~29ns, negligable (measured)
i follow the pre-exisitng example set by https://github.com/sgl-project/sglang/blob/v0.5.14/python/sglang/srt/layers/mhc.py#L1446-L1449
|
Hello @YAMY1234, I believe this bug was introduced in #26238 - can you help take a look / find a maintainer to review? Small change (missing import) that breaks dsv4 for certain configs. This is my first contribution, so unclear to me how to get help here. Fixed in my downstream fork but would like to get fixed upstream so I don't have to keep re-patching it. |
|
Let me take a look |
|
Hi @Fridge003, I see you approved the original PR - can I get a stamp here / help merging this change? |
2 participants
Motivation
Fixes #29738.
Modifications
tf32_hc_prenorm_gemmreferenced the module-leveldeep_gemm, which is only imported underENABLE_JIT_DEEPGEMM. The HC-prenorm path reaches it via the independentSGLANG_OPT_DEEPGEMM_HC_PRENORMgate, soSGLANG_ENABLE_JIT_DEEPGEMM=0raisedNameError. Added a localimport deep_gemm, matching the existing pattern inmhc.py.CI States
Latest PR Test (Base): ❌ Run #31571002834
Latest PR Test (Extra): ❌ Run #31571002654