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

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 BufFloat, typename InFloat, typename OutFloat>
+template <typename Trait, bool kUsePDL, typename InFloat, typename OutFloat>
 SGL_DEVICE void c128_forward(
-    const BufFloat* kv_buf,  // [128n, 128n + 127]
+    const InFloat* 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,15 +101,11 @@ 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. kv_score_buffer (fp32, runtime state pool)
+  /// NOTE: part 1: load kv + score
-  /// 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};
-  float kv[kElementsPerWarp][kTileElements];
+  StorageIn kv[kElementsPerWarp];
-  float score[kElementsPerWarp][kTileElements];
+  StorageIn score[kElementsPerWarp];
   StorageIn bias[kElementsPerWarp];
   const int32_t warp_offset = warp_id * kElementsPerWarp;
 
@@ -125,23 +121,9 @@ SGL_DEVICE void c128_forward(
   for (int32_t i = 0; i < kElementsPerWarp; ++i) {
     const int32_t j = i + warp_offset;
     __builtin_assume(j < 128);
-    if (j < buffer_len) {
+    const auto src = j < buffer_len ? kv_buf : kv_start;
-      const auto k = gmem_buf.load(kv_buf + j * Trait::kElementSize);
+    kv[i] = gmem_in.load(src + j * Trait::kElementSize);
-      const auto s = gmem_buf.load(kv_buf + j * Trait::kElementSize + Trait::kScoreOffset);
+    score[i] = gmem_in.load(src + 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
@@ -156,11 +138,11 @@ SGL_DEVICE void c128_forward(
 
   float score_fp32[kTileElements][kElementsPerWarp];
 
-  // kv/score already fp32 (converted at load); just add the bias
+  // convert to fp32 and apply bias first
 #pragma unroll
   for (int32_t i = 0; i < kTileElements; ++i) {
     for (int32_t j = 0; j < kElementsPerWarp; ++j) {
-      score_fp32[i][j] = score[j][i] + cast<float>(bias[j][i]);
+      score_fp32[i][j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
     }
   }
 
@@ -181,7 +163,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 += kv[j][i] * exp_score;
+      sum_product += cast<float>(kv[j][i]) * exp_score;
       sum_exp_value += exp_score;
     }
 
@@ -233,27 +215,25 @@ SGL_DEVICE void c128_forward(
   }
 }
 
-template <typename Trait, typename BufFloat, typename InFloat>
+template <typename Trait, typename InFloat>
-SGL_DEVICE void c128_write_decode(BufFloat* kv_buf, const InFloat* kv_src) {
+SGL_DEVICE void c128_write_decode(InFloat* kv_buf, const InFloat* kv_src) {
   using namespace device;
 
-  using StorageIn = AlignedVector<InFloat, kTileElements>;
+  using Storage = AlignedVector<InFloat, kTileElements>;
-  using StorageBuf = AlignedVector<BufFloat, kTileElements>;
+  const auto gmem = tile::Memory<Storage>::warp();
-  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) {
-    const auto d = gmem_in.load(kv_src + Trait::kHeadDim * i);
+    data[i] = gmem.load(kv_src + Trait::kHeadDim * i);
-    StorageBuf o;
+  }
 #pragma unroll
-    for (int32_t t = 0; t < kTileElements; ++t)
+  for (int32_t i = 0; i < 2; ++i) {
-      o[t] = cast<BufFloat>(d[t]);
+    gmem.store(kv_buf + Trait::kHeadDim * i, data[i]);
-    gmem_buf.store(kv_buf + Trait::kHeadDim * i, o);
   }
 }
 
-template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
 C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodeParams params) {
   using namespace device;
   using Trait = C128Trait<kHeadDim>;
@@ -267,7 +247,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<BufFloat*>(params.kv_buffer) + split_offset;
+  const auto kv_buffer = static_cast<InFloat*>(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;
@@ -278,15 +258,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, BufFloat, InFloat>(kv_dst, kv_src);
+    c128_write_decode<Trait>(kv_dst, kv_src);
   }
   if (plan.write_loc % 128 == 127) {
-    c128_forward<Trait, kUsePDL, BufFloat, InFloat, OutFloat>(kv_buf, kv_src, kv_out, score_bias, 128);
+    c128_forward<Trait, kUsePDL>(kv_buf, kv_src, kv_out, score_bias, 128);
   }
 }
 
 // compress kernel
-template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
 C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillParams params) {
   using namespace device;
   using Trait = C128Trait<kHeadDim>;
@@ -299,7 +279,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<BufFloat*>(params.kv_buffer) + split_offset;
+  const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
   const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
   if (plan.is_invalid()) return;
 
@@ -308,15 +288,14 @@ 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, BufFloat, InFloat, OutFloat>(kv_buf, kv_src, kv_out, score_bias, plan.buffer_len);
+  c128_forward<Trait, kUsePDL>(kv_buf, kv_src, kv_out, score_bias, plan.buffer_len);
 }
 
-template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, 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
@@ -329,37 +308,32 @@ 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<BufFloat*>(params.kv_buffer) + split_offset;
+  const auto kv_buffer = static_cast<InFloat*>(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_in = tile::Memory<StorageIn>::warp();
+  const auto gmem = 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_in.load(kv_src, i);
+    data[i] = gmem.load(kv_src, i);
   }
   PDLTriggerSecondary<kUsePDL>();
 #pragma unroll
   for (int32_t i = 0; i < 2; ++i) {
-    StorageBuf o;
+    gmem.store(kv_buf, data[i], i);
-#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 BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
 struct FlashCompress128Kernel {
-  static constexpr auto decode_kernel = flash_c128_decode<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
+  static constexpr auto decode_kernel = flash_c128_decode<kHeadDim, InFloat, OutFloat, kUsePDL>;
-  static constexpr auto prefill_c_kernel = flash_c128_prefill<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
+  static constexpr auto prefill_c_kernel = flash_c128_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
-  static constexpr auto prefill_w_kernel = write_c128_prefill<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
+  static constexpr auto prefill_w_kernel = write_c128_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
   static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads;  // 64
   static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
   using Trait = C128Trait<kHeadDim>;
@@ -377,7 +351,7 @@ struct FlashCompress128Kernel {
     device_.set_options<kDLGPU>();
 
     TensorMatcher({-1, 128, Trait::kElementSize})  // kv score
-        .with_dtype<BufFloat>()
+        .with_dtype<InFloat>()
         .with_device(device_)
         .verify(kv_buffer);
     TensorMatcher({N, Trait::kElementSize})  // kv score input
@@ -424,7 +398,7 @@ struct FlashCompress128Kernel {
     device_.set_options<kDLGPU>();
 
     TensorMatcher({-1, 128, Trait::kElementSize})  // kv score
-        .with_dtype<BufFloat>()
+        .with_dtype<InFloat>()
         .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,27 +74,23 @@ struct C4Trait {
   static_assert(kHeadDim % kTileDim == 0);
 };
 
-template <typename Trait, bool kUsePDL, typename BufFloat, typename InFloat, typename OutFloat>
+template <typename Trait, bool kUsePDL, typename InFloat, typename OutFloat>
 SGL_DEVICE void c4_forward(
-    const BufFloat* kv_buf_0,  // overlap [4n - 4, 4n - 1]
+    const InFloat* kv_buf_0,  // overlap [4n - 4, 4n - 1]
-    const BufFloat* kv_buf_1,  // normal [4n + 0, 4n + 3]
+    const InFloat* 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. kv_score_buffer (fp32, runtime state pool)
+  /// NOTE: part 1: load kv + score
-  /// 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();
-  float kv[8][kTileElements];
+  StorageIn kv[8];
-  float score[8][kTileElements];
+  StorageIn score[8];
   StorageIn bias[8];
 
 #pragma unroll
@@ -106,61 +102,28 @@ 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) {
-      if (i < buffer_len) {
+      const auto src = i < buffer_len ? kv_buf_0 : kv_start;
-        const auto base = kv_buf_0 + i * Trait::kElementSize;
+      const auto base = src + i * Trait::kElementSize;
-        const auto k = gmem_buf.load(base);
+      kv[i] = gmem_in.load(base);
-        const auto s = gmem_buf.load(base + Trait::kScoreOffset);
+      score[i] = 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 {
-        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) {
-#pragma unroll
+      kv[i].fill(cast<InFloat>(0.0f));
-      for (int32_t t = 0; t < kTileElements; ++t) {
+      score[i].fill(cast<InFloat>(kFloatNegInf));
-        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) {
-    if (i + 4 < buffer_len) {
+    const auto src = i + 4 < buffer_len ? kv_buf_1 : kv_start;
-      const auto base = kv_buf_1 + i * Trait::kElementSize + Trait::kOverlapOffset;
+    const auto base = src + i * Trait::kElementSize + Trait::kOverlapOffset;
-      const auto k = gmem_buf.load(base);
+    kv[i + 4] = gmem_in.load(base);
-      const auto s = gmem_buf.load(base + Trait::kScoreOffset);
+    score[i + 4] = 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]);
-      }
-    } 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
@@ -171,11 +134,11 @@ SGL_DEVICE void c4_forward(
   // consume 32 fp registers
   float score_fp32[kTileElements][8];
 
-  // kv/score already fp32 (converted at load); just add the bias
+  // convert to fp32 and apply bias first
 #pragma unroll
   for (int32_t i = 0; i < kTileElements; ++i) {
     for (int32_t j = 0; j < 8; ++j) {
-      score_fp32[i][j] = score[j][i] + cast<float>(bias[j][i]);
+      score_fp32[i][j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
     }
   }
 
@@ -196,7 +159,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 += kv[j][i] * exp_score;
+      sum_product += cast<float>(kv[j][i]) * exp_score;
       sum_exp_value += exp_score;
     }
 
@@ -208,27 +171,25 @@ SGL_DEVICE void c4_forward(
   gmem_out.store(kv_out, result);
 }
 
-template <typename Trait, typename BufFloat, typename InFloat>
+template <typename Trait, typename InFloat>
-SGL_DEVICE void c4_write_decode(BufFloat* kv_buf, const InFloat* kv_src) {
+SGL_DEVICE void c4_write_decode(InFloat* kv_buf, const InFloat* kv_src) {
   using namespace device;
 
   using StorageIn = AlignedVector<InFloat, kTileElements>;
-  using StorageBuf = AlignedVector<BufFloat, kTileElements>;
+  const auto gmem = tile::Memory<StorageIn>::warp();
-  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) {
-    const auto d = gmem_in.load(kv_src + Trait::kHeadDim * i);
+    data[i] = gmem.load(kv_src + Trait::kHeadDim * i);
-    StorageBuf o;
+  }
 #pragma unroll
-    for (int32_t t = 0; t < kTileElements; ++t)
+  for (int32_t i = 0; i < 4; ++i) {
-      o[t] = cast<BufFloat>(d[t]);
+    gmem.store(kv_buf + Trait::kHeadDim * i, data[i]);
-    gmem_buf.store(kv_buf + Trait::kHeadDim * i, o);
   }
 }
 
-template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
 C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams params) {
   using namespace device;
   using Trait = C4Trait<kHeadDim>;
@@ -243,7 +204,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<BufFloat*>(params.kv_buffer) + split_offset;
+  const auto kv_buffer = static_cast<InFloat*>(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;
@@ -253,15 +214,14 @@ 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, BufFloat, InFloat>(kv_dst, kv_src);
+  c4_write_decode<Trait>(kv_dst, kv_src);
   if (plan.seq_len % 4 == 0) {
     const auto need_overlap = plan.seq_len > 4;
-    c4_forward<Trait, kUsePDL, BufFloat, InFloat, OutFloat>(
+    c4_forward<Trait, kUsePDL>(kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, 8);
-        kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, 8);
   }
 }
 
-template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
 C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams params) {
   using namespace device;
   using Trait = C4Trait<kHeadDim>;
@@ -276,7 +236,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<BufFloat*>(params.kv_buffer) + split_offset;
+  const auto kv_buffer = static_cast<InFloat*>(params.kv_buffer) + split_offset;
   const auto score_bias = static_cast<const InFloat*>(params.score_bias) + split_offset;
   if (plan.is_invalid()) return;
 
@@ -287,16 +247,14 @@ 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, BufFloat, InFloat, OutFloat>(
+  c4_forward<Trait, kUsePDL>(kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, plan.buffer_len);
-      kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, plan.buffer_len);
 }
 
-template <int64_t kHeadDim, typename BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, 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
@@ -309,37 +267,32 @@ 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<BufFloat*>(params.kv_buffer) + split_offset;
+  const auto kv_buffer = static_cast<InFloat*>(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_in = tile::Memory<StorageIn>::warp();
+  const auto gmem = 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_in.load(kv_src, i);
+    data[i] = gmem.load(kv_src, i);
   }
   PDLTriggerSecondary<kUsePDL>();
 #pragma unroll
   for (int32_t i = 0; i < 4; ++i) {
-    StorageBuf o;
+    gmem.store(kv_buf, data[i], i);
-#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 BufFloat, typename InFloat, typename OutFloat, bool kUsePDL>
+template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
 struct FlashCompress4Kernel {
-  static constexpr auto decode_kernel = flash_c4_decode<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
+  static constexpr auto decode_kernel = flash_c4_decode<kHeadDim, InFloat, OutFloat, kUsePDL>;
-  static constexpr auto prefill_c_kernel = flash_c4_prefill<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
+  static constexpr auto prefill_c_kernel = flash_c4_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
-  static constexpr auto prefill_w_kernel = write_c4_prefill<kHeadDim, BufFloat, InFloat, OutFloat, kUsePDL>;
+  static constexpr auto prefill_w_kernel = write_c4_prefill<kHeadDim, InFloat, OutFloat, kUsePDL>;
   static constexpr uint32_t kBlockSize = 128;
   static constexpr uint32_t kTileDim = kTileElements * device::kWarpThreads;
   static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
@@ -359,7 +312,7 @@ struct FlashCompress4Kernel {
     device_.set_options<kDLGPU>();
 
     TensorMatcher({-1, 4, Trait::kElementSize})  // kv score
-        .with_dtype<BufFloat>()
+        .with_dtype<InFloat>()
         .with_device(device_)
         .verify(kv_buffer);
     TensorMatcher({N, Trait::kElementSize})  // kv score input
@@ -406,7 +359,7 @@ struct FlashCompress4Kernel {
     device_.set_options<kDLGPU>();
 
     TensorMatcher({-1, 4, Trait::kElementSize})  // kv score
-        .with_dtype<BufFloat>()
+        .with_dtype<InFloat>()
         .with_device(device_)
         .verify(kv_buffer);
     TensorMatcher({N, Trait::kElementSize})  // kv score input (ragged)

python/sglang/jit_kernel/dsv4/compress.py

@@ -44,14 +44,11 @@ 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(
+    args = make_cpp_args(head_dim, dtype_in, dtype_out, is_arch_support_pdl())
-        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"),
@@ -327,17 +324,8 @@ def compress_forward(
         assert compress_ratio == 128 and head_dim == 512
         module = _jit_compress_128_online_module(512)
     else:
-        # kv_score_buffer (fp32 runtime state pool) may differ from input/ape, so
+        dtype_in, dtype_out = kv_score_input.dtype, out.dtype
-        # the kernel keeps a BufFloat template and casts the buffer to fp32 at
+        module = _jit_compress_module(head_dim, dtype_in, dtype_out, compress_ratio)
-        # 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,10 +395,6 @@ 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)