From 19f55c0e6d6fe5cf9a0f9ab6fc148cc004facb7b Mon Sep 17 00:00:00 2001 From: DarkSharpness <76582120+DarkSharpness@users.noreply.github.com> Date: Thu, 21 May 2026 16:14:31 +0800 Subject: [PATCH] [Refactor] major JIT kernel clean up for dsv4 (#25884) Co-authored-by: Claude --- .../jit_kernel/csrc/deepseek_v4/topk_1024.cuh | 336 ------ .../deepseek_v4/{topk.cuh => topk_v1.cuh} | 22 +- python/sglang/jit_kernel/deepseek_v4.py | 1036 ----------------- python/sglang/jit_kernel/dsv4/__init__.py | 49 +- python/sglang/jit_kernel/dsv4/attn.py | 195 ++++ python/sglang/jit_kernel/dsv4/compress_old.py | 308 +++++ python/sglang/jit_kernel/dsv4/elementwise.py | 158 +++ python/sglang/jit_kernel/dsv4/gemm.py | 24 + python/sglang/jit_kernel/dsv4/hisparse.py | 27 + python/sglang/jit_kernel/dsv4/moe.py | 216 ++++ python/sglang/jit_kernel/dsv4/topk.py | 88 ++ .../srt/layers/attention/dsv4/compressor.py | 5 +- .../srt/layers/attention/dsv4/indexer.py | 2 +- .../srt/layers/attention/dsv4/metadata.py | 4 +- python/sglang/srt/layers/moe/hash_topk.py | 2 +- python/sglang/srt/layers/moe/mega_moe.py | 2 +- .../srt/layers/moe/moe_runner/deep_gemm.py | 4 +- .../moe/moe_runner/triton_utils/fused_moe.py | 2 +- python/sglang/srt/layers/moe/topk.py | 2 +- .../srt/mem_cache/deepseek_v4_memory_pool.py | 2 +- .../srt/mem_cache/hisparse_memory_pool.py | 2 +- python/sglang/srt/models/deepseek_v2.py | 4 +- python/sglang/srt/models/deepseek_v4.py | 2 +- 23 files changed, 1093 insertions(+), 1399 deletions(-) delete mode 100644 python/sglang/jit_kernel/csrc/deepseek_v4/topk_1024.cuh rename python/sglang/jit_kernel/csrc/deepseek_v4/{topk.cuh => topk_v1.cuh} (95%) delete mode 100644 python/sglang/jit_kernel/deepseek_v4.py create mode 100644 python/sglang/jit_kernel/dsv4/attn.py create mode 100644 python/sglang/jit_kernel/dsv4/compress_old.py create mode 100644 python/sglang/jit_kernel/dsv4/elementwise.py create mode 100644 python/sglang/jit_kernel/dsv4/gemm.py create mode 100644 python/sglang/jit_kernel/dsv4/hisparse.py create mode 100644 python/sglang/jit_kernel/dsv4/moe.py create mode 100644 python/sglang/jit_kernel/dsv4/topk.py diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/topk_1024.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/topk_1024.cuh deleted file mode 100644 index 6774734ec..000000000 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/topk_1024.cuh +++ /dev/null @@ -1,336 +0,0 @@ -#include -#include - -#include - -#include -#include - -#include -#include - -namespace { - -constexpr uint32_t kTopK = 1024; -constexpr uint32_t kTopKBlockSize = 1024; -constexpr uint32_t kSMEM = 16 * 1024 * sizeof(uint32_t); // 64KB (bytes) - -struct TopK1024Params { - const float* __restrict__ scores; - const int32_t* __restrict__ seq_lens; - const int32_t* __restrict__ page_table; - int32_t* __restrict__ page_indices; - int32_t* __restrict__ raw_indices; // optional: output raw abs position indices before page transform - const int64_t score_stride; - const int64_t page_table_stride; - uint32_t page_bits; -}; - -SGL_DEVICE uint8_t convert_to_uint8(float x) { - __half h = __float2half_rn(x); - uint16_t bits = __half_as_ushort(h); - uint16_t key = (bits & 0x8000) ? static_cast(~bits) : static_cast(bits | 0x8000); - return static_cast(key >> 8); -} - -SGL_DEVICE uint32_t convert_to_uint32(float x) { - uint32_t bits = __float_as_uint(x); - return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); -} - -SGL_DEVICE int32_t page_to_indices(const int32_t* __restrict__ page_table, uint32_t i, uint32_t page_bits) { - const uint32_t mask = (1u << page_bits) - 1u; - return (page_table[i >> page_bits] << page_bits) | (i & mask); -} - -[[maybe_unused]] -SGL_DEVICE void naive_transform( - const float* __restrict__, // unused - const int32_t* __restrict__ page_table, - int32_t* __restrict__ indices, - int32_t* __restrict__ raw_indices, // optional: output raw abs position indices - const uint32_t length, - const uint32_t page_bits) { - static_assert(kTopK <= kTopKBlockSize); - if (const auto tx = threadIdx.x; tx < length) { - indices[tx] = page_to_indices(page_table, tx, page_bits); - if (raw_indices != nullptr) { - raw_indices[tx] = tx; - } - } else if (kTopK == kTopKBlockSize || tx < kTopK) { - indices[tx] = -1; // fill invalid indices to -1 - if (raw_indices != nullptr) { - raw_indices[tx] = -1; - } - } -} - -[[maybe_unused]] -SGL_DEVICE void radix_topk(const float* __restrict__ input, int32_t* __restrict__ output, const uint32_t length) { - constexpr uint32_t RADIX = 256; - constexpr uint32_t BLOCK_SIZE = kTopKBlockSize; - constexpr uint32_t SMEM_INPUT_SIZE = kSMEM / (2 * sizeof(int32_t)); - - alignas(128) __shared__ uint32_t _s_histogram_buf[2][RADIX + 32]; - alignas(128) __shared__ uint32_t s_counter; - alignas(128) __shared__ uint32_t s_threshold_bin_id; - alignas(128) __shared__ uint32_t s_num_input[2]; - alignas(128) __shared__ int32_t s_last_remain; - - extern __shared__ uint32_t s_input_idx[][kSMEM / (2 * sizeof(int32_t))]; - - const uint32_t tx = threadIdx.x; - uint32_t remain_topk = kTopK; - auto& s_histogram = _s_histogram_buf[0]; - - const auto run_cumsum = [&] { -#pragma unroll 8 - for (int32_t i = 0; i < 8; ++i) { - static_assert(1 << 8 == RADIX); - if (tx < RADIX) { - const auto j = 1 << i; - const auto k = i & 1; - auto value = _s_histogram_buf[k][tx]; - if (tx + j < RADIX) { - value += _s_histogram_buf[k][tx + j]; - } - _s_histogram_buf[k ^ 1][tx] = value; - } - __syncthreads(); - } - }; - - // stage 1: 8bit coarse histogram - if (tx < RADIX + 1) s_histogram[tx] = 0; - __syncthreads(); - for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) { - const auto bin = convert_to_uint8(input[idx]); - ::atomicAdd(&s_histogram[bin], 1); - } - __syncthreads(); - run_cumsum(); - if (tx < RADIX && s_histogram[tx] > remain_topk && s_histogram[tx + 1] <= remain_topk) { - s_threshold_bin_id = tx; - s_num_input[0] = 0; - s_counter = 0; - } - __syncthreads(); - - const auto threshold_bin = s_threshold_bin_id; - remain_topk -= s_histogram[threshold_bin + 1]; - if (remain_topk == 0) { - for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) { - const uint32_t bin = convert_to_uint8(input[idx]); - if (bin > threshold_bin) { - const auto pos = ::atomicAdd(&s_counter, 1); - output[pos] = idx; - } - } - __syncthreads(); - return; - } else { - __syncthreads(); - if (tx < RADIX + 1) { - s_histogram[tx] = 0; - } - __syncthreads(); - - for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) { - const float raw_input = input[idx]; - const uint32_t bin = convert_to_uint8(raw_input); - if (bin > threshold_bin) { - const auto pos = ::atomicAdd(&s_counter, 1); - output[pos] = idx; - } else if (bin == threshold_bin) { - const auto pos = ::atomicAdd(&s_num_input[0], 1); - if (pos < SMEM_INPUT_SIZE) { - [[likely]] s_input_idx[0][pos] = idx; - const auto bin = convert_to_uint32(raw_input); - const auto sub_bin = (bin >> 24) & 0xFF; - ::atomicAdd(&s_histogram[sub_bin], 1); - } - } - } - __syncthreads(); - } - - // stage 2: refine with 8bit radix passes -#pragma unroll 4 - for (int round = 0; round < 4; ++round) { - const auto r_idx = round % 2; - - // clip here to prevent overflow - const auto raw_num_input = s_num_input[r_idx]; - const auto num_input = raw_num_input < SMEM_INPUT_SIZE ? raw_num_input : SMEM_INPUT_SIZE; - - run_cumsum(); - if (tx < RADIX && s_histogram[tx] > remain_topk && s_histogram[tx + 1] <= remain_topk) { - s_threshold_bin_id = tx; - s_num_input[r_idx ^ 1] = 0; - s_last_remain = remain_topk - s_histogram[tx + 1]; - } - __syncthreads(); - - const auto threshold_bin = s_threshold_bin_id; - remain_topk -= s_histogram[threshold_bin + 1]; - - if (remain_topk == 0) { - for (uint32_t i = tx; i < num_input; i += BLOCK_SIZE) { - const auto idx = s_input_idx[r_idx][i]; - const auto offset = 24 - round * 8; - const auto bin = (convert_to_uint32(input[idx]) >> offset) & 0xFF; - if (bin > threshold_bin) { - const auto pos = ::atomicAdd(&s_counter, 1); - output[pos] = idx; - } - } - __syncthreads(); - break; - } else { - __syncthreads(); - if (tx < RADIX + 1) { - s_histogram[tx] = 0; - } - __syncthreads(); - for (uint32_t i = tx; i < num_input; i += BLOCK_SIZE) { - const auto idx = s_input_idx[r_idx][i]; - const auto raw_input = input[idx]; - const auto offset = 24 - round * 8; - const auto bin = (convert_to_uint32(raw_input) >> offset) & 0xFF; - if (bin > threshold_bin) { - const auto pos = ::atomicAdd(&s_counter, 1); - output[pos] = idx; - } else if (bin == threshold_bin) { - if (round == 3) { - const auto pos = ::atomicAdd(&s_last_remain, -1); - if (pos > 0) { - output[kTopK - pos] = idx; - } - } else { - const auto pos = ::atomicAdd(&s_num_input[r_idx ^ 1], 1); - if (pos < SMEM_INPUT_SIZE) { - /// NOTE: (dark) fuse the histogram computation here - [[likely]] s_input_idx[r_idx ^ 1][pos] = idx; - const auto bin = convert_to_uint32(raw_input); - const auto sub_bin = (bin >> (offset - 8)) & 0xFF; - ::atomicAdd(&s_histogram[sub_bin], 1); - } - } - } - } - __syncthreads(); - } - } -} - -template -__global__ void topk_1024_transform(const __grid_constant__ TopK1024Params params) { - const auto &[ - scores, seq_lens, page_table, page_indices, raw_indices, // pointers - score_stride, page_table_stride, page_bits // sizes - ] = params; - const uint32_t work_id = blockIdx.x; - - /// NOTE: dangerous prefetch seq_len before PDL wait - const uint32_t seq_len = seq_lens[work_id]; - const auto score_ptr = scores + work_id * score_stride; - const auto page_ptr = page_table + work_id * page_table_stride; - const auto indices_ptr = page_indices + work_id * kTopK; - const auto raw_indices_ptr = raw_indices != nullptr ? raw_indices + work_id * kTopK : nullptr; - - device::PDLWaitPrimary(); - - if (seq_len <= kTopK) { - naive_transform(score_ptr, page_ptr, indices_ptr, raw_indices_ptr, seq_len, page_bits); - } else { - __shared__ int32_t s_topk_indices[kTopK]; - radix_topk(score_ptr, s_topk_indices, seq_len); - static_assert(kTopK <= kTopKBlockSize); - const auto tx = threadIdx.x; - if (kTopK == kTopKBlockSize || tx < kTopK) { - indices_ptr[tx] = page_to_indices(page_ptr, s_topk_indices[tx], page_bits); - if (raw_indices_ptr != nullptr) { - raw_indices_ptr[tx] = s_topk_indices[tx]; - } - } - } - - device::PDLTriggerSecondary(); -} - -template -void setup_kernel_smem_once(host::DebugInfo where = {}) { - [[maybe_unused]] - static const auto result = [] { - const auto fptr = std::bit_cast(f); - return ::cudaFuncSetAttribute(fptr, ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM); - }(); - host::RuntimeDeviceCheck(result, where); -} - -template -struct TopK1024Kernel { - static constexpr auto kernel = topk_1024_transform; - - static void transform( - const tvm::ffi::TensorView scores, - const tvm::ffi::TensorView seq_lens, - const tvm::ffi::TensorView page_table, - const tvm::ffi::TensorView page_indices, - const uint32_t page_size, - const tvm::ffi::Optional raw_indices) { - using namespace host; - auto B = SymbolicSize{"batch_size"}; - auto S = SymbolicSize{"score_stride"}; - auto P = SymbolicSize{"page_table_stride"}; - auto device = SymbolicDevice{}; - device.set_options(); - - TensorMatcher({B, -1}) // strided scores - .with_strides({S, 1}) - .with_dtype() - .with_device(device) - .verify(scores); - TensorMatcher({B}) // seq_lens, must be contiguous - .with_dtype() - .with_device(device) - .verify(seq_lens); - TensorMatcher({B, -1}) // strided page table - .with_strides({P, 1}) - .with_dtype() - .with_device(device) - .verify(page_table); - TensorMatcher({B, 1024}) // output, must be contiguous - .with_dtype() - .with_device(device) - .verify(page_indices); - - int32_t* raw_indices_ptr = nullptr; - if (raw_indices.has_value()) { - TensorMatcher({B, 1024}) // optional raw indices output, must be contiguous - .with_dtype() - .with_device(device) - .verify(raw_indices.value()); - raw_indices_ptr = static_cast(raw_indices.value().data_ptr()); - } - - RuntimeCheck(std::has_single_bit(page_size), "page_size must be power of 2"); - const auto page_bits = static_cast(std::countr_zero(page_size)); - const auto batch_size = static_cast(B.unwrap()); - const auto params = TopK1024Params{ - .scores = static_cast(scores.data_ptr()), - .seq_lens = static_cast(seq_lens.data_ptr()), - .page_table = static_cast(page_table.data_ptr()), - .page_indices = static_cast(page_indices.data_ptr()), - .raw_indices = raw_indices_ptr, - .score_stride = S.unwrap(), - .page_table_stride = P.unwrap(), - .page_bits = page_bits, - }; - constexpr auto kSMEM_ = kSMEM + sizeof(int32_t); // align up a little - setup_kernel_smem_once(); - LaunchKernel(batch_size, kTopKBlockSize, device.unwrap(), kSMEM_).enable_pdl(kUsePDL)(kernel, params); - } -}; - -} // namespace diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/topk.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/topk_v1.cuh similarity index 95% rename from python/sglang/jit_kernel/csrc/deepseek_v4/topk.cuh rename to python/sglang/jit_kernel/csrc/deepseek_v4/topk_v1.cuh index ef2be43c0..b1ccd24b2 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/topk.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/topk_v1.cuh @@ -11,11 +11,15 @@ namespace { -constexpr uint32_t kTopK = 512; -constexpr uint32_t kTopKBlockSize = 512; +#ifndef SGL_TOPK +#define SGL_TOPK 512 +#endif + +constexpr uint32_t kTopK = SGL_TOPK; +constexpr uint32_t kTopKBlockSize = SGL_TOPK; constexpr uint32_t kSMEM = 16 * 1024 * sizeof(uint32_t); // 64KB (bytes) -struct TopK512Params { +struct TopKParams { const float* __restrict__ scores; const int32_t* __restrict__ seq_lens; const int32_t* __restrict__ page_table; @@ -224,7 +228,7 @@ SGL_DEVICE void radix_topk(const float* __restrict__ input, int32_t* __restrict_ } template -__global__ void topk_512_transform(const __grid_constant__ TopK512Params params) { +__global__ void topk_transform_kernel(const __grid_constant__ TopKParams params) { const auto &[ scores, seq_lens, page_table, page_indices, raw_indices, // pointers score_stride, page_table_stride, page_bits // sizes @@ -269,8 +273,8 @@ void setup_kernel_smem_once(host::DebugInfo where = {}) { } template -struct TopK512Kernel { - static constexpr auto kernel = topk_512_transform; +struct TopKKernel { + static constexpr auto kernel = topk_transform_kernel; static void transform( const tvm::ffi::TensorView scores, @@ -300,14 +304,14 @@ struct TopK512Kernel { .with_dtype() .with_device(device) .verify(page_table); - TensorMatcher({B, 512}) // output, must be contiguous + TensorMatcher({B, kTopK}) // output, must be contiguous .with_dtype() .with_device(device) .verify(page_indices); int32_t* raw_indices_ptr = nullptr; if (raw_indices.has_value()) { - TensorMatcher({B, 512}) // optional raw indices output, must be contiguous + TensorMatcher({B, kTopK}) // optional raw indices output, must be contiguous .with_dtype() .with_device(device) .verify(raw_indices.value()); @@ -317,7 +321,7 @@ struct TopK512Kernel { RuntimeCheck(std::has_single_bit(page_size), "page_size must be power of 2"); const auto page_bits = static_cast(std::countr_zero(page_size)); const auto batch_size = static_cast(B.unwrap()); - const auto params = TopK512Params{ + const auto params = TopKParams{ .scores = static_cast(scores.data_ptr()), .seq_lens = static_cast(seq_lens.data_ptr()), .page_table = static_cast(page_table.data_ptr()), diff --git a/python/sglang/jit_kernel/deepseek_v4.py b/python/sglang/jit_kernel/deepseek_v4.py deleted file mode 100644 index 1216dcc04..000000000 --- a/python/sglang/jit_kernel/deepseek_v4.py +++ /dev/null @@ -1,1036 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Optional, Tuple, Union - -import torch -import triton -import triton.language as tl - -from sglang.jit_kernel.utils import ( - cache_once, - is_arch_support_pdl, - load_jit, - make_cpp_args, -) -from sglang.srt.environ import envs -from sglang.srt.utils import get_bool_env_var, is_hip - -_is_hip = is_hip() -_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip - -if _use_aiter: - from aiter.tuned_gemm import tgemm - -from sglang.srt.layers import deep_gemm_wrapper - -if TYPE_CHECKING: - from tvm_ffi.module import Module - -_linear_bf16_fp32_algo = envs.SGLANG_OPT_BF16_FP32_GEMM_ALGO.get() - - -def make_name(name: str) -> str: - return f"dpsk_v4_{name}" - - -@cache_once -def _jit_common_module() -> Module: - return load_jit( - make_name("common"), - cuda_files=["deepseek_v4/common.cuh"], - cuda_wrappers=[("plan_compress_prefill", "plan_compress_prefill")], - ) - - -@cache_once -def _jit_compress_128_online_plan_module() -> Module: - """Host-side plan generator for online compress 128 (no template args).""" - return load_jit( - make_name("compress_128_online_plan"), - cuda_files=["deepseek_v4/c128_online.cuh"], - cuda_wrappers=[ - ("plan_compress_online_prefill", "plan_compress_online_prefill"), - ], - ) - - -@cache_once -def _jit_compress_128_online_module(head_dim: int) -> Module: - """Online compress 128 kernel: ring_size=1, per-index (max, sum, kv) state.""" - args = make_cpp_args(head_dim, is_arch_support_pdl()) - kernel_class = f"FlashCompress128OnlineKernel<{args}>" - return load_jit( - make_name("compress_128_online"), - *args, - cuda_files=["deepseek_v4/c128_online.cuh"], - cuda_wrappers=[ - ("decode", f"{kernel_class}::run_decode"), - ("prefill", f"{kernel_class}::run_prefill"), - ], - extra_cuda_cflags=["-use_fast_math"], - ) - - -@cache_once -def _jit_topk_module() -> Module: - args = make_cpp_args(is_arch_support_pdl()) - return load_jit( - make_name("topk"), - *args, - cuda_files=["deepseek_v4/topk.cuh"], - cuda_wrappers=[("topk_transform", f"TopK512Kernel<{args}>::transform")], - ) - - -@cache_once -def _jit_topk1024_module() -> Module: - args = make_cpp_args(is_arch_support_pdl()) - return load_jit( - make_name("topk1024"), - *args, - cuda_files=["deepseek_v4/topk_1024.cuh"], - cuda_wrappers=[("topk_transform", f"TopK1024Kernel<{args}>::transform")], - ) - - -@cache_once -def _jit_topk_v2_module(topk: int) -> Module: - return load_jit( - make_name("topk_v2"), - str(topk), - cuda_files=["deepseek_v4/topk_v2.cuh"], - cuda_wrappers=[ - ("topk_transform", "CombinedTopKKernel::transform"), - ("topk_plan", "CombinedTopKKernel::plan"), - ], - extra_cuda_cflags=[f"-DSGL_TOPK={topk}"], - ) - - -@cache_once -def _jit_mask_topk_module() -> Module: - return load_jit( - make_name("mask_topk"), - cuda_files=["deepseek_v4/hash_topk.cuh"], - cuda_wrappers=[("run", "MaskKernel::run")], - ) - - -@cache_once -def _jit_hash_topk_module() -> Module: - args = make_cpp_args("act_sqrt_softplus", is_arch_support_pdl()) - return load_jit( - make_name("hash_topk"), - *args, - cuda_files=["deepseek_v4/hash_topk.cuh"], - cuda_wrappers=[("hash_topk", f"HashTopKKernel<{args}>::run")], - ) - - -@cache_once -def _jit_compress_module( - head_dim: int, - dtype_in: torch.dtype, - dtype_out: torch.dtype, - ratio: Literal[4, 128], -) -> Module: - args = make_cpp_args(head_dim, dtype_in, dtype_out, is_arch_support_pdl()) - kernel_class = f"FlashCompress{ratio}Kernel<{args}>" - return load_jit( - make_name(f"compress_{ratio}"), - *args, - cuda_files=[f"deepseek_v4/c{ratio}.cuh"], - cuda_wrappers=[ - ("decode", f"{kernel_class}::run_decode"), - ("prefill", f"{kernel_class}::run_prefill"), - ], - extra_cuda_cflags=["-use_fast_math"], - ) - - -@cache_once -def _jit_fused_rope_module() -> Module: - args = make_cpp_args(is_arch_support_pdl()) - return load_jit( - make_name("fused_rope"), - *args, - cuda_files=["deepseek_v4/rope.cuh"], - cuda_wrappers=[("forward", f"FusedQKRopeKernel<{args}>::forward")], - ) - - -@cache_once -def _jit_norm_rope_module( - dtype: torch.dtype, - head_dim: int, - rope_dim: int, -) -> Module: - args = make_cpp_args(dtype, head_dim, rope_dim, is_arch_support_pdl()) - return load_jit( - make_name("fused_norm_rope"), - *args, - cuda_files=["deepseek_v4/fused_norm_rope.cuh"], - cuda_wrappers=[ - ("forward", f"FusedNormRopeKernel<{args}>::forward"), - ], - ) - - -@cache_once -def _jit_fused_store_module( - name: Literal["flashmla", "indexer"], - input_dtype: torch.dtype, - index_dtype: torch.dtype, - page_size: int, -) -> Module: - args = make_cpp_args(input_dtype, index_dtype, page_size, is_arch_support_pdl()) - cname = "FlashMLA" if name == "flashmla" else "Indexer" - kernel_class = f"FusedStoreCache{cname}Kernel<{args}>" - return load_jit( - make_name("store_" + name), - *args, - cuda_files=["deepseek_v4/store.cuh"], - cuda_wrappers=[("run", f"{kernel_class}::run")], - ) - - -@cache_once -def _jit_main_q_norm_rope_module( - dtype: torch.dtype, head_dim: int, rope_dim: int -) -> Module: - """Main MLA path Q kernel: rmsnorm-self + RoPE, warp per (token, head).""" - args = make_cpp_args(dtype, head_dim, rope_dim, is_arch_support_pdl()) - return load_jit( - make_name("main_q_norm_rope"), - *args, - cuda_files=["deepseek_v4/main_norm_rope.cuh"], - cuda_wrappers=[ - ("forward", f"FusedQNormRopeKernel<{args}>::forward"), - ], - ) - - -@cache_once -def _jit_main_k_norm_rope_flashmla_module( - dtype: torch.dtype, head_dim: int, rope_dim: int, page_size: int -) -> Module: - """Main MLA path K kernel: rmsnorm + RoPE + write to FlashMLA paged cache.""" - args = make_cpp_args(dtype, head_dim, rope_dim, page_size, is_arch_support_pdl()) - return load_jit( - make_name("main_k_norm_rope_flashmla"), - *args, - cuda_files=["deepseek_v4/main_norm_rope.cuh"], - cuda_wrappers=[ - ("forward", f"FusedKNormRopeFlashMLAKernel<{args}>::forward"), - ], - ) - - -@cache_once -def _jit_main_q_indexer_rope_hadamard_quant_module(dtype: torch.dtype) -> Module: - """C4 indexer Q kernel: RoPE + 128-pt Hadamard + fp8 act-quant (no norm).""" - args = make_cpp_args(dtype, is_arch_support_pdl()) - return load_jit( - make_name("main_q_indexer_rope_hadamard_quant"), - *args, - cuda_files=["deepseek_v4/main_norm_rope.cuh"], - cuda_wrappers=[ - ("forward", f"FusedQIndexerRopeHadamardQuantKernel<{args}>::forward"), - ], - ) - - -@cache_once -def _jit_metadata_module(): - return load_jit( - make_name("metadata"), - cuda_files=["deepseek_v4/paged_mqa_metadata.cuh"], - cuda_wrappers=[("run", "IndexerMetadataKernel::run")], - ) - - -@cache_once -def _jit_silu_mul_quant_varlen_module( - quant_group_size: int, - scale_ue8m0: bool, - swizzle: bool, - apply_swiglu_limit: bool, -) -> Module: - args = make_cpp_args( - quant_group_size, - scale_ue8m0, - swizzle, - is_arch_support_pdl(), - apply_swiglu_limit, - ) - return load_jit( - make_name("silu_mul_quant_varlen"), - *args, - cuda_files=["deepseek_v4/silu_and_mul_masked_post_quant.cuh"], - cuda_wrappers=[("run", f"SiluAndMulMaskedPostQuantKernel<{args}>::run")], - extra_cuda_cflags=["-use_fast_math"], - ) - - -@cache_once -def _jit_silu_mul_quant_contig_module( - quant_group_size: int, - scale_ue8m0: bool, - swizzle: bool, - apply_swiglu_limit: bool, -) -> Module: - args = make_cpp_args( - quant_group_size, - scale_ue8m0, - swizzle, - is_arch_support_pdl(), - apply_swiglu_limit, - ) - return load_jit( - make_name("silu_mul_quant_contig"), - *args, - cuda_files=["deepseek_v4/silu_and_mul_masked_post_quant.cuh"], - cuda_wrappers=[("run", f"SiluAndMulContigPostQuantKernel<{args}>::run")], - extra_cuda_cflags=["-use_fast_math"], - ) - - -@cache_once -def _jit_silu_and_mul_clamp_module(dtype: torch.dtype) -> Module: - args = make_cpp_args(dtype, is_arch_support_pdl()) - return load_jit( - make_name("silu_and_mul_clamp"), - *args, - cuda_files=["deepseek_v4/silu_and_mul_masked_post_quant.cuh"], - cuda_wrappers=[("run", f"SiluAndMulClampKernel<{args}>::run")], - extra_cuda_cflags=["-use_fast_math"], - ) - - -@cache_once -def _jit_mega_moe_pre_dispatch_module(quant_group_size: int) -> Module: - args = make_cpp_args(quant_group_size, is_arch_support_pdl()) - return load_jit( - make_name("mega_moe_pre_dispatch"), - *args, - cuda_files=["deepseek_v4/mega_moe_pre_dispatch.cuh"], - cuda_wrappers=[("run", f"MegaMoEPreDispatchKernel<{args}>::run")], - ) - - -@cache_once -def _jit_hisparse_transfer_module() -> Module: - return load_jit( - make_name("hisparse_transfer"), - cuda_files=["deepseek_v4/hisparse_transfer.cuh"], - cuda_wrappers=[("hisparse_transfer", "hisparse_transfer")], - ) - - -def hisparse_offload_to_host( - gpu_ptrs: torch.Tensor, - cpu_ptrs: torch.Tensor, - gpu_indices: torch.Tensor, - cpu_indices: torch.Tensor, -) -> None: - module = _jit_hisparse_transfer_module() - module.hisparse_transfer(gpu_ptrs, cpu_ptrs, gpu_indices, cpu_indices) - - -def topk_transform_512( - scores: torch.Tensor, - seq_lens: torch.Tensor, - page_tables: torch.Tensor, - out_page_indices: torch.Tensor, - page_size: int, - out_raw_indices: Optional[torch.Tensor] = None, -) -> None: - if out_page_indices.shape[1] == 512: - module = _jit_topk_module() - else: - module = _jit_topk1024_module() - module.topk_transform( - scores, seq_lens, page_tables, out_page_indices, page_size, out_raw_indices - ) - - -_WORKSPACE_INTS_PER_BATCH = 2 + 1024 * 2 -_PLAN_METADATA_INTS_PER_BATCH = 4 - - -def plan_topk_v2(seq_lens: torch.Tensor, static_threshold: int = 0) -> torch.Tensor: - module = _jit_topk_v2_module(512) # does not matter - bs = seq_lens.shape[0] - metadata = seq_lens.new_empty(bs + 1, _PLAN_METADATA_INTS_PER_BATCH) - module.topk_plan(seq_lens, metadata, static_threshold) - return metadata - - -def topk_transform_512_v2( - scores: torch.Tensor, - seq_lens: torch.Tensor, - page_tables: torch.Tensor, - out_page_indices: torch.Tensor, - page_size: int, - metadata: torch.Tensor, -) -> None: - module = _jit_topk_v2_module(out_page_indices.shape[1]) - bs = scores.shape[0] - workspace = seq_lens.new_empty(bs, _WORKSPACE_INTS_PER_BATCH) - module.topk_transform( - scores, - seq_lens, - page_tables, - out_page_indices, - page_size, - workspace, - metadata, - ) - - -def hash_topk( - router_logits: torch.Tensor, - input_ids: torch.Tensor, - tid2eid: torch.Tensor, - num_fused_shared_experts: int = 0, - routed_scaling_factor: float = 1.0, - scoring_func: str = "sqrtsoftplus", -) -> Tuple[torch.Tensor, torch.Tensor]: - assert scoring_func == "sqrtsoftplus" - num_tokens = router_logits.size(0) - topk_routed = tid2eid.size(1) - topk_fused = topk_routed + num_fused_shared_experts - topk_ids = torch.empty( - (num_tokens, topk_fused), dtype=torch.int32, device=router_logits.device - ) - topk_weights = torch.empty( - (num_tokens, topk_fused), dtype=torch.float32, device=router_logits.device - ) - module = _jit_hash_topk_module() - module.hash_topk( - router_logits, - input_ids, - tid2eid, - topk_weights, - topk_ids, - routed_scaling_factor, - ) - return topk_weights, topk_ids - - -def mask_topk_ids(topk_ids: torch.Tensor, num_token_non_padded: torch.Tensor): - return _jit_mask_topk_module().run(topk_ids, num_token_non_padded) - - -class CompressorPrefillPlan(NamedTuple): - compress_ratio: int - compress_plan: torch.Tensor - write_plan: torch.Tensor - - def copy_(self, other: CompressorPrefillPlan) -> None: - assert self.compress_ratio == other.compress_ratio - self.compress_plan.copy_(other.compress_plan) - self.write_plan.copy_(other.write_plan) - - @staticmethod - def generate( - compress_ratio: Literal[4, 128], - num_q_tokens: int, - seq_lens: torch.Tensor, - extend_lens: torch.Tensor, - device: torch.device, - use_cuda_graph: bool = False, - ) -> CompressorPrefillPlan: - from sglang.srt.environ import envs - - # Online c128 keeps the same NamedTuple shape (compress_plan, write_plan) - # so call sites that splat `*plan[1:]` continue to work, but the C++ - # plan struct semantics differ (last-token coords + window_len). - if compress_ratio == 128 and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get(): - return CompressorPrefillPlan._generate_online( - num_q_tokens=num_q_tokens, - seq_lens=seq_lens, - extend_lens=extend_lens, - device=device, - use_cuda_graph=use_cuda_graph, - ) - assert seq_lens.device == extend_lens.device - seq_lens = seq_lens.to(torch.int64) - extend_lens = extend_lens.to(torch.int64) - plan_tensor = torch.empty( - (2, num_q_tokens, 16), - dtype=torch.uint8, - device=seq_lens.device, - pin_memory=seq_lens.is_cpu, - ) - module = _jit_common_module() - is_overlap = compress_ratio == 4 - plan_lens = module.plan_compress_prefill( - extend_lens, - seq_lens, - plan_tensor[0], - plan_tensor[1], - compress_ratio, - is_overlap, - use_cuda_graph, - ) - return CompressorPrefillPlan( - compress_ratio, - plan_tensor[0, : plan_lens[0]].to(device, non_blocking=True), - plan_tensor[1, : plan_lens[1]].to(device, non_blocking=True), - ) - - @staticmethod - def _generate_online( - num_q_tokens: int, - seq_lens: torch.Tensor, - extend_lens: torch.Tensor, - device: torch.device, - use_cuda_graph: bool, - ) -> CompressorPrefillPlan: - # Online plan host-side path: only CPU/cuda-host implemented today. - # Move inputs to CPU pinned memory then bounce the result to device. - seq_lens_cpu = seq_lens.detach().to(torch.int64).cpu() - extend_lens_cpu = extend_lens.detach().to(torch.int64).cpu() - plan_tensor = torch.empty( - (2, num_q_tokens, 16), - dtype=torch.uint8, - device="cpu", - pin_memory=True, - ) - module = _jit_compress_128_online_plan_module() - plan_lens = module.plan_compress_online_prefill( - extend_lens_cpu, - seq_lens_cpu, - plan_tensor[0], - plan_tensor[1], - use_cuda_graph, - ) - return CompressorPrefillPlan( - 128, - plan_tensor[0, : plan_lens[0]].to(device, non_blocking=True), - plan_tensor[1, : plan_lens[1]].to(device, non_blocking=True), - ) - - @property - def is_decode(self) -> bool: - return False - - -class CompressorDecodePlan(NamedTuple): - compress_ratio: int - seq_lens: torch.Tensor - - def copy_(self, other: CompressorDecodePlan) -> None: - assert self.compress_ratio == other.compress_ratio - self.seq_lens.copy_(other.seq_lens) - - @property - def is_decode(self) -> bool: - return True - - -def compress_plan( - compress_ratio: Literal[4, 128], - num_q_tokens: int, - seq_lens: torch.Tensor, - extend_lens: Optional[torch.Tensor], - device: torch.device, -) -> Union[CompressorDecodePlan, CompressorPrefillPlan]: - if extend_lens is not None: - return CompressorPrefillPlan.generate( - compress_ratio, - num_q_tokens, - seq_lens, - extend_lens, - device, - ) - else: - assert num_q_tokens == len(seq_lens) - seq_lens = seq_lens.to(device, non_blocking=True) - return CompressorDecodePlan(compress_ratio, seq_lens) - - -def compress_forward( - kv_score_buffer: torch.Tensor, - kv_score_input: torch.Tensor, - ape: torch.Tensor, - indices: torch.Tensor, - plan: Union[CompressorDecodePlan, CompressorPrefillPlan, None] = None, - extra_data: Optional[torch.Tensor] = None, - *, - head_dim: int, - compress_ratio: Literal[4, 128], - out: Optional[torch.Tensor] = None, - seq_lens: Optional[torch.Tensor] = None, - extend_lens: Optional[torch.Tensor] = None, -) -> torch.Tensor: - assert head_dim % 128 == 0 - num_q_tokens = kv_score_input.shape[0] - if out is None: - out = kv_score_input.new_empty((num_q_tokens, head_dim)) - if plan is None: - assert seq_lens is not None - plan = compress_plan( - compress_ratio, - num_q_tokens, - seq_lens, - extend_lens, - kv_score_input.device, - ) - assert plan.compress_ratio == compress_ratio, "Mismatched compress ratio in plan!" - # Online c128: separate JIT module, fp32 state, no compile-time dtypes. - if compress_ratio == 128 and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get(): - online_module = _jit_compress_128_online_module(head_dim=head_dim) - F = online_module.decode if plan.is_decode else online_module.prefill - F(kv_score_buffer, kv_score_input, out, ape, indices, *plan[1:], extra_data) - return out - module = _jit_compress_module( - head_dim, - kv_score_input.dtype, - out.dtype, - compress_ratio, - ) - F = module.decode if plan.is_decode else module.prefill - F(kv_score_buffer, kv_score_input, out, ape, indices, *plan[1:], extra_data) - return out - - -def compress_fused_norm_rope_inplace( - kv: torch.Tensor, - weight: torch.Tensor, - eps: float, - freq_cis: torch.Tensor, - plan: Union[CompressorDecodePlan, CompressorPrefillPlan], -) -> None: - freq_cis = torch.view_as_real(freq_cis).flatten(-2) - module = _jit_norm_rope_module(kv.dtype, kv.shape[-1], freq_cis.shape[-1]) - module.forward( - kv, - weight, - plan[1], - freq_cis, - int(plan.is_decode), - eps, - plan.compress_ratio, - ) - - -def fused_norm_rope_inplace( - kv: torch.Tensor, - weight: torch.Tensor, - eps: float, - freq_cis: torch.Tensor, - positions: torch.Tensor, -) -> None: - freq_cis = torch.view_as_real(freq_cis).flatten(-2) - module = _jit_norm_rope_module(kv.dtype, kv.shape[-1], freq_cis.shape[-1]) - module.forward( - kv, - weight, - positions, - freq_cis, - 2, - eps, - 0, - ) - - -def fused_rope( - q: torch.Tensor, - k: Optional[torch.Tensor], - freqs_cis: torch.Tensor, - positions: torch.Tensor, - inverse: bool = False, -) -> None: - """Apply rotary embeddings to both Q and K in a single fused CUDA kernel. - - Args: - q: [batch_size, num_q_heads, rope_dim] bfloat16 - k: [batch_size, num_k_heads, rope_dim] bfloat16 or None - freqs_cis: [max_seq_len, rope_dim // 2] complex64 (full table) - positions: [batch_size] int32 or int64, indices into freqs_cis - inverse: if True, apply inverse rotation (conjugate freqs) - """ - if _is_hip: - from sglang.srt.layers.deepseek_v4_rope import apply_rotary_emb_triton - - apply_rotary_emb_triton(q, freqs_cis, positions=positions, inverse=inverse) - if k is not None: - apply_rotary_emb_triton(k, freqs_cis, positions=positions, inverse=inverse) - return - - freqs_real = torch.view_as_real(freqs_cis).flatten(-2).contiguous() - module = _jit_fused_rope_module() - module.forward(q, k, freqs_real, positions, inverse) - - -# Alias for V2 code paths -fused_rope_inplace = fused_rope - - -def fused_q_norm_rope( - q_input: torch.Tensor, - q_output: torch.Tensor, - eps: float, - freqs_cis: torch.Tensor, - positions: torch.Tensor, -) -> None: - freqs_real = torch.view_as_real(freqs_cis).flatten(-2) - head_dim = q_input.shape[-1] - rope_dim = freqs_real.shape[-1] - module = _jit_main_q_norm_rope_module(q_input.dtype, head_dim, rope_dim) - module.forward(q_input, q_output, freqs_real, positions, eps) - - -def fused_q_indexer_rope_hadamard_quant( - q_input: torch.Tensor, - weight: torch.Tensor, - weight_scale: float, - freqs_cis: torch.Tensor, - positions: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - freqs_real = torch.view_as_real(freqs_cis).flatten(-2) - q_fp8 = torch.empty(q_input.shape, dtype=torch.float8_e4m3fn, device=q_input.device) - weights_out = torch.empty( - (*q_input.shape[:-1], 1), dtype=torch.float32, device=q_input.device - ) - module = _jit_main_q_indexer_rope_hadamard_quant_module(q_input.dtype) - module.forward( - q_input, q_fp8, weight, weights_out, float(weight_scale), freqs_real, positions - ) - return q_fp8, weights_out - - -def fused_k_norm_rope_flashmla( - kv: torch.Tensor, - kv_weight: torch.Tensor, - eps: float, - freqs_cis: torch.Tensor, - positions: torch.Tensor, - out_loc: torch.Tensor, - kvcache: torch.Tensor, - page_size: int, -) -> None: - freqs_real = torch.view_as_real(freqs_cis).flatten(-2) - head_dim = kv.shape[-1] - rope_dim = freqs_real.shape[-1] - module = _jit_main_k_norm_rope_flashmla_module( - kv.dtype, head_dim, rope_dim, page_size - ) - module.forward(kv, kv_weight, freqs_real, positions, out_loc, kvcache, eps) - - -@triton.jit -def create_paged_compress_data_kernel( - req_pool_indices_ptr, - seq_lens_ptr, - extend_seq_lens_ptr, - req_to_token_ptr, - full_to_swa_index_mapping_ptr, - out_0_ptr, - out_1_ptr, - batch_size, - stride_req_to_token_0, - stride_req_to_token_1: tl.constexpr, - stride_out_1_0, - stride_out_1_1: tl.constexpr, - compress_ratio: tl.constexpr, - is_overlap: tl.constexpr, - swa_page_size: tl.constexpr, - ring_size: tl.constexpr, - BLOCK: tl.constexpr, -) -> None: - pid = tl.program_id(0) - offs = pid * BLOCK + tl.arange(0, BLOCK) - mask = offs < batch_size - - rid = tl.load(req_pool_indices_ptr + offs, mask=mask, other=0).to(tl.int32) - seq_len = tl.load(seq_lens_ptr + offs, mask=mask, other=0).to(tl.int32) - extend_len = tl.load(extend_seq_lens_ptr + offs, mask=mask, other=0).to(tl.int32) - prefix_len = seq_len - extend_len - - cr = compress_ratio - write_pos = ((seq_len - 1) // cr) * cr - load_pos = ((prefix_len - 1) // cr) * cr - write_overlap_pos = write_pos - cr - load_overlap_pos = load_pos - cr - v0 = tl.zeros([BLOCK], tl.int32) - v1 = tl.zeros([BLOCK], tl.int32) - v2 = tl.zeros([BLOCK], tl.int32) - v3 = tl.zeros([BLOCK], tl.int32) - - for i in tl.static_range(4): - if i == 0: - pos = load_pos - elif i == 1: - pos = write_pos - elif i == 2: - pos = load_overlap_pos - else: - pos = write_overlap_pos - pos = tl.maximum(pos, 0) - loc = tl.load( - req_to_token_ptr - + rid.to(tl.int64) * stride_req_to_token_0 - + pos.to(tl.int64) * stride_req_to_token_1, - mask=mask, - other=0, - ).to(tl.int32) - swa_loc = tl.load(full_to_swa_index_mapping_ptr + loc, mask=mask, other=0).to( - tl.int32 - ) - swa_page = swa_loc // swa_page_size - state_loc = swa_page * ring_size + (swa_loc % ring_size) - state_loc = state_loc // cr - if i == 0: - v0 = state_loc - elif i == 1: - v1 = state_loc - elif i == 2: - v2 = state_loc - else: - v3 = state_loc - - tl.store(out_0_ptr + offs, v1, mask=mask) - - if is_overlap: - base = out_1_ptr + offs * stride_out_1_0 - tl.store(base + 0 * stride_out_1_1, v2, mask=mask) - tl.store(base + 1 * stride_out_1_1, v0, mask=mask) - tl.store(base + 2 * stride_out_1_1, v3, mask=mask) - tl.store(base + 3 * stride_out_1_1, write_pos.to(tl.int32), mask=mask) - else: - base = out_1_ptr + offs * stride_out_1_0 - tl.store(base + 0 * stride_out_1_1, v0, mask=mask) - - -def triton_create_paged_compress_data( - *, - compress_ratio: int, - is_overlap: bool, - swa_page_size: int, - ring_size: int, - req_pool_indices: torch.Tensor, - seq_lens: torch.Tensor, - extend_seq_lens: torch.Tensor, - req_to_token: torch.Tensor, - full_to_swa_index_mapping: torch.Tensor, - block: int = 128, -) -> Tuple[torch.Tensor, torch.Tensor]: - batch_size = req_pool_indices.shape[0] - out_dim = 4 if is_overlap else 1 - device_args: dict = dict(device=req_pool_indices.device, dtype=torch.int32) - out_0 = torch.empty((batch_size,), **device_args) - out_1 = torch.empty((batch_size, out_dim), **device_args) - grid = (triton.cdiv(batch_size, block),) - create_paged_compress_data_kernel[grid]( - req_pool_indices, - seq_lens, - extend_seq_lens, - req_to_token, - full_to_swa_index_mapping, - out_0, - out_1, - batch_size=batch_size, - stride_req_to_token_0=req_to_token.stride(0), - stride_req_to_token_1=req_to_token.stride(1), - stride_out_1_0=out_1.stride(0), - stride_out_1_1=out_1.stride(1), - compress_ratio=compress_ratio, - is_overlap=1 if is_overlap else 0, - swa_page_size=swa_page_size, - ring_size=ring_size, - BLOCK=block, - ) - - if not is_overlap: - out_1.squeeze_(1) - return out_0, out_1 - - -def fused_store_cache( - input: torch.Tensor, - cache: torch.Tensor, - indices: torch.Tensor, - *, - page_size: int, - type: Literal["flashmla", "indexer"], -) -> None: - module = _jit_fused_store_module( - name=type, - input_dtype=input.dtype, - index_dtype=indices.dtype, - page_size=page_size, - ) - module.run(input, cache, indices) - - -def silu_and_mul_clamp( - input: torch.Tensor, - output: torch.Tensor, - swiglu_limit: float, -) -> None: - module = _jit_silu_and_mul_clamp_module(input.dtype) - module.run(input, output, float(swiglu_limit)) - - -def silu_and_mul_masked_post_quant( - input: torch.Tensor, - output: torch.Tensor, - output_scale: torch.Tensor, - quant_group_size: int, - masked_m: torch.Tensor, - scale_ue8m0: bool = False, - topk: int = 8, - transposed: bool = False, - swiglu_limit: Optional[float] = None, - swizzle: bool = False, -) -> None: - apply_swiglu_limit = swiglu_limit is not None - module = _jit_silu_mul_quant_varlen_module( - quant_group_size, scale_ue8m0, swizzle, apply_swiglu_limit - ) - module.run( - input, - output, - output_scale, - masked_m, - topk, - transposed, - float(swiglu_limit) if apply_swiglu_limit else 0.0, - ) - - -def silu_and_mul_contig_post_quant( - input: torch.Tensor, - output: torch.Tensor, - output_scale: torch.Tensor, - quant_group_size: int, - scale_ue8m0: bool = False, - transposed: bool = False, - swiglu_limit: Optional[float] = None, - swizzle: bool = False, -) -> None: - apply_swiglu_limit = swiglu_limit is not None - module = _jit_silu_mul_quant_contig_module( - quant_group_size, scale_ue8m0, swizzle, apply_swiglu_limit - ) - module.run( - input, - output, - output_scale, - transposed, - float(swiglu_limit) if apply_swiglu_limit else 0.0, - ) - - -def mega_moe_pre_dispatch( - x: torch.Tensor, - topk_idx: torch.Tensor, - topk_weights: torch.Tensor, - buf_x: torch.Tensor, - buf_x_sf: torch.Tensor, - buf_topk_idx: torch.Tensor, - buf_topk_weights: torch.Tensor, - quant_group_size: int = 32, -) -> None: - module = _jit_mega_moe_pre_dispatch_module(quant_group_size) - module.run( - x, - topk_idx, - topk_weights, - buf_x, - buf_x_sf, - buf_topk_idx, - buf_topk_weights, - ) - - -def get_paged_mqa_logits_metadata(seq_lens: torch.Tensor, page_size: int, num_sm: int): - assert page_size == 64 - seq_lens = seq_lens.view(-1).to(torch.int32) - metadata = seq_lens.new_empty(num_sm + 1, 2) - module = _jit_metadata_module() - module.run(seq_lens, metadata) - return metadata - - -@cache_once -def _jit_torch_cublas_bf16_fp32() -> Any: - import torch.utils.cpp_extension - - source = """ -#include -#include -#include - -torch::Tensor linear_bf16_fp32( - torch::Tensor X, - torch::Tensor W) -{ - int batch = X.size(0); - int in_features = X.size(1); - int out_features = W.size(0); - - auto Y = torch::empty( - {batch, out_features}, - torch::dtype(torch::kFloat32).device(X.device())); - - cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle(); - - float alpha = 1.0f; - float beta = 0.0f; - - cublasGemmEx( - handle, - CUBLAS_OP_T, - CUBLAS_OP_N, - out_features, - batch, - in_features, - &alpha, - W.data_ptr(), CUDA_R_16BF, in_features, - X.data_ptr(), CUDA_R_16BF, in_features, - &beta, - Y.data_ptr(), CUDA_R_32F, out_features, - CUBLAS_COMPUTE_32F, - CUBLAS_GEMM_DEFAULT_TENSOR_OP - ); - - return Y; -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("linear_bf16_fp32", &linear_bf16_fp32, "BF16xBF16 -> FP32 linear (no bias)"); -} -""" - module = torch.utils.cpp_extension.load_inline( - name="linear_bf16_fp32", - cpp_sources="", - cuda_sources=source, - extra_cflags=["-O3"], - extra_cuda_cflags=["-O3"], - verbose=False, - ) - return module - - -def linear_bf16_fp32(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - return _dispatch_bf16_fp32_backend(x, y, algo=_linear_bf16_fp32_algo) - - -def _dispatch_bf16_fp32_backend( - x: torch.Tensor, y: torch.Tensor, *, algo: str -) -> torch.Tensor: - if algo == "cublas": - module = _jit_torch_cublas_bf16_fp32() - return module.linear_bf16_fp32(x, y) - elif algo == "deep_gemm": - z = torch.empty(x.size(0), y.size(0), dtype=torch.float32, device=x.device) - deep_gemm_wrapper.gemm_nt_bf16bf16f32(x, y, z) - return z - elif _use_aiter: - return tgemm.mm(x, y, otype=torch.float32) - else: - return torch.nn.functional.linear(x.float(), y.float()) diff --git a/python/sglang/jit_kernel/dsv4/__init__.py b/python/sglang/jit_kernel/dsv4/__init__.py index 0dcf60827..4939a73e5 100644 --- a/python/sglang/jit_kernel/dsv4/__init__.py +++ b/python/sglang/jit_kernel/dsv4/__init__.py @@ -1,4 +1,32 @@ -from .compress import * +from .attn import ( + fused_store_cache, + get_paged_mqa_logits_metadata, + triton_create_paged_compress_data, +) +from .compress import ( + CompressorDecodePlan, + CompressorPrefillPlan, + compress_forward, + compress_norm_rope_store, +) +from .compress_old import fused_norm_rope_inplace +from .elementwise import ( + fused_k_norm_rope_flashmla, + fused_q_indexer_rope_hadamard_quant, + fused_q_norm_rope, + fused_rope_inplace, +) +from .gemm import linear_bf16_fp32 +from .hisparse import hisparse_offload_to_host +from .moe import ( + hash_topk, + mask_topk_ids, + mega_moe_pre_dispatch, + silu_and_mul_clamp, + silu_and_mul_contig_post_quant, + silu_and_mul_masked_post_quant, +) +from .topk import plan_topk_v2, topk_transform_512, topk_transform_512_v2 from .utils import make_name __all__ = [ @@ -6,5 +34,24 @@ __all__ = [ "CompressorPrefillPlan", "compress_forward", "compress_norm_rope_store", + "fused_norm_rope_inplace", + "fused_store_cache", + "fused_rope_inplace", + "fused_q_norm_rope", + "fused_q_indexer_rope_hadamard_quant", + "fused_k_norm_rope_flashmla", "make_name", + "linear_bf16_fp32", + "hisparse_offload_to_host", + "get_paged_mqa_logits_metadata", + "triton_create_paged_compress_data", + "topk_transform_512", + "topk_transform_512_v2", + "plan_topk_v2", + "hash_topk", + "mega_moe_pre_dispatch", + "mask_topk_ids", + "silu_and_mul_clamp", + "silu_and_mul_masked_post_quant", + "silu_and_mul_contig_post_quant", ] diff --git a/python/sglang/jit_kernel/dsv4/attn.py b/python/sglang/jit_kernel/dsv4/attn.py new file mode 100644 index 000000000..aa8f24b4f --- /dev/null +++ b/python/sglang/jit_kernel/dsv4/attn.py @@ -0,0 +1,195 @@ +from typing import Literal, Tuple + +import torch +import triton +import triton.language as tl + +from sglang.jit_kernel.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + +from .utils import make_name + + +@cache_once +def _jit_metadata_module(): + return load_jit( + make_name("metadata"), + cuda_files=["deepseek_v4/paged_mqa_metadata.cuh"], + cuda_wrappers=[("run", "IndexerMetadataKernel::run")], + ) + + +@cache_once +def _jit_fused_store_module( + name: Literal["flashmla", "indexer"], + input_dtype: torch.dtype, + index_dtype: torch.dtype, + page_size: int, +): + args = make_cpp_args(input_dtype, index_dtype, page_size, is_arch_support_pdl()) + cname = "FlashMLA" if name == "flashmla" else "Indexer" + kernel_class = f"FusedStoreCache{cname}Kernel<{args}>" + return load_jit( + make_name("store_" + name), + *args, + cuda_files=["deepseek_v4/store.cuh"], + cuda_wrappers=[("run", f"{kernel_class}::run")], + ) + + +def get_paged_mqa_logits_metadata(seq_lens: torch.Tensor, page_size: int, num_sm: int): + assert page_size == 64 + seq_lens = seq_lens.view(-1).to(torch.int32) + metadata = seq_lens.new_empty(num_sm + 1, 2) + module = _jit_metadata_module() + module.run(seq_lens, metadata) + return metadata + + +def fused_store_cache( + input: torch.Tensor, + cache: torch.Tensor, + indices: torch.Tensor, + *, + page_size: int, + type: Literal["flashmla", "indexer"], +) -> None: + module = _jit_fused_store_module( + name=type, + input_dtype=input.dtype, + index_dtype=indices.dtype, + page_size=page_size, + ) + module.run(input, cache, indices) + + +@triton.jit +def create_paged_compress_data_kernel( + req_pool_indices_ptr, + seq_lens_ptr, + extend_seq_lens_ptr, + req_to_token_ptr, + full_to_swa_index_mapping_ptr, + out_0_ptr, + out_1_ptr, + batch_size, + stride_req_to_token_0, + stride_req_to_token_1: tl.constexpr, + stride_out_1_0, + stride_out_1_1: tl.constexpr, + compress_ratio: tl.constexpr, + is_overlap: tl.constexpr, + swa_page_size: tl.constexpr, + ring_size: tl.constexpr, + BLOCK: tl.constexpr, +) -> None: + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < batch_size + + rid = tl.load(req_pool_indices_ptr + offs, mask=mask, other=0).to(tl.int32) + seq_len = tl.load(seq_lens_ptr + offs, mask=mask, other=0).to(tl.int32) + extend_len = tl.load(extend_seq_lens_ptr + offs, mask=mask, other=0).to(tl.int32) + prefix_len = seq_len - extend_len + + cr = compress_ratio + write_pos = ((seq_len - 1) // cr) * cr + load_pos = ((prefix_len - 1) // cr) * cr + write_overlap_pos = write_pos - cr + load_overlap_pos = load_pos - cr + v0 = tl.zeros([BLOCK], tl.int32) + v1 = tl.zeros([BLOCK], tl.int32) + v2 = tl.zeros([BLOCK], tl.int32) + v3 = tl.zeros([BLOCK], tl.int32) + + for i in tl.static_range(4): + if i == 0: + pos = load_pos + elif i == 1: + pos = write_pos + elif i == 2: + pos = load_overlap_pos + else: + pos = write_overlap_pos + pos = tl.maximum(pos, 0) + loc = tl.load( + req_to_token_ptr + + rid.to(tl.int64) * stride_req_to_token_0 + + pos.to(tl.int64) * stride_req_to_token_1, + mask=mask, + other=0, + ).to(tl.int32) + swa_loc = tl.load(full_to_swa_index_mapping_ptr + loc, mask=mask, other=0).to( + tl.int32 + ) + swa_page = swa_loc // swa_page_size + state_loc = swa_page * ring_size + (swa_loc % ring_size) + state_loc = state_loc // cr + if i == 0: + v0 = state_loc + elif i == 1: + v1 = state_loc + elif i == 2: + v2 = state_loc + else: + v3 = state_loc + + tl.store(out_0_ptr + offs, v1, mask=mask) + + if is_overlap: + base = out_1_ptr + offs * stride_out_1_0 + tl.store(base + 0 * stride_out_1_1, v2, mask=mask) + tl.store(base + 1 * stride_out_1_1, v0, mask=mask) + tl.store(base + 2 * stride_out_1_1, v3, mask=mask) + tl.store(base + 3 * stride_out_1_1, write_pos.to(tl.int32), mask=mask) + else: + base = out_1_ptr + offs * stride_out_1_0 + tl.store(base + 0 * stride_out_1_1, v0, mask=mask) + + +def triton_create_paged_compress_data( + *, + compress_ratio: int, + is_overlap: bool, + swa_page_size: int, + ring_size: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + req_to_token: torch.Tensor, + full_to_swa_index_mapping: torch.Tensor, + block: int = 128, +) -> Tuple[torch.Tensor, torch.Tensor]: + batch_size = req_pool_indices.shape[0] + out_dim = 4 if is_overlap else 1 + device_args: dict = dict(device=req_pool_indices.device, dtype=torch.int32) + out_0 = torch.empty((batch_size,), **device_args) + out_1 = torch.empty((batch_size, out_dim), **device_args) + grid = (triton.cdiv(batch_size, block),) + create_paged_compress_data_kernel[grid]( + req_pool_indices, + seq_lens, + extend_seq_lens, + req_to_token, + full_to_swa_index_mapping, + out_0, + out_1, + batch_size=batch_size, + stride_req_to_token_0=req_to_token.stride(0), + stride_req_to_token_1=req_to_token.stride(1), # type: ignore + stride_out_1_0=out_1.stride(0), + stride_out_1_1=out_1.stride(1), # type: ignore + compress_ratio=compress_ratio, # type: ignore + is_overlap=1 if is_overlap else 0, # type: ignore + swa_page_size=swa_page_size, # type: ignore + ring_size=ring_size, # type: ignore + BLOCK=block, # type: ignore + ) + + if not is_overlap: + out_1.squeeze_(1) + return out_0, out_1 diff --git a/python/sglang/jit_kernel/dsv4/compress_old.py b/python/sglang/jit_kernel/dsv4/compress_old.py new file mode 100644 index 000000000..9bf96a964 --- /dev/null +++ b/python/sglang/jit_kernel/dsv4/compress_old.py @@ -0,0 +1,308 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union + +import torch + +from sglang.jit_kernel.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) +from sglang.srt.environ import envs + +from .utils import make_name + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@cache_once +def _jit_common_module() -> Module: + return load_jit( + make_name("common"), + cuda_files=["deepseek_v4/common.cuh"], + cuda_wrappers=[("plan_compress_prefill", "plan_compress_prefill")], + ) + + +@cache_once +def _jit_compress_128_online_plan_module() -> Module: + """Host-side plan generator for online compress 128 (no template args).""" + return load_jit( + make_name("compress_128_online_plan"), + cuda_files=["deepseek_v4/c128_online.cuh"], + cuda_wrappers=[ + ("plan_compress_online_prefill", "plan_compress_online_prefill"), + ], + ) + + +@cache_once +def _jit_compress_128_online_module(head_dim: int) -> Module: + """Online compress 128 kernel: ring_size=1, per-index (max, sum, kv) state.""" + args = make_cpp_args(head_dim, is_arch_support_pdl()) + kernel_class = f"FlashCompress128OnlineKernel<{args}>" + return load_jit( + make_name("compress_128_online"), + *args, + cuda_files=["deepseek_v4/c128_online.cuh"], + cuda_wrappers=[ + ("decode", f"{kernel_class}::run_decode"), + ("prefill", f"{kernel_class}::run_prefill"), + ], + extra_cuda_cflags=["-use_fast_math"], + ) + + +@cache_once +def _jit_norm_rope_module( + dtype: torch.dtype, + head_dim: int, + rope_dim: int, +) -> Module: + args = make_cpp_args(dtype, head_dim, rope_dim, is_arch_support_pdl()) + return load_jit( + make_name("fused_norm_rope"), + *args, + cuda_files=["deepseek_v4/fused_norm_rope.cuh"], + cuda_wrappers=[ + ("forward", f"FusedNormRopeKernel<{args}>::forward"), + ], + ) + + +@cache_once +def _jit_compress_module( + head_dim: int, + dtype_in: torch.dtype, + dtype_out: torch.dtype, + ratio: Literal[4, 128], +) -> Module: + args = make_cpp_args(head_dim, dtype_in, dtype_out, is_arch_support_pdl()) + kernel_class = f"FlashCompress{ratio}Kernel<{args}>" + return load_jit( + make_name(f"compress_{ratio}"), + *args, + cuda_files=[f"deepseek_v4/c{ratio}.cuh"], + cuda_wrappers=[ + ("decode", f"{kernel_class}::run_decode"), + ("prefill", f"{kernel_class}::run_prefill"), + ], + extra_cuda_cflags=["-use_fast_math"], + ) + + +class CompressorPrefillPlan(NamedTuple): + compress_ratio: int + compress_plan: torch.Tensor + write_plan: torch.Tensor + + def copy_(self, other: CompressorPrefillPlan) -> None: + assert self.compress_ratio == other.compress_ratio + self.compress_plan.copy_(other.compress_plan) + self.write_plan.copy_(other.write_plan) + + @staticmethod + def generate( + compress_ratio: Literal[4, 128], + num_q_tokens: int, + seq_lens: torch.Tensor, + extend_lens: torch.Tensor, + device: torch.device, + use_cuda_graph: bool = False, + ) -> CompressorPrefillPlan: + from sglang.srt.environ import envs + + # Online c128 keeps the same NamedTuple shape (compress_plan, write_plan) + # so call sites that splat `*plan[1:]` continue to work, but the C++ + # plan struct semantics differ (last-token coords + window_len). + if compress_ratio == 128 and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get(): + return CompressorPrefillPlan._generate_online( + num_q_tokens=num_q_tokens, + seq_lens=seq_lens, + extend_lens=extend_lens, + device=device, + use_cuda_graph=use_cuda_graph, + ) + assert seq_lens.device == extend_lens.device + seq_lens = seq_lens.to(torch.int64) + extend_lens = extend_lens.to(torch.int64) + plan_tensor = torch.empty( + (2, num_q_tokens, 16), + dtype=torch.uint8, + device=seq_lens.device, + pin_memory=seq_lens.is_cpu, + ) + module = _jit_common_module() + is_overlap = compress_ratio == 4 + plan_lens = module.plan_compress_prefill( + extend_lens, + seq_lens, + plan_tensor[0], + plan_tensor[1], + compress_ratio, + is_overlap, + use_cuda_graph, + ) + return CompressorPrefillPlan( + compress_ratio, + plan_tensor[0, : plan_lens[0]].to(device, non_blocking=True), + plan_tensor[1, : plan_lens[1]].to(device, non_blocking=True), + ) + + @staticmethod + def _generate_online( + num_q_tokens: int, + seq_lens: torch.Tensor, + extend_lens: torch.Tensor, + device: torch.device, + use_cuda_graph: bool, + ) -> CompressorPrefillPlan: + # Online plan host-side path: only CPU/cuda-host implemented today. + # Move inputs to CPU pinned memory then bounce the result to device. + seq_lens_cpu = seq_lens.detach().to(torch.int64).cpu() + extend_lens_cpu = extend_lens.detach().to(torch.int64).cpu() + plan_tensor = torch.empty( + (2, num_q_tokens, 16), + dtype=torch.uint8, + device="cpu", + pin_memory=True, + ) + module = _jit_compress_128_online_plan_module() + plan_lens = module.plan_compress_online_prefill( + extend_lens_cpu, + seq_lens_cpu, + plan_tensor[0], + plan_tensor[1], + use_cuda_graph, + ) + return CompressorPrefillPlan( + 128, + plan_tensor[0, : plan_lens[0]].to(device, non_blocking=True), + plan_tensor[1, : plan_lens[1]].to(device, non_blocking=True), + ) + + @property + def is_decode(self) -> bool: + return False + + +class CompressorDecodePlan(NamedTuple): + compress_ratio: int + seq_lens: torch.Tensor + + def copy_(self, other: CompressorDecodePlan) -> None: + assert self.compress_ratio == other.compress_ratio + self.seq_lens.copy_(other.seq_lens) + + @property + def is_decode(self) -> bool: + return True + + +def compress_plan( + compress_ratio: Literal[4, 128], + num_q_tokens: int, + seq_lens: torch.Tensor, + extend_lens: Optional[torch.Tensor], + device: torch.device, +) -> Union[CompressorDecodePlan, CompressorPrefillPlan]: + if extend_lens is not None: + return CompressorPrefillPlan.generate( + compress_ratio, + num_q_tokens, + seq_lens, + extend_lens, + device, + ) + else: + assert num_q_tokens == len(seq_lens) + seq_lens = seq_lens.to(device, non_blocking=True) + return CompressorDecodePlan(compress_ratio, seq_lens) + + +def compress_forward( + kv_score_buffer: torch.Tensor, + kv_score_input: torch.Tensor, + ape: torch.Tensor, + indices: torch.Tensor, + plan: Union[CompressorDecodePlan, CompressorPrefillPlan, None] = None, + extra_data: Optional[torch.Tensor] = None, + *, + head_dim: int, + compress_ratio: Literal[4, 128], + out: Optional[torch.Tensor] = None, + seq_lens: Optional[torch.Tensor] = None, + extend_lens: Optional[torch.Tensor] = None, +) -> torch.Tensor: + assert head_dim % 128 == 0 + num_q_tokens = kv_score_input.shape[0] + if out is None: + out = kv_score_input.new_empty((num_q_tokens, head_dim)) + if plan is None: + assert seq_lens is not None + plan = compress_plan( + compress_ratio, + num_q_tokens, + seq_lens, + extend_lens, + kv_score_input.device, + ) + assert plan.compress_ratio == compress_ratio, "Mismatched compress ratio in plan!" + # Online c128: separate JIT module, fp32 state, no compile-time dtypes. + if compress_ratio == 128 and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get(): + online_module = _jit_compress_128_online_module(head_dim=head_dim) + F = online_module.decode if plan.is_decode else online_module.prefill + F(kv_score_buffer, kv_score_input, out, ape, indices, *plan[1:], extra_data) + return out + module = _jit_compress_module( + head_dim, + kv_score_input.dtype, + out.dtype, + compress_ratio, + ) + F = module.decode if plan.is_decode else module.prefill + F(kv_score_buffer, kv_score_input, out, ape, indices, *plan[1:], extra_data) + return out + + +def compress_fused_norm_rope_inplace( + kv: torch.Tensor, + weight: torch.Tensor, + eps: float, + freq_cis: torch.Tensor, + plan: Union[CompressorDecodePlan, CompressorPrefillPlan], +) -> None: + freq_cis = torch.view_as_real(freq_cis).flatten(-2) + module = _jit_norm_rope_module(kv.dtype, kv.shape[-1], freq_cis.shape[-1]) + module.forward( + kv, + weight, + plan[1], + freq_cis, + int(plan.is_decode), + eps, + plan.compress_ratio, + ) + + +def fused_norm_rope_inplace( + kv: torch.Tensor, + weight: torch.Tensor, + eps: float, + freq_cis: torch.Tensor, + positions: torch.Tensor, +) -> None: + freq_cis = torch.view_as_real(freq_cis).flatten(-2) + module = _jit_norm_rope_module(kv.dtype, kv.shape[-1], freq_cis.shape[-1]) + module.forward( + kv, + weight, + positions, + freq_cis, + 2, + eps, + 0, + ) diff --git a/python/sglang/jit_kernel/dsv4/elementwise.py b/python/sglang/jit_kernel/dsv4/elementwise.py new file mode 100644 index 000000000..d2bbb8f3c --- /dev/null +++ b/python/sglang/jit_kernel/dsv4/elementwise.py @@ -0,0 +1,158 @@ +from typing import Optional, Tuple + +import torch + +from sglang.jit_kernel.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) +from sglang.srt.utils import is_hip + +from .utils import make_name + +_is_hip = is_hip() + + +@cache_once +def _jit_fused_rope_module(): + args = make_cpp_args(is_arch_support_pdl()) + return load_jit( + make_name("fused_rope"), + *args, + cuda_files=["deepseek_v4/rope.cuh"], + cuda_wrappers=[("forward", f"FusedQKRopeKernel<{args}>::forward")], + ) + + +@cache_once +def _jit_main_q_norm_rope_module( + dtype: torch.dtype, + head_dim: int, + rope_dim: int, +): + """Main MLA path Q kernel: rmsnorm-self + RoPE, warp per (token, head).""" + args = make_cpp_args(dtype, head_dim, rope_dim, is_arch_support_pdl()) + return load_jit( + make_name("main_q_norm_rope"), + *args, + cuda_files=["deepseek_v4/main_norm_rope.cuh"], + cuda_wrappers=[ + ("forward", f"FusedQNormRopeKernel<{args}>::forward"), + ], + ) + + +@cache_once +def _jit_main_k_norm_rope_flashmla_module( + dtype: torch.dtype, + head_dim: int, + rope_dim: int, + page_size: int, +): + """Main MLA path K kernel: rmsnorm + RoPE + write to FlashMLA paged cache.""" + args = make_cpp_args(dtype, head_dim, rope_dim, page_size, is_arch_support_pdl()) + return load_jit( + make_name("main_k_norm_rope_flashmla"), + *args, + cuda_files=["deepseek_v4/main_norm_rope.cuh"], + cuda_wrappers=[ + ("forward", f"FusedKNormRopeFlashMLAKernel<{args}>::forward"), + ], + ) + + +@cache_once +def _jit_main_q_indexer_rope_hadamard_quant_module(dtype: torch.dtype): + """C4 indexer Q kernel: RoPE + 128-pt Hadamard + fp8 act-quant""" + args = make_cpp_args(dtype, is_arch_support_pdl()) + return load_jit( + make_name("main_q_indexer_rope_hadamard_quant"), + *args, + cuda_files=["deepseek_v4/main_norm_rope.cuh"], + cuda_wrappers=[ + ("forward", f"FusedQIndexerRopeHadamardQuantKernel<{args}>::forward"), + ], + ) + + +def fused_rope_inplace( + q: torch.Tensor, + k: Optional[torch.Tensor], + freqs_cis: torch.Tensor, + positions: torch.Tensor, + inverse: bool = False, +) -> None: + """Apply rotary embeddings to both Q and K in a single fused CUDA kernel. + + Args: + q: [batch_size, num_q_heads, rope_dim] bfloat16 + k: [batch_size, num_k_heads, rope_dim] bfloat16 or None + freqs_cis: [max_seq_len, rope_dim // 2] complex64 (full table) + positions: [batch_size] int32 or int64, indices into freqs_cis + inverse: if True, apply inverse rotation (conjugate freqs) + """ + if _is_hip: + from sglang.srt.layers.deepseek_v4_rope import apply_rotary_emb_triton + + apply_rotary_emb_triton(q, freqs_cis, positions=positions, inverse=inverse) + if k is not None: + apply_rotary_emb_triton(k, freqs_cis, positions=positions, inverse=inverse) + return + + freqs_real = torch.view_as_real(freqs_cis).flatten(-2).contiguous() + module = _jit_fused_rope_module() + module.forward(q, k, freqs_real, positions, inverse) + + +def fused_q_norm_rope( + q_input: torch.Tensor, + q_output: torch.Tensor, + eps: float, + freqs_cis: torch.Tensor, + positions: torch.Tensor, +) -> None: + freqs_real = torch.view_as_real(freqs_cis).flatten(-2) + head_dim = q_input.shape[-1] + rope_dim = freqs_real.shape[-1] + module = _jit_main_q_norm_rope_module(q_input.dtype, head_dim, rope_dim) + module.forward(q_input, q_output, freqs_real, positions, eps) + + +def fused_q_indexer_rope_hadamard_quant( + q_input: torch.Tensor, + weight: torch.Tensor, + weight_scale: float, + freqs_cis: torch.Tensor, + positions: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + freqs_real = torch.view_as_real(freqs_cis).flatten(-2) + q_fp8 = torch.empty(q_input.shape, dtype=torch.float8_e4m3fn, device=q_input.device) + weights_out = torch.empty( + (*q_input.shape[:-1], 1), dtype=torch.float32, device=q_input.device + ) + module = _jit_main_q_indexer_rope_hadamard_quant_module(q_input.dtype) + module.forward( + q_input, q_fp8, weight, weights_out, float(weight_scale), freqs_real, positions + ) + return q_fp8, weights_out + + +def fused_k_norm_rope_flashmla( + kv: torch.Tensor, + kv_weight: torch.Tensor, + eps: float, + freqs_cis: torch.Tensor, + positions: torch.Tensor, + out_loc: torch.Tensor, + kvcache: torch.Tensor, + page_size: int, +) -> None: + freqs_real = torch.view_as_real(freqs_cis).flatten(-2) + head_dim = kv.shape[-1] + rope_dim = freqs_real.shape[-1] + module = _jit_main_k_norm_rope_flashmla_module( + kv.dtype, head_dim, rope_dim, page_size + ) + module.forward(kv, kv_weight, freqs_real, positions, out_loc, kvcache, eps) diff --git a/python/sglang/jit_kernel/dsv4/gemm.py b/python/sglang/jit_kernel/dsv4/gemm.py new file mode 100644 index 000000000..c89ca32a1 --- /dev/null +++ b/python/sglang/jit_kernel/dsv4/gemm.py @@ -0,0 +1,24 @@ +import torch + +from sglang.srt.environ import envs +from sglang.srt.layers import deep_gemm_wrapper +from sglang.srt.utils import get_bool_env_var, is_hip + +_is_hip = is_hip() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip + +if _use_aiter: + from aiter.tuned_gemm import tgemm + +_linear_bf16_fp32_algo = envs.SGLANG_OPT_BF16_FP32_GEMM_ALGO.get() + + +def linear_bf16_fp32(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + if _linear_bf16_fp32_algo == "deep_gemm": + z = torch.empty(x.size(0), y.size(0), dtype=torch.float32, device=x.device) + deep_gemm_wrapper.gemm_nt_bf16bf16f32(x, y, z) + return z + elif _use_aiter: + return tgemm.mm(x, y, otype=torch.float32) + else: + return torch.mm(x, y.t(), out_dtype=torch.float32) diff --git a/python/sglang/jit_kernel/dsv4/hisparse.py b/python/sglang/jit_kernel/dsv4/hisparse.py new file mode 100644 index 000000000..c1fdda2b0 --- /dev/null +++ b/python/sglang/jit_kernel/dsv4/hisparse.py @@ -0,0 +1,27 @@ +import torch + +from sglang.jit_kernel.utils import ( + cache_once, + load_jit, +) + +from .utils import make_name + + +@cache_once +def _jit_hisparse_transfer_module(): + return load_jit( + make_name("hisparse_transfer"), + cuda_files=["deepseek_v4/hisparse_transfer.cuh"], + cuda_wrappers=[("hisparse_transfer", "hisparse_transfer")], + ) + + +def hisparse_offload_to_host( + gpu_ptrs: torch.Tensor, + cpu_ptrs: torch.Tensor, + gpu_indices: torch.Tensor, + cpu_indices: torch.Tensor, +) -> None: + module = _jit_hisparse_transfer_module() + module.hisparse_transfer(gpu_ptrs, cpu_ptrs, gpu_indices, cpu_indices) diff --git a/python/sglang/jit_kernel/dsv4/moe.py b/python/sglang/jit_kernel/dsv4/moe.py new file mode 100644 index 000000000..a5372c781 --- /dev/null +++ b/python/sglang/jit_kernel/dsv4/moe.py @@ -0,0 +1,216 @@ +from typing import Optional, Tuple + +import torch + +from sglang.jit_kernel.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + +from .utils import make_name + + +@cache_once +def _jit_mask_topk_module(): + return load_jit( + make_name("mask_topk"), + cuda_files=["deepseek_v4/hash_topk.cuh"], + cuda_wrappers=[("run", "MaskKernel::run")], + ) + + +@cache_once +def _jit_hash_topk_module(): + args = make_cpp_args("act_sqrt_softplus", is_arch_support_pdl()) + return load_jit( + make_name("hash_topk"), + *args, + cuda_files=["deepseek_v4/hash_topk.cuh"], + cuda_wrappers=[("hash_topk", f"HashTopKKernel<{args}>::run")], + ) + + +@cache_once +def _jit_mega_moe_pre_dispatch_module(quant_group_size: int): + args = make_cpp_args(quant_group_size, is_arch_support_pdl()) + return load_jit( + make_name("mega_moe_pre_dispatch"), + *args, + cuda_files=["deepseek_v4/mega_moe_pre_dispatch.cuh"], + cuda_wrappers=[("run", f"MegaMoEPreDispatchKernel<{args}>::run")], + ) + + +@cache_once +def _jit_silu_mul_quant_varlen_module( + quant_group_size: int, + scale_ue8m0: bool, + swizzle: bool, + apply_swiglu_limit: bool, +): + args = make_cpp_args( + quant_group_size, + scale_ue8m0, + swizzle, + is_arch_support_pdl(), + apply_swiglu_limit, + ) + return load_jit( + make_name("silu_mul_quant_varlen"), + *args, + cuda_files=["deepseek_v4/silu_and_mul_masked_post_quant.cuh"], + cuda_wrappers=[("run", f"SiluAndMulMaskedPostQuantKernel<{args}>::run")], + extra_cuda_cflags=["-use_fast_math"], + ) + + +@cache_once +def _jit_silu_mul_quant_contig_module( + quant_group_size: int, + scale_ue8m0: bool, + swizzle: bool, + apply_swiglu_limit: bool, +): + args = make_cpp_args( + quant_group_size, + scale_ue8m0, + swizzle, + is_arch_support_pdl(), + apply_swiglu_limit, + ) + return load_jit( + make_name("silu_mul_quant_contig"), + *args, + cuda_files=["deepseek_v4/silu_and_mul_masked_post_quant.cuh"], + cuda_wrappers=[("run", f"SiluAndMulContigPostQuantKernel<{args}>::run")], + extra_cuda_cflags=["-use_fast_math"], + ) + + +@cache_once +def _jit_silu_and_mul_clamp_module(dtype: torch.dtype): + args = make_cpp_args(dtype, is_arch_support_pdl()) + return load_jit( + make_name("silu_and_mul_clamp"), + *args, + cuda_files=["deepseek_v4/silu_and_mul_masked_post_quant.cuh"], + cuda_wrappers=[("run", f"SiluAndMulClampKernel<{args}>::run")], + extra_cuda_cflags=["-use_fast_math"], + ) + + +def mask_topk_ids(topk_ids: torch.Tensor, num_token_non_padded: torch.Tensor): + return _jit_mask_topk_module().run(topk_ids, num_token_non_padded) + + +def hash_topk( + router_logits: torch.Tensor, + input_ids: torch.Tensor, + tid2eid: torch.Tensor, + num_fused_shared_experts: int = 0, + routed_scaling_factor: float = 1.0, + scoring_func: str = "sqrtsoftplus", +) -> Tuple[torch.Tensor, torch.Tensor]: + assert scoring_func == "sqrtsoftplus" + num_tokens = router_logits.size(0) + topk_routed = tid2eid.size(1) + topk_fused = topk_routed + num_fused_shared_experts + topk_ids = torch.empty( + (num_tokens, topk_fused), dtype=torch.int32, device=router_logits.device + ) + topk_weights = torch.empty( + (num_tokens, topk_fused), dtype=torch.float32, device=router_logits.device + ) + module = _jit_hash_topk_module() + module.hash_topk( + router_logits, + input_ids, + tid2eid, + topk_weights, + topk_ids, + routed_scaling_factor, + ) + return topk_weights, topk_ids + + +def mega_moe_pre_dispatch( + x: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + buf_x: torch.Tensor, + buf_x_sf: torch.Tensor, + buf_topk_idx: torch.Tensor, + buf_topk_weights: torch.Tensor, + quant_group_size: int = 32, +) -> None: + module = _jit_mega_moe_pre_dispatch_module(quant_group_size) + module.run( + x, + topk_idx, + topk_weights, + buf_x, + buf_x_sf, + buf_topk_idx, + buf_topk_weights, + ) + + +def silu_and_mul_clamp( + input: torch.Tensor, + output: torch.Tensor, + swiglu_limit: float, +) -> None: + module = _jit_silu_and_mul_clamp_module(input.dtype) + module.run(input, output, float(swiglu_limit)) + + +def silu_and_mul_masked_post_quant( + input: torch.Tensor, + output: torch.Tensor, + output_scale: torch.Tensor, + quant_group_size: int, + masked_m: torch.Tensor, + scale_ue8m0: bool = False, + topk: int = 8, + transposed: bool = False, + swiglu_limit: Optional[float] = None, + swizzle: bool = False, +) -> None: + apply_swiglu_limit = swiglu_limit is not None + module = _jit_silu_mul_quant_varlen_module( + quant_group_size, scale_ue8m0, swizzle, apply_swiglu_limit + ) + module.run( + input, + output, + output_scale, + masked_m, + topk, + transposed, + float(swiglu_limit) if apply_swiglu_limit else 0.0, + ) + + +def silu_and_mul_contig_post_quant( + input: torch.Tensor, + output: torch.Tensor, + output_scale: torch.Tensor, + quant_group_size: int, + scale_ue8m0: bool = False, + transposed: bool = False, + swiglu_limit: Optional[float] = None, + swizzle: bool = False, +) -> None: + apply_swiglu_limit = swiglu_limit is not None + module = _jit_silu_mul_quant_contig_module( + quant_group_size, scale_ue8m0, swizzle, apply_swiglu_limit + ) + module.run( + input, + output, + output_scale, + transposed, + float(swiglu_limit) if apply_swiglu_limit else 0.0, + ) diff --git a/python/sglang/jit_kernel/dsv4/topk.py b/python/sglang/jit_kernel/dsv4/topk.py new file mode 100644 index 000000000..3c7d14769 --- /dev/null +++ b/python/sglang/jit_kernel/dsv4/topk.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +from typing import Optional + +import torch + +from sglang.jit_kernel.utils import ( + cache_once, + is_arch_support_pdl, + load_jit, + make_cpp_args, +) + +from .utils import make_name + + +@cache_once +def _jit_topk_v1_module(topk: int): + args = make_cpp_args(is_arch_support_pdl()) + assert topk in (512, 1024), "Only support topk=512 or 1024" + return load_jit( + make_name(f"topk_v1_{topk}"), + *args, + cuda_files=["deepseek_v4/topk_v1.cuh"], + cuda_wrappers=[("topk_transform", f"TopKKernel<{args}>::transform")], + extra_cuda_cflags=[f"-DSGL_TOPK={topk}"], + ) + + +@cache_once +def _jit_topk_v2_module(topk: int): + return load_jit( + make_name(f"topk_v2_{topk}"), + cuda_files=["deepseek_v4/topk_v2.cuh"], + cuda_wrappers=[ + ("topk_transform", "CombinedTopKKernel::transform"), + ("topk_plan", "CombinedTopKKernel::plan"), + ], + extra_cuda_cflags=[f"-DSGL_TOPK={topk}"], + ) + + +def topk_transform_512( + scores: torch.Tensor, + seq_lens: torch.Tensor, + page_tables: torch.Tensor, + out_page_indices: torch.Tensor, + page_size: int, + out_raw_indices: Optional[torch.Tensor] = None, +) -> None: + module = _jit_topk_v1_module(out_page_indices.shape[1]) + module.topk_transform( + scores, seq_lens, page_tables, out_page_indices, page_size, out_raw_indices + ) + + +_WORKSPACE_INTS_PER_BATCH = 2 + 1024 * 2 +_PLAN_METADATA_INTS_PER_BATCH = 4 + + +def plan_topk_v2(seq_lens: torch.Tensor, static_threshold: int = 0) -> torch.Tensor: + module = _jit_topk_v2_module(512) # does not matter + bs = seq_lens.shape[0] + metadata = seq_lens.new_empty(bs + 1, _PLAN_METADATA_INTS_PER_BATCH) + module.topk_plan(seq_lens, metadata, static_threshold) + return metadata + + +def topk_transform_512_v2( + scores: torch.Tensor, + seq_lens: torch.Tensor, + page_tables: torch.Tensor, + out_page_indices: torch.Tensor, + page_size: int, + metadata: torch.Tensor, +) -> None: + module = _jit_topk_v2_module(out_page_indices.shape[1]) + bs = scores.shape[0] + workspace = seq_lens.new_empty(bs, _WORKSPACE_INTS_PER_BATCH) + module.topk_transform( + scores, + seq_lens, + page_tables, + out_page_indices, + page_size, + workspace, + metadata, + ) diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index 428b5af14..092b98e2c 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor.py @@ -5,13 +5,12 @@ from typing import TYPE_CHECKING, List, Literal, NamedTuple, Optional, Union import torch import torch.nn as nn -from sglang.jit_kernel.deepseek_v4 import ( +from sglang.jit_kernel.dsv4 import linear_bf16_fp32, triton_create_paged_compress_data +from sglang.jit_kernel.dsv4.compress_old import ( CompressorDecodePlan, CompressorPrefillPlan, compress_forward, compress_fused_norm_rope_inplace, - linear_bf16_fp32, - triton_create_paged_compress_data, ) from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config from sglang.srt.environ import envs diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index facf5ef3d..f8899264d 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -8,7 +8,7 @@ import torch.nn.functional as F import triton import triton.language as tl -from sglang.jit_kernel.deepseek_v4 import ( +from sglang.jit_kernel.dsv4 import ( fused_q_indexer_rope_hadamard_quant, topk_transform_512, topk_transform_512_v2, diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py index 7995dbd95..c3e804103 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata.py @@ -109,7 +109,7 @@ class PagedIndexerMetadata: import deep_gemm if envs.SGLANG_OPT_USE_JIT_INDEXER_METADATA.get(): - from sglang.jit_kernel.deepseek_v4 import get_paged_mqa_logits_metadata + from sglang.jit_kernel.dsv4 import get_paged_mqa_logits_metadata else: from deep_gemm import get_paged_mqa_logits_metadata @@ -124,7 +124,7 @@ class PagedIndexerMetadata: assert isinstance(self.deep_gemm_metadata, torch.Tensor) - from sglang.jit_kernel.deepseek_v4 import plan_topk_v2 + from sglang.jit_kernel.dsv4 import plan_topk_v2 if envs.SGLANG_OPT_USE_TOPK_V2.get(): self.topk_metadata = plan_topk_v2(self.c4_seq_lens) diff --git a/python/sglang/srt/layers/moe/hash_topk.py b/python/sglang/srt/layers/moe/hash_topk.py index cdeca35f7..959880f84 100644 --- a/python/sglang/srt/layers/moe/hash_topk.py +++ b/python/sglang/srt/layers/moe/hash_topk.py @@ -127,7 +127,7 @@ class HashTopK(nn.Module): ), f"{input_ids.shape=} {hidden_states.shape=} {router_logits.shape=}" if envs.SGLANG_OPT_USE_FUSED_HASH_TOPK.get(): - from sglang.jit_kernel.deepseek_v4 import hash_topk + from sglang.jit_kernel.dsv4 import hash_topk topk_weights, topk_ids = hash_topk( router_logits=router_logits, diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py index 9574bd2da..31eecf8a9 100644 --- a/python/sglang/srt/layers/moe/mega_moe.py +++ b/python/sglang/srt/layers/moe/mega_moe.py @@ -21,7 +21,7 @@ from typing import TYPE_CHECKING, Optional import torch -from sglang.jit_kernel.deepseek_v4 import mega_moe_pre_dispatch +from sglang.jit_kernel.dsv4 import mega_moe_pre_dispatch from sglang.srt.environ import envs from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo from sglang.srt.layers.dp_attention import get_dp_global_num_tokens diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index 61af5533f..f31366868 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, Tuple import einops import torch -from sglang.jit_kernel.deepseek_v4 import silu_and_mul_masked_post_quant +from sglang.jit_kernel.dsv4 import silu_and_mul_masked_post_quant from sglang.srt.environ import envs from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.moe.moe_runner.base import ( @@ -167,7 +167,7 @@ class DeepGemmRunnerCore(MoeRunnerCore): quant_info: DeepGemmMoeQuantInfo, running_state: dict, ) -> torch.Tensor: - from sglang.jit_kernel.deepseek_v4 import silu_and_mul_contig_post_quant + from sglang.jit_kernel.dsv4 import silu_and_mul_contig_post_quant from sglang.srt.layers.moe.ep_moe.kernels import tma_align_input_scale from sglang.srt.layers.quantization.fp8_kernel import ( create_per_token_group_quant_fp8_output_scale, diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py index 8062e93c4..a2f739b2d 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py @@ -569,7 +569,7 @@ def _fused_moe_kernel_sequence( if not filter_expert: if swiglu_limit_for_silu_and_mul_clamp is not None: - from sglang.jit_kernel.deepseek_v4 import silu_and_mul_clamp + from sglang.jit_kernel.dsv4 import silu_and_mul_clamp silu_and_mul_clamp( intermediate_cache1.view(-1, N), diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 732ed7ff7..ca716dc33 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -79,7 +79,7 @@ try: except ImportError: pass -from sglang.jit_kernel.deepseek_v4 import mask_topk_ids +from sglang.jit_kernel.dsv4 import mask_topk_ids from sglang.srt.distributed import ( get_moe_expert_parallel_rank, get_moe_expert_parallel_world_size, diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 8a5a6247d..5e77fa614 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -6,7 +6,7 @@ from typing import List, Literal, NamedTuple, Optional, Tuple import torch -from sglang.jit_kernel.deepseek_v4 import fused_k_norm_rope_flashmla, fused_store_cache +from sglang.jit_kernel.dsv4 import fused_k_norm_rope_flashmla, fused_store_cache from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.environ import envs from sglang.srt.layers.attention.dsa import index_buf_accessor diff --git a/python/sglang/srt/mem_cache/hisparse_memory_pool.py b/python/sglang/srt/mem_cache/hisparse_memory_pool.py index 679c0f46d..b08fe0db1 100644 --- a/python/sglang/srt/mem_cache/hisparse_memory_pool.py +++ b/python/sglang/srt/mem_cache/hisparse_memory_pool.py @@ -463,7 +463,7 @@ class DeepSeekV4SingleKVPoolHost(HiSparseHostPoolMixin): if io_backend != "kernel": raise ValueError(f"Unsupported IO backend: {io_backend}") - from sglang.jit_kernel.deepseek_v4 import hisparse_offload_to_host + from sglang.jit_kernel.dsv4 import hisparse_offload_to_host if host_indices.device != device_indices.device: host_indices = host_indices.to(device=device_indices.device) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 8e8459eb4..98944726f 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -29,7 +29,7 @@ import torch.nn.functional as F from torch import nn from transformers import PretrainedConfig -from sglang.jit_kernel.deepseek_v4 import ( +from sglang.jit_kernel.dsv4 import ( silu_and_mul_clamp, silu_and_mul_contig_post_quant, ) @@ -420,7 +420,7 @@ class MoEGate(nn.Module): logits = aiter_dsv3_router_gemm(hidden_states, self.weight) else: if self.is_deepseek_v4: - from sglang.jit_kernel.deepseek_v4 import linear_bf16_fp32 + from sglang.jit_kernel.dsv4 import linear_bf16_fp32 logits = linear_bf16_fp32(hidden_states, self.weight) else: diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index cb57ea400..fdf6d557c 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -21,7 +21,7 @@ import triton import triton.language as tl import sglang.srt.models.deepseek_v2 as deepseek_v2 -from sglang.jit_kernel.deepseek_v4 import ( +from sglang.jit_kernel.dsv4 import ( fused_norm_rope_inplace, fused_q_norm_rope, fused_rope_inplace,