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

Co-authored-by: Brayden Zhong <brayden@radixark.ai> Co-authored-by: Kaixi <kaiximatteoc@nvidia.com>
2026-06-27 14:29:56 -07:00
parent e4253b39e2
commit 073de15053
13 changed files with 1196 additions and 69 deletions
@@ -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
@@ -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,
    };
@@ -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",
+]
@@ -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),
+    )
@@ -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",
@@ -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,
@@ -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)
@@ -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,
@@ -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:
@@ -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: