[JIT] Inject target architecture flag into JIT compilation by xingsy97 · Pull Request #20103 · sgl-project/sglang

xingsy97 · 2026-03-07T17:28:55Z

Motivation

Follow-up from #19794 (comment by @DarkSharpness): __CUDA_ARCH__ is only available in the device pass, so arch macros can't be used for host-side dispatch. This PR injects -DSGL_CUDA_ARCH=<value> via load_jit() so that arch-dependent macros work in both host and device passes.

Modifications

utils.py: Inject -DSGL_CUDA_ARCH={major*100+minor*10} into JIT cuda cflags based on torch.cuda.get_device_capability().
utils.cuh: Add SGL_TARGET_CUDA_ARCH (prefers __CUDA_ARCH__ in device pass, falls back to SGL_CUDA_ARCH in host pass). Define SGL_ARCH_HOPPER_OR_GREATER and SGL_ARCH_BLACKWELL_OR_GREATER. Refactor PDL guards to use them.

Accuracy Tests

Compile-time only change. No kernel logic modified.

Benchmarking and Profiling

No performance impact.

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

gemini-code-assist · 2026-03-07T17:29:09Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the JIT compilation process by providing a consistent mechanism to detect and utilize the target CUDA architecture across both host and device compilation passes. This change is crucial for enabling architecture-specific optimizations and features, such as PDL, to function correctly regardless of the compilation stage, ultimately improving the flexibility and robustness of the JIT kernel system.

Highlights

JIT Compilation Flag Injection: Injected the -DSGL_CUDA_ARCH=<value> flag into JIT CUDA compilation cflags. This makes the target architecture information available during both host and device compilation passes, addressing a limitation where __CUDA_ARCH__ is only defined in the device pass.
Architecture Detection Macros: Introduced SGL_TARGET_CUDA_ARCH to dynamically select between __CUDA_ARCH__ (for device pass) and the newly injected SGL_CUDA_ARCH (for host pass). New macros SGL_ARCH_HOPPER_OR_GREATER and SGL_ARCH_BLACKWELL_OR_GREATER were defined for simplified architecture-specific checks.
PDL Guard Refactoring: Refactored the PDL (Program Dependent Launch) guards in PDLWaitPrimary and PDLTriggerSecondary to utilize the new SGL_ARCH_HOPPER_OR_GREATER macro, ensuring consistent architecture-specific behavior.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

python/sglang/jit_kernel/include/sgl_kernel/utils.cuh
- Added architecture detection macros SGL_TARGET_CUDA_ARCH, SGL_ARCH_HOPPER_OR_GREATER, and SGL_ARCH_BLACKWELL_OR_GREATER.
- Updated PDLWaitPrimary to use SGL_ARCH_HOPPER_OR_GREATER for architecture checks.
- Updated PDLTriggerSecondary to use SGL_ARCH_HOPPER_OR_GREATER for architecture checks.
python/sglang/jit_kernel/utils.py
- Modified load_jit to inject -DSGL_CUDA_ARCH into extra_cuda_cflags for non-ROCM builds.
- Added a new helper function _get_cuda_arch_value to compute the CUDA architecture value from torch.cuda.get_device_capability().

Activity

The pull request was motivated by a comment from @DarkSharpness in Add compile-time 256-bit vector guard for pre-Blackwell #19794 regarding the availability of __CUDA_ARCH__.
The author has marked the code as formatted according to pre-commit hooks and compliant with SGLang code style guidance.
The author noted that this is a compile-time only change with no kernel logic modified, and no performance impact, with PDL codegen remaining identical.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces a mechanism to inject the target CUDA architecture as a compile-time flag, which is a great improvement for enabling host-side dispatch based on GPU capabilities. The implementation is clean, with a clear fallback from __CUDA_ARCH__ to the new SGL_CUDA_ARCH macro. The refactoring of the PDL guards to use the new SGL_ARCH_HOPPER_OR_GREATER macro is also well done.

