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

[xingsy97]

Motivation

Follow-up from #19770 review discussion. On pre-Blackwell GPUs (SM < 100), CUDA only supports 128-bit vector load/store. The 32-byte AlignedVector path compiles but should never run. This PR adds a compile-time guard so invalid 256-bit instantiations are caught at compile time rather than relying solely on runtime dispatch.

Modifications

• utils.cuh: Add SGL_ARCH_IS_HOPPER_PLUS and SGL_ARCH_IS_BLACKWELL_PLUS macros. Refactor existing PDL __CUDA_ARCH__ checks to use them.
• vec.cuh: Split AlignedVector static_assert by arch — 16 bytes max on pre-Blackwell, 32 bytes on Blackwell+.
• fused_add_rmsnorm.cuh / qknorm_across_heads.cuh: Add if constexpr early return in 32B kernel template on pre-Blackwell, preventing the body (which instantiates 32-byte AlignedVector) from being compiled.

Accuracy Tests

No model output changes. This is a compile-time guard only — kernel logic is unchanged.

Benchmarking and Profiling

No performance impact. The guard is resolved at compile time (if constexpr / #if preprocessor), producing identical codegen for all existing paths.

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]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[DarkSharpness]

SGL_ARCH_IS_HOPPER_PLUS / SGL_ARCH_IS_BLACKWELL_PLUS . Can we improve the naming? Maybe we can refer to other open-source libraries as example.

[xingsy97]

What about SGL_ARCH_HOPPER_OR_GREATER / SGL_ARCH_BLACKWELL_OR_GREATER, inspired by .NET macro naming. See https://learn.microsoft.com/en-us/dotnet/standard/frameworks#preprocessor-symbols ,it looks pretty clear

[DarkSharpness]

I'm not sure here (though SGL_ARCH_HOPPER_OR_GREATER do look better to me). cc @HydraQYH @BBuf if you have any idea

[DarkSharpness]

Given that we are testing the architecture from compile-time macros, we should remove this line and change the kernel dispatch logic based on the compile-time macro values. We should never instantiate invalid kernel (e.g. VEC_SIZE_IN_BYTE = 32 on hopper), i.e. the constexpr guard in device side should be moved to host side.

[xingsy97]

found a blocker for such change

1. SGL_ARCH_IS_BLACKWELL_PLUS relies on __CUDA_ARCH__, which is only defined on the device side. In host-side run() (qknorm_across_heads.cuh:L168), it always evaluates to 0, even on Blackwell.

2. cc_major is a runtime value. It gives the correct answer but can't prevent template instantiation. Both kernel<16> and kernel<32> are always instantiated (qknorm_across_heads.cuh:L216).

3. To achieve zero-instantiation on the host side, we'd need the JIT compilation pipeline to pass a compiler flag (e.g. -DSGL_CC_MAJOR=10) so that host code can use #if to select the kernel template. This requires changes to load_jit().

The current device-side if constexpr approach is correct, though the invalid kernel gets instantiated with an empty body. Considering the additional complexity involved (also somewhat out of scope for this PR), I'd suggest keeping it as-is and maybe addressing it in a follow-up PR if needed. What do you think?

[DarkSharpness]

Actually, in Python-side, we can get the architecture information from torch. We may add some macros which indicate the target architecture (which works for both host/device code) in default compile flags.

[xingsy97]

Makes sense. I'll add the default architecture macro in load_jit() and update the host-side dispatch in a separate PR, since it will have a broader impact. This PR will keep the current approach as-is.

[DarkSharpness]

I think that a value that differs in host/device code is highly error prune. BTW this PR seems to break some JIT kernel unit-tests https://github.com/sgl-project/sglang/actions/runs/22722498873/job/65904747412?pr=19794

PTAL

[BBuf]

Why we not do like this:

void run_kernel(...) {
    // ...
    
    #if defined(SGL_ARCH_IS_BLACKWELL_PLUS)
      launch_kernel<32>(...);
    #else
      launch_kernel<16>(...);
    #endif
}

[DarkSharpness]

Why we not do like this:

void run_kernel(...) {
    // ...
    
    #if defined(SGL_ARCH_IS_BLACKWELL_PLUS)
      launch_kernel<32>(...);
    #else
      launch_kernel<16>(...);
    #endif
}

That's because "SGL_ARCH_IS_BLACKWELL_PLUS relies on CUDA_ARCH, which is only defined on the device side. In host-side run() (qknorm_across_heads.cuh:L168), it always evaluates to 0, even on Blackwell."

That's why I suggest we introduce new compile flags to specify the target architecture in both host/device code.

[BBuf]

Make sense.

[xingsy97]

I opened a follow-up PR #20103 to inject target arch flag into JIT compilation so that arch-dependent macros work in both host and device passes, as @DarkSharpness suggested. cc @BBuf

After that follow-up PR merged, this PR will be much simpler.

[DarkSharpness]

/tag-and-rerun-ci

[xingsy97]

Hi @DarkSharpness, I rewrote this PR, rebased on #20103. Added inline constexpr kMaxVecBytes (derived from SGL_ARCH_BLACKWELL_OR_GREATER, consistent in host/device). Kernel dispatch now uses kernel<DType, device::kMaxVecBytes> directly, so only the valid kernel is instantiated. PTAL.

[DarkSharpness]

LGTM. cc @BBuf @HydraQYH

[BBuf]

LGTM