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

[DSA] Use cos_sin_cache for DSA indexer fusion (#29613)

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2026-06-29 23:49:00 -07:00
committed by GitHub
parent dbb51c46ac
commit b6fceaa789
6 changed files with 160 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/deepseek_v32/indexer_k.cuh
+38 -24
View File
@@ -1,7 +1,7 @@
// 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
// kRopeDim dims, and fp8-quantizes the 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`
@@ -33,14 +33,28 @@ constexpr uint32_t kFusedKIndexerNumWarps = kFusedKIndexerBlockSize / device::kW
#define K_INDEXER_KERNEL __global__ __launch_bounds__(kFusedKIndexerBlockSize, 16)
template <int64_t kRopeDim>
SGL_DEVICE device::AlignedVector<float, 4>
load_rope_first_cos_sin(const float* __restrict__ cos_sin_cache, int32_t lane_id) {
constexpr int64_t kHalfRopeDim = kRopeDim / 2;
const int32_t pair0 = lane_id * 2;
const int32_t pair1 = pair0 + 1;
device::AlignedVector<float, 4> freq;
freq[0] = cos_sin_cache[pair0];
freq[1] = cos_sin_cache[kHalfRopeDim + pair0];
freq[2] = cos_sin_cache[pair1];
freq[3] = cos_sin_cache[kHalfRopeDim + pair1];
return freq;
}
// 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
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__ cos_sin_cache; // (max_pos, 64) fp32 [cos..., sin...]
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;
@@ -71,7 +85,7 @@ K_INDEXER_KERNEL void fused_k_indexer_norm_rope(const __grid_constant__ FusedKIn
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;
const auto cos_sin_cache = params.cos_sin_cache + position * kRopeDim;
PDLWaitPrimary<kUsePDL>();
Float4 data, freq, gamma, beta;
@@ -82,7 +96,7 @@ K_INDEXER_KERNEL void fused_k_indexer_norm_rope(const __grid_constant__ FusedKIn
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);
if (is_rope_lane) freq = load_rope_first_cos_sin<kRopeDim>(cos_sin_cache, lane_id);
float sum = 0.0f;
#pragma unroll
@@ -144,7 +158,7 @@ struct FusedKIndexerNormRopeKernel {
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 cos_sin_cache,
const tvm::ffi::TensorView positions,
double eps) {
using namespace host;
@@ -176,7 +190,7 @@ struct FusedKIndexerNormRopeKernel {
TensorMatcher({-1, kRopeDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(freqs_cis);
.verify(cos_sin_cache);
auto pos_dtype = SymbolicDType{};
TensorMatcher({B}) //
.with_dtype<int32_t, int64_t>(pos_dtype)
@@ -191,7 +205,7 @@ struct FusedKIndexerNormRopeKernel {
.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()),
.cos_sin_cache = static_cast<const float*>(cos_sin_cache.data_ptr()),
.positions = positions.data_ptr(),
.k_input_stride_batch = k_input.stride(0),
.batch_size = batch_size,
@@ -210,13 +224,13 @@ struct FusedKIndexerNormRopeKernel {
// 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
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__ cos_sin_cache; // (max_pos, 64) fp32 [cos..., sin...]
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;
@@ -249,7 +263,7 @@ K_INDEXER_KERNEL void fused_k_indexer_norm_rope_store(const __grid_constant__ Fu
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;
const auto cos_sin_cache = params.cos_sin_cache + position * kRopeDim;
PDLWaitPrimary<kUsePDL>();
Float4 data, freq, gamma, beta;
@@ -260,7 +274,7 @@ K_INDEXER_KERNEL void fused_k_indexer_norm_rope_store(const __grid_constant__ Fu
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);
if (is_rope_lane) freq = load_rope_first_cos_sin<kRopeDim>(cos_sin_cache, lane_id);
float sum = 0.0f;
#pragma unroll
@@ -345,7 +359,7 @@ struct FusedKIndexerNormRopeStoreKernel {
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 cos_sin_cache,
const tvm::ffi::TensorView positions,
double eps) {
using namespace host;
@@ -381,7 +395,7 @@ struct FusedKIndexerNormRopeStoreKernel {
TensorMatcher({-1, kRopeDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(freqs_cis);
.verify(cos_sin_cache);
auto pos_dtype = SymbolicDType{};
TensorMatcher({B}) //
.with_dtype<int32_t, int64_t>(pos_dtype)
@@ -397,7 +411,7 @@ struct FusedKIndexerNormRopeStoreKernel {
.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()),
.cos_sin_cache = static_cast<const float*>(cos_sin_cache.data_ptr()),
.positions = positions.data_ptr(),
.k_input_stride_batch = k_input.stride(0),
.batch_size = batch_size,
python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh
+33 -10
View File
@@ -46,6 +46,20 @@ constexpr uint32_t kFusedKNumWarps = kFusedKBlockSize / device::kWarpThreads;
#define Q_KERNEL __global__ __launch_bounds__(kFusedQBlockSize, 16)
#define K_KERNEL __global__ __launch_bounds__(kFusedKBlockSize, 8)
template <int64_t kRopeDim>
SGL_DEVICE device::AlignedVector<float, 4>
load_rope_first_cos_sin(const float* __restrict__ cos_sin_cache, int32_t lane_id) {
constexpr int64_t kHalfRopeDim = kRopeDim / 2;
const int32_t pair0 = lane_id * 2;
const int32_t pair1 = pair0 + 1;
device::AlignedVector<float, 4> freq;
freq[0] = cos_sin_cache[pair0];
freq[1] = cos_sin_cache[kHalfRopeDim + pair0];
freq[2] = cos_sin_cache[pair1];
freq[3] = cos_sin_cache[kHalfRopeDim + pair1];
return freq;
}
// ============================================================================
// Q kernel: warp-per-(token, head) rmsnorm-self + RoPE + write to q_out.
// ============================================================================
@@ -420,11 +434,14 @@ struct FusedQIndexerRopeHadamardQuantParams {
// weights_out[b, h] = weight[b, h] * weight_scale * q_scale[b, h].
// q_scale is computed internally and not exposed -- the only consumer of
// it is `weights_out`.
const void* __restrict__ weight; // (B, num_heads) DType
float* __restrict__ weights_out; // (B, num_heads) fp32 (== (B, H, 1) flat)
float weight_scale; // scalar c4_indexer.weight_scale
const float* __restrict__ freqs_cis; // (max_pos, 64) fp32
const void* __restrict__ positions; // (B,) PosT
const void* __restrict__ weight; // (B, num_heads) DType
float* __restrict__ weights_out; // (B, num_heads) fp32 (== (B, H, 1) flat)
float weight_scale; // scalar c4_indexer.weight_scale
// Template-dependent layout:
// kRopeFirst=false: (max_pos, 64) fp32 interleaved [cos0, sin0, ...]
// kRopeFirst=true : (max_pos, 64) fp32 halves [cos..., sin...]
const float* __restrict__ rope_cache;
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;
@@ -460,7 +477,7 @@ Q_KERNEL void fused_q_indexer_rope_hadamard_quant(const __grid_constant__ FusedQ
const uint32_t batch_id = work_id / params.num_heads;
const auto input_ptr = static_cast<const DType*>(params.q_input) + work_id * kHeadDim;
const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[batch_id]);
const auto freqs_cis = params.freqs_cis + position * kRopeDim;
const auto rope_cache = params.rope_cache + position * kRopeDim;
// Lane 0 prefetches the weight scalar for this (token, head) work item.
// Weight is (B, num_heads) DType; we need one scalar per warp -- offload
@@ -477,7 +494,13 @@ Q_KERNEL void fused_q_indexer_rope_hadamard_quant(const __grid_constant__ FusedQ
{
Storage input_vec;
input_vec.load(input_ptr, lane_id);
if (is_rope_lane) freq.load(freqs_cis, kRopeFirst ? lane_id : (lane_id - (kWarpThreads - kRopeSize)));
if (is_rope_lane) {
if constexpr (kRopeFirst) {
freq = load_rope_first_cos_sin<kRopeDim>(rope_cache, lane_id);
} else {
freq.load(rope_cache, lane_id - (kWarpThreads - kRopeSize));
}
}
#pragma unroll
for (int i = 0; i < kVecSize; ++i) {
data[i] = cast<float>(input_vec[i]);
@@ -567,7 +590,7 @@ struct FusedQIndexerRopeHadamardQuantKernel {
const tvm::ffi::TensorView weight,
const tvm::ffi::TensorView weights_out,
double weight_scale,
const tvm::ffi::TensorView freqs_cis,
const tvm::ffi::TensorView rope_cache,
const tvm::ffi::TensorView positions) {
using namespace host;
constexpr int64_t kHeadDim = 128;
@@ -603,7 +626,7 @@ struct FusedQIndexerRopeHadamardQuantKernel {
TensorMatcher({-1, kRopeDim}) //
.with_dtype<float>()
.with_device(device_)
.verify(freqs_cis);
.verify(rope_cache);
auto pos_dtype = SymbolicDType{};
TensorMatcher({B}) //
.with_dtype<int32_t, int64_t>(pos_dtype)
@@ -632,7 +655,7 @@ struct FusedQIndexerRopeHadamardQuantKernel {
.weight = weight.data_ptr(),
.weights_out = static_cast<float*>(weights_out.data_ptr()),
.weight_scale = static_cast<float>(weight_scale),
.freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
.rope_cache = static_cast<const float*>(rope_cache.data_ptr()),
.positions = positions.data_ptr(),
.weight_stride_batch = weight.stride(0),
.batch_size = batch_size,
python/sglang/jit_kernel/dsv32/elementwise.py
+4 -6
View File
@@ -43,11 +43,10 @@ def fused_k_indexer_norm_rope(
weight: torch.Tensor,
bias: torch.Tensor,
eps: float,
freqs_cis: torch.Tensor,
cos_sin_cache: 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)
@@ -56,7 +55,7 @@ def fused_k_indexer_norm_rope(
k_out,
weight,
bias,
freqs_real,
cos_sin_cache,
positions,
float(eps),
)
@@ -70,13 +69,12 @@ def fused_k_indexer_norm_rope_store(
weight: torch.Tensor,
bias: torch.Tensor,
eps: float,
freqs_cis: torch.Tensor,
cos_sin_cache: 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)
@@ -86,7 +84,7 @@ def fused_k_indexer_norm_rope_store(
out_cache_loc,
weight,
bias,
freqs_real,
cos_sin_cache,
positions,
float(eps),
)
python/sglang/jit_kernel/dsv4/elementwise.py
+2 -3
View File
@@ -188,11 +188,10 @@ def fused_q_indexer_rope_first_quant(
q_input: torch.Tensor,
weight: torch.Tensor,
weight_scale: float,
freqs_cis: torch.Tensor,
cos_sin_cache: 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
@@ -204,7 +203,7 @@ def fused_q_indexer_rope_first_quant(
weight,
weights_out,
float(weight_scale),
freqs_real,
cos_sin_cache,
positions,
)
return q_fp8, weights_out
python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+8 -18
View File
@@ -327,16 +327,6 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor:
return hadamard_transform(x, scale=hidden_size**-0.5)
def _shared_indexer_freqs_cis(rotary_emb: torch.nn.Module) -> torch.Tensor:
cached = getattr(rotary_emb, "_dsa_indexer_freqs_cis", None)
if cached is None:
c = rotary_emb.cos_sin_cache.to(torch.float32)
half = c.shape[-1] // 2
cached = torch.complex(c[:, :half].contiguous(), c[:, half:].contiguous())
rotary_emb._dsa_indexer_freqs_cis = cached
return cached
class Indexer(MultiPlatformOp):
_MQA_LOGITS_BYTES_PER_ELEM = 4
_MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000
@@ -437,10 +427,6 @@ class Indexer(MultiPlatformOp):
self.scale_fmt = scale_fmt
self.softmax_scale = self.head_dim**-0.5
self._indexer_freqs_cis: Optional[torch.Tensor] = None
if _use_dsa_indexer_fusion:
self._indexer_freqs_cis = _shared_indexer_freqs_cis(self.rotary_emb)
@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
@@ -456,6 +442,10 @@ class Indexer(MultiPlatformOp):
else:
yield
@property
def _indexer_cos_sin_cache(self) -> torch.Tensor:
return self.rotary_emb.cos_sin_cache
def _weights_proj_bf16_in_fp32_out(
self, x: Union[torch.Tensor, Tuple[torch.Tensor, ...]]
) -> torch.Tensor:
@@ -662,7 +652,7 @@ class Indexer(MultiPlatformOp):
self.k_norm.weight,
self.k_norm.bias,
self.k_norm.variance_epsilon,
self._indexer_freqs_cis,
self._indexer_cos_sin_cache,
positions,
page_size,
)
@@ -674,7 +664,7 @@ class Indexer(MultiPlatformOp):
self.k_norm.weight,
self.k_norm.bias,
self.k_norm.variance_epsilon,
self._indexer_freqs_cis,
self._indexer_cos_sin_cache,
positions,
)
self._store_index_k_cache(
@@ -726,7 +716,7 @@ class Indexer(MultiPlatformOp):
q.contiguous(),
weights_raw,
q_scale_gate,
self._indexer_freqs_cis,
self._indexer_cos_sin_cache,
positions,
)
@@ -763,7 +753,7 @@ class Indexer(MultiPlatformOp):
q.contiguous(),
weights_raw,
q_scale_gate,
self._indexer_freqs_cis,
self._indexer_cos_sin_cache,
positions,
)