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

Enable PDL for various kernels in DSV32/GLM5 (#23965)

Browse Source 
Co-authored-by: b8zhong <b8zhong@users.noreply.github.com>
This commit is contained in:
Brayden Zhong
2026-05-09 03:42:56 -07:00
committed by GitHub
co-authored by b8zhong
parent d5564c2a96
commit 05d1ab51e8
8 changed files with 71 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+7 -1
View File
@@ -299,6 +299,12 @@ class Indexer(MultiPlatformOp):
weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
return weights return weights
@torch.compile(dynamic=True)
def _apply_q_scale_and_softmax_scale(
self, weights: torch.Tensor, q_scale: torch.Tensor
):
return weights.unsqueeze(-1) * q_scale * self.softmax_scale
def _get_q_k_bf16( def _get_q_k_bf16(
self, self,
q_lora: torch.Tensor, q_lora: torch.Tensor,
@@ -1161,7 +1167,7 @@ class Indexer(MultiPlatformOp):
act_quant=act_quant, act_quant=act_quant,
) )
current_stream.wait_stream(self.alt_stream) current_stream.wait_stream(self.alt_stream)
weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale weights = self._apply_q_scale_and_softmax_scale(weights, q_scale)
else: else:
query, key = self._get_q_k_bf16( query, key = self._get_q_k_bf16(
q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
python/sglang/srt/layers/attention/utils.py
+3
View File
@@ -12,6 +12,8 @@ _is_cuda = is_cuda()
if _is_cuda: if _is_cuda:
from sgl_kernel import concat_mla_absorb_q from sgl_kernel import concat_mla_absorb_q
from sglang.jit_kernel.utils import is_arch_support_pdl
@triton.jit @triton.jit
def create_flashinfer_kv_indices_triton( def create_flashinfer_kv_indices_triton(
@@ -462,6 +464,7 @@ def mla_quantize_and_rope_for_fp8(
# Quantization scales (set to 1.0 for no additional scaling) # Quantization scales (set to 1.0 for no additional scaling)
quant_scale_q=1.0, quant_scale_q=1.0,
quant_scale_kv=1.0, quant_scale_kv=1.0,
enable_pdl=is_arch_support_pdl(),
) )
return q_out, k_nope_out, k_rope_out return q_out, k_nope_out, k_rope_out
python/sglang/srt/mem_cache/utils.py
+21
View File
@@ -20,6 +20,7 @@ import torch
import triton import triton
import triton.language as tl import triton.language as tl
from sglang.jit_kernel.utils import is_arch_support_pdl
from sglang.srt.environ import envs from sglang.srt.environ import envs
@@ -35,6 +36,7 @@ def set_mla_kv_buffer_kernel(
nope_dim: tl.constexpr, nope_dim: tl.constexpr,
rope_dim: tl.constexpr, rope_dim: tl.constexpr,
BLOCK: tl.constexpr, BLOCK: tl.constexpr,
USE_GDC: tl.constexpr = False,
): ):
pid_loc = tl.program_id(0) pid_loc = tl.program_id(0)
pid_blk = tl.program_id(1) pid_blk = tl.program_id(1)
@@ -44,6 +46,9 @@ def set_mla_kv_buffer_kernel(
total_dim = nope_dim + rope_dim total_dim = nope_dim + rope_dim
mask = offs < total_dim mask = offs < total_dim
if USE_GDC:
tl.extra.cuda.gdc_wait()
loc = tl.load(loc_ptr + pid_loc).to(tl.int64) loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs
@@ -82,6 +87,9 @@ def set_mla_kv_buffer_kernel(
tl.store(dst_ptr, src, mask=mask) tl.store(dst_ptr, src, mask=mask)
if USE_GDC:
tl.extra.cuda.gdc_launch_dependents()
def set_mla_kv_buffer_triton( def set_mla_kv_buffer_triton(
kv_buffer: torch.Tensor, kv_buffer: torch.Tensor,
@@ -96,6 +104,8 @@ def set_mla_kv_buffer_triton(
n_loc = loc.numel() n_loc = loc.numel()
grid = (n_loc, triton.cdiv(total_dim, BLOCK)) grid = (n_loc, triton.cdiv(total_dim, BLOCK))
pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
set_mla_kv_buffer_kernel[grid]( set_mla_kv_buffer_kernel[grid](
kv_buffer, kv_buffer,
cache_k_nope, cache_k_nope,
@@ -107,6 +117,7 @@ def set_mla_kv_buffer_triton(
nope_dim, nope_dim,
rope_dim, rope_dim,
BLOCK=BLOCK, BLOCK=BLOCK,
**pdl_kwargs,
) )
@@ -122,6 +133,7 @@ def set_mla_kv_buffer_fp8_quant_kernel(
nope_dim: tl.constexpr, nope_dim: tl.constexpr,
rope_dim: tl.constexpr, rope_dim: tl.constexpr,
BLOCK: tl.constexpr, BLOCK: tl.constexpr,
USE_GDC: tl.constexpr = False,
): ):
"""Fuse BF16/FP16->FP8 cast with paged KV write.""" """Fuse BF16/FP16->FP8 cast with paged KV write."""
pid_loc = tl.program_id(0) pid_loc = tl.program_id(0)
@@ -132,6 +144,9 @@ def set_mla_kv_buffer_fp8_quant_kernel(
total_dim = nope_dim + rope_dim total_dim = nope_dim + rope_dim
mask = offs < total_dim mask = offs < total_dim
if USE_GDC:
tl.extra.cuda.gdc_wait()
loc = tl.load(loc_ptr + pid_loc).to(tl.int64) loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
dst_ptr = kv_buffer_fp8_ptr + loc * buffer_stride + offs dst_ptr = kv_buffer_fp8_ptr + loc * buffer_stride + offs
@@ -165,6 +180,9 @@ def set_mla_kv_buffer_fp8_quant_kernel(
# Destination pointer is FP8-typed view; tl.store performs downcast. # Destination pointer is FP8-typed view; tl.store performs downcast.
tl.store(dst_ptr, src, mask=mask) tl.store(dst_ptr, src, mask=mask)
if USE_GDC:
tl.extra.cuda.gdc_launch_dependents()
def set_mla_kv_buffer_triton_fp8_quant( def set_mla_kv_buffer_triton_fp8_quant(
kv_buffer: torch.Tensor, kv_buffer: torch.Tensor,
@@ -183,6 +201,8 @@ def set_mla_kv_buffer_triton_fp8_quant(
n_loc = loc.numel() n_loc = loc.numel()
grid = (n_loc, triton.cdiv(total_dim, BLOCK)) grid = (n_loc, triton.cdiv(total_dim, BLOCK))
pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
set_mla_kv_buffer_fp8_quant_kernel[grid]( set_mla_kv_buffer_fp8_quant_kernel[grid](
kv_buffer_fp8, kv_buffer_fp8,
cache_k_nope, cache_k_nope,
@@ -194,6 +214,7 @@ def set_mla_kv_buffer_triton_fp8_quant(
nope_dim, nope_dim,
rope_dim, rope_dim,
BLOCK=BLOCK, BLOCK=BLOCK,
**pdl_kwargs,
) )
sgl-kernel/csrc/elementwise/pos_enc.cuh
+4 -4
View File
@@ -105,7 +105,7 @@ __global__ void BatchQKApplyRotaryPosIdsCosSinCacheEnhancedHeadParallelismKernel
const uint32_t bdy = blockDim.y; const uint32_t bdy = blockDim.y;
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.wait;"); cudaGridDependencySynchronize();
#endif #endif
vec_t<float, vec_size> cos, sin; vec_t<float, vec_size> cos, sin;
@@ -184,7 +184,7 @@ __global__ void BatchQKApplyRotaryPosIdsCosSinCacheEnhancedHeadParallelismKernel
} }
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.launch_dependents;"); cudaTriggerProgrammaticLaunchCompletion();
#endif #endif
} }
@@ -229,7 +229,7 @@ __global__ void BatchQKApplyRotaryPosIdsCosSinCacheEnhancedKernel(
const uint32_t bdy = blockDim.y; const uint32_t bdy = blockDim.y;
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.wait;"); cudaGridDependencySynchronize();
#endif #endif
vec_t<float, vec_size> cos, sin; vec_t<float, vec_size> cos, sin;
@@ -310,7 +310,7 @@ __global__ void BatchQKApplyRotaryPosIdsCosSinCacheEnhancedKernel(
} }
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.launch_dependents;"); cudaTriggerProgrammaticLaunchCompletion();
#endif #endif
} }
sgl-kernel/csrc/gemm/dsv3_fused_a_gemm.cu
+2 -2
View File
@@ -285,7 +285,7 @@ struct GmemLoaderB {
__device__ void issue_mainloop() { __device__ void issue_mainloop() {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
asm volatile("griddepcontrol.wait;"); cudaGridDependencySynchronize();
#pragma unroll 1 #pragma unroll 1
for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) { for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
if (need_wait) { if (need_wait) {
@@ -571,7 +571,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
mma_computer.issue_mainloop(); mma_computer.issue_mainloop();
mma_computer.epi(); mma_computer.epi();
} }
asm volatile("griddepcontrol.launch_dependents;"); cudaTriggerProgrammaticLaunchCompletion();
#endif #endif
} }
sgl-kernel/csrc/gemm/dsv3_router_gemm_bf16_out.cu
+2 -2
View File
@@ -75,7 +75,7 @@ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output(__nv_bfloat16* out
} }
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.wait;"); cudaGridDependencySynchronize();
#endif #endif
// Process the GEMM in chunks // Process the GEMM in chunks
@@ -159,7 +159,7 @@ __launch_bounds__(128, 1) void router_gemm_kernel_bf16_output(__nv_bfloat16* out
} }
} }
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.launch_dependents;"); cudaTriggerProgrammaticLaunchCompletion();
#endif #endif
} }
sgl-kernel/csrc/gemm/dsv3_router_gemm_float_out.cu
+2 -2
View File
@@ -74,7 +74,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output(float*
} }
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.wait;"); cudaGridDependencySynchronize();
#endif #endif
// Process the GEMM in chunks // Process the GEMM in chunks
@@ -158,7 +158,7 @@ __global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output(float*
} }
} }
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.launch_dependents;"); cudaTriggerProgrammaticLaunchCompletion();
#endif #endif
} }
sgl-kernel/csrc/gemm/per_token_group_quant_8bit_v2.cu
+30 -11
View File
@@ -271,6 +271,10 @@ __global__ void per_token_group_quant_8bit_kernel(
using scale_element_t = std::conditional_t<SCALE_UE8M0, uint8_t, float>; using scale_element_t = std::conditional_t<SCALE_UE8M0, uint8_t, float>;
static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0); static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0);
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
cudaGridDependencySynchronize();
#endif
SCHEDULER::execute<FUSE_SILU_AND_MUL, GROUP_SIZE, THREADS_PER_SUBWARP>( SCHEDULER::execute<FUSE_SILU_AND_MUL, GROUP_SIZE, THREADS_PER_SUBWARP>(
subwarps_per_block, subwarps_per_block,
hidden_dim_num_groups, hidden_dim_num_groups,
@@ -398,6 +402,10 @@ __global__ void per_token_group_quant_8bit_kernel(
reinterpret_cast<int4*>(output_q + offset_num_groups * GROUP_SIZE + lane_id * INPUT_PRIMARY_VEC_SIZE), reinterpret_cast<int4*>(output_q + offset_num_groups * GROUP_SIZE + lane_id * INPUT_PRIMARY_VEC_SIZE),
output_buf); output_buf);
}); });
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
cudaTriggerProgrammaticLaunchCompletion();
#endif
} }
void sgl_per_token_group_quant_8bit_v2( void sgl_per_token_group_quant_8bit_v2(
@@ -445,17 +453,28 @@ void sgl_per_token_group_quant_8bit_v2(
SCHEDULER::compute_exec_config( \ SCHEDULER::compute_exec_config( \
THREADS_PER_SUBWARP, num_local_experts, hidden_dim_num_groups, num_groups, subwarps_per_block, grid, block); \ THREADS_PER_SUBWARP, num_local_experts, hidden_dim_num_groups, num_groups, subwarps_per_block, grid, block); \
\ \
per_token_group_quant_8bit_kernel<SCHEDULER, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, __VA_ARGS__> \ cudaLaunchConfig_t config; \
<<<grid, block, 0, stream>>>( \ config.gridDim = grid; \
static_cast<T*>(input.data_ptr()), \ config.blockDim = block; \
static_cast<DST_DTYPE*>(output_q.data_ptr()), \ config.dynamicSmemBytes = 0; \
static_cast<output_s_dtype*>(output_s.data_ptr()), \ config.stream = stream; \
static_cast<int32_t*>(masked_m.has_value() ? masked_m->data_ptr() : 0), \ cudaLaunchAttribute attrs[1]; \
subwarps_per_block, \ attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \
hidden_dim_num_groups, \ attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL(); \
scale_expert_stride, \ config.numAttrs = 1; \
scale_hidden_stride, \ config.attrs = attrs; \
num_tokens_per_expert); \ cudaLaunchKernelEx( \
&config, \
per_token_group_quant_8bit_kernel<SCHEDULER, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, __VA_ARGS__>, \
static_cast<T*>(input.data_ptr()), \
static_cast<DST_DTYPE*>(output_q.data_ptr()), \
static_cast<output_s_dtype*>(output_s.data_ptr()), \
static_cast<int32_t*>(masked_m.has_value() ? masked_m->data_ptr() : 0), \
subwarps_per_block, \
hidden_dim_num_groups, \
scale_expert_stride, \
scale_hidden_stride, \
num_tokens_per_expert); \
} while (0) } while (0)
#define LAUNCH_KERNEL(GROUP_SIZE, T, DST_DTYPE) \ #define LAUNCH_KERNEL(GROUP_SIZE, T, DST_DTYPE) \