[AMD] Fix DeepSeek V4 Pro c128 state tensor dtype mismatch error and c4_sparse_raw_indices attribute error in cuda graph phase (#27529)

Co-authored-by: YC Yen-Ching Tseng <yctseng@amd.com>
Co-authored-by: Thomas Wang <thomawan@amd.com>

python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh

@@ -89,10 +89,10 @@ struct C128Trait {
   static_assert(kHeadDim % kTileDim == 0);
 };
 
-template <typename Trait, bool kUsePDL, typename InFloat, typename OutFloat>
+template <typename Trait, bool kUsePDL, typename BufFloat, typename InFloat, typename OutFloat>
 SGL_DEVICE void c128_forward(
-    const InFloat* kv_buf,  // [128n, 128n + 127]
+    const BufFloat* kv_buf,  // [128n, 128n + 127]
-    const InFloat* kv_src,  // ragged pointer at position = 128n + 127
+    const InFloat* kv_src,   // ragged pointer at position = 128n + 127
     OutFloat* kv_out,
     const InFloat* score_bias,
     const int32_t buffer_len) {
@@ -101,11 +101,15 @@ SGL_DEVICE void c128_forward(
   const auto warp_id = threadIdx.x / kWarpThreads;
   const auto lane_id = threadIdx.x % kWarpThreads;
 
-  /// NOTE: part 1: load kv + score
+  /// NOTE: part 1: load kv + score. kv_score_buffer (fp32, runtime state pool)
+  /// keeps its own BufFloat dtype; input/ape share InFloat (ape is cast to bf16
+  /// at load). Every value is converted to fp32 right after load.
+  using StorageBuf = AlignedVector<BufFloat, kTileElements>;
   using StorageIn = AlignedVector<InFloat, kTileElements>;
+  const auto gmem_buf = tile::Memory<StorageBuf>{lane_id, kWarpThreads};
   const auto gmem_in = tile::Memory<StorageIn>{lane_id, kWarpThreads};
-  StorageIn kv[kElementsPerWarp];
+  float kv[kElementsPerWarp][kTileElements];
-  StorageIn score[kElementsPerWarp];
+  float score[kElementsPerWarp][kTileElements];
   StorageIn bias[kElementsPerWarp];
   const int32_t warp_offset = warp_id * kElementsPerWarp;
 
@@ -121,9 +125,23 @@ SGL_DEVICE void c128_forward(
   for (int32_t i = 0; i < kElementsPerWarp; ++i) {
     const int32_t j = i + warp_offset;
     __builtin_assume(j < 128);
-    const auto src = j < buffer_len ? kv_buf : kv_start;
+    if (j < buffer_len) {
-    kv[i] = gmem_in.load(src + j * Trait::kElementSize);
+      const auto k = gmem_buf.load(kv_buf + j * Trait::kElementSize);
-    score[i] = gmem_in.load(src + j * Trait::kElementSize + Trait::kScoreOffset);
+      const auto s = gmem_buf.load(kv_buf + j * Trait::kElementSize + Trait::kScoreOffset);
+#pragma unroll
+      for (int32_t t = 0; t < kTileElements; ++t) {
+        kv[i][t] = cast<float>(k[t]);
+        score[i][t] = cast<float>(s[t]);
+      }
+    } else {
+      const auto k = gmem_in.load(kv_start + j * Trait::kElementSize);
+      const auto s = gmem_in.load(kv_start + j * Trait::kElementSize + Trait::kScoreOffset);
+#pragma unroll
+      for (int32_t t = 0; t < kTileElements; ++t) {
+        kv[i][t] = cast<float>(k[t]);
+        score[i][t] = cast<float>(s[t]);
+      }
+    }
   }
 
   /// NOTE: part 2: safe online softmax + weighted sum
@@ -138,11 +156,11 @@ SGL_DEVICE void c128_forward(
 
   float score_fp32[kTileElements][kElementsPerWarp];
 
-  // convert to fp32 and apply bias first
+  // kv/score already fp32 (converted at load); just add the bias
 #pragma unroll
   for (int32_t i = 0; i < kTileElements; ++i) {
     for (int32_t j = 0; j < kElementsPerWarp; ++j) {
-      score_fp32[i][j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
+      score_fp32[i][j] = score[j][i] + cast<float>(bias[j][i]);
     }
   }
 
@@ -163,7 +181,7 @@ SGL_DEVICE void c128_forward(
     for (int32_t j = 0; j < 8; ++j) {
       const auto fp32_score = score[j];
       const auto exp_score = expf(fp32_score - max_value);
-      sum_product += cast<float>(kv[j][i]) * exp_score;
+      sum_product += kv[j][i] * exp_score;
       sum_exp_value += exp_score;
     }
 
@@ -215,25 +233,27 @@ SGL_DEVICE void c128_forward(
   }
 }
 
-template <typename Trait, typename InFloat>
+template <typename Trait, typename BufFloat, typename InFloat>
-SGL_DEVICE void c128_write_decode(InFloat* kv_buf, const InFloat* kv_src) {
+SGL_DEVICE void c128_write_decode(BufFloat* kv_buf, const InFloat* kv_src) {
   using namespace device;
 
-  using Storage = AlignedVector<InFloat, kTileElements>;
+  using StorageIn = AlignedVector<InFloat, kTileElements>;
-  const auto gmem = tile::Memory<Storage>::warp();
+  using StorageBuf = AlignedVector<BufFloat, kTileElements>;
+  const auto gmem_in = tile::Memory<StorageIn>::warp();
+  const auto gmem_buf = tile::Memory<StorageBuf>::warp();
 
-  Storage data[2];
 #pragma unroll
   for (int32_t i = 0; i < 2; ++i) {
-    data[i] = gmem.load(kv_src + Trait::kHeadDim * i);
+    const auto d = gmem_in.load(kv_src + Trait::kHeadDim * i);
-  }
+    StorageBuf o;
 #pragma unroll
-  for (int32_t i = 0; i < 2; ++i) {
+    for (int32_t t = 0; t < kTileElements; ++t)
-    gmem.store(kv_buf + Trait::kHeadDim * i, data[i]);
+      o[t] = cast<BufFloat>(d[t]);
+    gmem_buf.store(kv_buf + Trait::kHeadDim * i, o);
   }
 }
 
-template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
 C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodeParams params) {
   using namespace device;
   using Trait = C128Trait<kHeadDim>;
@@ -247,7 +267,7 @@ C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodePara
   const auto plan = params.plan_d[global_bid];
   const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
   const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
-  const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
+  const auto kv_buffer = static_cast<BufFloat*>(params.kv_buffer) + split_offset;
   const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
 
   const auto kv_src = kv_input + global_bid * Trait::kElementSize;
@@ -258,15 +278,15 @@ C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodePara
   PDLWaitPrimary<kUsePDL>();
   // the write warp must match the load warp in the following `c128_forward`
   if (warp_id == kNumWarps - 1) {
-    c128_write_decode<Trait>(kv_dst, kv_src);
+    c128_write_decode<Trait, BufFloat, InFloat>(kv_dst, kv_src);
   }
   if (plan.write_loc % 128 == 127) {
-    c128_forward<Trait, kUsePDL>(kv_buf, kv_src, kv_out, score_bias, 128);
+    c128_forward<Trait, kUsePDL, BufFloat, InFloat, OutFloat>(kv_buf, kv_src, kv_out, score_bias, 128);
   }
 }
 
 // compress kernel
-template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
 C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillParams params) {
   using namespace device;
   using Trait = C128Trait<kHeadDim>;
@@ -279,7 +299,7 @@ C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillPa
   const auto plan = params.plan_c[global_pid];
   const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
   const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
-  const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
+  const auto kv_buffer = static_cast<BufFloat*>(params.kv_buffer) + split_offset;
   const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
   if (plan.is_invalid()) return;
 
@@ -288,14 +308,15 @@ C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillPa
   const auto kv_out = kv_output + global_pid * Trait::kHeadDim;
   const auto kv_buf = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
   PDLWaitPrimary<kUsePDL>();
-  c128_forward<Trait, kUsePDL>(kv_buf, kv_src, kv_out, score_bias, plan.buffer_len);
+  c128_forward<Trait, kUsePDL, BufFloat, InFloat, OutFloat>(kv_buf, kv_src, kv_out, score_bias, plan.buffer_len);
 }
 
-template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
 WRITE_KERNEL void write_c128_prefill(const __grid_constant__ Compress128PrefillParams params) {
   using namespace device;
   using Trait = C128Trait<kHeadDim>;
   using StorageIn = AlignedVector<InFloat, kTileElements>;
+  using StorageBuf = AlignedVector<BufFloat, kTileElements>;
 
   const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
   const uint32_t global_wid = global_tid / kWarpThreads;      // warp id
@@ -308,32 +329,37 @@ WRITE_KERNEL void write_c128_prefill(const __grid_constant__ Compress128PrefillP
 
   const auto plan = params.plan_w[global_pid];
   const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
-  const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
+  const auto kv_buffer = static_cast<BufFloat*>(params.kv_buffer) + split_offset;
   if (plan.is_invalid()) return;
 
   // each warp will handle a contiguous region
   const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
   const auto kv_buf = kv_buffer + plan.write_loc * Trait::kElementSize;
-  const auto gmem = tile::Memory<StorageIn>::warp();
+  const auto gmem_in = tile::Memory<StorageIn>::warp();
+  const auto gmem_buf = tile::Memory<StorageBuf>::warp();
 
   PDLWaitPrimary<kUsePDL>();
   StorageIn data[2];
 #pragma unroll
   for (int32_t i = 0; i < 2; ++i) {
-    data[i] = gmem.load(kv_src, i);
+    data[i] = gmem_in.load(kv_src, i);
   }
   PDLTriggerSecondary<kUsePDL>();
 #pragma unroll
   for (int32_t i = 0; i < 2; ++i) {
-    gmem.store(kv_buf, data[i], i);
+    StorageBuf o;
+#pragma unroll
+    for (int32_t t = 0; t < kTileElements; ++t)
+      o[t] = cast<BufFloat>(data[i][t]);
+    gmem_buf.store(kv_buf, o, i);
   }
 }
 
-template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
 struct FlashCompress128Kernel {
-  static constexpr auto decode_kernel = flash_c128_decode<kHeadDim, InFloat, OutFloat, kUsePDL>;
+  static constexpr auto decode_kernel = flash_c128_decode<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
-  static constexpr auto prefill_c_kernel = flash_c128_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
+  static constexpr auto prefill_c_kernel = flash_c128_prefill<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
-  static constexpr auto prefill_w_kernel = write_c128_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
+  static constexpr auto prefill_w_kernel = write_c128_prefill<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
   static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads;  // 64
   static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
   using Trait = C128Trait<kHeadDim>;
@@ -351,7 +377,7 @@ struct FlashCompress128Kernel {
     device_.set_options<kDLGPU>();
 
     TensorMatcher({-1, 128, Trait::kElementSize})  // kv score
-        .with_dtype<InFloat>()
+        .with_dtype<BufFloat>()
         .with_device(device_)
         .verify(kv_buffer);
     TensorMatcher({N, Trait::kElementSize})  // kv score input
@@ -398,7 +424,7 @@ struct FlashCompress128Kernel {
     device_.set_options<kDLGPU>();
 
     TensorMatcher({-1, 128, Trait::kElementSize})  // kv score
-        .with_dtype<InFloat>()
+        .with_dtype<BufFloat>()
         .with_device(device_)
         .verify(kv_buffer);
     TensorMatcher({N, Trait::kElementSize})  // kv score input (ragged)

python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh

@@ -74,23 +74,27 @@ struct C4Trait {
   static_assert(kHeadDim % kTileDim == 0);
 };
 
-template <typename Trait, bool kUsePDL, typename InFloat, typename OutFloat>
+template <typename Trait, bool kUsePDL, typename BufFloat, typename InFloat, typename OutFloat>
 SGL_DEVICE void c4_forward(
-    const InFloat* kv_buf_0,  // overlap [4n - 4, 4n - 1]
+    const BufFloat* kv_buf_0,  // overlap [4n - 4, 4n - 1]
-    const InFloat* kv_buf_1,  // normal [4n + 0, 4n + 3]
+    const BufFloat* kv_buf_1,  // normal [4n + 0, 4n + 3]
-    const InFloat* kv_src,    // ragged pointer at position = 4n + 3
+    const InFloat* kv_src,     // ragged pointer at position = 4n + 3
     OutFloat* kv_out,
     const InFloat* score_bias,
     const bool should_overlap,
     const int32_t buffer_len) {
   using namespace device;
 
-  /// NOTE: part 1: load kv + score
+  /// NOTE: part 1: load kv + score. kv_score_buffer (fp32, runtime state pool)
+  /// keeps its own BufFloat dtype; input/ape share InFloat (ape is cast to bf16
+  /// at load). Values are converted to fp32 at load.
+  using StorageBuf = AlignedVector<BufFloat, kTileElements>;
   using StorageIn = AlignedVector<InFloat, kTileElements>;
   /// NOTE: load one tile_dim (< head_dim) at at time
+  const auto gmem_buf = tile::Memory<StorageBuf>::warp();
   const auto gmem_in = tile::Memory<StorageIn>::warp();
-  StorageIn kv[8];
+  float kv[8][kTileElements];
-  StorageIn score[8];
+  float score[8][kTileElements];
   StorageIn bias[8];
 
 #pragma unroll
@@ -102,28 +106,61 @@ SGL_DEVICE void c4_forward(
     const auto kv_start = kv_src - 7 * Trait::kElementSize;  // point to start
 #pragma unroll
     for (int32_t i = 0; i < 4; ++i) {
-      const auto src = i < buffer_len ? kv_buf_0 : kv_start;
+      if (i < buffer_len) {
-      const auto base = src + i * Trait::kElementSize;
+        const auto base = kv_buf_0 + i * Trait::kElementSize;
-      kv[i] = gmem_in.load(base);
+        const auto k = gmem_buf.load(base);
-      score[i] = gmem_in.load(base + Trait::kScoreOffset);
+        const auto s = gmem_buf.load(base + Trait::kScoreOffset);
+#pragma unroll
+        for (int32_t t = 0; t < kTileElements; ++t) {
+          kv[i][t] = cast<float>(k[t]);
+          score[i][t] = cast<float>(s[t]);
+        }
+      } else {
+        const auto base = kv_start + i * Trait::kElementSize;
+        const auto k = gmem_in.load(base);
+        const auto s = gmem_in.load(base + Trait::kScoreOffset);
+#pragma unroll
+        for (int32_t t = 0; t < kTileElements; ++t) {
+          kv[i][t] = cast<float>(k[t]);
+          score[i][t] = cast<float>(s[t]);
+        }
+      }
     }
   } else {
     [[unlikely]];
     constexpr float kFloatNegInf = -FLT_MAX;
 #pragma unroll
     for (int32_t i = 0; i < 4; ++i) {
-      kv[i].fill(cast<InFloat>(0.0f));
+#pragma unroll
-      score[i].fill(cast<InFloat>(kFloatNegInf));
+      for (int32_t t = 0; t < kTileElements; ++t) {
+        kv[i][t] = 0.0f;
+        score[i][t] = kFloatNegInf;
+      }
     }
   }
 
   const auto kv_start = kv_src - 3 * Trait::kElementSize;  // point to start
 #pragma unroll
   for (int32_t i = 0; i < 4; ++i) {
-    const auto src = i + 4 < buffer_len ? kv_buf_1 : kv_start;
+    if (i + 4 < buffer_len) {
-    const auto base = src + i * Trait::kElementSize + Trait::kOverlapOffset;
+      const auto base = kv_buf_1 + i * Trait::kElementSize + Trait::kOverlapOffset;
-    kv[i + 4] = gmem_in.load(base);
+      const auto k = gmem_buf.load(base);
-    score[i + 4] = gmem_in.load(base + Trait::kScoreOffset);
+      const auto s = gmem_buf.load(base + Trait::kScoreOffset);
+#pragma unroll
+      for (int32_t t = 0; t < kTileElements; ++t) {
+        kv[i + 4][t] = cast<float>(k[t]);
+        score[i + 4][t] = cast<float>(s[t]);
+      }
+    } else {
+      const auto base = kv_start + i * Trait::kElementSize + Trait::kOverlapOffset;
+      const auto k = gmem_in.load(base);
+      const auto s = gmem_in.load(base + Trait::kScoreOffset);
+#pragma unroll
+      for (int32_t t = 0; t < kTileElements; ++t) {
+        kv[i + 4][t] = cast<float>(k[t]);
+        score[i + 4][t] = cast<float>(s[t]);
+      }
+    }
   }
 
   /// NOTE: part 2: safe online softmax + weighted sum
@@ -134,11 +171,11 @@ SGL_DEVICE void c4_forward(
   // consume 32 fp registers
   float score_fp32[kTileElements][8];
 
-  // convert to fp32 and apply bias first
+  // kv/score already fp32 (converted at load); just add the bias
 #pragma unroll
   for (int32_t i = 0; i < kTileElements; ++i) {
     for (int32_t j = 0; j < 8; ++j) {
-      score_fp32[i][j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
+      score_fp32[i][j] = score[j][i] + cast<float>(bias[j][i]);
     }
   }
 
@@ -159,7 +196,7 @@ SGL_DEVICE void c4_forward(
     for (int32_t j = 0; j < 8; ++j) {
       const auto fp32_score = score[j];
       const auto exp_score = expf(fp32_score - max_value);
-      sum_product += cast<float>(kv[j][i]) * exp_score;
+      sum_product += kv[j][i] * exp_score;
       sum_exp_value += exp_score;
     }
 
@@ -171,25 +208,27 @@ SGL_DEVICE void c4_forward(
   gmem_out.store(kv_out, result);
 }
 
-template <typename Trait, typename InFloat>
+template <typename Trait, typename BufFloat, typename InFloat>
-SGL_DEVICE void c4_write_decode(InFloat* kv_buf, const InFloat* kv_src) {
+SGL_DEVICE void c4_write_decode(BufFloat* kv_buf, const InFloat* kv_src) {
   using namespace device;
 
   using StorageIn = AlignedVector<InFloat, kTileElements>;
-  const auto gmem = tile::Memory<StorageIn>::warp();
+  using StorageBuf = AlignedVector<BufFloat, kTileElements>;
+  const auto gmem_in = tile::Memory<StorageIn>::warp();
+  const auto gmem_buf = tile::Memory<StorageBuf>::warp();
 
-  StorageIn data[4];
 #pragma unroll
   for (int32_t i = 0; i < 4; ++i) {
-    data[i] = gmem.load(kv_src + Trait::kHeadDim * i);
+    const auto d = gmem_in.load(kv_src + Trait::kHeadDim * i);
-  }
+    StorageBuf o;
 #pragma unroll
-  for (int32_t i = 0; i < 4; ++i) {
+    for (int32_t t = 0; t < kTileElements; ++t)
-    gmem.store(kv_buf + Trait::kHeadDim * i, data[i]);
+      o[t] = cast<BufFloat>(d[t]);
+    gmem_buf.store(kv_buf + Trait::kHeadDim * i, o);
   }
 }
 
-template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
 C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams params) {
   using namespace device;
   using Trait = C4Trait<kHeadDim>;
@@ -204,7 +243,7 @@ C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams par
   const auto plan = params.plan_d[global_bid];
   const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
   const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
-  const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
+  const auto kv_buffer = static_cast<BufFloat*>(params.kv_buffer) + split_offset;
   const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
 
   const auto kv_src = kv_input + global_bid * Trait::kElementSize;
@@ -214,14 +253,15 @@ C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams par
   const auto kv_dst = kv_buffer + plan.write_loc * Trait::kElementSize;
 
   PDLWaitPrimary<kUsePDL>();
-  c4_write_decode<Trait>(kv_dst, kv_src);
+  c4_write_decode<Trait, BufFloat, InFloat>(kv_dst, kv_src);
   if (plan.seq_len % 4 == 0) {
     const auto need_overlap = plan.seq_len > 4;
-    c4_forward<Trait, kUsePDL>(kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, 8);
+    c4_forward<Trait, kUsePDL, BufFloat, InFloat, OutFloat>(
+        kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, 8);
   }
 }
 
-template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
 C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams params) {
   using namespace device;
   using Trait = C4Trait<kHeadDim>;
@@ -236,7 +276,7 @@ C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams p
   const auto plan = params.plan_c[global_pid];
   const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
   const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
-  const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
+  const auto kv_buffer = static_cast<BufFloat*>(params.kv_buffer) + split_offset;
   const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
   if (plan.is_invalid()) return;
 
@@ -247,14 +287,16 @@ C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams p
   const auto kv_buf_1 = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
   const bool need_overlap = plan.seq_len > 4;
   PDLWaitPrimary<kUsePDL>();
-  c4_forward<Trait, kUsePDL>(kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, plan.buffer_len);
+  c4_forward<Trait, kUsePDL, BufFloat, InFloat, OutFloat>(
+      kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, plan.buffer_len);
 }
 
-template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
 WRITE_KERNEL void write_c4_prefill(const __grid_constant__ Compress4PrefillParams params) {
   using namespace device;
   using Trait = C4Trait<kHeadDim>;
   using StorageIn = AlignedVector<InFloat, kTileElements>;
+  using StorageBuf = AlignedVector<BufFloat, kTileElements>;
 
   const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
   const uint32_t global_wid = global_tid / kWarpThreads;      // warp id
@@ -267,32 +309,37 @@ WRITE_KERNEL void write_c4_prefill(const __grid_constant__ Compress4PrefillParam
 
   const auto plan = params.plan_w[global_pid];
   const auto kv_input = static_cast<const InFloat*>(params.kv_input) + split_offset;
-  const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
+  const auto kv_buffer = static_cast<BufFloat*>(params.kv_buffer) + split_offset;
   if (plan.is_invalid()) return;
 
   // each warp will handle a contiguous region
   const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
   const auto kv_buf = kv_buffer + plan.write_loc * Trait::kElementSize;
-  const auto gmem = tile::Memory<StorageIn>::warp();
+  const auto gmem_in = tile::Memory<StorageIn>::warp();
+  const auto gmem_buf = tile::Memory<StorageBuf>::warp();
 
   PDLWaitPrimary<kUsePDL>();
   StorageIn data[4];
 #pragma unroll
   for (int32_t i = 0; i < 4; ++i) {
-    data[i] = gmem.load(kv_src, i);
+    data[i] = gmem_in.load(kv_src, i);
   }
   PDLTriggerSecondary<kUsePDL>();
 #pragma unroll
   for (int32_t i = 0; i < 4; ++i) {
-    gmem.store(kv_buf, data[i], i);
+    StorageBuf o;
+#pragma unroll
+    for (int32_t t = 0; t < kTileElements; ++t)
+      o[t] = cast<BufFloat>(data[i][t]);
+    gmem_buf.store(kv_buf, o, i);
   }
 }
 
-template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
 struct FlashCompress4Kernel {
-  static constexpr auto decode_kernel = flash_c4_decode<kHeadDim, InFloat, OutFloat, kUsePDL>;
+  static constexpr auto decode_kernel = flash_c4_decode<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
-  static constexpr auto prefill_c_kernel = flash_c4_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
+  static constexpr auto prefill_c_kernel = flash_c4_prefill<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
-  static constexpr auto prefill_w_kernel = write_c4_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
+  static constexpr auto prefill_w_kernel = write_c4_prefill<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
   static constexpr uint32_t kBlockSize = 128;
   static constexpr uint32_t kTileDim = kTileElements * device::kWarpThreads;
   static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
@@ -312,7 +359,7 @@ struct FlashCompress4Kernel {
     device_.set_options<kDLGPU>();
 
     TensorMatcher({-1, 4, Trait::kElementSize})  // kv score
-        .with_dtype<InFloat>()
+        .with_dtype<BufFloat>()
         .with_device(device_)
         .verify(kv_buffer);
     TensorMatcher({N, Trait::kElementSize})  // kv score input
@@ -359,7 +406,7 @@ struct FlashCompress4Kernel {
     device_.set_options<kDLGPU>();
 
     TensorMatcher({-1, 4, Trait::kElementSize})  // kv score
-        .with_dtype<InFloat>()
+        .with_dtype<BufFloat>()
         .with_device(device_)
         .verify(kv_buffer);
     TensorMatcher({N, Trait::kElementSize})  // kv score input (ragged)

python/sglang/jit_kernel/dsv4/compress.py

@@ -44,11 +44,14 @@ def _jit_compress_norm_rope_module(
 @cache_once
 def _jit_compress_module(
     head_dim: int,
+    dtype_buf: torch.dtype,
     dtype_in: torch.dtype,
     dtype_out: torch.dtype,
     ratio: Literal[4, 128],
 ) -> Module:
-    args = make_cpp_args(head_dim, dtype_in, dtype_out, is_arch_support_pdl())
+    args = make_cpp_args(
+        head_dim, dtype_buf, dtype_in, dtype_out, is_arch_support_pdl()
+    )
     kernel_class = f"FlashCompress{ratio}Kernel<{args}>"
     return load_jit(
         make_name(f"compress_{ratio}_v2"),
@@ -324,8 +327,17 @@ def compress_forward(
         assert compress_ratio == 128 and head_dim == 512
         module = _jit_compress_128_online_module(512)
     else:
-        dtype_in, dtype_out = kv_score_input.dtype, out.dtype
+        # kv_score_buffer (fp32 runtime state pool) may differ from input/ape, so
-        module = _jit_compress_module(head_dim, dtype_in, dtype_out, compress_ratio)
+        # the kernel keeps a BufFloat template and casts the buffer to fp32 at
+        # load. ape/weight are cast to bf16 at load (apply_ape_hotfix), matching
+        # kv_score_input's dtype, so they need no extra template params.
+        module = _jit_compress_module(
+            head_dim,
+            kv_score_buffer.dtype,
+            kv_score_input.dtype,
+            out.dtype,
+            compress_ratio,
+        )
     fn = module.decode if plan.is_decode else module.prefill
     fn(kv_score_buffer, kv_score_input, out, ape, *plan[1:3])
     return out

python/sglang/srt/layers/attention/dsv4/compressor.py

@@ -395,6 +395,10 @@ class Compressor(nn.Module):
             ape = torch.cat([ape[0], ape[1]], dim=0)
             self.ape.data.copy_(ape.view(self.ratio, -1))
 
+        if _use_aiter:
+            self.ape.data = self.ape.data.to(torch.bfloat16)
+            self.norm.weight.data = self.norm.weight.data.to(torch.bfloat16)
+
     def get_state_pool(self, attn_backend: AttentionBackend) -> CompressStatePool:
         token_to_kv_pool = attn_backend.token_to_kv_pool
         assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)