minke.yu/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
python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh
+3 -7
View File
@@ -1,7 +1,6 @@
#include <sgl_kernel/tensor.h> #include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h> #include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/tile.cuh> #include <sgl_kernel/tile.cuh>
#include <sgl_kernel/type.cuh> #include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh> #include <sgl_kernel/utils.cuh>
@@ -150,11 +149,9 @@ struct FusedAddRMSNormKernel {
.with_device(device) .with_device(device)
.verify(residual); .verify(residual);
auto cc_major = host::runtime::get_cc_major(device.unwrap().device_id);
int hidden_size = static_cast<int>(D.unwrap()); int hidden_size = static_cast<int>(D.unwrap());
if ((cc_major <= 9 && hidden_size <= 8192) || (cc_major >= 10 && hidden_size <= 12288)) { if (hidden_size <= (device::kMaxVecBytes == 32 ? 12288 : 8192)) {
int max_vec_size_byte = cc_major >= 10 ? 32 : 16; int elements_in_vec = device::kMaxVecBytes / sizeof(DType);
int elements_in_vec = max_vec_size_byte / sizeof(DType);
int vec_hidden_size = hidden_size / elements_in_vec; int vec_hidden_size = hidden_size / elements_in_vec;
uint threads = (vec_hidden_size + 31) / 32 * 32; uint threads = (vec_hidden_size + 31) / 32 * 32;
@@ -167,8 +164,7 @@ struct FusedAddRMSNormKernel {
elements_in_vec); elements_in_vec);
// Launch kernel // Launch kernel
auto kernel = auto kernel = fused_add_rmsnorm_reg_kernel<DType, device::kMaxVecBytes>;
max_vec_size_byte == 32 ? fused_add_rmsnorm_reg_kernel<DType, 32> : fused_add_rmsnorm_reg_kernel<DType, 16>;
LaunchKernel(static_cast<uint>(N.unwrap()), threads, device.unwrap()) LaunchKernel(static_cast<uint>(N.unwrap()), threads, device.unwrap())
.enable_pdl(false)( .enable_pdl(false)(
kernel, kernel,
python/sglang/jit_kernel/csrc/elementwise/qknorm_across_heads.cuh
+3 -7
View File
@@ -1,7 +1,6 @@
#include <sgl_kernel/tensor.h> #include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h> #include <sgl_kernel/utils.h>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/tile.cuh> #include <sgl_kernel/tile.cuh>
#include <sgl_kernel/type.cuh> #include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh> #include <sgl_kernel/utils.cuh>
@@ -194,11 +193,9 @@ struct QKNormAcrossHeadsKernel {
.with_device(device) .with_device(device)
.verify(k_weight); .verify(k_weight);
auto cc_major = host::runtime::get_cc_major(device.unwrap().device_id);
int hidden_size = static_cast<int>(D.unwrap()); int hidden_size = static_cast<int>(D.unwrap());
if ((cc_major <= 9 && hidden_size <= 8192) || (cc_major >= 10 && hidden_size <= 12288)) { if (hidden_size <= (device::kMaxVecBytes == 32 ? 12288 : 8192)) {
int max_vec_size_byte = cc_major >= 10 ? 32 : 16; int elements_in_vec = device::kMaxVecBytes / sizeof(DType);
int elements_in_vec = max_vec_size_byte / sizeof(DType);
int vec_hidden_size = hidden_size / elements_in_vec; int vec_hidden_size = hidden_size / elements_in_vec;
uint threads = (vec_hidden_size + 31) / 32 * 32; uint threads = (vec_hidden_size + 31) / 32 * 32;
@@ -211,8 +208,7 @@ struct QKNormAcrossHeadsKernel {
elements_in_vec); elements_in_vec);
// Launch single kernel for both q and k // Launch single kernel for both q and k
auto kernel = max_vec_size_byte == 32 ? qknorm_across_heads_reg_kernel<DType, 32> auto kernel = qknorm_across_heads_reg_kernel<DType, device::kMaxVecBytes>;
: qknorm_across_heads_reg_kernel<DType, 16>;
LaunchKernel(static_cast<uint>(N.unwrap()), threads, device.unwrap()) LaunchKernel(static_cast<uint>(N.unwrap()), threads, device.unwrap())
.enable_pdl(false)( .enable_pdl(false)(
python/sglang/jit_kernel/include/sgl_kernel/utils.cuh
+5
View File
@@ -106,6 +106,11 @@ static_assert(
#define SGL_ARCH_BLACKWELL_OR_GREATER 0 #define SGL_ARCH_BLACKWELL_OR_GREATER 0
#endif #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). /// \brief Number of threads per warp (always 32 on NVIDIA/AMD GPUs).
inline constexpr auto kWarpThreads = 32u; inline constexpr auto kWarpThreads = 32u;
/// \brief Full warp active mask (all 32 lanes). /// \brief Full warp active mask (all 32 lanes).
python/sglang/jit_kernel/include/sgl_kernel/vec.cuh
+4 -2
View File
@@ -73,8 +73,10 @@ struct alignas(sizeof(T) * N) AlignedStorage {
template <typename T, std::size_t N> template <typename T, std::size_t N>
struct AlignedVector { struct AlignedVector {
private: private:
/// NOTE: N must be a power of two and sizeof(T) * N <= 32 bytes (256 bits) static_assert(
static_assert((N > 0 && (N & (N - 1)) == 0) && sizeof(T) * N <= 32, "CUDA only supports at most 256-bit vector op"); (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 element_t = typename details::sized_int<T>;
using storage_t = AlignedStorage<element_t, N>; using storage_t = AlignedStorage<element_t, N>;