Skip to content

[CI] fix bench_jit_per_token_group_quant_8bit sglang import cascade - #453

Open
arathi-hlab wants to merge 3 commits into
sgl-project:mainfrom
arathi-hlab:ci/fix-jit-8bit-bench-sglang-import
Open

arathi-hlab wants to merge 3 commits into
sgl-project:mainfrom
arathi-hlab:ci/fix-jit-8bit-bench-sglang-import

Conversation

@arathi-hlab

Copy link
Copy Markdown
Collaborator

Summary

Follow-up to #443. The Run Sglang Kernel Benchmarks step is now failing on bench_jit_per_token_group_quant_8bit.py with:

ModuleNotFoundError: No module named 'orjson'
  File ".../sgl_kernel/gemm.py", line 99, in sgl_per_token_group_quant_8bit
    from sglang.srt.utils import get_bool_env_var
  File "/root/sglang/python/sglang/__init__.py", line 21, in <module>
    ...
  File "/root/sglang/python/sglang/srt/utils/common.py", line 82, in <module>
    import orjson

Root cause

sgl_kernel.gemm.sgl_per_token_group_quant_8bit only imports sglang.srt.utils.get_bool_env_var when enable_v2 is None, to read the SGLANG_PER_TOKEN_GROUP_QUANT_8BIT_V2 env var. The bench was leaving enable_v2 unset, so PYTHONPATH=/root/sglang/python (added in #443 to unblock this) made sglang findable — but sglang/__init__.py cascades into srt/utils/common.py which pulls orjson, psutil, pybase64, requests, starlette, torchvision, … none installed in the CI container. Chasing those with pip installs is fragile.

Fix

  • Pin enable_v2=False in the AOT calls in bench_jit_per_token_group_quant_8bit.py so sgl_kernel never reaches the sglang-lookup branch. v2 is already covered by bench_jit_per_token_group_quant_8bit_v2.py.
  • Drop the export PYTHONPATH=/root/sglang/python from the bench step; nothing else in that step needs it.

@arathi-hlab
arathi-hlab force-pushed the ci/fix-jit-8bit-bench-sglang-import branch from 8aabfc4 to a1e16f3 Compare August 27, 2026 17:35
PR sgl-project#450 skipped this bench because it raised
`ModuleNotFoundError: No module named 'orjson'`. Root cause is that
sgl_kernel.gemm.sgl_per_token_group_quant_8bit only reads
`SGLANG_PER_TOKEN_GROUP_QUANT_8BIT_V2` via
`sglang.srt.utils.get_bool_env_var` when `enable_v2 is None`; that path
imports the full sglang package, whose common.py needs orjson, psutil,
pybase64, etc. — none installed in the bench container.

Passing enable_v2=False from the bench short-circuits the sglang lookup
entirely. v2 is already covered by
bench_jit_per_token_group_quant_8bit_v2.py, so this bench legitimately
wants the v1 kernel.

- bench: pin enable_v2=False in aot_per_token_group_quant_8bit.
- CI: run bench_jit_per_token_group_quant_8bit.py again instead of the
  skip echo added in sgl-project#450.
@arathi-hlab
arathi-hlab force-pushed the ci/fix-jit-8bit-bench-sglang-import branch from a1e16f3 to fd910ae Compare August 27, 2026 17:35

This branch has not been deployed

No deployments
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.

1 participant