minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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/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,
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/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)(