chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
xingsy97
2026-03-20 18:25:12 +08:00
committed by GitHub
parent 2dd9196079
commit f41832795e
4 changed files with 15 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh
View File

@@ -1,7 +1,6 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/tile.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
@@ -150,11 +149,9 @@ struct FusedAddRMSNormKernel {
.with_device(device)
.verify(residual);
auto cc_major = host::runtime::get_cc_major(device.unwrap().device_id);
int hidden_size = static_cast<int>(D.unwrap());
if ((cc_major <= 9 && hidden_size <= 8192) || (cc_major >= 10 && hidden_size <= 12288)) {
int max_vec_size_byte = cc_major >= 10 ? 32 : 16;
int elements_in_vec = max_vec_size_byte / sizeof(DType);
if (hidden_size <= (device::kMaxVecBytes == 32 ? 12288 : 8192)) {
int elements_in_vec = device::kMaxVecBytes / sizeof(DType);
int vec_hidden_size = hidden_size / elements_in_vec;
uint threads = (vec_hidden_size + 31) / 32 * 32;
@@ -167,8 +164,7 @@ struct FusedAddRMSNormKernel {
elements_in_vec);
// Launch kernel
auto kernel =
max_vec_size_byte == 32 ? fused_add_rmsnorm_reg_kernel<DType, 32> : fused_add_rmsnorm_reg_kernel<DType, 16>;
auto kernel = fused_add_rmsnorm_reg_kernel<DType, device::kMaxVecBytes>;
LaunchKernel(static_cast<uint>(N.unwrap()), threads, device.unwrap())
.enable_pdl(false)(
kernel,

10
python/sglang/jit_kernel/csrc/elementwise/qknorm_across_heads.cuh
View File

@@ -1,7 +1,6 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/tile.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
@@ -194,11 +193,9 @@ struct QKNormAcrossHeadsKernel {
.with_device(device)
.verify(k_weight);
auto cc_major = host::runtime::get_cc_major(device.unwrap().device_id);
int hidden_size = static_cast<int>(D.unwrap());
if ((cc_major <= 9 && hidden_size <= 8192) || (cc_major >= 10 && hidden_size <= 12288)) {
int max_vec_size_byte = cc_major >= 10 ? 32 : 16;
int elements_in_vec = max_vec_size_byte / sizeof(DType);
if (hidden_size <= (device::kMaxVecBytes == 32 ? 12288 : 8192)) {
int elements_in_vec = device::kMaxVecBytes / sizeof(DType);
int vec_hidden_size = hidden_size / elements_in_vec;
uint threads = (vec_hidden_size + 31) / 32 * 32;
@@ -211,8 +208,7 @@ struct QKNormAcrossHeadsKernel {
elements_in_vec);
// Launch single kernel for both q and k
auto kernel = max_vec_size_byte == 32 ? qknorm_across_heads_reg_kernel<DType, 32>
: qknorm_across_heads_reg_kernel<DType, 16>;
auto kernel = qknorm_across_heads_reg_kernel<DType, device::kMaxVecBytes>;
LaunchKernel(static_cast<uint>(N.unwrap()), threads, device.unwrap())
.enable_pdl(false)(

5
python/sglang/jit_kernel/include/sgl_kernel/utils.cuh
View File

@@ -106,6 +106,11 @@ static_assert(
#define SGL_ARCH_BLACKWELL_OR_GREATER 0
#endif
// Maximum vector size in bytes supported by current architecture.
// Pre-Blackwell / AMD: 128-bit (16 bytes)
// Blackwell or greater: 256-bit (32 bytes)
inline constexpr std::size_t kMaxVecBytes = SGL_ARCH_BLACKWELL_OR_GREATER ? 32 : 16;
/// \brief Number of threads per warp (always 32 on NVIDIA/AMD GPUs).
inline constexpr auto kWarpThreads = 32u;
/// \brief Full warp active mask (all 32 lanes).

6
python/sglang/jit_kernel/include/sgl_kernel/vec.cuh
View File

@@ -73,8 +73,10 @@ struct alignas(sizeof(T) * N) AlignedStorage {
template <typename T, std::size_t N>
struct AlignedVector {
private:
/// NOTE: N must be a power of two and sizeof(T) * N <= 32 bytes (256 bits)
static_assert((N > 0 && (N & (N - 1)) == 0) && sizeof(T) * N <= 32, "CUDA only supports at most 256-bit vector op");
static_assert(
(N > 0 && (N & (N - 1)) == 0) && sizeof(T) * N <= kMaxVecBytes,
"CUDA vector size exceeds arch limit: max 16 bytes on pre-Blackwell/AMD, "
"32 bytes on Blackwell or greater");
using element_t = typename details::sized_int<T>;
using storage_t = AlignedStorage<element_t, N>;