diff --git a/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh b/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh index 7750d1266..1d2774f6f 100644 --- a/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh +++ b/python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh @@ -55,7 +55,7 @@ SGL_DEVICE CacheDType load_cache_value(const CacheDType* ptr, int64_t idx) { template SGL_DEVICE T rotary_mul_rn(T lhs, T rhs) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1200 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1030 || __CUDA_ARCH__ >= 1200) uint16_t lhs_bits; uint16_t rhs_bits; if constexpr (std::is_same_v) { @@ -83,9 +83,10 @@ SGL_DEVICE T rotary_mul_rn(T lhs, T rhs) { template SGL_DEVICE T rotary_add(T x, T cos, T y, T sin) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1200 - // nvcc may contract the packed local expression on SM120 even though the - // reference RoPE kernel rounds both products to the activation dtype first. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1030 || __CUDA_ARCH__ >= 1200) + // nvcc may contract the packed local expression on Blackwell SM103/SM120 + // even though the reference RoPE kernel rounds both products to the + // activation dtype first. const T lhs = rotary_mul_rn(x, cos); const T rhs = rotary_mul_rn(y, sin); uint16_t lhs_bits; @@ -115,7 +116,7 @@ SGL_DEVICE T rotary_add(T x, T cos, T y, T sin) { template SGL_DEVICE T rotary_sub(T x, T cos, T y, T sin) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1200 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1030 || __CUDA_ARCH__ >= 1200) const T lhs = rotary_mul_rn(x, cos); const T rhs = rotary_mul_rn(y, sin); uint16_t lhs_bits;