diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh index 43525c7db..31353e6a1 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh @@ -89,10 +89,10 @@ struct C128Trait { static_assert(kHeadDim % kTileDim == 0); }; -template +template SGL_DEVICE void c128_forward( - const BufFloat* kv_buf, // [128n, 128n + 127] - const InFloat* kv_src, // ragged pointer at position = 128n + 127 + const InFloat* kv_buf, // [128n, 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) - /// 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; + /// NOTE: part 1: load kv + score using StorageIn = AlignedVector; - const auto gmem_buf = tile::Memory{lane_id, kWarpThreads}; const auto gmem_in = tile::Memory{lane_id, kWarpThreads}; - float kv[kElementsPerWarp][kTileElements]; - float score[kElementsPerWarp][kTileElements]; + StorageIn kv[kElementsPerWarp]; + 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 k = gmem_buf.load(kv_buf + j * Trait::kElementSize); - 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(k[t]); - score[i][t] = cast(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(k[t]); - score[i][t] = cast(s[t]); - } - } + const auto src = j < buffer_len ? kv_buf : kv_start; + kv[i] = gmem_in.load(src + j * Trait::kElementSize); + score[i] = gmem_in.load(src + j * Trait::kElementSize + Trait::kScoreOffset); } /// 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(bias[j][i]); + score_fp32[i][j] = cast(score[j][i]) + cast(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(kv[j][i]) * exp_score; sum_exp_value += exp_score; } @@ -233,27 +215,25 @@ SGL_DEVICE void c128_forward( } } -template -SGL_DEVICE void c128_write_decode(BufFloat* kv_buf, const InFloat* kv_src) { +template +SGL_DEVICE void c128_write_decode(InFloat* kv_buf, const InFloat* kv_src) { using namespace device; - using StorageIn = AlignedVector; - using StorageBuf = AlignedVector; - const auto gmem_in = tile::Memory::warp(); - const auto gmem_buf = tile::Memory::warp(); + using Storage = AlignedVector; + const auto gmem = tile::Memory::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); - StorageBuf o; + data[i] = gmem.load(kv_src + Trait::kHeadDim * i); + } #pragma unroll - for (int32_t t = 0; t < kTileElements; ++t) - o[t] = cast(d[t]); - gmem_buf.store(kv_buf + Trait::kHeadDim * i, o); + for (int32_t i = 0; i < 2; ++i) { + gmem.store(kv_buf + Trait::kHeadDim * i, data[i]); } } -template +template C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodeParams params) { using namespace device; using Trait = C128Trait; @@ -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(params.kv_input) + split_offset; const auto kv_output = static_cast(params.kv_output) + split_offset; - const auto kv_buffer = static_cast(params.kv_buffer) + split_offset; + const auto kv_buffer = static_cast(params.kv_buffer) + split_offset; const auto score_bias = static_cast(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(); // the write warp must match the load warp in the following `c128_forward` if (warp_id == kNumWarps - 1) { - c128_write_decode(kv_dst, kv_src); + c128_write_decode(kv_dst, kv_src); } if (plan.write_loc % 128 == 127) { - c128_forward(kv_buf, kv_src, kv_out, score_bias, 128); + c128_forward(kv_buf, kv_src, kv_out, score_bias, 128); } } // compress kernel -template +template C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillParams params) { using namespace device; using Trait = C128Trait; @@ -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(params.kv_input) + split_offset; const auto kv_output = static_cast(params.kv_output) + split_offset; - const auto kv_buffer = static_cast(params.kv_buffer) + split_offset; + const auto kv_buffer = static_cast(params.kv_buffer) + split_offset; const auto score_bias = static_cast(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(); - c128_forward(kv_buf, kv_src, kv_out, score_bias, plan.buffer_len); + c128_forward(kv_buf, kv_src, kv_out, score_bias, plan.buffer_len); } -template +template WRITE_KERNEL void write_c128_prefill(const __grid_constant__ Compress128PrefillParams params) { using namespace device; using Trait = C128Trait; using StorageIn = AlignedVector; - using StorageBuf = AlignedVector; 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(params.kv_input) + split_offset; - const auto kv_buffer = static_cast(params.kv_buffer) + split_offset; + const auto kv_buffer = static_cast(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::warp(); - const auto gmem_buf = tile::Memory::warp(); + const auto gmem = tile::Memory::warp(); PDLWaitPrimary(); 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(); #pragma unroll for (int32_t i = 0; i < 2; ++i) { - StorageBuf o; -#pragma unroll - for (int32_t t = 0; t < kTileElements; ++t) - o[t] = cast(data[i][t]); - gmem_buf.store(kv_buf, o, i); + gmem.store(kv_buf, data[i], i); } } -template +template struct FlashCompress128Kernel { - static constexpr auto decode_kernel = flash_c128_decode; - static constexpr auto prefill_c_kernel = flash_c128_prefill; - static constexpr auto prefill_w_kernel = write_c128_prefill; + static constexpr auto decode_kernel = flash_c128_decode; + static constexpr auto prefill_c_kernel = flash_c128_prefill; + static constexpr auto prefill_w_kernel = write_c128_prefill; static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads; // 64 static constexpr uint32_t kNumSplit = kHeadDim / kTileDim; using Trait = C128Trait; @@ -377,7 +351,7 @@ struct FlashCompress128Kernel { device_.set_options(); TensorMatcher({-1, 128, Trait::kElementSize}) // kv score - .with_dtype() + .with_dtype() .with_device(device_) .verify(kv_buffer); TensorMatcher({N, Trait::kElementSize}) // kv score input @@ -424,7 +398,7 @@ struct FlashCompress128Kernel { device_.set_options(); TensorMatcher({-1, 128, Trait::kElementSize}) // kv score - .with_dtype() + .with_dtype() .with_device(device_) .verify(kv_buffer); TensorMatcher({N, Trait::kElementSize}) // kv score input (ragged) diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh index e844e1031..efa9f0510 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh @@ -74,27 +74,23 @@ struct C4Trait { static_assert(kHeadDim % kTileDim == 0); }; -template +template SGL_DEVICE void c4_forward( - const BufFloat* kv_buf_0, // overlap [4n - 4, 4n - 1] - const BufFloat* kv_buf_1, // normal [4n + 0, 4n + 3] - const InFloat* kv_src, // ragged pointer at position = 4n + 3 + const InFloat* kv_buf_0, // overlap [4n - 4, 4n - 1] + const InFloat* kv_buf_1, // normal [4n + 0, 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) - /// 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; + /// NOTE: part 1: load kv + score using StorageIn = AlignedVector; /// NOTE: load one tile_dim (< head_dim) at at time - const auto gmem_buf = tile::Memory::warp(); const auto gmem_in = tile::Memory::warp(); - float kv[8][kTileElements]; - float score[8][kTileElements]; + StorageIn kv[8]; + 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 base = kv_buf_0 + i * Trait::kElementSize; - const auto k = gmem_buf.load(base); - const auto s = gmem_buf.load(base + Trait::kScoreOffset); -#pragma unroll - for (int32_t t = 0; t < kTileElements; ++t) { - kv[i][t] = cast(k[t]); - score[i][t] = cast(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(k[t]); - score[i][t] = cast(s[t]); - } - } + const auto src = i < buffer_len ? kv_buf_0 : kv_start; + const auto base = src + i * Trait::kElementSize; + kv[i] = gmem_in.load(base); + score[i] = gmem_in.load(base + Trait::kScoreOffset); } } else { [[unlikely]]; constexpr float kFloatNegInf = -FLT_MAX; #pragma unroll for (int32_t i = 0; i < 4; ++i) { -#pragma unroll - for (int32_t t = 0; t < kTileElements; ++t) { - kv[i][t] = 0.0f; - score[i][t] = kFloatNegInf; - } + kv[i].fill(cast(0.0f)); + score[i].fill(cast(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 base = kv_buf_1 + i * Trait::kElementSize + Trait::kOverlapOffset; - const auto k = gmem_buf.load(base); - const auto s = gmem_buf.load(base + Trait::kScoreOffset); -#pragma unroll - for (int32_t t = 0; t < kTileElements; ++t) { - kv[i + 4][t] = cast(k[t]); - score[i + 4][t] = cast(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(k[t]); - score[i + 4][t] = cast(s[t]); - } - } + const auto src = i + 4 < buffer_len ? kv_buf_1 : kv_start; + const auto base = src + i * Trait::kElementSize + Trait::kOverlapOffset; + kv[i + 4] = gmem_in.load(base); + score[i + 4] = gmem_in.load(base + Trait::kScoreOffset); } /// 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(bias[j][i]); + score_fp32[i][j] = cast(score[j][i]) + cast(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(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 -SGL_DEVICE void c4_write_decode(BufFloat* kv_buf, const InFloat* kv_src) { +template +SGL_DEVICE void c4_write_decode(InFloat* kv_buf, const InFloat* kv_src) { using namespace device; using StorageIn = AlignedVector; - using StorageBuf = AlignedVector; - const auto gmem_in = tile::Memory::warp(); - const auto gmem_buf = tile::Memory::warp(); + const auto gmem = tile::Memory::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); - StorageBuf o; + data[i] = gmem.load(kv_src + Trait::kHeadDim * i); + } #pragma unroll - for (int32_t t = 0; t < kTileElements; ++t) - o[t] = cast(d[t]); - gmem_buf.store(kv_buf + Trait::kHeadDim * i, o); + for (int32_t i = 0; i < 4; ++i) { + gmem.store(kv_buf + Trait::kHeadDim * i, data[i]); } } -template +template C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams params) { using namespace device; using Trait = C4Trait; @@ -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(params.kv_input) + split_offset; const auto kv_output = static_cast(params.kv_output) + split_offset; - const auto kv_buffer = static_cast(params.kv_buffer) + split_offset; + const auto kv_buffer = static_cast(params.kv_buffer) + split_offset; const auto score_bias = static_cast(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(); - c4_write_decode(kv_dst, kv_src); + c4_write_decode(kv_dst, kv_src); if (plan.seq_len % 4 == 0) { const auto need_overlap = plan.seq_len > 4; - c4_forward( - kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, 8); + c4_forward(kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, 8); } } -template +template C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams params) { using namespace device; using Trait = C4Trait; @@ -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(params.kv_input) + split_offset; const auto kv_output = static_cast(params.kv_output) + split_offset; - const auto kv_buffer = static_cast(params.kv_buffer) + split_offset; + const auto kv_buffer = static_cast(params.kv_buffer) + split_offset; const auto score_bias = static_cast(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(); - c4_forward( - kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, plan.buffer_len); + c4_forward(kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, plan.buffer_len); } -template +template WRITE_KERNEL void write_c4_prefill(const __grid_constant__ Compress4PrefillParams params) { using namespace device; using Trait = C4Trait; using StorageIn = AlignedVector; - using StorageBuf = AlignedVector; 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(params.kv_input) + split_offset; - const auto kv_buffer = static_cast(params.kv_buffer) + split_offset; + const auto kv_buffer = static_cast(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::warp(); - const auto gmem_buf = tile::Memory::warp(); + const auto gmem = tile::Memory::warp(); PDLWaitPrimary(); 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(); #pragma unroll for (int32_t i = 0; i < 4; ++i) { - StorageBuf o; -#pragma unroll - for (int32_t t = 0; t < kTileElements; ++t) - o[t] = cast(data[i][t]); - gmem_buf.store(kv_buf, o, i); + gmem.store(kv_buf, data[i], i); } } -template +template struct FlashCompress4Kernel { - static constexpr auto decode_kernel = flash_c4_decode; - static constexpr auto prefill_c_kernel = flash_c4_prefill; - static constexpr auto prefill_w_kernel = write_c4_prefill; + static constexpr auto decode_kernel = flash_c4_decode; + static constexpr auto prefill_c_kernel = flash_c4_prefill; + static constexpr auto prefill_w_kernel = write_c4_prefill; 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(); TensorMatcher({-1, 4, Trait::kElementSize}) // kv score - .with_dtype() + .with_dtype() .with_device(device_) .verify(kv_buffer); TensorMatcher({N, Trait::kElementSize}) // kv score input @@ -406,7 +359,7 @@ struct FlashCompress4Kernel { device_.set_options(); TensorMatcher({-1, 4, Trait::kElementSize}) // kv score - .with_dtype() + .with_dtype() .with_device(device_) .verify(kv_buffer); TensorMatcher({N, Trait::kElementSize}) // kv score input (ragged) diff --git a/python/sglang/jit_kernel/dsv4/compress.py b/python/sglang/jit_kernel/dsv4/compress.py index 86c0c06ea..ea3d8f729 100644 --- a/python/sglang/jit_kernel/dsv4/compress.py +++ b/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( - head_dim, dtype_buf, dtype_in, dtype_out, is_arch_support_pdl() - ) + args = make_cpp_args(head_dim, 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 - # 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, - ) + dtype_in, dtype_out = kv_score_input.dtype, out.dtype + module = _jit_compress_module(head_dim, dtype_in, dtype_out, 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 diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index 5d986bb4b..73a58e52f 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/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)