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

[Refactor] major JIT kernel clean up for dsv4 (#25884)

Browse Source 
Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
DarkSharpness
2026-05-21 01:14:31 -07:00
committed by GitHub
co-authored by Claude
parent 1b3d8da827
commit 19f55c0e6d
23 changed files with 1093 additions and 1399 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/deepseek_v4/topk_1024.cuh
-336
View File
@@ -1,336 +0,0 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh>
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
#include <bit>
#include <cstdint>
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<uint16_t>(~bits) : static_cast<uint16_t>(bits | 0x8000);
return static_cast<uint8_t>(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 <bool kUsePDL>
__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<kUsePDL>();
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<kUsePDL>();
}
template <auto* f, size_t kMaxDynamicSMEM>
void setup_kernel_smem_once(host::DebugInfo where = {}) {
[[maybe_unused]]
static const auto result = [] {
const auto fptr = std::bit_cast<const void*>(f);
return ::cudaFuncSetAttribute(fptr, ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM);
}();
host::RuntimeDeviceCheck(result, where);
}
template <bool kUsePDL>
struct TopK1024Kernel {
static constexpr auto kernel = topk_1024_transform<kUsePDL>;
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<tvm::ffi::TensorView> 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<kDLCUDA>();
TensorMatcher({B, -1}) // strided scores
.with_strides({S, 1})
.with_dtype<float>()
.with_device(device)
.verify(scores);
TensorMatcher({B}) // seq_lens, must be contiguous
.with_dtype<int32_t>()
.with_device(device)
.verify(seq_lens);
TensorMatcher({B, -1}) // strided page table
.with_strides({P, 1})
.with_dtype<int32_t>()
.with_device(device)
.verify(page_table);
TensorMatcher({B, 1024}) // output, must be contiguous
.with_dtype<int32_t>()
.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<int32_t>()
.with_device(device)
.verify(raw_indices.value());
raw_indices_ptr = static_cast<int32_t*>(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<uint32_t>(std::countr_zero(page_size));
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto params = TopK1024Params{
.scores = static_cast<float*>(scores.data_ptr()),
.seq_lens = static_cast<int32_t*>(seq_lens.data_ptr()),
.page_table = static_cast<int32_t*>(page_table.data_ptr()),
.page_indices = static_cast<int32_t*>(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<kernel, kSMEM_>();
LaunchKernel(batch_size, kTopKBlockSize, device.unwrap(), kSMEM_).enable_pdl(kUsePDL)(kernel, params);
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/topk.cuh → python/sglang/jit_kernel/csrc/deepseek_v4/topk_v1.cuh
+13 -9
View File
@@ -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 <bool kUsePDL>
__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 <bool kUsePDL>
struct TopK512Kernel {
static constexpr auto kernel = topk_512_transform<kUsePDL>;
struct TopKKernel {
static constexpr auto kernel = topk_transform_kernel<kUsePDL>;
static void transform(
const tvm::ffi::TensorView scores,
@@ -300,14 +304,14 @@ struct TopK512Kernel {
.with_dtype<int32_t>()
.with_device(device)
.verify(page_table);
TensorMatcher({B, 512}) // output, must be contiguous
TensorMatcher({B, kTopK}) // output, must be contiguous
.with_dtype<int32_t>()
.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<int32_t>()
.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<uint32_t>(std::countr_zero(page_size));
const auto batch_size = static_cast<uint32_t>(B.unwrap());
const auto params = TopK512Params{
const auto params = TopKParams{
.scores = static_cast<float*>(scores.data_ptr()),
.seq_lens = static_cast<int32_t*>(seq_lens.data_ptr()),
.page_table = static_cast<int32_t*>(page_table.data_ptr()),
python/sglang/jit_kernel/deepseek_v4.py
-1036
View File
File diff suppressed because it is too large Load Diff
python/sglang/jit_kernel/dsv4/__init__.py
+48 -1
View File
@@ -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",
]
python/sglang/jit_kernel/dsv4/attn.py
+195
View File
@@ -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
python/sglang/jit_kernel/dsv4/compress_old.py
+308
View File
@@ -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,
)
python/sglang/jit_kernel/dsv4/elementwise.py
+158
View File
@@ -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)
python/sglang/jit_kernel/dsv4/gemm.py
+24
View File
@@ -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)
python/sglang/jit_kernel/dsv4/hisparse.py
+27
View File
@@ -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)
python/sglang/jit_kernel/dsv4/moe.py
+216
View File
@@ -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,
)
python/sglang/jit_kernel/dsv4/topk.py
+88
View File
@@ -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,
)
python/sglang/srt/layers/attention/dsv4/compressor.py
+2 -3
View File
@@ -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
python/sglang/srt/layers/attention/dsv4/indexer.py
+1 -1
View File
@@ -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,
python/sglang/srt/layers/attention/dsv4/metadata.py
+2 -2
View File
@@ -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)
python/sglang/srt/layers/moe/hash_topk.py
+1 -1
View File
@@ -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,
python/sglang/srt/layers/moe/mega_moe.py
+1 -1
View File
@@ -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
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py
+2 -2
View File
@@ -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,
python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py
+1 -1
View File
@@ -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),
python/sglang/srt/layers/moe/topk.py
+1 -1
View File
@@ -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,
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
+1 -1
View File
@@ -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
python/sglang/srt/mem_cache/hisparse_memory_pool.py
+1 -1
View File
@@ -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)
python/sglang/srt/models/deepseek_v2.py
+2 -2
View File
@@ -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:
python/sglang/srt/models/deepseek_v4.py
+1 -1
View File
@@ -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,