From 0e4d1b49d301b741fe64108f88c34aaf11a3f1ff Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Tue, 19 May 2026 20:04:32 +0800 Subject: [PATCH] [Codex] Remove stale DeepSeek V4 JIT kernels (#25764) --- .../jit_kernel/csrc/deepseek_v4/rmsnorm.cuh | 134 ------- .../silu_and_mul_masked_post_quant_tmp.cuh | 371 ------------------ python/sglang/jit_kernel/deepseek_v4.py | 22 -- 3 files changed, 527 deletions(-) delete mode 100644 python/sglang/jit_kernel/csrc/deepseek_v4/rmsnorm.cuh delete mode 100644 python/sglang/jit_kernel/csrc/deepseek_v4/silu_and_mul_masked_post_quant_tmp.cuh diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/rmsnorm.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/rmsnorm.cuh deleted file mode 100644 index a69642c5a..000000000 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/rmsnorm.cuh +++ /dev/null @@ -1,134 +0,0 @@ -#include -#include - -#include -#include -#include -#include -#include -#include - -#include - -namespace { - -constexpr uint32_t kBlockSize = 128; -constexpr uint32_t kNumWarps = kBlockSize / device::kWarpThreads; - -struct RMSNormSelfParams { - const void* __restrict__ input; - void* __restrict__ output; - int64_t stride_batch_bytes_0; - int64_t stride_head_bytes_0; - int64_t stride_batch_bytes_1; - int64_t stride_head_bytes_1; - uint32_t batch_size; - uint32_t num_head; - float eps; -}; - -template -__global__ __launch_bounds__(kBlockSize, 20) // - void rmsnorm_self(const __grid_constant__ RMSNormSelfParams params) { - using namespace device; - constexpr int64_t kVecSize = 16 / sizeof(DType); - constexpr uint32_t kNumLoop = kHeadDim / (kVecSize * kWarpThreads); - static_assert(kHeadDim % (kWarpThreads * kVecSize) == 0); - using DType2 = packed_t; - using Vec = AlignedVector; - - const auto warp_id = blockIdx.x * kNumWarps + threadIdx.x / kWarpThreads; - const auto batch_id = warp_id / params.num_head; - const auto head_id = warp_id % params.num_head; - const auto gmem = tile::Memory::warp(); - if (batch_id >= params.batch_size) return; - const auto input_ptr = pointer::offset( // - params.input, - batch_id * params.stride_batch_bytes_0, - head_id * params.stride_head_bytes_0); - const auto output_ptr = pointer::offset( // - params.output, - batch_id * params.stride_batch_bytes_1, - head_id * params.stride_head_bytes_1); - PDLWaitPrimary(); // wait for primary kernel - - Vec inputs[kNumLoop]; -#pragma unroll - for (uint32_t i = 0; i < kNumLoop; ++i) { - inputs[i] = gmem.load(input_ptr, i); - } - - // compute sum of squares - float local_sum = 0; -#pragma unroll - for (uint32_t i = 0; i < kNumLoop; ++i) { -#pragma unroll - for (uint32_t j = 0; j < kVecSize / 2; ++j) { - const auto [x, y] = cast(inputs[i][j]); - local_sum += x * x + y * y; - } - } - - const auto sum_of_squares = warp::reduce_sum(local_sum); - const auto factor = math::rsqrt(sum_of_squares / kHeadDim + params.eps); - - // weight must be identity (null, not used) -#pragma unroll - for (uint32_t i = 0; i < kNumLoop; ++i) { -#pragma unroll - for (uint32_t j = 0; j < kVecSize / 2; ++j) { - const auto [x, y] = cast(inputs[i][j]); - inputs[i][j] = cast(fp32x2_t{x * factor, y * factor}); - } - gmem.store(output_ptr, inputs[i], i); - } - - PDLTriggerSecondary(); // launch secondary kernel -} - -template -struct RMSNormKernel { - static constexpr auto kernel_self = rmsnorm_self; - - static void run_self(tvm::ffi::TensorView input, tvm::ffi::TensorView output, float eps) { - using namespace host; - - auto N = SymbolicSize{"batch_size"}; - auto H = SymbolicSize{"num_heads"}; - constexpr auto D = kHeadDim; - auto device = SymbolicDevice{}; - device.set_options(); - - TensorMatcher({N, H, D}) // input - .with_strides({-1, -1, 1}) - .with_dtype() - .with_device(device) - .verify(input); - TensorMatcher({N, H, D}) // output - .with_strides({-1, -1, 1}) - .with_dtype() - .with_device(device) - .verify(output); - - const auto batch_size = static_cast(N.unwrap()); - const auto num_head = static_cast(H.unwrap()); - const auto params = RMSNormSelfParams{ - .input = input.data_ptr(), - .output = output.data_ptr(), - .stride_batch_bytes_0 = static_cast(input.stride(0) * sizeof(DType)), - .stride_head_bytes_0 = static_cast(input.stride(1) * sizeof(DType)), - .stride_batch_bytes_1 = static_cast(output.stride(0) * sizeof(DType)), - .stride_head_bytes_1 = static_cast(output.stride(1) * sizeof(DType)), - .batch_size = batch_size, - .num_head = num_head, - .eps = eps, - }; - if (batch_size == 0 || num_head == 0) return; - const auto needed_warps = batch_size * num_head; - const auto num_blocks = div_ceil(needed_warps, kNumWarps); - LaunchKernel(num_blocks, kBlockSize, device.unwrap()) // - .enable_pdl(kUsePDL)(kernel_self, params); - } -}; - -} // namespace diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/silu_and_mul_masked_post_quant_tmp.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/silu_and_mul_masked_post_quant_tmp.cuh deleted file mode 100644 index 3e2bd9258..000000000 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/silu_and_mul_masked_post_quant_tmp.cuh +++ /dev/null @@ -1,371 +0,0 @@ -#include -#include - -#include -#include -#include -#include -#include -#include - -#include - -#include -#include - -namespace { - -using deepseek_v4::fp8::cast_to_ue8m0; -using deepseek_v4::fp8::pack_fp8; - -struct SiluMulQuantParams { - const bf16_t* __restrict__ input; - fp8_e4m3_t* __restrict__ output; - float* __restrict__ output_scale; - const int32_t* __restrict__ masked_m; - float swiglu_limit; // only read when kApplySwigluLimit=true - int64_t hidden_dim; - uint32_t num_tokens; - uint32_t num_experts; -}; - -constexpr uint32_t kMaxExperts = 256; - -struct alignas(16) CTAWork { - uint32_t expert_id; - uint32_t expert_token_id; - bool valid; -}; - -SGL_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) { - static_assert(device::kWarpThreads == 32); -#pragma unroll - for (uint32_t offset = 1; offset < 32; offset *= 2) { - uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset); - if (lane_id >= offset) val += n; - } - return val; -} - -[[maybe_unused]] -SGL_DEVICE CTAWork get_work(const SiluMulQuantParams& params) { - // Preconditions: - // 1. blockDim.x >= params.num_experts - // 2. params.num_experts <= kMaxExperts - using namespace device; - static_assert(kWarpThreads == 32); - - static __shared__ uint32_t s_warp_sum[32]; - static __shared__ CTAWork result; - - result.valid = false; - - const uint32_t tx = threadIdx.x; - const uint32_t lane_id = tx % kWarpThreads; - const uint32_t warp_id = tx / kWarpThreads; - - const uint32_t val = tx < params.num_experts ? params.masked_m[tx] : 0u; - - // Per-warp inclusive scan of masked_m. - const uint32_t warp_inclusive = warp_inclusive_sum(lane_id, val); - const uint32_t warp_exclusive = warp_inclusive - val; - - // Write each warp total. - if (lane_id == kWarpThreads - 1) s_warp_sum[warp_id] = warp_inclusive; - __syncthreads(); - const auto tmp_val = lane_id < warp_id ? s_warp_sum[lane_id] : 0u; - const auto prefix_exclusive = warp::reduce_sum(tmp_val) + warp_exclusive; - const auto bx = blockIdx.x; - if (prefix_exclusive <= bx && bx < prefix_exclusive + val) { - result = {tx, bx - prefix_exclusive, true}; - } - __syncthreads(); - return result; -} - -template -__global__ __launch_bounds__(1024, 2) void // maximize occupancy - silu_mul_quant_kernel(const SiluMulQuantParams __grid_constant__ params) { - using namespace device; - - constexpr uint32_t kGroupSize = 128u; - constexpr uint32_t kWorkThreads = 16u; - // each thread will handle 8 elements - using InputVec = AlignedVector; - using OutputVec = AlignedVector; - static_assert(8 * kWorkThreads == 128, "Invalid tiling"); - static_assert(!(kTransposed && !kScaleUE8M0), "transposed layout only supports ue8m0"); - - const auto [expert_id, token_id, valid] = get_work(params); - - if (!valid) return; - - const auto work_id = threadIdx.x / kWorkThreads; - - const auto offset = expert_id * params.num_tokens + token_id; - const auto input = params.input + offset * params.hidden_dim * 2; - const auto output = params.output + offset * params.hidden_dim; - [[maybe_unused]] - const auto output_scale = [&] { - const auto num_groups = params.hidden_dim / kGroupSize; - if constexpr (kTransposed) { - const auto base = reinterpret_cast(params.output_scale); - // Physical layout is [E, G//4, N] int32. Each int32 packs 4 consecutive - // group scales for the same token, so the byte address is: - // expert_offset + (group/4)*N*4 + token*4 + group%4 - return base + expert_id * num_groups * params.num_tokens + (work_id / 4u) * (params.num_tokens * 4u) + - token_id * 4u + (work_id % 4u); - } else { - return params.output_scale + offset * num_groups + work_id; - } - }(); - - PDLWaitPrimary(); - - InputVec gate_vec, up_vec; - gate_vec.load(input, threadIdx.x); - up_vec.load(input, threadIdx.x + blockDim.x); - - float local_max = 0.0f; - float results[8]; - -#pragma unroll - for (uint32_t i = 0; i < 4; ++i) { - if constexpr (kApplySwigluLimit) { - // Fused fp32 path: bf16 load ??? fp32 clamp ??? fp32 silu ??? fp32 mul ??? fp32 result. - // Avoids the silu???bf16???mul???fp32 round-trip of the non-fused path since we already - // have gate/up in fp32 registers after clamp. - const float limit = params.swiglu_limit; - - const auto [g0_raw, g1_raw] = cast(gate_vec[i]); - const float g0 = fminf(g0_raw, limit); - const float g1 = fminf(g1_raw, limit); - - const float silu0 = g0 / (1.0f + expf(-g0)); - const float silu1 = g1 / (1.0f + expf(-g1)); - - const auto [u0_raw, u1_raw] = cast(up_vec[i]); - const float u0 = fmaxf(fminf(u0_raw, limit), -limit); - const float u1 = fmaxf(fminf(u1_raw, limit), -limit); - - const float val0 = u0 * silu0; - const float val1 = u1 * silu1; - results[2 * i + 0] = val0; - results[2 * i + 1] = val1; - local_max = fmaxf(local_max, fmaxf(fabsf(val0), fabsf(val1))); - } else { - // original code path ??? must stay byte-equal to pre-fusion kernel. - const auto [g0, g1] = cast(gate_vec[i]); - - float silu0 = g0 / (1.0f + expf(-g0)); - float silu1 = g1 / (1.0f + expf(-g1)); - - bf16x2_t silu_d = cast(fp32x2_t{silu0, silu1}); - auto [val0, val1] = cast(up_vec[i] * silu_d); - results[2 * i + 0] = val0; - results[2 * i + 1] = val1; - local_max = fmaxf(local_max, fmaxf(fabsf(val0), fabsf(val1))); - } - } - - local_max = warp::reduce_max(local_max); - - const float absmax = fmaxf(local_max, 1e-10f); - float scale; - uint32_t ue8m0_exp; - - if constexpr (kScaleUE8M0) { - const float raw_scale = absmax / math::FP8_E4M3_MAX; - ue8m0_exp = cast_to_ue8m0(raw_scale); - scale = __uint_as_float(ue8m0_exp << 23); - } else { - scale = absmax / math::FP8_E4M3_MAX; - } - const auto inv_scale = 1.0f / scale; - - OutputVec out_vec; -#pragma unroll - for (uint32_t i = 0; i < 4; ++i) { - const float scaled_val0 = results[2 * i + 0] * inv_scale; - const float scaled_val1 = results[2 * i + 1] * inv_scale; - out_vec[i] = pack_fp8(scaled_val0, scaled_val1); - } - - PDLTriggerSecondary(); - - out_vec.store(output, threadIdx.x); - if constexpr (kTransposed) { - *output_scale = ue8m0_exp; - } else { - *output_scale = scale; - } -} - -struct SiluAndMulClampParams { - const void* __restrict__ input; - void* __restrict__ output; - float swiglu_limit; -}; - -template -__global__ __launch_bounds__(1024, 2) void // maximize occupancy - silu_mul_clamp_kernel(const SiluAndMulClampParams __grid_constant__ params) { - using namespace device; - static_assert(sizeof(DType) == 2, "only fp16/bf16 supported"); - using DType2 = packed_t; - constexpr auto kVecSize = 16 / sizeof(DType); - static_assert(kVecSize % 2 == 0 && kVecSize > 0); - using Vec = AlignedVector; - const auto bid = blockIdx.x; - const auto tile = tile::Memory::cta(); - const float limit = params.swiglu_limit; - - PDLWaitPrimary(); - const auto gate = tile.load(params.input, bid * 2 + 0); - const auto up = tile.load(params.input, bid * 2 + 1); - Vec out; - -#pragma unroll - for (uint32_t i = 0; i < kVecSize / 2; ++i) { - const auto [g0_raw, g1_raw] = cast(gate[i]); - const float g0 = fminf(g0_raw, limit); - const float g1 = fminf(g1_raw, limit); - const float silu0 = g0 / (1.0f + expf(-g0)); - const float silu1 = g1 / (1.0f + expf(-g1)); - const auto [u0_raw, u1_raw] = cast(up[i]); - const float u0 = fmaxf(fminf(u0_raw, limit), -limit); - const float u1 = fmaxf(fminf(u1_raw, limit), -limit); - const float val0 = u0 * silu0; - const float val1 = u1 * silu1; - out[i] = cast(fp32x2_t{val0, val1}); - } - - tile.store(params.output, out, bid); - PDLTriggerSecondary(); -} - -// ---- Host wrapper -// ------------------------------------------------------------------------------------------------------------------------ - -template -struct SiluAndMulMaskedPostQuantKernel { - static_assert(kGroupSize == 128); - static constexpr auto kernel_normal = silu_mul_quant_kernel; - static constexpr auto kernel_transposed = silu_mul_quant_kernel; - - static void - run(const tvm::ffi::TensorView input, - const tvm::ffi::TensorView output, - const tvm::ffi::TensorView output_scale, - const tvm::ffi::TensorView masked_m, - const uint32_t topk, - const bool transposed, - const double swiglu_limit) { - using namespace host; - - auto device = SymbolicDevice{}; - auto E = SymbolicSize{"num_experts"}; - auto T = SymbolicSize{"num_tokens_padded"}; - auto D = SymbolicSize{"hidden_dim x 2"}; - auto N = SymbolicSize{"hidden_dim"}; - auto G = SymbolicSize{"num_groups"}; - device.set_options(); - - TensorMatcher({E, T, D}) // input - .with_dtype() - .with_device(device) - .verify(input); - TensorMatcher({E, T, N}) // output - .with_dtype() - .with_device(device) - .verify(output); - if (!transposed) { - TensorMatcher({E, T, G}) // - .with_dtype() - .with_device(device) - .verify(output_scale); - } else { - RuntimeCheck(kScaleUE8M0, "transposed layout only supports scale_ue8m0=true"); - auto G_ = SymbolicSize{"G // 4"}; - TensorMatcher({E, G_, T}) // - .with_dtype() - .with_device(device) - .verify(output_scale); - G.set_value(G_.unwrap() * 4); - } - TensorMatcher({E}) // - .with_dtype() - .with_device(device) - .verify(masked_m); - - const auto num_experts = static_cast(E.unwrap()); - const auto num_tokens = static_cast(T.unwrap()); - const auto num_groups = static_cast(G.unwrap()); - const auto hidden_dim = N.unwrap(); - - RuntimeCheck(D.unwrap() == 2 * hidden_dim, "invalid dimension"); - RuntimeCheck(hidden_dim % kGroupSize == 0); - RuntimeCheck(num_experts <= kMaxExperts, "num_experts exceeds maximum (256)"); - RuntimeCheck(num_groups * kGroupSize == hidden_dim, "invalid num_groups"); - - const auto params = SiluMulQuantParams{ - .input = static_cast(input.data_ptr()), - .output = static_cast(output.data_ptr()), - .output_scale = static_cast(output_scale.data_ptr()), - .masked_m = static_cast(masked_m.data_ptr()), - .swiglu_limit = static_cast(swiglu_limit), - .hidden_dim = hidden_dim, - .num_tokens = num_tokens, - .num_experts = num_experts, - }; - - const auto num_threads = hidden_dim / 8; - RuntimeCheck(num_threads % device::kWarpThreads == 0); - RuntimeCheck(num_threads >= num_experts); - const auto kernel = transposed ? kernel_transposed : kernel_normal; - LaunchKernel(num_tokens * topk, num_threads, device.unwrap()) // - .enable_pdl(kUsePDL)(kernel, params); - } -}; - -template -struct SiluAndMulClampKernel { - static constexpr auto kernel = silu_mul_clamp_kernel; - - static void run(const tvm::ffi::TensorView input, const tvm::ffi::TensorView output, const double swiglu_limit) { - using namespace host; - - auto device = SymbolicDevice{}; - auto M = SymbolicSize{"num_tokens"}; - auto D = SymbolicSize{"gate_up_dim"}; // 2 * out_dim - auto H = SymbolicSize{"out_dim"}; - device.set_options(); - - TensorMatcher({M, D}) // input (gate || up) - .with_dtype() - .with_device(device) - .verify(input); - TensorMatcher({M, H}) // output - .with_dtype() - .with_device(device) - .verify(output); - RuntimeCheck(D.unwrap() == 2 * H.unwrap(), "input last dim must be 2 * output last dim"); - - constexpr uint32_t kVecSize = 16 / sizeof(DType); - const auto out_dim = static_cast(H.unwrap()); - const auto num_tokens = static_cast(M.unwrap()); - RuntimeCheck(out_dim % kVecSize == 0, "out_dim must be divisible by vector size"); - const auto num_threads = out_dim / kVecSize; - RuntimeCheck(num_threads <= 1024, "out_dim too large for single-block-per-row launch"); - - const auto params = SiluAndMulClampParams{ - .input = input.data_ptr(), - .output = output.data_ptr(), - .swiglu_limit = static_cast(swiglu_limit), - }; - LaunchKernel(num_tokens, num_threads, device.unwrap()) // - .enable_pdl(kUsePDL)(kernel, params); - } -}; - -} // namespace diff --git a/python/sglang/jit_kernel/deepseek_v4.py b/python/sglang/jit_kernel/deepseek_v4.py index 5ff07e88a..9f64d4ae2 100644 --- a/python/sglang/jit_kernel/deepseek_v4.py +++ b/python/sglang/jit_kernel/deepseek_v4.py @@ -144,18 +144,6 @@ def _jit_compress_module( ) -@cache_once -def _jit_rmsnorm_head_module(head_dim: int, dtype: torch.dtype): - args = make_cpp_args(head_dim, dtype, is_arch_support_pdl()) - kernel_class = f"RMSNormKernel<{args}>" - return load_jit( - make_name("rmsnorm_head"), - *args, - cuda_files=["deepseek_v4/rmsnorm.cuh"], - cuda_wrappers=[("run_self", f"{kernel_class}::run_self")], - ) - - @cache_once def _jit_fused_rope_module() -> Module: args = make_cpp_args(is_arch_support_pdl()) @@ -964,16 +952,6 @@ def get_paged_mqa_logits_metadata(seq_lens: torch.Tensor, page_size: int, num_sm return metadata -def rmsnorm_self( - q: torch.Tensor, eps: float, out: Optional[torch.Tensor] = None -) -> torch.Tensor: - module = _jit_rmsnorm_head_module(q.shape[-1], q.dtype) - if out is None: - out = q.new_empty(q.shape) - module.run_self(q, out, eps) - return out - - @cache_once def _jit_torch_cublas_bf16_fp32() -> Any: import torch.utils.cpp_extension