diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh index 534527ba3..2055712e3 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh @@ -20,6 +20,7 @@ #include #include #include +#include namespace { @@ -40,7 +41,7 @@ struct Compress128OnlineDecodeParams { uint32_t batch_size; }; -template +template __global__ void flash_c128_online_decode_v2(const __grid_constant__ Compress128OnlineDecodeParams params) { using namespace device; constexpr uint32_t kVecSize = 4; @@ -59,7 +60,7 @@ __global__ void flash_c128_online_decode_v2(const __grid_constant__ Compress128O const auto plan = params.plan_d[batch_id]; const auto pos_in_chunk = (plan.seq_len - 1) % 128; - const auto kv_score_buffer = static_cast(params.kv_score_buffer); + const auto kv_score_buffer = static_cast(params.kv_score_buffer); const auto kv_score_input = static_cast(params.kv_score_input); const auto kv_load_buf = kv_score_buffer + plan.read_page_0 * (kHeadDim * 3); const auto kv_store_buf = kv_score_buffer + plan.write_loc * (kHeadDim * 3); @@ -75,9 +76,24 @@ __global__ void flash_c128_online_decode_v2(const __grid_constant__ Compress128O Vec out_sum_vec; if (pos_in_chunk != 0) { // Mid-chunk: combine prior partial state with the new token. - const auto max_score_vec = gmem.load(kv_load_buf, 0); - const auto sum_score_vec = gmem.load(kv_load_buf, 1); - const auto old_kv_vec = gmem.load(kv_load_buf, 2); + Vec max_score_vec, sum_score_vec, old_kv_vec; + if constexpr (std::is_same_v) { + max_score_vec = gmem.load(kv_load_buf, 0); + sum_score_vec = gmem.load(kv_load_buf, 1); + old_kv_vec = gmem.load(kv_load_buf, 2); + } else { + using BufferVec = AlignedVector; + const auto gmem_buffer = tile::Memory::cta(kBlockSize); + const auto max_score_tmp = gmem_buffer.load(kv_load_buf, 0); + const auto sum_score_tmp = gmem_buffer.load(kv_load_buf, 1); + const auto old_kv_tmp = gmem_buffer.load(kv_load_buf, 2); +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + max_score_vec[i] = cast(max_score_tmp[i]); + sum_score_vec[i] = cast(sum_score_tmp[i]); + old_kv_vec[i] = cast(old_kv_tmp[i]); + } + } #pragma unroll for (uint32_t i = 0; i < kVecSize; ++i) { const auto old_max = max_score_vec[i]; @@ -107,9 +123,24 @@ __global__ void flash_c128_online_decode_v2(const __grid_constant__ Compress128O const auto kv_out = static_cast(params.kv_compressed_output) + batch_id * kHeadDim; gmem.store(kv_out, out_kv_vec); } else { - gmem.store(kv_store_buf, out_max_vec, 0); - gmem.store(kv_store_buf, out_sum_vec, 1); - gmem.store(kv_store_buf, out_kv_vec, 2); + if constexpr (std::is_same_v) { + gmem.store(kv_store_buf, out_max_vec, 0); + gmem.store(kv_store_buf, out_sum_vec, 1); + gmem.store(kv_store_buf, out_kv_vec, 2); + } else { + using BufferVec = AlignedVector; + const auto gmem_buffer = tile::Memory::cta(kBlockSize); + BufferVec out_max_tmp, out_sum_tmp, out_kv_tmp; +#pragma unroll + for (uint32_t i = 0; i < kVecSize; ++i) { + out_max_tmp[i] = cast(out_max_vec[i]); + out_sum_tmp[i] = cast(out_sum_vec[i]); + out_kv_tmp[i] = cast(out_kv_vec[i]); + } + gmem_buffer.store(kv_store_buf, out_max_tmp, 0); + gmem_buffer.store(kv_store_buf, out_sum_tmp, 1); + gmem_buffer.store(kv_store_buf, out_kv_tmp, 2); + } } } @@ -150,8 +181,7 @@ struct Compress128SharedBuffer { /// \brief Sentinel score for padded positions in a 128-segment. constexpr float kPadScore = -FLT_MAX; -[[maybe_unused]] -SGL_DEVICE void c128_prefill_segment_softmax( +[[maybe_unused]] SGL_DEVICE void c128_prefill_segment_softmax( const PrefillStorage (&kv)[kElementsPerWarp], const PrefillStorage (&score)[kElementsPerWarp], float* seg_kv, @@ -239,7 +269,7 @@ SGL_DEVICE void c128_prefill_segment_softmax( /// `kWrite=true` (write pass) : handles trailing partial segments. /// Reads optional prior state from `read_page_1` (-1 = fallback to /// `read_page_0`), writes new running state to `read_page_0`. -template +template __global__ __launch_bounds__(kPrefillBlockSize, 2) // void flash_c128_online_prefill_v2(const __grid_constant__ Compress128OnlinePrefillParams params) { using namespace device; @@ -272,7 +302,7 @@ __global__ __launch_bounds__(kPrefillBlockSize, 2) // if (plan.is_invalid()) [[unlikely]] return; - const auto kv_score_buffer = static_cast(params.kv_score_buffer); + const auto kv_score_buffer = static_cast(params.kv_score_buffer); const auto kv_score_input = static_cast(params.kv_score_input); const auto kv_compressed_output = static_cast(params.kv_compressed_output); const auto score_bias_base = static_cast(params.score_bias); @@ -340,9 +370,23 @@ __global__ __launch_bounds__(kPrefillBlockSize, 2) // // Combine with prior partial state for this slot. const auto buf_load = kv_score_buffer + read_page * (kHeadDim * 3) + split_offset; PrefillStorage buf_max_vec, buf_sum_vec, buf_kv_vec; - buf_max_vec.load(buf_load + 0 * kHeadDim, lane_id); - buf_sum_vec.load(buf_load + 1 * kHeadDim, lane_id); - buf_kv_vec.load(buf_load + 2 * kHeadDim, lane_id); + if constexpr (std::is_same_v) { + buf_max_vec.load(buf_load + 0 * kHeadDim, lane_id); + buf_sum_vec.load(buf_load + 1 * kHeadDim, lane_id); + buf_kv_vec.load(buf_load + 2 * kHeadDim, lane_id); + } else { + using BufferPrefillStorage = AlignedVector; + BufferPrefillStorage buf_max_tmp, buf_sum_tmp, buf_kv_tmp; + buf_max_tmp.load(buf_load + 0 * kHeadDim, lane_id); + buf_sum_tmp.load(buf_load + 1 * kHeadDim, lane_id); + buf_kv_tmp.load(buf_load + 2 * kHeadDim, lane_id); +#pragma unroll + for (uint32_t ii = 0; ii < kTileElements; ++ii) { + buf_max_vec[ii] = cast(buf_max_tmp[ii]); + buf_sum_vec[ii] = cast(buf_sum_tmp[ii]); + buf_kv_vec[ii] = cast(buf_kv_tmp[ii]); + } + } #pragma unroll for (uint32_t ii = 0; ii < kTileElements; ++ii) { const float m1 = buf_max_vec[ii]; @@ -367,9 +411,23 @@ __global__ __launch_bounds__(kPrefillBlockSize, 2) // // segment's own chunk slot (the request keeps a single in-progress // chunk's running state at any time), so we reuse `read_page_0`. const auto buf_store = kv_score_buffer + plan.read_page_0 * (kHeadDim * 3) + split_offset; - reinterpret_cast(buf_store + 0 * kHeadDim)[lane_id] = out_max_vec; - reinterpret_cast(buf_store + 1 * kHeadDim)[lane_id] = out_sum_vec; - reinterpret_cast(buf_store + 2 * kHeadDim)[lane_id] = out_kv_vec; + if constexpr (std::is_same_v) { + reinterpret_cast(buf_store + 0 * kHeadDim)[lane_id] = out_max_vec; + reinterpret_cast(buf_store + 1 * kHeadDim)[lane_id] = out_sum_vec; + reinterpret_cast(buf_store + 2 * kHeadDim)[lane_id] = out_kv_vec; + } else { + using BufferPrefillStorage = AlignedVector; + BufferPrefillStorage out_max_tmp, out_sum_tmp, out_kv_tmp; +#pragma unroll + for (uint32_t ii = 0; ii < kTileElements; ++ii) { + out_max_tmp[ii] = cast(out_max_vec[ii]); + out_sum_tmp[ii] = cast(out_sum_vec[ii]); + out_kv_tmp[ii] = cast(out_kv_vec[ii]); + } + reinterpret_cast(buf_store + 0 * kHeadDim)[lane_id] = out_max_tmp; + reinterpret_cast(buf_store + 1 * kHeadDim)[lane_id] = out_sum_tmp; + reinterpret_cast(buf_store + 2 * kHeadDim)[lane_id] = out_kv_tmp; + } } else { // Compact output: one row per compress plan, indexed by `global_pid`. const auto out_ptr = kv_compressed_output + global_pid * kHeadDim + split_offset; @@ -381,14 +439,15 @@ __global__ __launch_bounds__(kPrefillBlockSize, 2) // // --------------------------------------------------------------------------- // Host wrapper: matches the c128_v2 / c4_v2 host API style (run_decode / // run_prefill methods on a kernel-class template). We only expose `kHeadDim` -// + `kUsePDL`; the dtype is fixed to fp32 for the online state pool. +// + state-buffer dtype + `kUsePDL`. Inputs, outputs, and APE remain fp32; only +// the online C128 running state buffer can be fp32/bf16. // --------------------------------------------------------------------------- -template +template struct FlashCompress128OnlineKernel { - static constexpr auto decode_kernel = flash_c128_online_decode_v2; + static constexpr auto decode_kernel = flash_c128_online_decode_v2; template - static constexpr auto prefill_kernel = flash_c128_online_prefill_v2; + static constexpr auto prefill_kernel = flash_c128_online_prefill_v2; static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads; // 64 static constexpr uint32_t kNumSplit = kHeadDim / kTileDim; static constexpr uint32_t kDecodeBlockSize = kHeadDim / 4; @@ -406,7 +465,7 @@ struct FlashCompress128OnlineKernel { device_.set_options(); TensorMatcher({-1, 1, kHeadDim * 3}) // kv score buffer (max, sum, kv) - .with_dtype() + .with_dtype() .with_device(device_) .verify(kv_score_buffer); TensorMatcher({B, kHeadDim * 2}) // kv score input @@ -453,7 +512,7 @@ struct FlashCompress128OnlineKernel { device_.set_options(); TensorMatcher({-1, 1, kHeadDim * 3}) // kv score buffer - .with_dtype() + .with_dtype() .with_device(device_) .verify(kv_score_buffer); TensorMatcher({N, kHeadDim * 2}) // kv score input (ragged) @@ -868,9 +927,7 @@ inline OnlinePrefillPlan plan_online_prefill( namespace { -[[maybe_unused]] -constexpr auto& plan_compress_128_online_decode = host::compress::plan_online_decode; -[[maybe_unused]] -constexpr auto& plan_compress_128_online_prefill = host::compress::plan_online_prefill; +[[maybe_unused]] constexpr auto& plan_compress_128_online_decode = host::compress::plan_online_decode; +[[maybe_unused]] constexpr auto& plan_compress_128_online_prefill = host::compress::plan_online_prefill; } // namespace diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh index 72eeea609..44d695ad2 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh @@ -4,11 +4,13 @@ #include #include +#include #include #include #include +#include namespace { @@ -17,14 +19,14 @@ SGL_DEVICE int64_t clamp_accept_len(int64_t delta, int64_t max_accept) { return delta < max_accept ? delta : max_accept; } -template +template struct OnlineC128MTPWritePrefixParams { const float* __restrict__ kv_score_input; const TSeq* __restrict__ seq_lens; const TReq* __restrict__ req_pool_indices; const int32_t* __restrict__ req_to_token; const float* __restrict__ ape; - float* __restrict__ state; + BufferFloat* __restrict__ state; int64_t kv_score_stride_b; int64_t req_to_token_stride_b; int64_t ape_stride_r; @@ -43,13 +45,13 @@ struct OnlineC128MTPMarkPendingParams { int64_t max_num_reqs; }; -template +template struct OnlineC128MTPCommitPendingParams { const TSeq* __restrict__ cur_seq_lens; const TReq* __restrict__ cur_req_pool_indices; const int32_t* __restrict__ req_to_token; const int64_t* __restrict__ pending_seq_lens; - float* __restrict__ state; + BufferFloat* __restrict__ state; int64_t cur_bs; int64_t req_to_token_stride_b; int64_t state_stride_b; @@ -73,8 +75,9 @@ __global__ void online_c128_mtp_mark_pending_kernel(const OnlineC128MTPMarkPendi } } -template -__global__ void online_c128_mtp_commit_pending_kernel(const OnlineC128MTPCommitPendingParams params) { +template +__global__ void +online_c128_mtp_commit_pending_kernel(const OnlineC128MTPCommitPendingParams params) { const int64_t bid = static_cast(blockIdx.x); if (bid >= params.cur_bs) return; @@ -91,16 +94,17 @@ __global__ void online_c128_mtp_commit_pending_kernel(const OnlineC128MTPCommitP if ((final_seq & 127) == 0) return; const int64_t slot = req; - const float* const src = params.state + (slot + accept * params.state_slot_stride) * params.state_stride_b; - float* const dst = params.state + slot * params.state_stride_b; + const BufferFloat* const src = params.state + (slot + accept * params.state_slot_stride) * params.state_stride_b; + BufferFloat* const dst = params.state + slot * params.state_stride_b; for (int64_t d = static_cast(threadIdx.x); d < kHeadDim * 3; d += blockDim.x) { dst[d] = src[d]; } } -template -__global__ void online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePrefixParams params) { +template +__global__ void +online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePrefixParams params) { const int64_t bid = static_cast(blockIdx.x); if (bid >= params.layer_bs) return; @@ -119,10 +123,17 @@ __global__ void online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePref float run_sum = 0.0f; float run_kv = 0.0f; if (has_partial) { - const float* const init = params.state + init_slot * params.state_stride_b; - run_max = init[d]; - run_sum = init[kHeadDim + d]; - run_kv = init[kHeadDim * 2 + d]; + if constexpr (std::is_same_v) { + const float* const init = params.state + init_slot * params.state_stride_b; + run_max = init[d]; + run_sum = init[kHeadDim + d]; + run_kv = init[kHeadDim * 2 + d]; + } else { + const BufferFloat* const init = params.state + init_slot * params.state_stride_b; + run_max = device::cast(init[d]); + run_sum = device::cast(init[kHeadDim + d]); + run_kv = device::cast(init[kHeadDim * 2 + d]); + } } constexpr int kMaxVerifyTokens = 8; @@ -163,10 +174,17 @@ __global__ void online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePref const int64_t final_seq = seq_before + step + 1; if ((final_seq & 127) != 0) { const int64_t slot = req_idx + (step + 1) * params.state_slot_stride; - float* const out = params.state + slot * params.state_stride_b; - out[d] = run_max; - out[kHeadDim + d] = run_sum; - out[kHeadDim * 2 + d] = run_kv; + if constexpr (std::is_same_v) { + float* const out = params.state + slot * params.state_stride_b; + out[d] = run_max; + out[kHeadDim + d] = run_sum; + out[kHeadDim * 2 + d] = run_kv; + } else { + BufferFloat* const out = params.state + slot * params.state_stride_b; + out[d] = device::cast(run_max); + out[kHeadDim + d] = device::cast(run_sum); + out[kHeadDim * 2 + d] = device::cast(run_kv); + } } if (pos == 127) { @@ -177,7 +195,7 @@ __global__ void online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePref } } -template +template struct OnlineC128MTPWritePrefixKernel { static void launch( tvm::ffi::TensorView kv_score_input, @@ -192,13 +210,13 @@ struct OnlineC128MTPWritePrefixKernel { DLDevice device) { using namespace host; - const auto params = OnlineC128MTPWritePrefixParams{ + const auto params = OnlineC128MTPWritePrefixParams{ .kv_score_input = static_cast(kv_score_input.data_ptr()), .seq_lens = static_cast(seq_lens.data_ptr()), .req_pool_indices = static_cast(req_pool_indices.data_ptr()), .req_to_token = static_cast(req_to_token.data_ptr()), .ape = static_cast(ape.data_ptr()), - .state = static_cast(state.data_ptr()), + .state = static_cast(state.data_ptr()), .kv_score_stride_b = kv_score_input.stride(0), .req_to_token_stride_b = req_to_token.stride(0), .ape_stride_r = ape.stride(0), @@ -211,7 +229,7 @@ struct OnlineC128MTPWritePrefixKernel { static_assert(kHeadDim == 512, "online c128 MTP write-prefix only supports head_dim=512"); constexpr uint32_t kThreads = static_cast(kHeadDim); LaunchKernel(static_cast(layer_bs), kThreads, device)( - online_c128_mtp_write_prefix_kernel, params); + online_c128_mtp_write_prefix_kernel, params); } static void @@ -234,7 +252,7 @@ struct OnlineC128MTPWritePrefixKernel { TensorMatcher({-1}).with_dtype().with_device(device).verify(req_pool_indices); TensorMatcher({-1, -1}).with_dtype().with_device(device).verify(req_to_token); TensorMatcher({128, kHeadDim}).with_dtype().with_device(device).verify(ape); - TensorMatcher({-1, kHeadDim * 3}).with_dtype().with_device(device).verify(state); + TensorMatcher({-1, kHeadDim * 3}).with_dtype().with_device(device).verify(state); if (layer_bs <= 0) return; RuntimeCheck(num_verify_tokens > 0 && num_verify_tokens <= 8, "unsupported num_verify_tokens=", num_verify_tokens); @@ -257,7 +275,7 @@ struct OnlineC128MTPWritePrefixKernel { } }; -template +template struct OnlineC128MTPMarkPendingKernel { static void launch( tvm::ffi::TensorView seq_lens, @@ -308,7 +326,7 @@ struct OnlineC128MTPMarkPendingKernel { } }; -template +template struct OnlineC128MTPCommitPendingKernel { static void launch( tvm::ffi::TensorView cur_seq_lens, @@ -323,12 +341,12 @@ struct OnlineC128MTPCommitPendingKernel { DLDevice device) { using namespace host; - const auto params = OnlineC128MTPCommitPendingParams{ + const auto params = OnlineC128MTPCommitPendingParams{ .cur_seq_lens = static_cast(cur_seq_lens.data_ptr()), .cur_req_pool_indices = static_cast(cur_req_pool_indices.data_ptr()), .req_to_token = static_cast(req_to_token.data_ptr()), .pending_seq_lens = static_cast(pending_seq_lens.data_ptr()), - .state = static_cast(state.data_ptr()), + .state = static_cast(state.data_ptr()), .cur_bs = cur_bs, .req_to_token_stride_b = req_to_token.stride(0), .state_stride_b = state.stride(0), @@ -339,7 +357,7 @@ struct OnlineC128MTPCommitPendingKernel { constexpr uint32_t kThreads = 256; LaunchKernel(static_cast(cur_bs), kThreads, device)( - online_c128_mtp_commit_pending_kernel, params); + online_c128_mtp_commit_pending_kernel, params); } static void @@ -361,7 +379,7 @@ struct OnlineC128MTPCommitPendingKernel { TensorMatcher({-1}).with_dtype().with_device(device).verify(cur_req_pool_indices); TensorMatcher({-1, -1}).with_dtype().with_device(device).verify(req_to_token); TensorMatcher({-1}).with_dtype().with_device(device).verify(pending_seq_lens); - TensorMatcher({-1, kHeadDim * 3}).with_dtype().with_device(device).verify(state); + TensorMatcher({-1, kHeadDim * 3}).with_dtype().with_device(device).verify(state); if (cur_bs <= 0) return; RuntimeCheck(num_verify_tokens > 0 && num_verify_tokens <= 8, "unsupported num_verify_tokens=", num_verify_tokens); diff --git a/python/sglang/jit_kernel/dsv4/compress.py b/python/sglang/jit_kernel/dsv4/compress.py index 7c41d5f74..d07636bdf 100644 --- a/python/sglang/jit_kernel/dsv4/compress.py +++ b/python/sglang/jit_kernel/dsv4/compress.py @@ -66,9 +66,11 @@ def _jit_compress_module( @cache_once -def _jit_compress_128_online_module(head_dim: int) -> Module: +def _jit_compress_128_online_module( + head_dim: int, dtype_buffer: torch.dtype = torch.float32 +) -> Module: assert head_dim == 512 - args = make_cpp_args(head_dim, is_arch_support_pdl()) + args = make_cpp_args(head_dim, dtype_buffer, is_arch_support_pdl()) kernel_class = f"FlashCompress128OnlineKernel<{args}>" return load_jit( make_name(f"compress_128_online_v2"), @@ -327,7 +329,7 @@ def compress_forward( assert plan.compress_ratio == compress_ratio if is_online: assert compress_ratio == 128 and head_dim == 512 - module = _jit_compress_128_online_module(512) + module = _jit_compress_128_online_module(512, kv_score_buffer.dtype) else: dtype_in, dtype_out = kv_score_input.dtype, out.dtype module = _jit_compress_module( diff --git a/python/sglang/jit_kernel/dsv4/online_c128_mtp.py b/python/sglang/jit_kernel/dsv4/online_c128_mtp.py index d04c34f13..5a901b4d9 100644 --- a/python/sglang/jit_kernel/dsv4/online_c128_mtp.py +++ b/python/sglang/jit_kernel/dsv4/online_c128_mtp.py @@ -15,9 +15,12 @@ if TYPE_CHECKING: @cache_once def _jit_online_c128_mtp_module( - head_dim: int, seq_dtype: torch.dtype, req_dtype: torch.dtype + head_dim: int, + seq_dtype: torch.dtype, + req_dtype: torch.dtype, + dtype_buffer: torch.dtype, ) -> Module: - args = make_cpp_args(head_dim, seq_dtype, req_dtype) + args = make_cpp_args(head_dim, seq_dtype, req_dtype, dtype_buffer) return load_jit( make_name(f"online_c128_mtp_{head_dim}"), *args, @@ -35,6 +38,7 @@ def _jit_online_c128_mtp_module( class _OnlineC128LayerRuntime: head_dim: int main_state: torch.Tensor + state_dtype: torch.dtype state_slot_offset: int @@ -76,11 +80,12 @@ class OnlineC128MTPController: seq_lens=seq_lens.detach(), ) head_dim = self._head_dim() - if head_dim is None or self._num_verify_tokens() == 0: + state_dtype = self._state_dtype() + if head_dim is None or state_dtype is None or self._num_verify_tokens() == 0: return token_to_kv_pool = self.backend.token_to_kv_pool _jit_online_c128_mtp_module( - head_dim, seq_lens.dtype, req_pool_indices.dtype + head_dim, seq_lens.dtype, req_pool_indices.dtype, state_dtype ).mark_pending( seq_lens, req_pool_indices, @@ -156,20 +161,21 @@ class OnlineC128MTPController: token_to_kv_pool = self.backend.token_to_kv_pool head_dim = compressor.head_dim state_pool = token_to_kv_pool.get_attention_compress_states(layer_id) + state = state_pool.kv_score_buffer.kv_score total_bs = kv_score_input.numel() // (num_verify_tokens * head_dim * 2) layer_bs = min(ctx.seq_lens.shape[0], ctx.req_pool_indices.shape[0], total_bs) if layer_bs <= 0: return _jit_online_c128_mtp_module( - head_dim, ctx.seq_lens.dtype, ctx.req_pool_indices.dtype + head_dim, ctx.seq_lens.dtype, ctx.req_pool_indices.dtype, state.dtype ).write_prefix_states( kv_score_input, ctx.seq_lens, ctx.req_pool_indices, self.backend.req_to_token, compressor.ape.reshape(128, head_dim), - state_pool.kv_score_buffer.kv_score, + state, layer_bs, num_verify_tokens, state_pool.online_mtp_state_slot_offset, @@ -200,7 +206,10 @@ class OnlineC128MTPController: for runtime in self._iter_layer_runtimes(): _jit_online_c128_mtp_module( - runtime.head_dim, seq_lens.dtype, req_pool_indices.dtype + runtime.head_dim, + seq_lens.dtype, + req_pool_indices.dtype, + runtime.state_dtype, ).commit_pending( seq_lens, req_pool_indices, @@ -239,6 +248,11 @@ class OnlineC128MTPController: return runtime.head_dim return None + def _state_dtype(self) -> Optional[torch.dtype]: + for runtime in self._iter_layer_runtimes(): + return runtime.state_dtype + return None + def _iter_layer_runtimes(self): if self._layer_runtimes is None: runtimes = [] @@ -255,6 +269,7 @@ class OnlineC128MTPController: _OnlineC128LayerRuntime( head_dim=compressor.head_dim, main_state=state_pool.kv_score_buffer.kv_score, + state_dtype=state_pool.kv_score_buffer.kv_score.dtype, state_slot_offset=state_pool.online_mtp_state_slot_offset, ) ) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index c5cf3a993..29db2d9f6 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -74,11 +74,6 @@ def _get_dsv4_compress_state_dtypes() -> tuple[torch.dtype, torch.dtype]: if dtype_name in ("float32", "fp32"): return torch.float32, torch.float32 if dtype_name in ("bfloat16", "bf16"): - if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get(): - raise ValueError( - "SGLANG_DSV4_COMPRESS_STATE_DTYPE=bf16 is not supported when " - "SGLANG_OPT_USE_ONLINE_COMPRESS=1; online c128 state must stay float32." - ) return torch.bfloat16, torch.bfloat16 raise ValueError( "Unsupported SGLANG_DSV4_COMPRESS_STATE_DTYPE=" diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 58ac874b7..e728a531f 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -78,11 +78,6 @@ def _get_dsv4_compress_state_dtype_sizes() -> tuple[int, int]: if dtype_name in ("float32", "fp32"): return 4, 4 if dtype_name in ("bfloat16", "bf16"): - if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get(): - raise ValueError( - "SGLANG_DSV4_COMPRESS_STATE_DTYPE=bf16 is not supported when " - "SGLANG_OPT_USE_ONLINE_COMPRESS=1; online c128 state must stay float32." - ) return 2, 2 raise ValueError( "Unsupported SGLANG_DSV4_COMPRESS_STATE_DTYPE=" diff --git a/test/registered/jit/benchmark/bench_online_c128_mtp.py b/test/registered/jit/benchmark/bench_online_c128_mtp.py index 80f722f98..9164f272a 100644 --- a/test/registered/jit/benchmark/bench_online_c128_mtp.py +++ b/test/registered/jit/benchmark/bench_online_c128_mtp.py @@ -148,7 +148,7 @@ def benchmark( ) -> tuple[float, float, float]: case = make_case(batch_size, num_verify_tokens) module = _jit_online_c128_mtp_module( - HEAD_DIM, case.seq_lens.dtype, case.req_pool_indices.dtype + HEAD_DIM, case.seq_lens.dtype, case.req_pool_indices.dtype, case.state.dtype ) fn = lambda: call_write_prefix(module, case)