minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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
sgl-kernel/csrc/elementwise/pos_enc.cuh
+4 -4
View File
@@ -105,7 +105,7 @@ __global__ void BatchQKApplyRotaryPosIdsCosSinCacheEnhancedHeadParallelismKernel
const uint32_t bdy = blockDim.y;
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.wait;");
cudaGridDependencySynchronize();
#endif
vec_t<float, vec_size> cos, sin;
@@ -184,7 +184,7 @@ __global__ void BatchQKApplyRotaryPosIdsCosSinCacheEnhancedHeadParallelismKernel
}
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.launch_dependents;");
cudaTriggerProgrammaticLaunchCompletion();
#endif
}
@@ -229,7 +229,7 @@ __global__ void BatchQKApplyRotaryPosIdsCosSinCacheEnhancedKernel(
const uint32_t bdy = blockDim.y;
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.wait;");
cudaGridDependencySynchronize();
#endif
vec_t<float, vec_size> cos, sin;
@@ -310,7 +310,7 @@ __global__ void BatchQKApplyRotaryPosIdsCosSinCacheEnhancedKernel(
}
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.launch_dependents;");
cudaTriggerProgrammaticLaunchCompletion();
#endif
}
sgl-kernel/csrc/gemm/dsv3_fused_a_gemm.cu
+2 -2
View File
@@ -285,7 +285,7 @@ struct GmemLoaderB {
__device__ void issue_mainloop() {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
asm volatile("griddepcontrol.wait;");
cudaGridDependencySynchronize();
#pragma unroll 1
for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
if (need_wait) {
@@ -571,7 +571,7 @@ __global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
mma_computer.issue_mainloop();
mma_computer.epi();
}
asm volatile("griddepcontrol.launch_dependents;");
cudaTriggerProgrammaticLaunchCompletion();
#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))
asm volatile("griddepcontrol.wait;");
cudaGridDependencySynchronize();
#endif
// 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))
asm volatile("griddepcontrol.launch_dependents;");
cudaTriggerProgrammaticLaunchCompletion();
#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))
asm volatile("griddepcontrol.wait;");
cudaGridDependencySynchronize();
#endif
// 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))
asm volatile("griddepcontrol.launch_dependents;");
cudaTriggerProgrammaticLaunchCompletion();
#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>;
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>(
subwarps_per_block,
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),
output_buf);
});
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
cudaTriggerProgrammaticLaunchCompletion();
#endif
}
void sgl_per_token_group_quant_8bit_v2(
@@ -445,17 +453,28 @@ void sgl_per_token_group_quant_8bit_v2(
SCHEDULER::compute_exec_config( \
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__> \
<<<grid, block, 0, stream>>>( \
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); \
cudaLaunchConfig_t config; \
config.gridDim = grid; \
config.blockDim = block; \
config.dynamicSmemBytes = 0; \
config.stream = stream; \
cudaLaunchAttribute attrs[1]; \
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \
attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL(); \
config.numAttrs = 1; \
config.attrs = attrs; \
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)
#define LAUNCH_KERNEL(GROUP_SIZE, T, DST_DTYPE) \