I have one point of feedback regarding the definition for Blackwell architecture support, which uses a speculative future CUDA toolkit version. Please see the specific comment for details.

xingsy97 · 2026-03-12T14:18:20Z

Hi @DarkSharpness , could you help take a look at this PR? It's a follow-up to #19794 based on your suggestion. Once this is merged, #19794 can be resolved as well. Thanks!

DarkSharpness

LGTM. cc @BBuf

DarkSharpness · 2026-03-12T17:14:36Z

/tag-and-rerun-ci

xingsy97 · 2026-03-14T08:49:52Z

/rerun-failed-ci

xingsy97 · 2026-03-14T09:31:22Z

Hi @BBuf , request merge for PR.

jit-kernel-benchmark-test fails because new compile flag -DSGL_CUDA_ARCH invalidates JIT cache. bench_gptq_marlin.py and bench_moe_wna16_marlin.py need > 120s to cold compile, hitting the CI timeout 120s. Thus I think it's not a code bug.

xingsy97 · 2026-03-17T10:06:23Z

Hi @BBuf , request merge for PR.

jit-kernel-benchmark-test fails because new compile flag -DSGL_CUDA_ARCH invalidates JIT cache. bench_gptq_marlin.py and bench_moe_wna16_marlin.py need > 120s to cold compile, hitting the CI timeout 120s. Thus I think it's not a code bug.

Hi @Fridge003 @BBuf could you help merge this PR? It's approved and I have other PRs depending on this one. Thanks!

b8zhong · 2026-03-18T16:04:46Z

@xingsy97 Hi can you take a look at this bug? I encountered this on SM120... thanks.

