Skip to content

Fix DeepGEMM import for independently enabled MHC prenorm - #38442

Open
superuct wants to merge 1 commit into
sgl-project:mainfrom
superuct:fix/mhc-prenorm-deepgemm-import
Open

superuct wants to merge 1 commit into
sgl-project:mainfrom
superuct:fix/mhc-prenorm-deepgemm-import

Conversation

@superuct

@superuct superuct commented Sep 8, 2026

Copy link
Copy Markdown

Motivation

DeepSeek V4 can reach tf32_hc_prenorm_gemm through SGLANG_OPT_DEEPGEMM_HC_PRENORM even when SGLANG_ENABLE_JIT_DEEPGEMM=0. In that configuration, the module-level import is skipped and nonempty inputs fail with NameError: name 'deep_gemm' is not defined, despite DeepGEMM being installed. This can also surface during MHC prewarming at model load time.

Fixes #29738. This uses the same local-import approach proposed in #29740 and adds CPU regression coverage against the current main branch.

Modifications

Import deep_gemm inside tf32_hc_prenorm_gemm, after the existing empty-input return. The independently enabled HC-prenorm path now loads its own dependency instead of relying on the JIT gate's module-level import. Empty inputs still return without requiring DeepGEMM, and a missing dependency raises its actual import error.

Add three registered CPU tests that execute the real entrypoint module with the JIT gate disabled and stubbed external dependencies. They check kernel argument forwarding, the empty-input fast path, and missing-dependency error propagation. Before the fix, the two nonempty-input tests reproduce the NameError; after the fix, all three pass.

Accuracy Tests

  • python3 test/registered/unit/layers/test_deep_gemm_hc_prenorm_import.py: 3 passed.
  • uvx pre-commit run --files python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py test/registered/unit/layers/test_deep_gemm_hc_prenorm_import.py: all applicable hooks passed, including CI registry validation.
  • No CUDA kernel implementation or tensor arithmetic changed. GPU numerical accuracy and full-model startup were not tested for this patch.

Speed Tests and Profiling

Not run. The change adds a function-local Python import (cached by Python after the first import); no measured throughput or latency claim is made.

Checklist

  • Add unit tests according to the contribution guide.
  • Follow the SGLang code style guidance.
  • Format your code with pre-commit (all applicable hooks passed).
  • Update documentation (no public API or configuration change).
  • Provide accuracy and speed benchmark results (not run; scope described above).

CI States

Latest PR Test (Base): ❌ Run #34197052129
Latest PR Test (Extra): ❌ Run #34197051961
Latest PR Test (AMD ROCm 7.2): ❌ Run #34197051914

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] NameError: name 'deep_gemm' is not defined in tf32_hc_prenorm_gemm when SGLANG_ENABLE_JIT_DEEPGEMM=0

1 participant