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

Fuse the DSA (V3.2, GLM-5.x) indexer Q/K paths into single kernels (#27705)

Browse Source 
Co-authored-by: Brayden Zhong <brayden@radixark.ai>
Co-authored-by: Kaixi <kaiximatteoc@nvidia.com>
This commit is contained in:
Brayden Zhong
2026-06-27 14:29:56 -07:00
committed by GitHub
co-authored by Brayden Zhong Kaixi
parent e4253b39e2
commit 073de15053
13 changed files with 1196 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/deepseek_v32/indexer_k.cuh
+415
View File
@@ -0,0 +1,415 @@
// DeepSeek-V3.2 only.
//
// DSA indexer K kernels: single-head LayerNorm (not RMS), ropes the leading
// kRopeDim dims, and fp8-quantizes the un-rotated activations. V3.2 drops the
// Hadamard incoherence rotation; it is logit-preserving (see main_norm_rope.cuh).
//
// Independent of the wk + weights_proj GEMM fusion (dsa_indexer.py): `k_input`
// here is the non-contiguous wk slice kw[:, :head_dim] read via
// k_input_stride_batch (no copy).
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/math.cuh>
#include <sgl_kernel/tile.cuh>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
#include <tvm/ffi/container/tensor.h>
#include <bit>
#include <cstdint>
namespace {
using deepseek_v4::fp8::pack_fp8;
constexpr uint32_t kFusedKIndexerBlockSize = 128;
constexpr uint32_t kFusedKIndexerNumWarps = kFusedKIndexerBlockSize / device::kWarpThreads;
#define K_INDEXER_KERNEL __global__ __launch_bounds__(kFusedKIndexerBlockSize, 16)
// Indexer K: LayerNorm + RoPE -> bf16.
struct FusedKIndexerNormRopeParams {
const void* __restrict__ k_input; // (B, 128) DType
void* __restrict__ k_out; // (B, 128) DType
const float* __restrict__ weight; // (128,) fp32 -- LayerNorm gamma
const float* __restrict__ bias; // (128,) fp32 -- LayerNorm beta
const float* __restrict__ freqs_cis; // (max_pos, 64) fp32
const void* __restrict__ positions; // (B,) PosT
// Row stride for `k_input` in elements (caller passes the wk slice directly).
int64_t k_input_stride_batch;
uint32_t batch_size;
float eps;
};
template <typename DType, typename PosT, bool kUsePDL>
K_INDEXER_KERNEL void fused_k_indexer_norm_rope(const __grid_constant__ FusedKIndexerNormRopeParams params) {
using namespace device;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kRopeDim = 64;
constexpr int64_t kVecSize = 4;
constexpr uint32_t kRopeSize = kRopeDim / kVecSize; // = 16
static_assert(kHeadDim == kWarpThreads * kVecSize);
static_assert(kRopeDim == kWarpThreads * 2);
static_assert(kRopeSize <= kWarpThreads);
using Storage = AlignedVector<DType, kVecSize>;
using Float4 = AlignedVector<float, kVecSize>;
const auto warp_id = threadIdx.x / kWarpThreads;
const auto lane_id = threadIdx.x % kWarpThreads;
const auto work_id = blockIdx.x * kFusedKIndexerNumWarps + warp_id;
const bool is_rope_lane = lane_id < kRopeSize;
if (work_id >= params.batch_size) return;
const auto input_ptr = static_cast<const DType*>(params.k_input) + work_id * params.k_input_stride_batch;
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[work_id]);
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
PDLWaitPrimary<kUsePDL>();
Float4 data, freq, gamma, beta;
// part 1: LayerNorm
{
Storage input_vec;
input_vec.load(input_ptr, lane_id);
gamma.load(params.weight, lane_id);
beta.load(params.bias, lane_id);
if (is_rope_lane) freq.load(freqs_cis, lane_id);
float sum = 0.0f;
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
data[i] = cast<float>(input_vec[i]);
sum += data[i];
}
const float mean = warp::reduce_sum(sum) / kHeadDim;
float var = 0.0f;
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
const float centered = data[i] - mean;
var += centered * centered;
}
const float inv_std = math::rsqrt(warp::reduce_sum(var) / kHeadDim + params.eps);
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
data[i] = (data[i] - mean) * inv_std * gamma[i] + beta[i];
}
}
// part 2: rope on rope lanes
if (is_rope_lane) {
const auto x_real = data[0];
const auto x_imag = data[1];
const auto y_real = data[2];
const auto y_imag = data[3];
const auto fxr = freq[0];
const auto fxi = freq[1];
const auto fyr = freq[2];
const auto fyi = freq[3];
data[0] = x_real * fxr - x_imag * fxi;
data[1] = x_real * fxi + x_imag * fxr;
data[2] = y_real * fyr - y_imag * fyi;
data[3] = y_real * fyi + y_imag * fyr;
}
PDLTriggerSecondary<kUsePDL>();
{
Storage out_vec;
#pragma unroll
for (int i = 0; i < kVecSize; ++i)
out_vec[i] = cast<DType>(data[i]);
auto out_row = static_cast<DType*>(params.k_out) + work_id * kHeadDim;
out_vec.store(out_row, lane_id);
}
}
template <typename DType, bool kUsePDL>
struct FusedKIndexerNormRopeKernel {
template <typename PosT>
static constexpr auto kernel = fused_k_indexer_norm_rope<DType, PosT, kUsePDL>;
static void forward(
const tvm::ffi::TensorView k_input,
const tvm::ffi::TensorView k_out,
const tvm::ffi::TensorView weight,
const tvm::ffi::TensorView bias,
const tvm::ffi::TensorView freqs_cis,
const tvm::ffi::TensorView positions,
double eps) {
using namespace host;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kRopeDim = 64;
auto B = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({B, kHeadDim}) //
.with_strides({-1, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(k_input);
TensorMatcher({B, kHeadDim}) //
.with_strides({kHeadDim, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(k_out);
TensorMatcher({kHeadDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(weight);
TensorMatcher({kHeadDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(bias);
TensorMatcher({-1, kRopeDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(freqs_cis);
auto pos_dtype = SymbolicDType{};
TensorMatcher({B}) //
.with_dtype<int32_t, int64_t>(pos_dtype)
.with_device(device_)
.verify(positions);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
if (batch_size == 0) return;
const auto params = FusedKIndexerNormRopeParams{
.k_input = k_input.data_ptr(),
.k_out = k_out.data_ptr(),
.weight = static_cast<const float*>(weight.data_ptr()),
.bias = static_cast<const float*>(bias.data_ptr()),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.positions = positions.data_ptr(),
.k_input_stride_batch = k_input.stride(0),
.batch_size = batch_size,
.eps = static_cast<float>(eps),
};
const auto num_blocks = div_ceil(batch_size, kFusedKIndexerNumWarps);
const auto k_int32 = kernel<int32_t>;
const auto k_int64 = kernel<int64_t>;
const auto k = pos_dtype.is_type<int32_t>() ? k_int32 : k_int64;
LaunchKernel(num_blocks, kFusedKIndexerBlockSize, device_.unwrap()) //
.enable_pdl(kUsePDL)(k, params);
}
};
// Indexer K + fused store: LayerNorm + RoPE + fp8 quant + paged store in one
// launch. Page layout matches fused_store_index_cache.cuh: each page is
// 132*page_size bytes (128*page_size fp8 keys, then 4*page_size fp32 scales).
struct FusedKIndexerNormRopeStoreParams {
const void* __restrict__ k_input; // (B, 128) DType
void* __restrict__ cache; // (num_pages, 132*page_size) uint8
const void* __restrict__ indices; // (B,) int64 -- out_cache_loc
const float* __restrict__ weight; // (128,) fp32 -- LayerNorm gamma
const float* __restrict__ bias; // (128,) fp32 -- LayerNorm beta
const float* __restrict__ freqs_cis; // (max_pos, 64) fp32
const void* __restrict__ positions; // (B,) PosT
// Row stride for `k_input` (caller passes the non-contiguous wk slice directly).
int64_t k_input_stride_batch;
uint32_t batch_size;
float eps;
};
template <typename DType, typename PosT, bool kUsePDL, int32_t kPageBits>
K_INDEXER_KERNEL void fused_k_indexer_norm_rope_store(const __grid_constant__ FusedKIndexerNormRopeStoreParams params) {
using namespace device;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kRopeDim = 64;
constexpr int64_t kVecSize = 4;
constexpr uint32_t kRopeSize = kRopeDim / kVecSize; // = 16
constexpr int64_t kPageBytes = 132ll << kPageBits;
static_assert(kHeadDim == kWarpThreads * kVecSize);
static_assert(kRopeDim == kWarpThreads * 2);
static_assert(kRopeSize <= kWarpThreads);
using Storage = AlignedVector<DType, kVecSize>;
using Float4 = AlignedVector<float, kVecSize>;
using OutStorage = AlignedVector<fp8x2_e4m3_t, 2>; // 4 fp8 / lane
const auto warp_id = threadIdx.x / kWarpThreads;
const auto lane_id = threadIdx.x % kWarpThreads;
const auto work_id = blockIdx.x * kFusedKIndexerNumWarps + warp_id;
const bool is_rope_lane = lane_id < kRopeSize;
if (work_id >= params.batch_size) return;
const auto input_ptr = static_cast<const DType*>(params.k_input) + work_id * params.k_input_stride_batch;
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[work_id]);
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
PDLWaitPrimary<kUsePDL>();
Float4 data, freq, gamma, beta;
// part 1: LayerNorm
{
Storage input_vec;
input_vec.load(input_ptr, lane_id);
gamma.load(params.weight, lane_id);
beta.load(params.bias, lane_id);
if (is_rope_lane) freq.load(freqs_cis, lane_id);
float sum = 0.0f;
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
data[i] = cast<float>(input_vec[i]);
sum += data[i];
}
const float mean = warp::reduce_sum(sum) / kHeadDim;
float var = 0.0f;
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
const float centered = data[i] - mean;
var += centered * centered;
}
const float inv_std = math::rsqrt(warp::reduce_sum(var) / kHeadDim + params.eps);
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
data[i] = (data[i] - mean) * inv_std * gamma[i] + beta[i];
}
}
// part 2: rope on rope lanes
if (is_rope_lane) {
const auto x_real = data[0];
const auto x_imag = data[1];
const auto y_real = data[2];
const auto y_imag = data[3];
const auto fxr = freq[0];
const auto fxi = freq[1];
const auto fyr = freq[2];
const auto fyi = freq[3];
data[0] = x_real * fxr - x_imag * fxi;
data[1] = x_real * fxi + x_imag * fxr;
data[2] = y_real * fyr - y_imag * fyi;
data[3] = y_real * fyi + y_imag * fyr;
}
PDLTriggerSecondary<kUsePDL>();
// part 3: fp8 act-quant + paged store. Round through bf16 first so the fp8
// scale matches the un-fused path.
#pragma unroll
for (int i = 0; i < kVecSize; ++i)
data[i] = cast<float>(cast<DType>(data[i]));
float local_max = math::abs(data[0]);
#pragma unroll
for (int i = 1; i < kVecSize; ++i)
local_max = math::max(local_max, math::abs(data[i]));
const auto abs_max = warp::reduce_max(local_max);
const auto scale = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX;
const auto inv_scale = 1.0f / scale;
const auto index = static_cast<const int64_t*>(params.indices)[work_id];
const int32_t page = static_cast<int32_t>(index >> kPageBits);
const int32_t offset = static_cast<int32_t>(index & ((1 << kPageBits) - 1));
const auto page_ptr = static_cast<uint8_t*>(params.cache) + page * kPageBytes;
const auto value_ptr = page_ptr + offset * kHeadDim;
const auto scale_ptr = page_ptr + (kHeadDim << kPageBits) + offset * 4;
OutStorage result;
result[0] = pack_fp8(data[0] * inv_scale, data[1] * inv_scale);
result[1] = pack_fp8(data[2] * inv_scale, data[3] * inv_scale);
reinterpret_cast<OutStorage*>(value_ptr)[lane_id] = result;
if (lane_id == 0) *reinterpret_cast<float*>(scale_ptr) = scale;
}
template <typename DType, bool kUsePDL, uint32_t kPageSize>
struct FusedKIndexerNormRopeStoreKernel {
static constexpr int32_t kPageBits = std::countr_zero(kPageSize);
static constexpr int64_t kPageBytes = 132ll * kPageSize;
static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2");
template <typename PosT>
static constexpr auto kernel = fused_k_indexer_norm_rope_store<DType, PosT, kUsePDL, kPageBits>;
static void forward(
const tvm::ffi::TensorView k_input,
const tvm::ffi::TensorView cache,
const tvm::ffi::TensorView indices,
const tvm::ffi::TensorView weight,
const tvm::ffi::TensorView bias,
const tvm::ffi::TensorView freqs_cis,
const tvm::ffi::TensorView positions,
double eps) {
using namespace host;
constexpr int64_t kHeadDim = 128;
constexpr int64_t kRopeDim = 64;
auto B = SymbolicSize{"batch_size"};
auto device_ = SymbolicDevice{};
device_.set_options<kDLCUDA>();
TensorMatcher({B, kHeadDim}) //
.with_strides({-1, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(k_input);
TensorMatcher({-1, -1}) //
.with_strides({kPageBytes, 1})
.with_dtype<uint8_t>()
.with_device(device_)
.verify(cache);
TensorMatcher({B}) //
.with_dtype<int64_t>()
.with_device(device_)
.verify(indices);
TensorMatcher({kHeadDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(weight);
TensorMatcher({kHeadDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(bias);
TensorMatcher({-1, kRopeDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(freqs_cis);
auto pos_dtype = SymbolicDType{};
TensorMatcher({B}) //
.with_dtype<int32_t, int64_t>(pos_dtype)
.with_device(device_)
.verify(positions);
const auto batch_size = static_cast<uint32_t>(B.unwrap());
if (batch_size == 0) return;
const auto params = FusedKIndexerNormRopeStoreParams{
.k_input = k_input.data_ptr(),
.cache = cache.data_ptr(),
.indices = indices.data_ptr(),
.weight = static_cast<const float*>(weight.data_ptr()),
.bias = static_cast<const float*>(bias.data_ptr()),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.positions = positions.data_ptr(),
.k_input_stride_batch = k_input.stride(0),
.batch_size = batch_size,
.eps = static_cast<float>(eps),
};
const auto num_blocks = div_ceil(batch_size, kFusedKIndexerNumWarps);
const auto k_int32 = kernel<int32_t>;
const auto k_int64 = kernel<int64_t>;
const auto k = pos_dtype.is_type<int32_t>() ? k_int32 : k_int64;
LaunchKernel(num_blocks, kFusedKIndexerBlockSize, device_.unwrap()) //
.enable_pdl(kUsePDL)(k, params);
}
};
} // namespace
python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh
+18 -10
View File
@@ -425,11 +425,13 @@ struct FusedQIndexerRopeHadamardQuantParams {
float weight_scale; // scalar c4_indexer.weight_scale
const float* __restrict__ freqs_cis; // (max_pos, 64) fp32
const void* __restrict__ positions; // (B,) PosT
// Row stride for `weight` (caller passes the non-contiguous wk slice directly).
int64_t weight_stride_batch;
uint32_t batch_size;
uint32_t num_heads;
};
template <typename DType, typename PosT, bool kUsePDL>
template <typename DType, typename PosT, bool kUsePDL, bool kRopeFirst = false, bool kHadamard = true>
Q_KERNEL void fused_q_indexer_rope_hadamard_quant(const __grid_constant__ FusedQIndexerRopeHadamardQuantParams params) {
using namespace device;
@@ -448,9 +450,9 @@ Q_KERNEL void fused_q_indexer_rope_hadamard_quant(const __grid_constant__ FusedQ
const auto warp_id = threadIdx.x / kWarpThreads;
const auto lane_id = threadIdx.x % kWarpThreads;
const auto work_id = blockIdx.x * kFusedQNumWarps + warp_id;
// Last `kRopeSize` lanes own the rope tail; their 4-elem packs cover the
// trailing kRopeDim elements.
const bool is_rope_lane = lane_id >= kWarpThreads - kRopeSize;
// V4 ropes the trailing kRopeDim dims (kRopeFirst=false); V3.2 ropes the
// leading kRopeDim dims (kRopeFirst=true). Select the owning lanes per layout.
const bool is_rope_lane = kRopeFirst ? (lane_id < kRopeSize) : (lane_id >= kWarpThreads - kRopeSize);
const uint32_t total_works = params.batch_size * params.num_heads;
if (work_id >= total_works) return;
@@ -467,13 +469,15 @@ Q_KERNEL void fused_q_indexer_rope_hadamard_quant(const __grid_constant__ FusedQ
PDLWaitPrimary<kUsePDL>();
Float4 data, freq;
const auto weight_val = cast<float>(static_cast<const DType*>(params.weight)[work_id]);
const uint32_t head_id = work_id - batch_id * params.num_heads;
const auto weight_val =
cast<float>(static_cast<const DType*>(params.weight)[batch_id * params.weight_stride_batch + head_id]);
// part 1: load (no norm). Each lane owns a 4-elem pack.
{
Storage input_vec;
input_vec.load(input_ptr, lane_id);
if (is_rope_lane) freq.load(freqs_cis, lane_id - (kWarpThreads - kRopeSize));
if (is_rope_lane) freq.load(freqs_cis, kRopeFirst ? lane_id : (lane_id - (kWarpThreads - kRopeSize)));
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
data[i] = cast<float>(input_vec[i]);
@@ -500,8 +504,10 @@ Q_KERNEL void fused_q_indexer_rope_hadamard_quant(const __grid_constant__ FusedQ
// part 3: 128-point Hadamard (2 local stages + 5 cross-lane shfl_xor stages).
// Same recipe as `fused_norm_rope_indexer`; see comments there for the
// butterfly invariants and the early-return safety argument.
{
// butterfly invariants and the early-return safety argument. V3.2 omits the
// rotation (kHadamard=false): it is logit-preserving (H orthonormal, applied
// to both q and k), so dropping it only trades fp8 quant accuracy.
if constexpr (kHadamard) {
{
const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
data[0] = a0 + a1;
@@ -550,10 +556,10 @@ Q_KERNEL void fused_q_indexer_rope_hadamard_quant(const __grid_constant__ FusedQ
}
}
template <typename DType, bool kUsePDL>
template <typename DType, bool kUsePDL, bool kRopeFirst = false, bool kHadamard = true>
struct FusedQIndexerRopeHadamardQuantKernel {
template <typename PosT>
static constexpr auto kernel = fused_q_indexer_rope_hadamard_quant<DType, PosT, kUsePDL>;
static constexpr auto kernel = fused_q_indexer_rope_hadamard_quant<DType, PosT, kUsePDL, kRopeFirst, kHadamard>;
static void forward(
const tvm::ffi::TensorView q_input,
@@ -586,6 +592,7 @@ struct FusedQIndexerRopeHadamardQuantKernel {
.with_device(device_)
.verify(q_fp8);
TensorMatcher({B, H}) //
.with_strides({-1, 1})
.with_dtype<DType>()
.with_device(device_)
.verify(weight);
@@ -627,6 +634,7 @@ struct FusedQIndexerRopeHadamardQuantKernel {
.weight_scale = static_cast<float>(weight_scale),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.positions = positions.data_ptr(),
.weight_stride_batch = weight.stride(0),
.batch_size = batch_size,
.num_heads = num_heads,
};
python/sglang/jit_kernel/dsv32/__init__.py
+11
View File
@@ -0,0 +1,11 @@
"""DSA only."""
from .elementwise import (
fused_k_indexer_norm_rope,
fused_k_indexer_norm_rope_store,
)
__all__ = [
"fused_k_indexer_norm_rope",
"fused_k_indexer_norm_rope_store",
]
python/sglang/jit_kernel/dsv32/elementwise.py
+92
View File
@@ -0,0 +1,92 @@
"""DSA only. Indexer K kernels (JIT)."""
import torch
from sglang.jit_kernel.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
_CUDA_FILE = "deepseek_v32/indexer_k.cuh"
@cache_once
def _jit_k_indexer_norm_rope_module(dtype: torch.dtype):
args = make_cpp_args(dtype, is_arch_support_pdl())
return load_jit(
"dpsk_v32_k_indexer_norm_rope",
*args,
cuda_files=[_CUDA_FILE],
cuda_wrappers=[
("forward", f"FusedKIndexerNormRopeKernel<{args}>::forward"),
],
)
@cache_once
def _jit_k_indexer_norm_rope_store_module(dtype: torch.dtype, page_size: int):
args = make_cpp_args(dtype, is_arch_support_pdl(), page_size)
return load_jit(
f"dpsk_v32_k_indexer_norm_rope_store_p{page_size}",
*args,
cuda_files=[_CUDA_FILE],
cuda_wrappers=[
("forward", f"FusedKIndexerNormRopeStoreKernel<{args}>::forward"),
],
)
def fused_k_indexer_norm_rope(
k_input: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor,
eps: float,
freqs_cis: torch.Tensor,
positions: torch.Tensor,
) -> torch.Tensor:
"""V3.2 indexer K: LayerNorm + RoPE on leading dims -> bf16. CUDA only."""
freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
# k_input may be a non-contiguous wk slice; output is always contiguous.
k_out = torch.empty(k_input.shape, dtype=k_input.dtype, device=k_input.device)
module = _jit_k_indexer_norm_rope_module(k_input.dtype)
module.forward(
k_input,
k_out,
weight,
bias,
freqs_real,
positions,
float(eps),
)
return k_out
def fused_k_indexer_norm_rope_store(
k_input: torch.Tensor,
cache: torch.Tensor,
out_cache_loc: torch.Tensor,
weight: torch.Tensor,
bias: torch.Tensor,
eps: float,
freqs_cis: torch.Tensor,
positions: torch.Tensor,
page_size: int,
) -> None:
"""V3.2 indexer K + fused store: LayerNorm + RoPE on leading dims + fp8
act-quant + paged index-k cache write, in one launch. CUDA only."""
freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
if not out_cache_loc.is_contiguous():
out_cache_loc = out_cache_loc.contiguous()
module = _jit_k_indexer_norm_rope_store_module(k_input.dtype, page_size)
module.forward(
k_input,
cache,
out_cache_loc,
weight,
bias,
freqs_real,
positions,
float(eps),
)
python/sglang/jit_kernel/dsv4/__init__.py
+2
View File
@@ -12,6 +12,7 @@ from .compress import (
from .compress_old import fused_norm_rope_inplace
from .elementwise import (
fused_k_norm_rope_flashmla,
fused_q_indexer_rope_first_quant,
fused_q_indexer_rope_hadamard_fp4_quant,
fused_q_indexer_rope_hadamard_quant,
fused_q_norm_rope,
@@ -38,6 +39,7 @@ __all__ = [
"fused_store_cache",
"fused_rope_inplace",
"fused_q_norm_rope",
"fused_q_indexer_rope_first_quant",
"fused_q_indexer_rope_hadamard_fp4_quant",
"fused_q_indexer_rope_hadamard_quant",
"fused_k_norm_rope_flashmla",
python/sglang/jit_kernel/dsv4/elementwise.py
+41
View File
@@ -77,6 +77,21 @@ def _jit_main_q_indexer_rope_hadamard_quant_module(dtype: torch.dtype):
)
# V3.2 lays q out as [rope | nope] (V4 is [nope | rope]) -> kRopeFirst=true, and
# drops the Hadamard rotation (kHadamard=false).
@cache_once
def _jit_main_q_indexer_rope_first_quant_module(dtype: torch.dtype):
args = make_cpp_args(dtype, is_arch_support_pdl(), True, False)
return load_jit(
make_name("main_q_indexer_rope_first_quant"),
*args,
cuda_files=["deepseek_v4/main_norm_rope.cuh"],
cuda_wrappers=[
("forward", f"FusedQIndexerRopeHadamardQuantKernel<{args}>::forward"),
],
)
@cache_once
def _jit_main_q_indexer_rope_hadamard_fp4_quant_module(dtype: torch.dtype):
args = make_cpp_args(dtype, is_arch_support_pdl())
@@ -169,6 +184,32 @@ def fused_q_indexer_rope_hadamard_quant(
return q_fp8, weights_out
def fused_q_indexer_rope_first_quant(
q_input: torch.Tensor,
weight: torch.Tensor,
weight_scale: float,
freqs_cis: torch.Tensor,
positions: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""DeepSeek-V3.2 only. Indexer Q: RoPE on the leading dims + fp8 act-quant. CUDA only."""
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_first_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_q_indexer_rope_hadamard_fp4_quant(
q_input: torch.Tensor,
weight: torch.Tensor,
python/sglang/srt/environ.py
+1
View File
@@ -649,6 +649,7 @@ class Envs:
SGLANG_ENABLE_PCG_DSV2_DUAL_STREAM = EnvBool(False)
SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True)
SGLANG_DSA_TOPK_BROADCAST = EnvBool(False)
SGLANG_DISABLE_DSA_INDEXER_FUSION = EnvBool(False)
# sgl-kernel
SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False)
python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+294 -20
View File
@@ -53,6 +53,7 @@ logger = logging.getLogger(__name__)
global _use_multi_stream
_is_cuda = is_cuda()
_use_dsa_indexer_fusion = _is_cuda and not envs.SGLANG_DISABLE_DSA_INDEXER_FUSION.get()
_is_hip = is_hip()
_is_npu = is_npu()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
@@ -147,6 +148,35 @@ def _uses_dsa_attention_backend(forward_batch: ForwardBatch) -> bool:
if _is_cuda:
from sglang.jit_kernel.dsv4 import fused_q_indexer_rope_first_quant
from sglang.jit_kernel.dsv32 import (
fused_k_indexer_norm_rope,
fused_k_indexer_norm_rope_store,
)
def _scale_head_gate_graph_fake_impl(
weights_raw: torch.Tensor,
n_heads_inv_sqrt: float,
softmax_scale: float,
q_scale: torch.Tensor,
) -> torch.Tensor:
return torch.empty(
(weights_raw.shape[0], weights_raw.shape[1], q_scale.shape[-1]),
dtype=torch.float32,
device=weights_raw.device,
)
# In-graph (PCG/BCG) head gate for the fused path: weights_proj is folded
# into wk_weights_proj, so weights_raw is precomputed and there is no GEMM.
@register_custom_op(fake_impl=_scale_head_gate_graph_fake_impl)
def scale_head_gate_graph(
weights_raw: torch.Tensor,
n_heads_inv_sqrt: float,
softmax_scale: float,
q_scale: torch.Tensor,
) -> torch.Tensor:
weights = weights_raw * n_heads_inv_sqrt
return weights.unsqueeze(-1) * q_scale * softmax_scale
def _logits_head_gate_graph_fake_impl(
x: torch.Tensor,
@@ -358,6 +388,15 @@ class Indexer(MultiPlatformOp):
prefix=add_prefix("wq_b", prefix),
)
if _use_dsa_indexer_fusion:
self.wk_weights_proj = ReplicatedLinear(
self.hidden_size,
self.head_dim + self.n_heads,
bias=False,
params_dtype=torch.bfloat16,
prefix=add_prefix("wk_weights_proj", prefix),
)
else:
self.wk = ReplicatedLinear(
self.hidden_size,
self.head_dim,
@@ -388,6 +427,15 @@ class Indexer(MultiPlatformOp):
self.scale_fmt = scale_fmt
self.softmax_scale = self.head_dim**-0.5
# freqs_cis is built from the fp32 cos/sin cache before any forward casts it to bf16.
self._indexer_freqs_cis: Optional[torch.Tensor] = None
if _use_dsa_indexer_fusion:
c = self.rotary_emb.cos_sin_cache.to(torch.float32)
half = c.shape[-1] // 2
self._indexer_freqs_cis = torch.complex(
c[:, :half].contiguous(), c[:, half:].contiguous()
)
@contextlib.contextmanager
def _with_real_sm_count(self):
# When pipeline parallelism is enabled, each PP rank initiates a recv operation after the _pp_launch_batch
@@ -443,6 +491,20 @@ class Indexer(MultiPlatformOp):
):
return weights.unsqueeze(-1) * q_scale * self.softmax_scale
@torch.compile(dynamic=True)
def _scale_head_gates(self, weights_raw: torch.Tensor, q_scale: torch.Tensor):
weights = weights_raw * self.n_heads**-0.5
return weights.unsqueeze(-1) * q_scale * self.softmax_scale
def _fused_k_weights(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
kw, _ = self.wk_weights_proj(x)
return kw.split([self.head_dim, self.n_heads], dim=-1)
def _maybe_rotate(self, x: torch.Tensor) -> torch.Tensor:
# Fusion drops the (logit-preserving) Hadamard rotation; without it the
# index-K cache here matches the fused path that decode reads back.
return x if _use_dsa_indexer_fusion else rotate_activation(x)
def _should_skip_logits_computation(self, forward_batch: ForwardBatch) -> bool:
if (
forward_batch.forward_mode.is_extend_without_speculative()
@@ -460,6 +522,7 @@ class Indexer(MultiPlatformOp):
enable_dual_stream: bool,
forward_batch: ForwardBatch,
):
weights_raw = None
if enable_dual_stream:
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
@@ -476,6 +539,9 @@ class Indexer(MultiPlatformOp):
)
with torch.cuda.stream(self.alt_stream):
# TODO we should also put DeepGEMM half SM here?
if _use_dsa_indexer_fusion:
key, weights_raw = self._fused_k_weights(x)
else:
key, _ = self.wk(x)
key = self.k_norm(key)
@@ -492,6 +558,9 @@ class Indexer(MultiPlatformOp):
q_rope, _ = torch.split(
query, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1
)
if _use_dsa_indexer_fusion:
key, weights_raw = self._fused_k_weights(x)
else:
key, _ = self.wk(x)
key = self.k_norm(key)
k_rope, _ = torch.split(
@@ -506,20 +575,20 @@ class Indexer(MultiPlatformOp):
if enable_dual_stream:
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
query = rotate_activation(query)
query = self._maybe_rotate(query)
with torch.cuda.stream(self.alt_stream):
key = rotate_activation(key)
key = self._maybe_rotate(key)
current_stream.wait_stream(self.alt_stream)
elif (
self.alt_stream is not None
and forward_batch.attn_cp_metadata is not None
and self.dsa_enable_prefill_cp
):
key = rotate_activation(key)
key = self._maybe_rotate(key)
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
query = rotate_activation(query)
query = self._maybe_rotate(query)
with torch.cuda.stream(self.alt_stream):
key = cp_all_gather_rerange_output(
@@ -529,10 +598,10 @@ class Indexer(MultiPlatformOp):
torch.cuda.current_stream(),
)
current_stream.wait_stream(self.alt_stream)
return query, key
return query, key, weights_raw
else:
query = rotate_activation(query)
key = rotate_activation(key)
query = self._maybe_rotate(query)
key = self._maybe_rotate(key)
# allgather+rerrange
if forward_batch.attn_cp_metadata is not None and self.dsa_enable_prefill_cp:
@@ -542,7 +611,7 @@ class Indexer(MultiPlatformOp):
forward_batch,
torch.cuda.current_stream(),
)
return query, key
return query, key, weights_raw
def _get_k_bf16(
self,
@@ -550,7 +619,7 @@ class Indexer(MultiPlatformOp):
positions: torch.Tensor,
enable_dual_stream: bool,
):
# Compute only key, skip query
# Non-fusion path only; self.wk does not exist when fusion is on.
key, _ = self.wk(x)
key = self.k_norm(key)
k_rope, _ = torch.split(
@@ -563,6 +632,139 @@ class Indexer(MultiPlatformOp):
return key
def _fused_k_prepare_and_store(
self,
key_raw: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
layer_id: int,
act_quant,
out_cache_loc: Optional[torch.Tensor] = None,
) -> None:
if out_cache_loc is None:
out_cache_loc = forward_batch.out_cache_loc
pool = get_token_to_kv_pool()
page_size = pool.page_size
if (
not _is_fp8_fnuz
and out_cache_loc is not None
and can_use_dsa_fused_store(torch.bfloat16, out_cache_loc.dtype, page_size)
):
fused_k_indexer_norm_rope_store(
key_raw,
pool.get_index_k_with_scale_buffer(layer_id=layer_id),
out_cache_loc,
self.k_norm.weight,
self.k_norm.bias,
self.k_norm.variance_epsilon,
self._indexer_freqs_cis,
positions,
page_size,
)
return
# Fallback: separate K kernel + store kernel.
key = fused_k_indexer_norm_rope(
key_raw,
self.k_norm.weight,
self.k_norm.bias,
self.k_norm.variance_epsilon,
self._indexer_freqs_cis,
positions,
)
self._store_index_k_cache(
forward_batch=forward_batch,
layer_id=layer_id,
key=key,
act_quant=act_quant,
out_cache_loc=out_cache_loc,
)
def _fused_q_prepare_and_store(
self,
x: torch.Tensor,
q_lora: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
layer_id: int,
act_quant,
*,
num_tokens: Optional[int] = None,
enable_dual_stream: bool = True,
) -> Tuple[torch.Tensor, torch.Tensor]:
# num_tokens (graph split-op contract) slices q/k/positions/out_cache_loc
# to the unpadded count; the returned q_fp8/weights are sliced to match.
q_scale_gate = self.softmax_scale * self.n_heads**-0.5
out_cache_loc = forward_batch.out_cache_loc
if num_tokens is not None:
positions = positions[:num_tokens]
out_cache_loc = out_cache_loc[:num_tokens]
if self.alt_stream is None or not enable_dual_stream:
kw, _ = self.wk_weights_proj(x)
key, weights_raw = kw.split([self.head_dim, self.n_heads], dim=-1)
if num_tokens is not None:
key = key[:num_tokens]
weights_raw = weights_raw[:num_tokens]
self._fused_k_prepare_and_store(
key,
positions,
forward_batch,
layer_id,
act_quant,
out_cache_loc=out_cache_loc,
)
q = self.wq_b(q_lora)[0].view(-1, self.n_heads, self.head_dim)
if num_tokens is not None:
q = q[:num_tokens]
return fused_q_indexer_rope_first_quant(
q.contiguous(),
weights_raw,
q_scale_gate,
self._indexer_freqs_cis,
positions,
)
# Two overlap stages: wq_b GEMM (alt) || wk_weights_proj GEMM (current),
# then fused Q kernel (current) || fused K kernel + cache store (alt).
# wait_stream calls are ordered by issue position so each side waits only
# on the GEMMs it consumes, not on the other side's fused kernel.
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
with torch.cuda.stream(self.alt_stream):
q = self.wq_b(q_lora)[0].view(-1, self.n_heads, self.head_dim)
if num_tokens is not None:
q = q[:num_tokens]
kw, _ = self.wk_weights_proj(x)
key, weights_raw = kw.split([self.head_dim, self.n_heads], dim=-1)
if num_tokens is not None:
key = key[:num_tokens]
weights_raw = weights_raw[:num_tokens]
current_stream.wait_stream(self.alt_stream)
self.alt_stream.wait_stream(current_stream)
with torch.cuda.stream(self.alt_stream):
self._fused_k_prepare_and_store(
key,
positions,
forward_batch,
layer_id,
act_quant,
out_cache_loc=out_cache_loc,
)
q_fp8, weights = fused_q_indexer_rope_first_quant(
q.contiguous(),
weights_raw,
q_scale_gate,
self._indexer_freqs_cis,
positions,
)
current_stream.wait_stream(self.alt_stream)
return q_fp8, weights
@staticmethod
def _update_rope_guarded(dst: torch.Tensor, src: torch.Tensor) -> None:
# On AMD with in-place RoPE kernels, self-aliasing can occur;
@@ -997,16 +1199,36 @@ class Indexer(MultiPlatformOp):
assert forward_batch.forward_mode.is_extend_without_speculative()
x_meta = x[0] if isinstance(x, tuple) else x
# Fast path: only compute and store k cache, skip all q and weights ops
key = self._get_k_bf16(x, positions, enable_dual_stream)
# Fast path: only compute and store k cache, skip all q and weights ops.
# num_tokens (graph contract) slices to the unpadded count.
out_cache_loc = None
if num_tokens is not None:
assert num_tokens <= key.shape[0]
assert num_tokens <= forward_batch.out_cache_loc.shape[0]
key = key[:num_tokens]
out_cache_loc = forward_batch.out_cache_loc[:num_tokens]
elif not forward_batch.out_cache_loc.is_contiguous():
forward_batch.out_cache_loc = forward_batch.out_cache_loc.contiguous()
# Write the same K representation the decode path reads back: fused
# (no-Hadamard) when fusion is on, else the legacy Hadamard path.
if _use_dsa_indexer_fusion:
key_raw, _ = self._fused_k_weights(x)
if num_tokens is not None:
assert num_tokens <= key_raw.shape[0]
key_raw = key_raw[:num_tokens]
positions = positions[:num_tokens]
self._fused_k_prepare_and_store(
key_raw,
positions,
forward_batch,
layer_id,
act_quant,
out_cache_loc=out_cache_loc,
)
else:
key = self._get_k_bf16(x, positions, enable_dual_stream)
if num_tokens is not None:
assert num_tokens <= key.shape[0]
key = key[:num_tokens]
self._store_index_k_cache(
forward_batch=forward_batch,
layer_id=layer_id,
@@ -1436,16 +1658,29 @@ class Indexer(MultiPlatformOp):
return maybe_capture_indexer_topk(layer_id, topk_result)
# When weights_proj is LoRA-wrapped, use an eager module call so the
# wrapper owns base+delta and no LoRA kernel runs under torch.compile
weights_proj_lora = getattr(self.weights_proj, "set_lora", False)
# wrapper owns base+delta and no LoRA kernel runs under torch.compile.
# Fusion folds weights_proj into wk_weights_proj, so weights_proj is
# absent then; short-circuit before touching it.
weights_proj_lora = not _use_dsa_indexer_fusion and getattr(
self.weights_proj, "set_lora", False
)
if (
_use_dsa_indexer_fusion
and not in_piecewise_or_breakable_cuda_graph
and forward_batch.attn_cp_metadata is None
):
q_fp8, weights = self._fused_q_prepare_and_store(
x, q_lora, positions, forward_batch, layer_id, act_quant
)
elif (
is_graph_dsa_split_op_surface(forward_batch)
and not self.dsa_enable_prefill_cp
):
# Default path for non-CP prefill under PCG/BCG: run the whole indexer
# (q/k proj, head gate, k-cache store, topk) as a single eager split op
# instead of capturing it piecemeal in the graph.
# instead of capturing it piecemeal in the graph. The split op is
# fusion-aware, so this also covers the fused path here.
if weights_proj_lora:
raise RuntimeError(GRAPH_WEIGHTS_PROJ_LORA_ERROR)
if return_indices:
@@ -1476,14 +1711,15 @@ class Indexer(MultiPlatformOp):
)
return maybe_capture_indexer_topk(layer_id, result)
if enable_dual_stream and forward_batch.forward_mode.is_decode_or_idle():
elif enable_dual_stream and forward_batch.forward_mode.is_decode_or_idle():
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
if not _use_dsa_indexer_fusion:
if weights_proj_lora:
weights = self.weights_proj(x)[0].float() * self.n_heads**-0.5
else:
weights = self._project_and_scale_head_gates(x)
query, key = self._get_q_k_bf16(
query, key, weights_raw = self._get_q_k_bf16(
q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
)
q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
@@ -1495,9 +1731,12 @@ class Indexer(MultiPlatformOp):
act_quant=act_quant,
)
current_stream.wait_stream(self.alt_stream)
if _use_dsa_indexer_fusion:
weights = self._scale_head_gates(weights_raw, q_scale)
else:
weights = self._apply_q_scale_and_softmax_scale(weights, q_scale)
else:
query, key = self._get_q_k_bf16(
query, key, weights_raw = self._get_q_k_bf16(
q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
)
@@ -1571,6 +1810,14 @@ class Indexer(MultiPlatformOp):
x_for_gate = x
if in_piecewise_or_breakable_cuda_graph:
if _use_dsa_indexer_fusion:
weights = scale_head_gate_graph(
weights_raw,
self.n_heads**-0.5,
self.softmax_scale,
q_scale,
)
else:
if weights_proj_lora:
raise RuntimeError(GRAPH_WEIGHTS_PROJ_LORA_ERROR)
weights = logits_head_gate_graph(
@@ -1580,6 +1827,8 @@ class Indexer(MultiPlatformOp):
self.softmax_scale,
q_scale,
)
elif _use_dsa_indexer_fusion:
weights = self._scale_head_gates(weights_raw, q_scale)
elif weights_proj_lora:
weights = self.weights_proj(x_for_gate)[0].float() * self.n_heads**-0.5
weights = self._apply_q_scale_and_softmax_scale(weights, q_scale)
@@ -2067,7 +2316,32 @@ def pcg_dsa_indexer_prefill_split(
)
return
query, key = indexer._get_q_k_bf16(
# Fused path stores K (no-Hadamard) and computes q_fp8 + head gate in the
# fused kernels, sliced to the unpadded count. Single stream: the split op is
# captured, so the dual-stream overlap is disabled.
if _use_dsa_indexer_fusion:
q_fp8, weights = indexer._fused_q_prepare_and_store(
x,
q_lora,
positions,
forward_batch,
layer_id,
act_quant,
num_tokens=extend_num_tokens,
enable_dual_stream=False,
)
indexer._get_topk_ragged(
False,
forward_batch,
layer_id,
q_fp8,
weights,
metadata,
topk_result,
)
return
query, key, _ = indexer._get_q_k_bf16(
q_lora,
x,
positions,
python/sglang/srt/lora/lora_manager.py
+16
View File
@@ -36,6 +36,7 @@ from sglang.srt.lora.lora_config import LoRAConfig
from sglang.srt.lora.lora_registry import LoRARef
from sglang.srt.lora.mem_pool import LoRAMemoryPool
from sglang.srt.lora.utils import (
DSA_INDEXER_LORA_NAMES,
EMBEDDING_NAMES,
LoRAType,
auto_detect_lora_target_modules,
@@ -605,6 +606,21 @@ class LoRAManager:
# Otherwise, infer target_modules from adapter configs.
self.target_modules.update(adapter_target_modules)
# Fusion folds wk + weights_proj into wk_weights_proj, so the modules
# LoRA wraps are absent and an indexer-targeted adapter is silently dropped.
indexer_targets = self.target_modules & DSA_INDEXER_LORA_NAMES
if indexer_targets:
from sglang.srt.layers.attention.dsa.dsa_indexer import (
_use_dsa_indexer_fusion,
)
if _use_dsa_indexer_fusion:
raise ValueError(
f"LoRA targets the DSA indexer ({sorted(indexer_targets)}), which is "
"incompatible with DSA indexer Q/K fusion. Set "
"SGLANG_DISABLE_DSA_INDEXER_FUSION=1 to disable fusion and use indexer LoRA."
)
if max_lora_rank is not None:
self.max_lora_rank = max_lora_rank
else:
python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py
+61
View File
@@ -76,6 +76,52 @@ def _clone_if_runai_streamed_tensor(tensor: torch.Tensor) -> torch.Tensor:
return tensor
def _load_fused_indexer_wk(
name: str,
loaded_weight: torch.Tensor,
params_dict: Dict[str, torch.Tensor],
pending: Dict[str, Dict[str, torch.Tensor]],
quant_config: Optional[QuantizationConfig],
) -> bool:
"""Load an indexer wk / weights_proj shard into the fused bf16 wk_weights_proj
param: wk fills the top head_dim rows (dequantized from block-fp8 if needed),
weights_proj the bottom n_heads rows.
Returns False when there is no fused param (non-CUDA, or CUDA with
SGLANG_DISABLE_DSA_INDEXER_FUSION set, where wk and weights_proj are
separate) so the caller falls through to per-tensor loading.
"""
fused_name = name.rsplit(".indexer.", 1)[0] + ".indexer.wk_weights_proj.weight"
fused_param = params_dict.get(fused_name)
if fused_param is None or fused_param.dtype != torch.bfloat16:
return False
if ".indexer.weights_proj." in name:
w = _clone_if_runai_streamed_tensor(loaded_weight)
fused_param.data[-w.shape[0] :].copy_(w)
return True
# wk: a bf16 checkpoint copies straight in; block-fp8 needs weight + scale.
is_scale = name.endswith(".weight_scale_inv")
if not is_scale and loaded_weight.dtype != torch.float8_e4m3fn:
w = _clone_if_runai_streamed_tensor(loaded_weight)
fused_param.data[: w.shape[0]].copy_(w)
return True
entry = pending.setdefault(fused_name, {})
entry["scale" if is_scale else "weight"] = _clone_if_runai_streamed_tensor(
loaded_weight
)
if "weight" in entry and "scale" in entry:
pending.pop(fused_name)
block_size = getattr(quant_config, "weight_block_size", None) or [128, 128]
wk_bf16 = block_quant_dequant(
entry["weight"], entry["scale"], block_size, torch.bfloat16
)
fused_param.data[: wk_bf16.shape[0]].copy_(wk_bf16)
return True
@dataclass(frozen=True)
class NextNEnabledConfig:
num_nextn_layers: int
@@ -147,6 +193,8 @@ class DeepseekV2WeightLoaderMixin:
)
cached_a_proj = {} if fuse_qkv_a_proj else None
pending_indexer_wk: Dict[str, Dict[str, torch.Tensor]] = {}
if self.num_fused_shared_experts > 0:
assert self.num_fused_shared_experts == 1
log_info_on_rank0(logger, "Shared experts fusion optimization enabled.")
@@ -207,6 +255,19 @@ class DeepseekV2WeightLoaderMixin:
if "rotary_emb.inv_freq" in name:
continue
# CUDA fuses wk + weights_proj into one bf16 wk_weights_proj; the
# helper returns True once it has consumed the shard.
if (
".indexer.wk." in name or ".indexer.weights_proj." in name
) and _load_fused_indexer_wk(
name,
loaded_weight,
params_dict,
pending_indexer_wk,
self.quant_config,
):
continue
for param_name, weight_name, shard_id in stacked_params_mapping:
# Skip non-stacked layers and experts (experts handled below).
if weight_name not in name:
test/registered/8-gpu-models/test_deepseek_v32_indexcache.py
+1 -1
View File
@@ -63,7 +63,7 @@ class TestDeepseekV32IndexTopkPattern(CustomTestCase):
write_github_step_summary(
f"### test_gsm8k (deepseek-v32)\n" f'{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.935)
self.assertGreater(metrics["accuracy"], 0.93)
class TestDeepseekV32IndexFreq(CustomTestCase):
test/registered/jit/test_dsv32_indexer_fusion.py
+206
View File
@@ -0,0 +1,206 @@
"""Correctness tests for the DeepSeek-V3.2 DSA indexer fused kernels.
Covers:
- fused_q_indexer_rope_first_quant (Q: rope-first + fp8 quant + head-gate fold)
- fused_k_indexer_norm_rope (K: LayerNorm + rope-first -> bf16)
- fused_k_indexer_norm_rope_store (K: the above + fp8 quant + paged index-k cache write)
The store kernel is checked for byte-exact equivalence against the un-fused path
(bf16 K kernel + standalone fused_store_index_k_cache), so it needs no fp8
reference. The Q/K math kernels are checked against torch references. Strided
inputs (the non-contiguous wk_weights_proj slices) are checked to match
contiguous inputs (the no-copy path).
"""
from __future__ import annotations
import pytest
import torch
from sglang.jit_kernel.dsv4 import fused_q_indexer_rope_first_quant
from sglang.jit_kernel.dsv32 import (
fused_k_indexer_norm_rope,
fused_k_indexer_norm_rope_store,
)
from sglang.jit_kernel.fused_store_index_cache import (
can_use_dsa_fused_store,
fused_store_index_k_cache,
)
from sglang.srt.utils import is_hip
from sglang.test.ci.ci_register import register_cuda_ci
_is_hip = is_hip()
register_cuda_ci(est_time=45, suite="base-b-kernel-unit-1-gpu-large")
register_cuda_ci(est_time=90, suite="nightly-kernel-1-gpu", nightly=True)
HEAD_DIM = 128
ROPE_DIM = 64
HALF = ROPE_DIM // 2
FP8_MAX = 448.0
PAGE_SIZE = 64
BYTES_PER_TOKEN = HEAD_DIM + 4 # 128 fp8 + 4-byte fp32 scale
EPS = 1e-6
MAX_POS = 8192
def _skip_if_unavailable():
if not torch.cuda.is_available():
pytest.skip("CUDA required")
if _is_hip:
pytest.skip("Indexer fused kernels are CUDA-specific")
def _make_inputs(B, seed=0, pos_dtype=torch.int32):
g = torch.Generator(device="cuda").manual_seed(seed)
dev = "cuda"
cos = torch.randn(MAX_POS, HALF, device=dev, generator=g)
sin = torch.randn(MAX_POS, HALF, device=dev, generator=g)
freqs_cis = torch.complex(cos, sin)
positions = torch.randint(0, 4096, (B,), device=dev, dtype=pos_dtype, generator=g)
return cos, sin, freqs_cis, positions
def _rope_first(x, cos_p, sin_p):
"""Interleaved complex rope on the leading ROPE_DIM dims (kRopeFirst)."""
x = x.clone()
xr = x[..., 0:ROPE_DIM:2].clone()
xi = x[..., 1:ROPE_DIM:2].clone()
x[..., 0:ROPE_DIM:2] = xr * cos_p - xi * sin_p
x[..., 1:ROPE_DIM:2] = xr * sin_p + xi * cos_p
return x
# ----------------------------------------------------------------------------
# K kernel (-> bf16): LayerNorm + rope-first
# ----------------------------------------------------------------------------
def test_k_norm_rope_matches_reference():
_skip_if_unavailable()
dev = "cuda"
B = 37
cos, sin, freqs_cis, positions = _make_inputs(B)
key = torch.randn(B, HEAD_DIM, dtype=torch.bfloat16, device=dev)
weight = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev)
bias = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev)
out = fused_k_indexer_norm_rope(key, weight, bias, EPS, freqs_cis, positions)
torch.cuda.synchronize()
normed = torch.nn.functional.layer_norm(
key.float(), (HEAD_DIM,), weight=weight, bias=bias, eps=EPS
)
cp, sp = cos[positions.long()], sin[positions.long()]
ref = _rope_first(normed, cp, sp)
torch.testing.assert_close(out.float(), ref, atol=0.06, rtol=0.0)
# ----------------------------------------------------------------------------
# K kernel + fused store == bf16 K kernel + standalone store (byte-exact).
# Also covers the strided (non-contiguous wk slice) no-copy input path.
# ----------------------------------------------------------------------------
@pytest.mark.parametrize("strided", [False, True])
def test_k_store_matches_unfused(strided):
_skip_if_unavailable()
if not can_use_dsa_fused_store(torch.bfloat16, torch.int64, PAGE_SIZE):
pytest.skip("fused store JIT unavailable")
dev = "cuda"
B, n_heads = 41, 64
cos, sin, freqs_cis, positions = _make_inputs(B)
weight = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev)
bias = torch.randn(HEAD_DIM, dtype=torch.float32, device=dev)
if strided:
# Mimic the fused wk_weights_proj GEMM output: key is kw[:, :head_dim].
kw = torch.randn(B, HEAD_DIM + n_heads, dtype=torch.bfloat16, device=dev)
key = kw[:, :HEAD_DIM]
assert not key.is_contiguous()
else:
key = torch.randn(B, HEAD_DIM, dtype=torch.bfloat16, device=dev)
loc = torch.randperm(B * 4, device=dev)[:B].to(torch.int64)
num_pages = int(loc.max().item()) // PAGE_SIZE + 2
buf_ref = torch.zeros(
num_pages, BYTES_PER_TOKEN * PAGE_SIZE, dtype=torch.uint8, device=dev
)
buf_fused = torch.zeros_like(buf_ref)
key_bf16 = fused_k_indexer_norm_rope(key, weight, bias, EPS, freqs_cis, positions)
fused_store_index_k_cache(key_bf16, buf_ref, loc, PAGE_SIZE)
fused_k_indexer_norm_rope_store(
key, buf_fused, loc, weight, bias, EPS, freqs_cis, positions, PAGE_SIZE
)
torch.cuda.synchronize()
assert torch.equal(buf_ref, buf_fused)
# ----------------------------------------------------------------------------
# Q kernel: rope-first + fp8 quant + head-gate fold
# ----------------------------------------------------------------------------
@pytest.mark.parametrize("pos_dtype", [torch.int32, torch.int64])
def test_q_rope_quant_matches_reference(pos_dtype):
_skip_if_unavailable()
dev = "cuda"
B, n_heads = 37, 64
cos, sin, freqs_cis, positions = _make_inputs(B, pos_dtype=pos_dtype)
q = torch.randn(B, n_heads, HEAD_DIM, dtype=torch.bfloat16, device=dev)
weight = torch.randn(B, n_heads, dtype=torch.bfloat16, device=dev)
weight_scale = 0.137
q_fp8, weights_out = fused_q_indexer_rope_first_quant(
q, weight, weight_scale, freqs_cis, positions
)
torch.cuda.synchronize()
cp = cos[positions.long()][:, None, :]
sp = sin[positions.long()][:, None, :]
ref = _rope_first(q.float(), cp, sp) # [B, n_heads, 128]
amax = ref.abs().amax(dim=-1, keepdim=True)
scale = torch.clamp(amax, min=1e-4) / FP8_MAX
# weights_out[b,h] = weight * weight_scale * scale
w_ref = weight.float() * weight_scale * scale.squeeze(-1)
torch.testing.assert_close(weights_out.squeeze(-1), w_ref, atol=1e-3, rtol=1e-3)
# dequantized q should match the rope result within fp8-e4m3 precision:
# round-to-nearest with 3 mantissa bits => <= 1/16 relative error, plus one
# scale step at the bottom of the range.
deq = q_fp8.float() * scale
err = (deq - ref).abs()
assert (
err <= 0.0625 * ref.abs() + scale
).all(), f"max fp8 dequant error {err.max().item()}"
# ----------------------------------------------------------------------------
# Strided weight (the wk_weights_proj slice) matches contiguous for the Q kernel
# ----------------------------------------------------------------------------
def test_q_strided_weight_matches_contiguous():
_skip_if_unavailable()
dev = "cuda"
B, n_heads = 29, 64
cos, sin, freqs_cis, positions = _make_inputs(B)
q = torch.randn(B, n_heads, HEAD_DIM, dtype=torch.bfloat16, device=dev)
# weights_raw = kw[:, head_dim:] is a non-contiguous slice.
kw = torch.randn(B, HEAD_DIM + n_heads, dtype=torch.bfloat16, device=dev)
w_strided = kw[:, HEAD_DIM:]
w_contig = w_strided.contiguous()
assert not w_strided.is_contiguous()
a_fp8, a_w = fused_q_indexer_rope_first_quant(
q, w_strided, 0.137, freqs_cis, positions
)
b_fp8, b_w = fused_q_indexer_rope_first_quant(
q, w_contig, 0.137, freqs_cis, positions
)
torch.cuda.synchronize()
assert torch.equal(a_fp8, b_fp8)
assert torch.equal(a_w, b_w)
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v"]))
test/registered/models_e2e/test_dsa_glm5_dp_mtp.py
+1 -1
View File
@@ -19,7 +19,7 @@ class TestGLM5DPMTP(
DsaMtpServerBase, DsaMtpEvalConfigDefaults, GSM8KMixin, SpecDecodingMixin
):
model = "zai-org/GLM-5-FP8"
mem_fraction_static = 0.8
mem_fraction_static = 0.88
enable_dp_attention = True
bs_1_speed_thres = 70