[2026-03-18 15:57:31 TP1] Scheduler hit an exception: Traceback (most recent call last):
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py", line 643, in __init__
    self.capture()
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py", line 799, in capture
    _capture_one_stream()
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py", line 786, in _capture_one_stream
    ) = self.capture_one_batch_size(bs, forward, stream_idx)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py", line 1008, in capture_one_batch_size
    run_once()
  File "/sgl-workspace/sglang/python/sglang/srt/model_executor/cuda_graph_runner.py", line 995, in run_once
    logits_output_or_pp_proxy_tensors = forward(
                                        ^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/models/minimax_m2.py", line 1005, in forward
    hidden_states = self.model(
                    ^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/models/minimax_m2.py", line 919, in forward
    hidden_states, residual = layer(
                              ^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/models/minimax_m2.py", line 763, in forward
    hidden_states = self.block_sparse_moe(hidden_states, forward_batch)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/models/minimax_m2.py", line 391, in forward
    return self.forward_normal(hidden_states)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/models/minimax_m2.py", line 401, in forward_normal
    final_hidden_states = self.experts(hidden_states, topk_output)
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/layer.py", line 994, in forward
    return self.forward_impl(hidden_states, topk_output)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/layer.py", line 1013, in forward_impl
    combine_input = self.run_moe_core(
                    ^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/layer.py", line 1034, in run_moe_core
    return self.quant_method.apply(
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/layers/quantization/modelopt_quant.py", line 2003, in apply
    output = cutlass_moe_fp4(
             ^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/srt/layers/moe/cutlass_moe.py", line 451, in cutlass_moe_fp4
    rep_a_fp4, rep_a_blockscale = scaled_fp4_experts_quant(
                                  ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/jit_kernel/nvfp4.py", line 433, in scaled_fp4_experts_quant
    _scaled_fp4_experts_quant_custom_op(
  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 1255, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/jit_kernel/nvfp4.py", line 374, in _scaled_fp4_experts_quant_custom_op
    module = _jit_nvfp4_expert_quant_module()
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/jit_kernel/utils.py", line 43, in wrapper
    result_map[key] = fn(*args, **kwargs)
                      ^^^^^^^^^^^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/jit_kernel/nvfp4.py", line 159, in _jit_nvfp4_expert_quant_module
    return load_jit(
           ^^^^^^^^^
  File "/sgl-workspace/sglang/python/sglang/jit_kernel/utils.py", line 198, in load_jit
    return load_inline(
           ^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tvm_ffi/cpp/extension.py", line 1035, in load_inline
    build_inline(
  File "/usr/local/lib/python3.12/dist-packages/tvm_ffi/cpp/extension.py", line 877, in build_inline
    return _build_impl(
           ^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tvm_ffi/cpp/extension.py", line 672, in _build_impl
    build_ninja(str(build_dir))
  File "/usr/local/lib/python3.12/dist-packages/tvm_ffi/cpp/extension.py", line 542, in build_ninja
    raise RuntimeError("\n".join(msg))
RuntimeError: ninja exited with status 1
stdout:
[1/2] /usr/local/cuda/bin/nvcc  --generate-dependencies-with-compile --dependency-output cuda_0.o.d -Xcompiler -fPIC -std=c++17 -O2 -gencode=arch=compute_120a,code=sm_120a -gencode=arch=compute_103a,code=sm_103a -std=c++20 -O3 --expt-relaxed-constexpr -DSGL_CUDA_ARCH=1200 -DNDEBUG -DFLASHINFER_ENABLE_F16 -DCUTE_USE_PACKED_TUPLE=1 -DCUTLASS_ENABLE_TENSOR_CORE_MMA=1 -DCUTLASS_VERSIONS_GENERATED -DCUTLASS_TEST_LEVEL=0 -DCUTLASS_TEST_ENABLE_CACHED_RESULTS=1 -DCUTLASS_DEBUG_TRACE_LEVEL=0 --expt-extended-lambda -I/usr/local/lib/python3.12/dist-packages/tvm_ffi/include -I/usr/local/lib/python3.12/dist-packages/tvm_ffi/include -I/sgl-workspace/sglang/python/sglang/jit_kernel/include -I/usr/local/lib/python3.12/dist-packages/flashinfer/data/cutlass/include -I/usr/local/lib/python3.12/dist-packages/flashinfer/data/cutlass/tools/util/include -I/usr/local/lib/python3.12/dist-packages/deep_gemm/include -c /root/.cache/tvm-ffi/sgl_kernel_jit_nvfp4_expert_quant_ae71782d466c305d/cuda.cu -o cuda_0.o
FAILED: [code=1] cuda_0.o 
/usr/local/cuda/bin/nvcc  --generate-dependencies-with-compile --dependency-output cuda_0.o.d -Xcompiler -fPIC -std=c++17 -O2 -gencode=arch=compute_120a,code=sm_120a -gencode=arch=compute_103a,code=sm_103a -std=c++20 -O3 --expt-relaxed-constexpr -DSGL_CUDA_ARCH=1200 -DNDEBUG -DFLASHINFER_ENABLE_F16 -DCUTE_USE_PACKED_TUPLE=1 -DCUTLASS_ENABLE_TENSOR_CORE_MMA=1 -DCUTLASS_VERSIONS_GENERATED -DCUTLASS_TEST_LEVEL=0 -DCUTLASS_TEST_ENABLE_CACHED_RESULTS=1 -DCUTLASS_DEBUG_TRACE_LEVEL=0 --expt-extended-lambda -I/usr/local/lib/python3.12/dist-packages/tvm_ffi/include -I/usr/local/lib/python3.12/dist-packages/tvm_ffi/include -I/sgl-workspace/sglang/python/sglang/jit_kernel/include -I/usr/local/lib/python3.12/dist-packages/flashinfer/data/cutlass/include -I/usr/local/lib/python3.12/dist-packages/flashinfer/data/cutlass/tools/util/include -I/usr/local/lib/python3.12/dist-packages/deep_gemm/include -c /root/.cache/tvm-ffi/sgl_kernel_jit_nvfp4_expert_quant_ae71782d466c305d/cuda.cu -o cuda_0.o
nvcc warning : incompatible redefinition for option 'std', the last value of this option was used
nvcc warning : incompatible redefinition for option 'optimize', the last value of this option was used
ptxas warning : Value of threads per SM for entry _Z15cvt_fp16_to_fp4I13__nv_bfloat16Lb0ELb1EEviiPKT_PKfPjS6_S6_S6_Piib is out of range. .minnctapersm will be ignored
ptxas warning : Value of threads per SM for entry _Z15cvt_fp16_to_fp4I13__nv_bfloat16Lb0ELb0EEviiPKT_PKfPjS6_S6_S6_Piib is out of range. .minnctapersm will be ignored
ptxas warning : Value of threads per SM for entry _Z15cvt_fp16_to_fp4I13__nv_bfloat16Lb0ELb1EEviiPKT_PKfPjS6_S6_S6_Pii is out of range. .minnctapersm will be ignored
ptxas warning : Value of threads per SM for entry _Z15cvt_fp16_to_fp4I13__nv_bfloat16Lb0ELb0EEviiPKT_PKfPjS6_S6_S6_Pii is out of range. .minnctapersm will be ignored
ptxas warning : Value of threads per SM for entry _Z22cvt_fp16_to_fp4_expertI13__nv_bfloat16Lb0EEviiPKT_PKfPjS6_Pibi is out of range. .minnctapersm will be ignored
ptxas warning : Value of threads per SM for entry _Z15cvt_fp16_to_fp4I6__halfLb0ELb1EEviiPKT_PKfPjS6_S6_S6_Piib is out of range. .minnctapersm will be ignored
ptxas warning : Value of threads per SM for entry _Z15cvt_fp16_to_fp4I6__halfLb0ELb0EEviiPKT_PKfPjS6_S6_S6_Piib is out of range. .minnctapersm will be ignored
ptxas warning : Value of threads per SM for entry _Z15cvt_fp16_to_fp4I6__halfLb0ELb1EEviiPKT_PKfPjS6_S6_S6_Pii is out of range. .minnctapersm will be ignored
ptxas warning : Value of threads per SM for entry _Z15cvt_fp16_to_fp4I6__halfLb0ELb0EEviiPKT_PKfPjS6_S6_S6_Pii is out of range. .minnctapersm will be ignored
ptxas warning : Value of threads per SM for entry _Z22cvt_fp16_to_fp4_expertI6__halfLb0EEviiPKT_PKfPjS6_Pibi is out of range. .minnctapersm will be ignored
/sgl-workspace/sglang/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh(99): error: static assertion failed with "SGL_CUDA_ARCH mismatch: injected arch flag does not match device target"
  static_assert(
  ^

1 error detected in the compilation of "/root/.cache/tvm-ffi/sgl_kernel_jit_nvfp4_expert_quant_ae71782d466c305d/cuda.cu".
ninja: build stopped: subcommand failed.

DarkSharpness · 2026-03-18T16:11:33Z

/usr/local/cuda/bin/nvcc --generate-dependencies-with-compile --dependency-output cuda_0.o.d -Xcompiler -fPIC -std=c++17 -O2 -gencode=arch=compute_120a,code=sm_120a -gencode=arch=compute_103a,code=sm_103a -std=c++20 -O3 --expt-relaxed-constexpr -DSGL_CUDA_ARCH=1200 -DNDEBUG -DFLASHINFER_ENABLE_F16 -DCUTE_USE_PACKED_TUPLE=1 -DCUTLASS_ENABLE_TENSOR_CORE_MMA=1 -DCUTLASS_VERSIONS_GENERATED -DCUTLASS_TEST_LEVEL=0 -DCUTLASS_TEST_ENABLE_CACHED_RESULTS=1 -DCUTLASS_DEBUG_TRACE_LEVEL=0 --expt-extended-lambda -I/usr/local/lib/python3.12/dist-packages/tvm_ffi/include -I/usr/local/lib/python3.12/dist-packages/tvm_ffi/include -I/sgl-workspace/sglang/python/sglang/jit_kernel/include -I/usr/local/lib/python3.12/dist-packages/flashinfer/data/cutlass/include -I/usr/local/lib/python3.12/dist-packages/flashinfer/data/cutlass/tools/util/include -I/usr/local/lib/python3.12/dist-packages/deep_gemm/include -c /root/.cache/tvm-ffi/sgl_kernel_jit_nvfp4_expert_quant_ae71782d466c305d/cuda.cu -o cuda_0.o

The compile commands looks strange @b8zhong :

/usr/local/cuda/bin/nvcc  --generate-dependencies-with-compile --dependency-output cuda_0.o.d -Xcompiler -fPIC -std=c++17 -O2 -gencode=arch=compute_120a,code=sm_120a -gencode=arch=compute_103a,code=sm_103a -std=c++20 -O3 --expt-relaxed-constexpr -DSGL_CUDA_ARCH=1200 -DNDEBUG -DFLASHINFER_ENABLE_F16 -DCUTE_USE_PACKED_TUPLE=1 -DCUTLASS_ENABLE_TENSOR_CORE_MMA=1 -DCUTLASS_VERSIONS_GENERATED -DCUTLASS_TEST_LEVEL=0 -DCUTLASS_TEST_ENABLE_CACHED_RESULTS=1 -DCUTLASS_DEBUG_TRACE_LEVEL=0 --expt-extended-lambda -I/usr/local/lib/python3.12/dist-packages/tvm_ffi/include -I/usr/local/lib/python3.12/dist-packages/tvm_ffi/include -I/sgl-workspace/sglang/python/sglang/jit_kernel/include -I/usr/local/lib/python3.12/dist-packages/flashinfer/data/cutlass/include -I/usr/local/lib/python3.12/dist-packages/flashinfer/data/cutlass/tools/util/include -I/usr/local/lib/python3.12/dist-packages/deep_gemm/include -c /root/.cache/tvm-ffi/sgl_kernel_jit_nvfp4_expert_quant_ae71782d466c305d/cuda.cu -o cuda_0.o

We somehow have both -gencode=arch=compute_120a,code=sm_120a -gencode=arch=compute_103a,code=sm_103a

Have you ever set TVM_FFI_CUDA_ARCH_LIST?

DarkSharpness · 2026-03-18T16:16:08Z

The root cause should be this: @b8zhong

sglang/python/sglang/jit_kernel/nvfp4.py

Lines 112 to 124 in 20a23e3

    
           @contextmanager 
        
           def _nvfp4_arch_env(): 
        
               key = "TVM_FFI_CUDA_ARCH_LIST" 
        
               old_val = os.environ.get(key) 
        
               os.environ[key] = _get_nvfp4_cuda_arch_list() 
        
               try: 
        
                   yield 
        
               finally: 
        
                   if old_val is None: 
        
                       os.environ.pop(key, None) 
        
                   else: 
        
                       os.environ[key] = old_val

These lines override the env var, which lead to error since we only assume one target in this PR.

mmangkad · 2026-03-18T16:47:42Z

Good catch on adding the static_assert here

xingsy97 · 2026-03-18T16:48:37Z

@b8zhong @DarkSharpness Thanks for report and quick investigation. I'm writing a fix for this.

DarkSharpness · 2026-03-18T16:49:20Z

@b8zhong @DarkSharpness Thanks for report and quick investigation. I'm writing a fix for this.

#20874 is fixing this. PTAL @xingsy97

b8zhong · 2026-03-18T16:50:02Z

Hi, thanks @DarkSharpness. I did not

xingsy97 · 2026-03-18T17:29:01Z

Conclusion for the bug reported by @b8zhong

Root cause: a JIT kernel built for multi-arch targets (seems copied from AOT logic), but load_jit() only injects a single SGL_CUDA_ARCH value, so the static_assert is triggered.
Being fixed by [JIT Kernel] Fix NVFP4 multi-arch compilation failure #20874 (@mmangkad).
This PR should be kept as-is. JIT only targets the current device, building kernel for multi-arch should be considered as unexpected usage. cc @DarkSharpness
Checked other JIT kernels, none of them have similar issue.

…ct#20103)

feat(jit): inject target arch flag into JIT compilation

2a0c584

gemini-code-assist Bot reviewed Mar 7, 2026

View reviewed changes

Comment thread python/sglang/jit_kernel/include/sgl_kernel/utils.cuh Outdated

xingsy97 mentioned this pull request Mar 7, 2026

Add compile-time 256-bit vector guard for pre-Blackwell #19794

Merged

5 tasks

xingsy97 marked this pull request as ready for review March 7, 2026 17:36

xingsy97 requested review from BBuf, DarkSharpness, HydraQYH, celve and yuan-luo as code owners March 7, 2026 17:36

DarkSharpness requested changes Mar 12, 2026

View reviewed changes

Comment thread python/sglang/jit_kernel/include/sgl_kernel/utils.cuh

add static assert

61b94bd

xingsy97 requested a review from DarkSharpness March 12, 2026 16:27

simplify

bdcd40e

DarkSharpness approved these changes Mar 12, 2026

View reviewed changes

github-actions Bot added the run-ci label Mar 12, 2026

fix lint error

455c9fe

github-actions Bot added the jit-kernel label Mar 14, 2026

Kangyan-Zhou merged commit d20e9a2 into sgl-project:main Mar 18, 2026
114 of 163 checks passed

xingsy97 deleted the wt/jit-inject-arch-flag branch March 18, 2026 06:33

mmangkad mentioned this pull request Mar 18, 2026

[JIT Kernel] Fix NVFP4 multi-arch compilation failure #20874

Merged

Wangzheee pushed a commit to Wangzheee/sglang that referenced this pull request Mar 21, 2026

[JIT] Inject target architecture flag into JIT compilation (sgl-proje…

2f74fde

…ct#20103)

JustinTong0323 pushed a commit to JustinTong0323/sglang that referenced this pull request Apr 7, 2026

[JIT] Inject target architecture flag into JIT compilation (sgl-proje…

092bce4

…ct#20103)

Chronostasys pushed a commit to MindLab-Research/sglang that referenced this pull request Aug 24, 2026

[JIT] Inject target architecture flag into JIT compilation (sgl-proje…

1bf4331

…ct#20103)

Uh oh!

Conversation

xingsy97 commented Mar 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Uh oh!

gemini-code-assist Bot commented Mar 7, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

xingsy97 commented Mar 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

DarkSharpness left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

DarkSharpness commented Mar 12, 2026

Uh oh!

xingsy97 commented Mar 14, 2026

Uh oh!

xingsy97 commented Mar 14, 2026

Uh oh!

xingsy97 commented Mar 17, 2026

Uh oh!

Uh oh!

b8zhong commented Mar 18, 2026

Uh oh!

DarkSharpness commented Mar 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

DarkSharpness commented Mar 18, 2026

Uh oh!

mmangkad commented Mar 18, 2026

Uh oh!

xingsy97 commented Mar 18, 2026

Uh oh!

DarkSharpness commented Mar 18, 2026

Uh oh!

b8zhong commented Mar 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

xingsy97 commented Mar 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

xingsy97 commented Mar 7, 2026 •

edited

Loading

xingsy97 commented Mar 12, 2026 •

edited

Loading

DarkSharpness commented Mar 18, 2026 •

edited

Loading

b8zhong commented Mar 18, 2026 •

edited

Loading

xingsy97 commented Mar 18, 2026 •

edited

Loading