From 073de150530ab27e0646492ec00016c727f0d7f7 Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Sat, 27 Jun 2026 14:29:56 -0700 Subject: [PATCH] Fuse the DSA (V3.2, GLM-5.x) indexer Q/K paths into single kernels (#27705) Co-authored-by: Brayden Zhong Co-authored-by: Kaixi --- .../csrc/deepseek_v32/indexer_k.cuh | 415 ++++++++++++++++++ .../csrc/deepseek_v4/main_norm_rope.cuh | 28 +- python/sglang/jit_kernel/dsv32/__init__.py | 11 + python/sglang/jit_kernel/dsv32/elementwise.py | 92 ++++ python/sglang/jit_kernel/dsv4/__init__.py | 2 + python/sglang/jit_kernel/dsv4/elementwise.py | 41 ++ python/sglang/srt/environ.py | 1 + .../srt/layers/attention/dsa/dsa_indexer.py | 388 +++++++++++++--- python/sglang/srt/lora/lora_manager.py | 16 + .../deepseek_common/deepseek_weight_loader.py | 61 +++ .../test_deepseek_v32_indexcache.py | 2 +- .../jit/test_dsv32_indexer_fusion.py | 206 +++++++++ .../models_e2e/test_dsa_glm5_dp_mtp.py | 2 +- 13 files changed, 1196 insertions(+), 69 deletions(-) create mode 100644 python/sglang/jit_kernel/csrc/deepseek_v32/indexer_k.cuh create mode 100644 python/sglang/jit_kernel/dsv32/__init__.py create mode 100644 python/sglang/jit_kernel/dsv32/elementwise.py create mode 100644 test/registered/jit/test_dsv32_indexer_fusion.py diff --git a/python/sglang/jit_kernel/csrc/deepseek_v32/indexer_k.cuh b/python/sglang/jit_kernel/csrc/deepseek_v32/indexer_k.cuh new file mode 100644 index 000000000..0cf4ea46d --- /dev/null +++ b/python/sglang/jit_kernel/csrc/deepseek_v32/indexer_k.cuh @@ -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 +#include + +#include +#include +#include +#include +#include +#include + +#include + +#include + +#include +#include + +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 +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; + using Float4 = AlignedVector; + + 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(params.k_input) + work_id * params.k_input_stride_batch; + const auto position = static_cast(static_cast(params.positions)[work_id]); + const auto freqs_cis = params.freqs_cis + position * kRopeDim; + + PDLWaitPrimary(); + 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(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(); + + { + Storage out_vec; +#pragma unroll + for (int i = 0; i < kVecSize; ++i) + out_vec[i] = cast(data[i]); + auto out_row = static_cast(params.k_out) + work_id * kHeadDim; + out_vec.store(out_row, lane_id); + } +} + +template +struct FusedKIndexerNormRopeKernel { + template + static constexpr auto kernel = fused_k_indexer_norm_rope; + + 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(); + + TensorMatcher({B, kHeadDim}) // + .with_strides({-1, 1}) + .with_dtype() + .with_device(device_) + .verify(k_input); + TensorMatcher({B, kHeadDim}) // + .with_strides({kHeadDim, 1}) + .with_dtype() + .with_device(device_) + .verify(k_out); + TensorMatcher({kHeadDim}) // + .with_dtype() + .with_device(device_) + .verify(weight); + TensorMatcher({kHeadDim}) // + .with_dtype() + .with_device(device_) + .verify(bias); + TensorMatcher({-1, kRopeDim}) // + .with_dtype() + .with_device(device_) + .verify(freqs_cis); + auto pos_dtype = SymbolicDType{}; + TensorMatcher({B}) // + .with_dtype(pos_dtype) + .with_device(device_) + .verify(positions); + + const auto batch_size = static_cast(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(weight.data_ptr()), + .bias = static_cast(bias.data_ptr()), + .freqs_cis = static_cast(freqs_cis.data_ptr()), + .positions = positions.data_ptr(), + .k_input_stride_batch = k_input.stride(0), + .batch_size = batch_size, + .eps = static_cast(eps), + }; + const auto num_blocks = div_ceil(batch_size, kFusedKIndexerNumWarps); + const auto k_int32 = kernel; + const auto k_int64 = kernel; + const auto k = pos_dtype.is_type() ? 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 +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; + using Float4 = AlignedVector; + using OutStorage = AlignedVector; // 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(params.k_input) + work_id * params.k_input_stride_batch; + const auto position = static_cast(static_cast(params.positions)[work_id]); + const auto freqs_cis = params.freqs_cis + position * kRopeDim; + + PDLWaitPrimary(); + 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(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(); + + // 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(cast(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(params.indices)[work_id]; + const int32_t page = static_cast(index >> kPageBits); + const int32_t offset = static_cast(index & ((1 << kPageBits) - 1)); + const auto page_ptr = static_cast(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(value_ptr)[lane_id] = result; + if (lane_id == 0) *reinterpret_cast(scale_ptr) = scale; +} + +template +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 + static constexpr auto kernel = fused_k_indexer_norm_rope_store; + + 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(); + + TensorMatcher({B, kHeadDim}) // + .with_strides({-1, 1}) + .with_dtype() + .with_device(device_) + .verify(k_input); + TensorMatcher({-1, -1}) // + .with_strides({kPageBytes, 1}) + .with_dtype() + .with_device(device_) + .verify(cache); + TensorMatcher({B}) // + .with_dtype() + .with_device(device_) + .verify(indices); + TensorMatcher({kHeadDim}) // + .with_dtype() + .with_device(device_) + .verify(weight); + TensorMatcher({kHeadDim}) // + .with_dtype() + .with_device(device_) + .verify(bias); + TensorMatcher({-1, kRopeDim}) // + .with_dtype() + .with_device(device_) + .verify(freqs_cis); + auto pos_dtype = SymbolicDType{}; + TensorMatcher({B}) // + .with_dtype(pos_dtype) + .with_device(device_) + .verify(positions); + + const auto batch_size = static_cast(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(weight.data_ptr()), + .bias = static_cast(bias.data_ptr()), + .freqs_cis = static_cast(freqs_cis.data_ptr()), + .positions = positions.data_ptr(), + .k_input_stride_batch = k_input.stride(0), + .batch_size = batch_size, + .eps = static_cast(eps), + }; + const auto num_blocks = div_ceil(batch_size, kFusedKIndexerNumWarps); + const auto k_int32 = kernel; + const auto k_int64 = kernel; + const auto k = pos_dtype.is_type() ? k_int32 : k_int64; + LaunchKernel(num_blocks, kFusedKIndexerBlockSize, device_.unwrap()) // + .enable_pdl(kUsePDL)(k, params); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh index 8fc8d0821..430914606 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh @@ -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 +template 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(); Float4 data, freq; - const auto weight_val = cast(static_cast(params.weight)[work_id]); + const uint32_t head_id = work_id - batch_id * params.num_heads; + const auto weight_val = + cast(static_cast(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(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 +template struct FusedQIndexerRopeHadamardQuantKernel { template - static constexpr auto kernel = fused_q_indexer_rope_hadamard_quant; + static constexpr auto kernel = fused_q_indexer_rope_hadamard_quant; 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() .with_device(device_) .verify(weight); @@ -627,6 +634,7 @@ struct FusedQIndexerRopeHadamardQuantKernel { .weight_scale = static_cast(weight_scale), .freqs_cis = static_cast(freqs_cis.data_ptr()), .positions = positions.data_ptr(), + .weight_stride_batch = weight.stride(0), .batch_size = batch_size, .num_heads = num_heads, }; diff --git a/python/sglang/jit_kernel/dsv32/__init__.py b/python/sglang/jit_kernel/dsv32/__init__.py new file mode 100644 index 000000000..a1b357c5f --- /dev/null +++ b/python/sglang/jit_kernel/dsv32/__init__.py @@ -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", +] diff --git a/python/sglang/jit_kernel/dsv32/elementwise.py b/python/sglang/jit_kernel/dsv32/elementwise.py new file mode 100644 index 000000000..39f48b357 --- /dev/null +++ b/python/sglang/jit_kernel/dsv32/elementwise.py @@ -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), + ) diff --git a/python/sglang/jit_kernel/dsv4/__init__.py b/python/sglang/jit_kernel/dsv4/__init__.py index e344599b0..a9952ebfe 100644 --- a/python/sglang/jit_kernel/dsv4/__init__.py +++ b/python/sglang/jit_kernel/dsv4/__init__.py @@ -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", diff --git a/python/sglang/jit_kernel/dsv4/elementwise.py b/python/sglang/jit_kernel/dsv4/elementwise.py index 91a5c501e..96bc8c1a2 100644 --- a/python/sglang/jit_kernel/dsv4/elementwise.py +++ b/python/sglang/jit_kernel/dsv4/elementwise.py @@ -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, diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index a69a7e6b6..9f0d03b89 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index d088c9243..8d70695a3 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -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,20 +388,29 @@ class Indexer(MultiPlatformOp): prefix=add_prefix("wq_b", prefix), ) - self.wk = ReplicatedLinear( - self.hidden_size, - self.head_dim, - bias=False, - quant_config=quant_config, - prefix=add_prefix("wk", prefix), - ) - self.weights_proj = ReplicatedLinear( - self.hidden_size, - self.n_heads, - bias=False, - params_dtype=torch.bfloat16, - prefix=add_prefix("weights_proj", 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, + bias=False, + quant_config=quant_config, + prefix=add_prefix("wk", prefix), + ) + self.weights_proj = ReplicatedLinear( + self.hidden_size, + self.n_heads, + bias=False, + params_dtype=torch.bfloat16, + prefix=add_prefix("weights_proj", prefix), + ) self.k_norm = LayerNorm( self.head_dim, dtype=torch.bfloat16 if _use_aiter else torch.float32 ) @@ -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,7 +539,10 @@ class Indexer(MultiPlatformOp): ) with torch.cuda.stream(self.alt_stream): # TODO we should also put DeepGEMM half SM here? - key, _ = self.wk(x) + 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( @@ -492,7 +558,10 @@ class Indexer(MultiPlatformOp): q_rope, _ = torch.split( query, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 ) - key, _ = self.wk(x) + 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( key, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 @@ -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,23 +1199,43 @@ 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() - 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, - ) + + # 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, + key=key, + act_quant=act_quant, + out_cache_loc=out_cache_loc, + ) # MHA doesn't need topk_indices if not return_indices: @@ -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 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( + 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, 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) - weights = self._apply_q_scale_and_softmax_scale(weights, q_scale) + 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,15 +1810,25 @@ class Indexer(MultiPlatformOp): x_for_gate = x if in_piecewise_or_breakable_cuda_graph: - if weights_proj_lora: - raise RuntimeError(GRAPH_WEIGHTS_PROJ_LORA_ERROR) - weights = logits_head_gate_graph( - x_for_gate, - self.weights_proj.weight, - self.n_heads**-0.5, - self.softmax_scale, - q_scale, - ) + 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( + x_for_gate, + self.weights_proj.weight, + self.n_heads**-0.5, + 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, diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index 91a9e48d5..7279114c5 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -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: diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index 754dc1bba..ad83266b7 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -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: diff --git a/test/registered/8-gpu-models/test_deepseek_v32_indexcache.py b/test/registered/8-gpu-models/test_deepseek_v32_indexcache.py index f710996b0..6b644e72d 100644 --- a/test/registered/8-gpu-models/test_deepseek_v32_indexcache.py +++ b/test/registered/8-gpu-models/test_deepseek_v32_indexcache.py @@ -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): diff --git a/test/registered/jit/test_dsv32_indexer_fusion.py b/test/registered/jit/test_dsv32_indexer_fusion.py new file mode 100644 index 000000000..596ad4d62 --- /dev/null +++ b/test/registered/jit/test_dsv32_indexer_fusion.py @@ -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"])) diff --git a/test/registered/models_e2e/test_dsa_glm5_dp_mtp.py b/test/registered/models_e2e/test_dsa_glm5_dp_mtp.py index b81f47495..ea3a3902b 100644 --- a/test/registered/models_e2e/test_dsa_glm5_dp_mtp.py +++ b/test/registered/models_e2e/test_dsa_glm5_dp_mtp.py @@ -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