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 31353e6a1..1e2f2051c 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh @@ -27,7 +27,9 @@ #include #include +#include #include +#include namespace { @@ -89,12 +91,12 @@ struct C128Trait { static_assert(kHeadDim % kTileDim == 0); }; -template +template SGL_DEVICE void c128_forward( - const InFloat* kv_buf, // [128n, 128n + 127] - const InFloat* kv_src, // ragged pointer at position = 128n + 127 + const BufferFloat* kv_buf, // [128n, 128n + 127] + const InputFloat* kv_src, // ragged pointer at position = 128n + 127 OutFloat* kv_out, - const InFloat* score_bias, + const InputFloat* score_bias, const int32_t buffer_len) { using namespace device; @@ -102,7 +104,7 @@ SGL_DEVICE void c128_forward( const auto lane_id = threadIdx.x % kWarpThreads; /// NOTE: part 1: load kv + score - using StorageIn = AlignedVector; + using StorageIn = AlignedVector; const auto gmem_in = tile::Memory{lane_id, kWarpThreads}; StorageIn kv[kElementsPerWarp]; StorageIn score[kElementsPerWarp]; @@ -117,13 +119,38 @@ SGL_DEVICE void c128_forward( const auto kv_start = kv_src - 127 * Trait::kElementSize; // point to start + if constexpr (std::is_same_v) { #pragma unroll - 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; - kv[i] = gmem_in.load(src + j * Trait::kElementSize); - score[i] = gmem_in.load(src + j * Trait::kElementSize + Trait::kScoreOffset); + 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; + kv[i] = gmem_in.load(src + j * Trait::kElementSize); + score[i] = gmem_in.load(src + j * Trait::kElementSize + Trait::kScoreOffset); + } + } else { // mixed dtype + using StorageBuffer = AlignedVector; + const auto gmem_buffer = tile::Memory{lane_id, kWarpThreads}; + +#pragma unroll + 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 = kv_buf + j * Trait::kElementSize; + const auto kv_tmp = gmem_buffer.load(src); + const auto score_tmp = gmem_buffer.load(src + Trait::kScoreOffset); +#pragma unroll + for (int32_t k = 0; k < kTileElements; ++k) { + kv[i][k] = cast(kv_tmp[k]); + score[i][k] = cast(score_tmp[k]); + } + } else { + const auto src = kv_start + j * Trait::kElementSize; + kv[i] = gmem_in.load(src); + score[i] = gmem_in.load(src + Trait::kScoreOffset); + } + } } /// NOTE: part 2: safe online softmax + weighted sum @@ -141,6 +168,7 @@ SGL_DEVICE void c128_forward( // convert to fp32 and apply bias first #pragma unroll for (int32_t i = 0; i < kTileElements; ++i) { +#pragma unroll for (int32_t j = 0; j < kElementsPerWarp; ++j) { score_fp32[i][j] = cast(score[j][i]) + cast(bias[j][i]); } @@ -215,25 +243,41 @@ SGL_DEVICE void c128_forward( } } -template -SGL_DEVICE void c128_write_decode(InFloat* kv_buf, const InFloat* kv_src) { +template +SGL_DEVICE void c128_write_decode(BufferFloat* kv_buf, const InputFloat* kv_src) { using namespace device; - using Storage = AlignedVector; - const auto gmem = tile::Memory::warp(); + using StorageInput = AlignedVector; + const auto gmem_input = tile::Memory::warp(); - Storage data[2]; + StorageInput data[2]; #pragma unroll for (int32_t i = 0; i < 2; ++i) { - data[i] = gmem.load(kv_src + Trait::kHeadDim * i); + data[i] = gmem_input.load(kv_src + Trait::kHeadDim * i); } + + if constexpr (std::is_same_v) { #pragma unroll - for (int32_t i = 0; i < 2; ++i) { - gmem.store(kv_buf + Trait::kHeadDim * i, data[i]); + for (int32_t i = 0; i < 2; ++i) { + gmem_input.store(kv_buf + Trait::kHeadDim * i, data[i]); + } + } else { + using StorageBuffer = AlignedVector; + const auto gmem_buffer = tile::Memory::warp(); + + StorageBuffer data_cast[2]; +#pragma unroll + for (int32_t i = 0; i < 2; ++i) { +#pragma unroll + for (int32_t j = 0; j < kTileElements; ++j) { + data_cast[i][j] = cast(data[i][j]); + } + gmem_buffer.store(kv_buf + Trait::kHeadDim * i, data_cast[i]); + } } } -template +template C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodeParams params) { using namespace device; using Trait = C128Trait; @@ -245,10 +289,10 @@ C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodePara if (global_bid >= params.batch_size) return; const auto plan = params.plan_d[global_bid]; - const auto kv_input = static_cast(params.kv_input) + split_offset; + 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 score_bias = static_cast(params.score_bias) + 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; const auto kv_out = kv_output + global_bid * Trait::kHeadDim; @@ -258,15 +302,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; @@ -277,10 +321,10 @@ C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillPa if (global_pid >= params.num_compress) return; const auto plan = params.plan_c[global_pid]; - const auto kv_input = static_cast(params.kv_input) + split_offset; + 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 score_bias = static_cast(params.score_bias) + 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; const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize; @@ -288,14 +332,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 StorageInput = AlignedVector; const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x; const uint32_t global_wid = global_tid / kWarpThreads; // warp id @@ -307,33 +351,53 @@ WRITE_KERNEL void write_c128_prefill(const __grid_constant__ Compress128PrefillP if (global_pid >= params.num_write) return; 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_input = static_cast(params.kv_input) + 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 = tile::Memory::warp(); + const auto gmem_input = tile::Memory::warp(); PDLWaitPrimary(); - StorageIn data[2]; + StorageInput data[2]; #pragma unroll for (int32_t i = 0; i < 2; ++i) { - data[i] = gmem.load(kv_src, i); + data[i] = gmem_input.load(kv_src, i); } - PDLTriggerSecondary(); + + if constexpr (std::is_same_v) { + PDLTriggerSecondary(); #pragma unroll - for (int32_t i = 0; i < 2; ++i) { - gmem.store(kv_buf, data[i], i); + for (int32_t i = 0; i < 2; ++i) { + gmem_input.store(kv_buf, data[i], i); + } + } else { + using StorageBuffer = AlignedVector; + const auto gmem_buffer = tile::Memory::warp(); + + StorageBuffer data_cast[2]; +#pragma unroll + for (int32_t i = 0; i < 2; ++i) { +#pragma unroll + for (int32_t j = 0; j < kTileElements; ++j) { + data_cast[i][j] = cast(data[i][j]); + } + } + PDLTriggerSecondary(); +#pragma unroll + for (int32_t i = 0; i < 2; ++i) { + gmem_buffer.store(kv_buf, data_cast[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; @@ -351,11 +415,11 @@ 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 - .with_dtype() + .with_dtype() .with_device(device_) .verify(kv_input); TensorMatcher({N, kHeadDim}) // kv compressed output @@ -363,7 +427,7 @@ struct FlashCompress128Kernel { .with_device(device_) .verify(kv_output); TensorMatcher({128, kHeadDim}) // ape - .with_dtype() + .with_dtype() .with_device(device_) .verify(ape); @@ -398,11 +462,11 @@ 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) - .with_dtype() + .with_dtype() .with_device(device_) .verify(kv_input); TensorMatcher({C, kHeadDim}) // kv compressed output (compact) @@ -410,7 +474,7 @@ struct FlashCompress128Kernel { .with_device(device_) .verify(kv_output); TensorMatcher({128, kHeadDim}) // ape - .with_dtype() + .with_dtype() .with_device(device_) .verify(ape); 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 efa9f0510..15c8e740a 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh @@ -28,6 +28,7 @@ #include #include +#include namespace { @@ -74,20 +75,18 @@ struct C4Trait { static_assert(kHeadDim % kTileDim == 0); }; -template +template SGL_DEVICE void c4_forward( - 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 + const BufferFloat* kv_buf_0, // overlap [4n - 4, 4n - 1] + const BufferFloat* kv_buf_1, // normal [4n + 0, 4n + 3] + const InputFloat* kv_src, // ragged pointer at position = 4n + 3 OutFloat* kv_out, - const InFloat* score_bias, + const InputFloat* score_bias, const bool should_overlap, const int32_t buffer_len) { using namespace device; - /// NOTE: part 1: load kv + score - using StorageIn = AlignedVector; - /// NOTE: load one tile_dim (< head_dim) at at time + using StorageIn = AlignedVector; const auto gmem_in = tile::Memory::warp(); StorageIn kv[8]; StorageIn score[8]; @@ -98,32 +97,80 @@ SGL_DEVICE void c4_forward( bias[i] = gmem_in.load(score_bias + i * Trait::kHeadDim); } - if (should_overlap) { - const auto kv_start = kv_src - 7 * Trait::kElementSize; // point to start + if constexpr (std::is_same_v) { + if (should_overlap) { + 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; - 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; + for (int32_t i = 0; i < 4; ++i) { + 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) { - kv[i].fill(cast(0.0f)); - score[i].fill(cast(kFloatNegInf)); + for (int32_t i = 0; i < 4; ++i) { + kv[i].fill(cast(0.0f)); + score[i].fill(cast(kFloatNegInf)); + } } - } - const auto kv_start = kv_src - 3 * Trait::kElementSize; // point to start + 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; - 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); + for (int32_t i = 0; i < 4; ++i) { + 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); + } + } else { // mixed dtype + using StorageBuffer = AlignedVector; + const auto gmem_buffer = tile::Memory::warp(); + const auto kv_start_0 = kv_src - 7 * Trait::kElementSize; // point to start + +#pragma unroll + for (int32_t i = 0; i < 4; ++i) { + if (should_overlap && i < buffer_len) { + const auto base = kv_buf_0 + i * Trait::kElementSize; + const auto kv_tmp = gmem_buffer.load(base); + const auto score_tmp = gmem_buffer.load(base + Trait::kScoreOffset); +#pragma unroll + for (int32_t j = 0; j < kTileElements; ++j) { + kv[i][j] = cast(kv_tmp[j]); + score[i][j] = cast(score_tmp[j]); + } + } else if (should_overlap) { + const auto base = kv_start_0 + i * Trait::kElementSize; + kv[i] = gmem_in.load(base); + score[i] = gmem_in.load(base + Trait::kScoreOffset); + } else { + [[unlikely]]; + constexpr float kFloatNegInf = -FLT_MAX; + 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 kv_tmp = gmem_buffer.load(base); + const auto score_tmp = gmem_buffer.load(base + Trait::kScoreOffset); +#pragma unroll + for (int32_t j = 0; j < kTileElements; ++j) { + kv[i + 4][j] = cast(kv_tmp[j]); + score[i + 4][j] = cast(score_tmp[j]); + } + } else { + const auto base = kv_start + 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 @@ -137,6 +184,7 @@ SGL_DEVICE void c4_forward( // convert to fp32 and apply bias first #pragma unroll for (int32_t i = 0; i < kTileElements; ++i) { +#pragma unroll for (int32_t j = 0; j < 8; ++j) { score_fp32[i][j] = cast(score[j][i]) + cast(bias[j][i]); } @@ -171,25 +219,41 @@ SGL_DEVICE void c4_forward( gmem_out.store(kv_out, result); } -template -SGL_DEVICE void c4_write_decode(InFloat* kv_buf, const InFloat* kv_src) { +template +SGL_DEVICE void c4_write_decode(BufferFloat* kv_buf, const InputFloat* kv_src) { using namespace device; - using StorageIn = AlignedVector; - const auto gmem = tile::Memory::warp(); + using StorageInput = AlignedVector; + const auto gmem_input = tile::Memory::warp(); - StorageIn data[4]; + StorageInput data[4]; #pragma unroll for (int32_t i = 0; i < 4; ++i) { - data[i] = gmem.load(kv_src + Trait::kHeadDim * i); + data[i] = gmem_input.load(kv_src + Trait::kHeadDim * i); } + + if constexpr (std::is_same_v) { #pragma unroll - for (int32_t i = 0; i < 4; ++i) { - gmem.store(kv_buf + Trait::kHeadDim * i, data[i]); + for (int32_t i = 0; i < 4; ++i) { + gmem_input.store(kv_buf + Trait::kHeadDim * i, data[i]); + } + } else { + using StorageBuffer = AlignedVector; + const auto gmem_buffer = tile::Memory::warp(); + + StorageBuffer data_cast[4]; +#pragma unroll + for (int32_t i = 0; i < 4; ++i) { +#pragma unroll + for (int32_t j = 0; j < kTileElements; ++j) { + data_cast[i][j] = cast(data[i][j]); + } + gmem_buffer.store(kv_buf + Trait::kHeadDim * i, data_cast[i]); + } } } -template +template C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams params) { using namespace device; using Trait = C4Trait; @@ -202,10 +266,10 @@ C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams par if (global_bid >= params.batch_size) return; const auto plan = params.plan_d[global_bid]; - const auto kv_input = static_cast(params.kv_input) + split_offset; + 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 score_bias = static_cast(params.score_bias) + 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; const auto kv_out = kv_output + global_bid * Trait::kHeadDim; @@ -214,14 +278,15 @@ 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; @@ -234,10 +299,10 @@ C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams p if (global_pid >= params.num_compress) return; const auto plan = params.plan_c[global_pid]; - const auto kv_input = static_cast(params.kv_input) + split_offset; + 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 score_bias = static_cast(params.score_bias) + 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; const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize; @@ -247,14 +312,15 @@ 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 StorageInput = AlignedVector; const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x; const uint32_t global_wid = global_tid / kWarpThreads; // warp id @@ -266,33 +332,53 @@ WRITE_KERNEL void write_c4_prefill(const __grid_constant__ Compress4PrefillParam if (global_pid >= params.num_write) return; 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_input = static_cast(params.kv_input) + 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 = tile::Memory::warp(); + const auto gmem_input = tile::Memory::warp(); PDLWaitPrimary(); - StorageIn data[4]; + StorageInput data[4]; #pragma unroll for (int32_t i = 0; i < 4; ++i) { - data[i] = gmem.load(kv_src, i); + data[i] = gmem_input.load(kv_src, i); } - PDLTriggerSecondary(); + + if constexpr (std::is_same_v) { + PDLTriggerSecondary(); #pragma unroll - for (int32_t i = 0; i < 4; ++i) { - gmem.store(kv_buf, data[i], i); + for (int32_t i = 0; i < 4; ++i) { + gmem_input.store(kv_buf, data[i], i); + } + } else { + using StorageBuffer = AlignedVector; + const auto gmem_buffer = tile::Memory::warp(); + + StorageBuffer data_cast[4]; +#pragma unroll + for (int32_t i = 0; i < 4; ++i) { +#pragma unroll + for (int32_t j = 0; j < kTileElements; ++j) { + data_cast[i][j] = cast(data[i][j]); + } + } + PDLTriggerSecondary(); +#pragma unroll + for (int32_t i = 0; i < 4; ++i) { + gmem_buffer.store(kv_buf, data_cast[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; @@ -312,11 +398,11 @@ 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 - .with_dtype() + .with_dtype() .with_device(device_) .verify(kv_input); TensorMatcher({N, kHeadDim}) // kv compressed output @@ -324,7 +410,7 @@ struct FlashCompress4Kernel { .with_device(device_) .verify(kv_output); TensorMatcher({8, kHeadDim}) // ape - .with_dtype() + .with_dtype() .with_device(device_) .verify(ape); @@ -359,11 +445,11 @@ 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) - .with_dtype() + .with_dtype() .with_device(device_) .verify(kv_input); TensorMatcher({C, kHeadDim}) // kv compressed output (compact) @@ -371,7 +457,7 @@ struct FlashCompress4Kernel { .with_device(device_) .verify(kv_output); TensorMatcher({8, kHeadDim}) // ape - .with_dtype() + .with_dtype() .with_device(device_) .verify(ape); const auto plan_c = compress::verify_plan_c(plan_c_, C, device_); diff --git a/python/sglang/jit_kernel/dsv4/compress.py b/python/sglang/jit_kernel/dsv4/compress.py index 610406932..ca8d2c342 100644 --- a/python/sglang/jit_kernel/dsv4/compress.py +++ b/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_buffer: 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_buffer, dtype_in, dtype_out, is_arch_support_pdl() + ) kernel_class = f"FlashCompress{ratio}Kernel<{args}>" return load_jit( make_name(f"compress_{ratio}_v2"), @@ -336,7 +339,9 @@ def compress_forward( module = _jit_compress_128_online_module(512) else: dtype_in, dtype_out = kv_score_input.dtype, out.dtype - module = _jit_compress_module(head_dim, dtype_in, dtype_out, compress_ratio) + module = _jit_compress_module( + head_dim, kv_score_buffer.dtype, 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/environ.py b/python/sglang/srt/environ.py index 1694e651c..f67f4da2b 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -788,6 +788,7 @@ class Envs: SGLANG_OPT_USE_JIT_INDEXER_METADATA = EnvBool(True) SGLANG_OPT_USE_ONLINE_COMPRESS = EnvBool(False) SGLANG_EXPERIMENTAL_ONLINE_C128_MTP = EnvBool(False) + SGLANG_DSV4_COMPRESS_STATE_DTYPE = EnvStr("float32") SGLANG_OPT_USE_COMPRESSOR_V2 = EnvBool(True) SGLANG_FP8_PAGED_MQA_LOGITS_TORCH = EnvBool(False) SGLANG_TOPK_TRANSFORM_512_TORCH = EnvBool(False) diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 08fed1e33..6d13fb6be 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -448,7 +448,8 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): page_size: int, swa_page_size: int, dtype: torch.dtype, - state_dtype: torch.dtype, + c4_state_dtype: torch.dtype, + c128_state_dtype: torch.dtype, qk_nope_head_dim: int, qk_rope_head_dim: int, indexer_head_dim: int, @@ -494,7 +495,8 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): self.c128_size = c128_size self.c4_state_pool_size = c4_state_pool_size self.c128_state_pool_size = c128_state_pool_size - self.state_dtype = state_dtype + self.c4_state_dtype = c4_state_dtype + self.c128_state_dtype = c128_state_dtype self.compression_ratios = compression_ratios self.online_mtp_max_draft_tokens = online_mtp_max_draft_tokens self.online_c128_mtp_pending_seq_lens: Optional[torch.Tensor] = None @@ -761,7 +763,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): ring_size=ring_size, overlap=overlap, head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim, - dtype=self.state_dtype, + dtype=self.c4_state_dtype if ratio == 4 else self.c128_state_dtype, device=self.device, enable_memory_saver=enable_memory_saver, ratio=ratio, @@ -779,7 +781,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): overlap=overlap, head_dim=self.indexer_head_dim, device=self.device, - dtype=self.state_dtype, + dtype=self.c4_state_dtype, enable_memory_saver=enable_memory_saver, ratio=ratio, swa_page_size=self.swa_page_size, diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 0127dd1ce..a86fd9a2c 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -58,6 +58,24 @@ MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP = 1 logger = logging.getLogger(__name__) + +def _get_dsv4_compress_state_dtypes() -> tuple[torch.dtype, torch.dtype]: + dtype_name = envs.SGLANG_DSV4_COMPRESS_STATE_DTYPE.get().strip().lower() + 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=" + f"{dtype_name!r}. Expected one of: float32, fp32, bfloat16, bf16." + ) + + _is_npu = is_npu() _is_hip = is_hip() @@ -418,7 +436,8 @@ class ModelRunnerKVCacheMixin: swa_page_size=swa_page_size, sliding_window=self.model_config.window_size, dtype=self.kv_cache_dtype, - state_dtype=self.state_dtype, + c4_state_dtype=self.c4_state_dtype, + c128_state_dtype=self.c128_state_dtype, qk_nope_head_dim=self.model_config.qk_nope_head_dim, qk_rope_head_dim=self.model_config.qk_rope_head_dim, indexer_head_dim=self.model_config.index_head_dim, @@ -947,12 +966,12 @@ class ModelRunnerKVCacheMixin: self.c4_state_pool_size = config.c4_state_pool_size self.c128_state_pool_size = config.c128_state_pool_size - # state_dtype is a DSV4 architectural constant (fp32 for c4/c128 - # state buffers); set unconditionally so draft workers have it before - # _init_pools reads it (target path also overwrites this in the - # configurator's resolve() for parity, harmless here). + # Draft worker does not own the compression-state pools, but keep the + # dtype attributes initialized so _init_pools can share one code path. if is_deepseek_v4(self.model_config.hf_config): - self.state_dtype = torch.float32 + self.c4_state_dtype, self.c128_state_dtype = ( + _get_dsv4_compress_state_dtypes() + ) self._init_pools() diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index b18a26b76..fb1a7f095 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -62,6 +62,23 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +def _get_dsv4_compress_state_dtype_sizes() -> tuple[int, int]: + dtype_name = envs.SGLANG_DSV4_COMPRESS_STATE_DTYPE.get().strip().lower() + 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=" + f"{dtype_name!r}. Expected one of: float32, fp32, bfloat16, bf16." + ) + + class MemoryPoolConfigurator: """Base class for memory pool configurators. @@ -419,16 +436,18 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): ) attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim - state_dtype_size = 4 - c4_state_bytes = 2 * 2 * attn_head_dim * state_dtype_size + c4_state_dtype_size, c128_state_dtype_size = ( + _get_dsv4_compress_state_dtype_sizes() + ) + c4_state_bytes = 2 * 2 * attn_head_dim * c4_state_dtype_size # Online c128 stores (max, sum, kv) per slot (3*head_dim) instead of # raw (kv, score) (2*head_dim). Combined with ring_size=1 this still # nets a large reduction (~3/256x) but the per-slot bytes go up. c128_online = envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get() c128_state_bytes = ( - (3 if c128_online else 2 * 1) * attn_head_dim * state_dtype_size + (3 if c128_online else 2 * 1) * attn_head_dim * c128_state_dtype_size ) - c4_indexer_state_bytes = 2 * 2 * self.indexer_head_dim * state_dtype_size + c4_indexer_state_bytes = 2 * 2 * self.indexer_head_dim * c4_state_dtype_size c4_state_ratio = self.c4_ring_size / self.swa_page_size c128_state_ratio = self.c128_ring_size / self.swa_page_size diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py index 5249b5126..6ed24703f 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py @@ -390,7 +390,8 @@ class MockDSV4ModelRunner: page_size=case.page_size, swa_page_size=DSV4_SWA_WINDOW, dtype=torch.float8_e4m3fn, - state_dtype=dtype, + c4_state_dtype=dtype, + c128_state_dtype=dtype, qk_nope_head_dim=DSV4_QK_NOPE_HEAD_DIM, qk_rope_head_dim=DSV4_QK_ROPE_HEAD_DIM, indexer_head_dim=128, diff --git a/test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py b/test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py new file mode 100644 index 000000000..11e01dff3 --- /dev/null +++ b/test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py @@ -0,0 +1,1082 @@ +# DeepSeek V4 compress-state test and benchmark entry. +# +# What this covers: +# - Synthetic Flash/Pro C4/C128 decode/prefill shapes for broad operator +# performance coverage. +# - Replays 84 compress shapes captured from zc01 DeepSeek-V4-Flash serving. +# - The capture used EAGLE, so runtime compress plans were prefill-style plans +# for both EXTEND and TARGET_VERIFY. Some short C128 cases have out_shape=0. +# - This file keeps the out_shape=0 diff handling local to the runtime replay +# benchmark path. +# +# Test command: +# python3 -m pytest -q \ +# test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py +# +# Runtime-shape benchmark command: +# python3 test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py \ +# --benchmark \ +# --shape-source runtime \ +# --warmup 20 \ +# --iters 100 \ +# --csv /data00/eval_results/operator_bench/runtime_shape_bench.csv +# +# Synthetic Flash/Pro shape benchmark command: +# python3 test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py \ +# --benchmark \ +# --shape-source preset \ +# --shape-presets all \ +# --shape-tier smoke \ +# --warmup 10 \ +# --iters 30 \ +# --csv /data00/eval_results/operator_bench/preset_shape_bench.csv +# +# Service-level scenario where BF16 state compression is more likely to help: +# - Disable speculative/EAGLE so the workload is not dominated by small +# TARGET_VERIFY shapes. +# - Prefer TP-only or DP=1 first; DP attention may reduce chunked prefill from +# 4096 to 512, which makes fixed kernel overhead dominate. +# - Use long random prompts and short outputs to make prefill dominate: +# SGLANG_DSV4_COMPRESS_STATE_DTYPE=bf16 \ +# SGLANG_SHARED_EXPERT_TP1=1 \ +# SGLANG_ENABLE_THINKING=1 \ +# SGLANG_DSV4_FP4_EXPERTS=1 \ +# SGLANG_JIT_DEEPGEMM_PRECOMPILE=1 \ +# sglang serve \ +# --trust-remote-code \ +# --model-path /data00/models/DeepSeek-V4-Flash \ +# --tp 8 \ +# --host 0.0.0.0 \ +# --port 8080 \ +# --mem-fraction-static 0.9 \ +# --moe-runner-backend marlin \ +# --chunked-prefill-size 4096 \ +# --max-prefill-tokens 16384 \ +# --max-running-requests 32 \ +# --cuda-graph-max-bs 16 \ +# --enable-metrics \ +# --disable-radix-cache +# +# Workload for the service-level scenario: +# HF_ENDPOINT=https://hf-mirror.com \ +# python3 -m sglang.bench_serving \ +# --host localhost \ +# --port 8080 \ +# --model /data00/models/DeepSeek-V4-Flash \ +# --dataset-name random \ +# --random-input-len 8192 \ +# --random-output-len 8 \ +# --random-range-ratio 1 \ +# --num-prompts 128 \ +# --max-concurrency 16 \ +# --request-rate 16 +# +# Compare with the same service command without +# SGLANG_DSV4_COMPRESS_STATE_DTYPE=bf16. A visible service-level gain requires +# large prefill shapes and a non-trivial compress-kernel share in the profile. + +from __future__ import annotations + +import argparse +import csv +import math +import sys +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace +from typing import Callable, Iterable, Literal, Optional + +import pytest +import torch + +from sglang.jit_kernel.dsv4 import ( + CompressorDecodePlan, + CompressorPrefillPlan, + compress_forward, +) +from sglang.jit_kernel.tests.deepseek_v4.common import ( + make_legacy_context, + to_seq_extend, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.utils import is_in_ci + +register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large") + +Mode = Literal["decode", "prefill"] +ShapeTier = Literal["ci", "smoke", "full"] +MAX_PRESET_PREFILL_Q_TOKENS = 32768 + + +@dataclass(frozen=True) +class ShapePreset: + name: str + hidden_size: int + num_attention_heads: int + index_topk: int + ratios: tuple[Literal[4, 128], ...] + decode_batch_sizes: dict[ShapeTier, tuple[int, ...]] + prefill_batch_sizes: dict[ShapeTier, tuple[int, ...]] + decode_seq_lens: dict[ShapeTier, tuple[int, ...]] + prefill_extend_lens: dict[ShapeTier, dict[Literal[4, 128], tuple[int, ...]]] + + +SHAPE_PRESETS: dict[str, ShapePreset] = { + # DeepSeek-V4-Flash config on zc01: + # hidden_size=4096, num_attention_heads=64, index_topk=512, + # compress_ratios=[0, 0, 4, 128, ... alternating ..., 0]. + # The compress kernel itself uses head_dim=512 + # (= qk_nope_head_dim 448 + qk_rope_head_dim 64). + "flash": ShapePreset( + name="flash", + hidden_size=4096, + num_attention_heads=64, + index_topk=512, + ratios=(4, 128), + decode_batch_sizes={ + "ci": (16,), + "smoke": (16, 128), + "full": (1, 2, 4, 8, 16, 32, 64, 128), + }, + prefill_batch_sizes={ + "ci": (1,), + "smoke": (1, 16), + "full": (1, 2, 4, 8, 16, 32), + }, + decode_seq_lens={ + "ci": (128,), + "smoke": (128,), + "full": (128, 256, 4096), + }, + prefill_extend_lens={ + "ci": {4: (128,), 128: (128,)}, + "smoke": {4: (16, 128, 4096), 128: (128, 4096)}, + "full": {4: (16, 128, 4096), 128: (128, 4096)}, + }, + ), + # DeepSeek-V4-Pro config on zc01: + # hidden_size=7168, num_attention_heads=128, index_topk=1024, + # compress_ratios=[128, 128, 4, 128, ... alternating ..., 0]. + "pro": ShapePreset( + name="pro", + hidden_size=7168, + num_attention_heads=128, + index_topk=1024, + ratios=(4, 128), + decode_batch_sizes={ + "ci": (16,), + "smoke": (16, 32), + "full": (1, 2, 4, 8, 16, 32), + }, + prefill_batch_sizes={ + "ci": (1,), + "smoke": (1, 16), + "full": (1, 2, 4, 8, 16, 32), + }, + decode_seq_lens={ + "ci": (128,), + "smoke": (128,), + "full": (128, 256, 4096), + }, + prefill_extend_lens={ + "ci": {4: (128,), 128: (128,)}, + "smoke": {4: (16, 128, 4096), 128: (128, 4096)}, + "full": {4: (16, 128, 4096), 128: (128, 4096)}, + }, + ), +} + + +@dataclass(frozen=True) +class BenchSpec: + shape_name: str + ratio: Literal[4, 128] + mode: Mode + batch_size: int + head_dim: int + tokens_per_req: int + + +@dataclass +class BenchInput: + state_pool: torch.Tensor + kv_score_input: torch.Tensor + ape: torch.Tensor + out: torch.Tensor + plan: CompressorDecodePlan | CompressorPrefillPlan + spec: BenchSpec + + @property + def effective_bytes(self) -> int: + tensors = (self.state_pool, self.kv_score_input, self.ape, self.out) + return sum(t.numel() * t.element_size() for t in tensors) + + def run(self) -> torch.Tensor: + return compress_forward( + kv_score_buffer=self.state_pool, + kv_score_input=self.kv_score_input, + ape=self.ape, + plan=self.plan, + head_dim=self.spec.head_dim, + compress_ratio=self.spec.ratio, + out=self.out, + ) + + +@dataclass +class BenchResult: + shape_name: str + ratio: int + mode: str + batch_size: int + tokens_per_req: int + head_dim: int + state_shape: str + input_shape: str + ape_shape: str + out_shape: str + fp32_state_mib: float + bf16_state_mib: float + input_mib: float + ape_mib: float + out_mib: float + fp32_effective_mib: float + bf16_effective_mib: float + fp32_us: float + fp32_gbps: float + bf16_us: float + bf16_gbps: float + out_diff: float + state_diff: float + + @property + def speedup(self) -> float: + return self.fp32_us / self.bf16_us + + +def _last_dim(ratio: int, head_dim: int) -> int: + return head_dim * (4 if ratio == 4 else 2) + + +def _ape_len(ratio: int) -> int: + return 8 if ratio == 4 else 128 + + +def _shape_str(tensor: torch.Tensor) -> str: + return "x".join(str(dim) for dim in tensor.shape) + + +def _mib(num_bytes: int) -> float: + return num_bytes / 1024 / 1024 + + +def _tensor_mib(tensor: torch.Tensor) -> float: + return _mib(tensor.numel() * tensor.element_size()) + + +def _make_plan(spec: BenchSpec): + ctx = make_legacy_context( + bs=spec.batch_size, + compress_ratio=spec.ratio, + head_dim=spec.head_dim, + ) + if spec.mode == "decode": + seq_lens = torch.full( + (spec.batch_size,), + spec.tokens_per_req, + dtype=torch.int64, + device="cuda", + ) + plan = ctx.make_decode_plan(seq_lens) + num_output_tokens = spec.batch_size + else: + seq_lens_cpu, extend_lens_cpu, num_q_tokens = to_seq_extend( + [(spec.tokens_per_req, spec.tokens_per_req)] * spec.batch_size + ) + plan = ctx.make_prefill_plan(seq_lens_cpu, extend_lens_cpu, num_q_tokens) + num_output_tokens = int(plan.plan_c.shape[0]) + return ctx, plan, num_output_tokens + + +def _make_bench_input( + spec: BenchSpec, + state_dtype: torch.dtype, + base_state: torch.Tensor, + base_input: torch.Tensor, + base_ape: torch.Tensor, + plan: CompressorDecodePlan | CompressorPrefillPlan, + num_q_tokens: int, +) -> BenchInput: + state_pool = base_state.to(state_dtype) + kv_score_input = base_input.clone() + ape = base_ape.clone() + out = torch.empty( + (num_q_tokens, spec.head_dim), + dtype=kv_score_input.dtype, + device=kv_score_input.device, + ) + return BenchInput( + state_pool=state_pool, + kv_score_input=kv_score_input, + ape=ape, + out=out, + plan=plan, + spec=spec, + ) + + +def _make_pair(spec: BenchSpec, seed: int) -> tuple[BenchInput, BenchInput]: + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + ctx, plan, num_q_tokens = _make_plan(spec) + last_dim = _last_dim(spec.ratio, spec.head_dim) + + if spec.mode == "decode": + base_state = torch.randn( + (ctx.num_pages, spec.ratio, last_dim), + dtype=torch.float32, + device="cuda", + ) + input_rows = spec.batch_size + else: + base_state = torch.zeros( + (ctx.num_pages, spec.ratio, last_dim), + dtype=torch.float32, + device="cuda", + ) + input_rows = spec.batch_size * spec.tokens_per_req + + base_input = torch.randn( + (input_rows, last_dim), + dtype=torch.float32, + device="cuda", + ) + base_ape = torch.randn( + (_ape_len(spec.ratio), spec.head_dim), + dtype=torch.float32, + device="cuda", + ) + + return ( + _make_bench_input( + spec, torch.float32, base_state, base_input, base_ape, plan, num_q_tokens + ), + _make_bench_input( + spec, torch.bfloat16, base_state, base_input, base_ape, plan, num_q_tokens + ), + ) + + +def _time_us(fn: Callable[[], torch.Tensor], warmup: int, iters: int) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) * 1000.0 / iters + + +def _gbps(effective_bytes: int, elapsed_us: float) -> float: + return effective_bytes / (elapsed_us * 1e-6) / 1e9 + + +def _benchmark_spec( + spec: BenchSpec, + warmup: int, + iters: int, + seed: int, +) -> BenchResult: + fp32_case, bf16_case = _make_pair(spec, seed) + fp32_case.run() + bf16_case.run() + torch.cuda.synchronize() + + fp32_us = _time_us(fp32_case.run, warmup=warmup, iters=iters) + bf16_us = _time_us(bf16_case.run, warmup=warmup, iters=iters) + + # Recreate once for diff so the timed loop's repeated in-place state writes + # do not affect the comparison. + fp32_diff_case, bf16_diff_case = _make_pair(spec, seed) + fp32_out = fp32_diff_case.run() + bf16_out = bf16_diff_case.run() + torch.cuda.synchronize() + + out_diff = (fp32_out.float() - bf16_out.float()).abs().max().item() + state_diff = ( + (fp32_diff_case.state_pool.float() - bf16_diff_case.state_pool.float()) + .abs() + .max() + .item() + ) + + return BenchResult( + shape_name=spec.shape_name, + ratio=spec.ratio, + mode=spec.mode, + batch_size=spec.batch_size, + tokens_per_req=spec.tokens_per_req, + head_dim=spec.head_dim, + state_shape=_shape_str(fp32_case.state_pool), + input_shape=_shape_str(fp32_case.kv_score_input), + ape_shape=_shape_str(fp32_case.ape), + out_shape=_shape_str(fp32_case.out), + fp32_state_mib=_tensor_mib(fp32_case.state_pool), + bf16_state_mib=_tensor_mib(bf16_case.state_pool), + input_mib=_tensor_mib(fp32_case.kv_score_input), + ape_mib=_tensor_mib(fp32_case.ape), + out_mib=_tensor_mib(fp32_case.out), + fp32_effective_mib=_mib(fp32_case.effective_bytes), + bf16_effective_mib=_mib(bf16_case.effective_bytes), + fp32_us=fp32_us, + fp32_gbps=_gbps(fp32_case.effective_bytes, fp32_us), + bf16_us=bf16_us, + bf16_gbps=_gbps(bf16_case.effective_bytes, bf16_us), + out_diff=out_diff, + state_diff=state_diff, + ) + + +def _expand_shape_names(shape_presets: list[str]) -> list[str]: + if "all" in shape_presets: + return ["flash", "pro"] + if "custom" in shape_presets and len(shape_presets) > 1: + raise ValueError("--shape-presets custom cannot be mixed with shape presets.") + return shape_presets + + +def _one_or_many( + value: Optional[list[int]], fallback: tuple[int, ...] +) -> tuple[int, ...]: + return tuple(value) if value is not None else fallback + + +def _validate_tokens(mode: Mode, ratio: int, tokens_per_req: int) -> None: + if tokens_per_req % ratio != 0: + raise ValueError( + f"{mode} tokens_per_req={tokens_per_req} must be a multiple " + f"of ratio={ratio}" + ) + + +def _make_shape_specs(args: argparse.Namespace) -> Iterable[BenchSpec]: + shape_names = _expand_shape_names(args.shape_presets) + for shape_name in shape_names: + if shape_name == "custom": + yield from _make_custom_specs(args) + continue + + preset = SHAPE_PRESETS[shape_name] + ratios = tuple(args.ratios) if args.ratios is not None else preset.ratios + modes = tuple(args.modes) if args.modes is not None else ("decode", "prefill") + decode_batch_sizes = _one_or_many( + args.decode_batch_sizes or args.batch_sizes, + preset.decode_batch_sizes[args.shape_tier], + ) + prefill_batch_sizes = _one_or_many( + args.prefill_batch_sizes or args.batch_sizes, + preset.prefill_batch_sizes[args.shape_tier], + ) + decode_seq_lens = _one_or_many( + args.decode_seq_lens, + preset.decode_seq_lens[args.shape_tier], + ) + + for ratio in ratios: + if ratio not in preset.ratios: + continue + if "decode" in modes: + for tokens_per_req in decode_seq_lens: + _validate_tokens("decode", ratio, tokens_per_req) + for batch_size in decode_batch_sizes: + yield BenchSpec( + shape_name=shape_name, + ratio=ratio, + mode="decode", + batch_size=batch_size, + head_dim=args.head_dim, + tokens_per_req=tokens_per_req, + ) + + if "prefill" in modes: + prefill_extend_lens = _one_or_many( + args.prefill_extend_lens, + preset.prefill_extend_lens[args.shape_tier][ratio], + ) + for tokens_per_req in prefill_extend_lens: + _validate_tokens("prefill", ratio, tokens_per_req) + for batch_size in prefill_batch_sizes: + # The legacy prefill planner rejects very large + # synthetic q-token grids. Keep long-prefill coverage, + # but do not form unsupported cross-product cases. + if batch_size * tokens_per_req > MAX_PRESET_PREFILL_Q_TOKENS: + continue + yield BenchSpec( + shape_name=shape_name, + ratio=ratio, + mode="prefill", + batch_size=batch_size, + head_dim=args.head_dim, + tokens_per_req=tokens_per_req, + ) + + +def _make_custom_specs(args: argparse.Namespace) -> Iterable[BenchSpec]: + ratios = tuple(args.ratios) if args.ratios is not None else (4, 128) + modes = tuple(args.modes) if args.modes is not None else ("decode", "prefill") + batch_sizes = tuple(args.batch_sizes or [16, 32, 64, 128, 256]) + decode_seq_lens = tuple(args.decode_seq_lens or [128]) + prefill_extend_lens = tuple(args.prefill_extend_lens or [128]) + + for ratio in ratios: + for mode in modes: + seq_lens = decode_seq_lens if mode == "decode" else prefill_extend_lens + for tokens_per_req in seq_lens: + _validate_tokens(mode, ratio, tokens_per_req) + for batch_size in batch_sizes: + yield BenchSpec( + shape_name="custom", + ratio=ratio, + mode=mode, + batch_size=batch_size, + head_dim=args.head_dim, + tokens_per_req=tokens_per_req, + ) + + +def _format_table(results: list[BenchResult]) -> str: + if not results: + return "No benchmark cases selected." + + headers = [ + "shape", + "ratio", + "mode", + "bs", + "tok/req", + "head", + "state", + "input", + "ape", + "out", + "state MiB", + "input MiB", + "ape MiB", + "out MiB", + "FP32 MiB", + "BF16 MiB", + "FP32 us", + "FP32 GB/s", + "BF16 us", + "BF16 GB/s", + "out diff", + "state diff", + "Speedup", + ] + rows = [ + [ + r.shape_name, + str(r.ratio), + r.mode, + str(r.batch_size), + str(r.tokens_per_req), + str(r.head_dim), + r.state_shape, + r.input_shape, + r.ape_shape, + r.out_shape, + f"{r.fp32_state_mib:.2f}/{r.bf16_state_mib:.2f}", + f"{r.input_mib:.2f}", + f"{r.ape_mib:.3f}", + f"{r.out_mib:.2f}", + f"{r.fp32_effective_mib:.2f}", + f"{r.bf16_effective_mib:.2f}", + f"{r.fp32_us:.2f}", + f"{r.fp32_gbps:.1f}", + f"{r.bf16_us:.2f}", + f"{r.bf16_gbps:.1f}", + f"{r.out_diff:.4g}", + f"{r.state_diff:.4g}", + f"{r.speedup:.2f}x", + ] + for r in results + ] + widths = [ + max(len(headers[i]), *(len(row[i]) for row in rows)) + for i in range(len(headers)) + ] + line = " ".join(headers[i].rjust(widths[i]) for i in range(len(headers))) + sep = " ".join("-" * widths[i] for i in range(len(headers))) + body = [ + " ".join(row[i].rjust(widths[i]) for i in range(len(headers))) for row in rows + ] + return "\n".join([line, sep, *body]) + + +def _write_csv(path: Path, results: list[BenchResult]) -> None: + with path.open("w", newline="") as f: + writer = csv.DictWriter( + f, + fieldnames=[ + "ratio", + "shape", + "mode", + "batch_size", + "tokens_per_req", + "head_dim", + "state_shape", + "input_shape", + "ape_shape", + "out_shape", + "fp32_state_mib", + "bf16_state_mib", + "input_mib", + "ape_mib", + "out_mib", + "fp32_effective_mib", + "bf16_effective_mib", + "fp32_us", + "fp32_gbps", + "bf16_us", + "bf16_gbps", + "out_diff", + "state_diff", + "speedup", + ], + ) + writer.writeheader() + for r in results: + writer.writerow( + { + "ratio": r.ratio, + "shape": r.shape_name, + "mode": r.mode, + "batch_size": r.batch_size, + "tokens_per_req": r.tokens_per_req, + "head_dim": r.head_dim, + "state_shape": r.state_shape, + "input_shape": r.input_shape, + "ape_shape": r.ape_shape, + "out_shape": r.out_shape, + "fp32_state_mib": r.fp32_state_mib, + "bf16_state_mib": r.bf16_state_mib, + "input_mib": r.input_mib, + "ape_mib": r.ape_mib, + "out_mib": r.out_mib, + "fp32_effective_mib": r.fp32_effective_mib, + "bf16_effective_mib": r.bf16_effective_mib, + "fp32_us": r.fp32_us, + "fp32_gbps": r.fp32_gbps, + "bf16_us": r.bf16_us, + "bf16_gbps": r.bf16_gbps, + "out_diff": r.out_diff, + "state_diff": r.state_diff, + "speedup": r.speedup, + } + ) + + +def _preset_args(**overrides): + args = dict( + shape_presets=["all"], + model_shapes=None, + shape_tier="ci", + ratios=None, + modes=None, + batch_sizes=None, + decode_batch_sizes=None, + prefill_batch_sizes=None, + decode_seq_lens=None, + prefill_extend_lens=None, + head_dim=512, + ) + args.update(overrides) + return SimpleNamespace(**args) + + +def test_flash_pro_shape_presets_cover_compress_paths() -> None: + specs = list(_make_shape_specs(_preset_args())) + keys = { + (s.shape_name, s.ratio, s.mode, s.batch_size, s.tokens_per_req) for s in specs + } + + for shape_name in ("flash", "pro"): + assert (shape_name, 4, "decode", 16, 128) in keys + assert (shape_name, 4, "prefill", 1, 128) in keys + assert (shape_name, 128, "decode", 16, 128) in keys + assert (shape_name, 128, "prefill", 1, 128) in keys + + assert all(s.head_dim == 512 for s in specs) + assert all(s.tokens_per_req % s.ratio == 0 for s in specs) + + +@pytest.mark.parametrize( + "spec", + [ + BenchSpec("flash", 4, "decode", 2, 512, 128), + BenchSpec("flash", 4, "prefill", 1, 512, 16), + BenchSpec("pro", 128, "decode", 2, 512, 128), + BenchSpec("pro", 128, "prefill", 1, 512, 128), + ], +) +def test_compress_state_benchmark_cuda_smoke(spec: BenchSpec) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for DeepSeek V4 compress benchmark smoke.") + + result = _benchmark_spec(spec, warmup=1, iters=1, seed=20260603) + assert math.isfinite(result.fp32_us) and result.fp32_us > 0 + assert math.isfinite(result.bf16_us) and result.bf16_us > 0 + assert result.out_diff < 0.1 + assert result.state_diff < 0.1 + + +# Captured from DeepSeek-V4-Flash runtime on zc01 with EAGLE enabled: +# /data00/eval_results/operator_bench/runtime_shape_20260603/shapes.jsonl +# The service used prefill-style compress plans for both EXTEND and TARGET_VERIFY. +CAPTURED_FLASH_RUNTIME_SPECS = [ + BenchSpec("flash-runtime-indexer", 4, "prefill", 1, 128, 1), + BenchSpec("flash-runtime-indexer", 4, "prefill", 1, 128, 4), + BenchSpec("flash-runtime-indexer", 4, "prefill", 1, 128, 6), + BenchSpec("flash-runtime-indexer", 4, "prefill", 1, 128, 28), + BenchSpec("flash-runtime-indexer", 4, "prefill", 1, 128, 29), + BenchSpec("flash-runtime-indexer", 4, "prefill", 1, 128, 30), + BenchSpec("flash-runtime-indexer", 4, "prefill", 1, 128, 284), + BenchSpec("flash-runtime-indexer", 4, "prefill", 1, 128, 285), + BenchSpec("flash-runtime-indexer", 4, "prefill", 1, 128, 286), + BenchSpec("flash-runtime-indexer", 4, "prefill", 1, 128, 287), + BenchSpec("flash-runtime-indexer", 4, "prefill", 1, 128, 288), + BenchSpec("flash-runtime-indexer", 4, "prefill", 1, 128, 289), + BenchSpec("flash-runtime-indexer", 4, "prefill", 1, 128, 512), + BenchSpec("flash-runtime-indexer", 4, "prefill", 2, 128, 4), + BenchSpec("flash-runtime-indexer", 4, "prefill", 2, 128, 142), + BenchSpec("flash-runtime-indexer", 4, "prefill", 2, 128, 143), + BenchSpec("flash-runtime-indexer", 4, "prefill", 2, 128, 144), + BenchSpec("flash-runtime-indexer", 4, "prefill", 2, 128, 145), + BenchSpec("flash-runtime-indexer", 4, "prefill", 3, 128, 4), + BenchSpec("flash-runtime-indexer", 4, "prefill", 4, 128, 4), + BenchSpec("flash-runtime-indexer", 4, "prefill", 5, 128, 4), + BenchSpec("flash-runtime-indexer", 4, "prefill", 6, 128, 4), + BenchSpec("flash-runtime-indexer", 4, "prefill", 7, 128, 4), + BenchSpec("flash-runtime-indexer", 4, "prefill", 8, 128, 4), + BenchSpec("flash-runtime-indexer", 4, "prefill", 10, 128, 4), + BenchSpec("flash-runtime-indexer", 4, "prefill", 12, 128, 4), + BenchSpec("flash-runtime-indexer", 4, "prefill", 14, 128, 4), + BenchSpec("flash-runtime-indexer", 4, "prefill", 16, 128, 4), + BenchSpec("flash-runtime-core", 4, "prefill", 1, 512, 1), + BenchSpec("flash-runtime-core", 4, "prefill", 1, 512, 4), + BenchSpec("flash-runtime-core", 4, "prefill", 1, 512, 6), + BenchSpec("flash-runtime-core", 4, "prefill", 1, 512, 28), + BenchSpec("flash-runtime-core", 4, "prefill", 1, 512, 29), + BenchSpec("flash-runtime-core", 4, "prefill", 1, 512, 30), + BenchSpec("flash-runtime-core", 4, "prefill", 1, 512, 284), + BenchSpec("flash-runtime-core", 4, "prefill", 1, 512, 285), + BenchSpec("flash-runtime-core", 4, "prefill", 1, 512, 286), + BenchSpec("flash-runtime-core", 4, "prefill", 1, 512, 287), + BenchSpec("flash-runtime-core", 4, "prefill", 1, 512, 288), + BenchSpec("flash-runtime-core", 4, "prefill", 1, 512, 289), + BenchSpec("flash-runtime-core", 4, "prefill", 1, 512, 512), + BenchSpec("flash-runtime-core", 4, "prefill", 2, 512, 4), + BenchSpec("flash-runtime-core", 4, "prefill", 2, 512, 142), + BenchSpec("flash-runtime-core", 4, "prefill", 2, 512, 143), + BenchSpec("flash-runtime-core", 4, "prefill", 2, 512, 144), + BenchSpec("flash-runtime-core", 4, "prefill", 2, 512, 145), + BenchSpec("flash-runtime-core", 4, "prefill", 3, 512, 4), + BenchSpec("flash-runtime-core", 4, "prefill", 4, 512, 4), + BenchSpec("flash-runtime-core", 4, "prefill", 5, 512, 4), + BenchSpec("flash-runtime-core", 4, "prefill", 6, 512, 4), + BenchSpec("flash-runtime-core", 4, "prefill", 7, 512, 4), + BenchSpec("flash-runtime-core", 4, "prefill", 8, 512, 4), + BenchSpec("flash-runtime-core", 4, "prefill", 10, 512, 4), + BenchSpec("flash-runtime-core", 4, "prefill", 12, 512, 4), + BenchSpec("flash-runtime-core", 4, "prefill", 14, 512, 4), + BenchSpec("flash-runtime-core", 4, "prefill", 16, 512, 4), + BenchSpec("flash-runtime-core", 128, "prefill", 1, 512, 1), + BenchSpec("flash-runtime-core", 128, "prefill", 1, 512, 4), + BenchSpec("flash-runtime-core", 128, "prefill", 1, 512, 6), + BenchSpec("flash-runtime-core", 128, "prefill", 1, 512, 28), + BenchSpec("flash-runtime-core", 128, "prefill", 1, 512, 29), + BenchSpec("flash-runtime-core", 128, "prefill", 1, 512, 30), + BenchSpec("flash-runtime-core", 128, "prefill", 1, 512, 284), + BenchSpec("flash-runtime-core", 128, "prefill", 1, 512, 285), + BenchSpec("flash-runtime-core", 128, "prefill", 1, 512, 286), + BenchSpec("flash-runtime-core", 128, "prefill", 1, 512, 287), + BenchSpec("flash-runtime-core", 128, "prefill", 1, 512, 288), + BenchSpec("flash-runtime-core", 128, "prefill", 1, 512, 289), + BenchSpec("flash-runtime-core", 128, "prefill", 1, 512, 512), + BenchSpec("flash-runtime-core", 128, "prefill", 2, 512, 4), + BenchSpec("flash-runtime-core", 128, "prefill", 2, 512, 142), + BenchSpec("flash-runtime-core", 128, "prefill", 2, 512, 143), + BenchSpec("flash-runtime-core", 128, "prefill", 2, 512, 144), + BenchSpec("flash-runtime-core", 128, "prefill", 2, 512, 145), + BenchSpec("flash-runtime-core", 128, "prefill", 3, 512, 4), + BenchSpec("flash-runtime-core", 128, "prefill", 4, 512, 4), + BenchSpec("flash-runtime-core", 128, "prefill", 5, 512, 4), + BenchSpec("flash-runtime-core", 128, "prefill", 6, 512, 4), + BenchSpec("flash-runtime-core", 128, "prefill", 7, 512, 4), + BenchSpec("flash-runtime-core", 128, "prefill", 8, 512, 4), + BenchSpec("flash-runtime-core", 128, "prefill", 10, 512, 4), + BenchSpec("flash-runtime-core", 128, "prefill", 12, 512, 4), + BenchSpec("flash-runtime-core", 128, "prefill", 14, 512, 4), + BenchSpec("flash-runtime-core", 128, "prefill", 16, 512, 4), +] + + +def _max_abs_diff(lhs: torch.Tensor, rhs: torch.Tensor) -> float: + if lhs.numel() == 0: + return 0.0 + return (lhs.float() - rhs.float()).abs().max().item() + + +def _benchmark_runtime_spec( + spec: BenchSpec, + warmup: int, + iters: int, + seed: int, +) -> BenchResult: + fp32_case, bf16_case = _make_pair(spec, seed) + fp32_case.run() + bf16_case.run() + torch.cuda.synchronize() + + fp32_us = _time_us(fp32_case.run, warmup=warmup, iters=iters) + bf16_us = _time_us(bf16_case.run, warmup=warmup, iters=iters) + + fp32_diff_case, bf16_diff_case = _make_pair(spec, seed) + fp32_out = fp32_diff_case.run() + bf16_out = bf16_diff_case.run() + torch.cuda.synchronize() + + out_diff = _max_abs_diff(fp32_out, bf16_out) + state_diff = _max_abs_diff( + fp32_diff_case.state_pool, + bf16_diff_case.state_pool, + ) + + return BenchResult( + shape_name=spec.shape_name, + ratio=spec.ratio, + mode=spec.mode, + batch_size=spec.batch_size, + tokens_per_req=spec.tokens_per_req, + head_dim=spec.head_dim, + state_shape=_shape_str(fp32_case.state_pool), + input_shape=_shape_str(fp32_case.kv_score_input), + ape_shape=_shape_str(fp32_case.ape), + out_shape=_shape_str(fp32_case.out), + fp32_state_mib=_tensor_mib(fp32_case.state_pool), + bf16_state_mib=_tensor_mib(bf16_case.state_pool), + input_mib=_tensor_mib(fp32_case.kv_score_input), + ape_mib=_tensor_mib(fp32_case.ape), + out_mib=_tensor_mib(fp32_case.out), + fp32_effective_mib=_mib(fp32_case.effective_bytes), + bf16_effective_mib=_mib(bf16_case.effective_bytes), + fp32_us=fp32_us, + fp32_gbps=_gbps(fp32_case.effective_bytes, fp32_us), + bf16_us=bf16_us, + bf16_gbps=_gbps(bf16_case.effective_bytes, bf16_us), + out_diff=out_diff, + state_diff=state_diff, + ) + + +def test_captured_flash_runtime_specs_are_unique() -> None: + keys = { + (s.shape_name, s.ratio, s.mode, s.batch_size, s.head_dim, s.tokens_per_req) + for s in CAPTURED_FLASH_RUNTIME_SPECS + } + assert len(keys) == len(CAPTURED_FLASH_RUNTIME_SPECS) + assert len(CAPTURED_FLASH_RUNTIME_SPECS) == 84 + assert any(s.head_dim == 128 and s.ratio == 4 for s in CAPTURED_FLASH_RUNTIME_SPECS) + assert any(s.head_dim == 512 and s.ratio == 4 for s in CAPTURED_FLASH_RUNTIME_SPECS) + assert any( + s.head_dim == 512 and s.ratio == 128 for s in CAPTURED_FLASH_RUNTIME_SPECS + ) + + +@pytest.mark.parametrize("spec", CAPTURED_FLASH_RUNTIME_SPECS) +def test_compress_state_runtime_shape_cuda_smoke(spec: BenchSpec) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for DeepSeek V4 runtime compress shape smoke.") + + result = _benchmark_runtime_spec(spec, warmup=1, iters=1, seed=20260603) + assert math.isfinite(result.fp32_us) and result.fp32_us > 0 + assert math.isfinite(result.bf16_us) and result.bf16_us > 0 + assert result.out_diff < 0.1 + assert result.state_diff < 0.1 + + +def _run_benchmark(args: argparse.Namespace) -> int: + ci = is_in_ci() + args.warmup = args.warmup if args.warmup is not None else (5 if ci else 10) + args.iters = args.iters if args.iters is not None else (20 if ci else 30) + + if not torch.cuda.is_available(): + print("[skip] CUDA is required for this benchmark.") + return 0 + if args.head_dim % 128 != 0: + raise ValueError("--head-dim must be a multiple of 128.") + + print("DeepSeek V4 compress state dtype benchmark") + print( + "effective GB/s = " + "state_pool + kv_score_input + ape + output footprint / kernel time" + ) + + if args.shape_source == "runtime": + specs = CAPTURED_FLASH_RUNTIME_SPECS + if args.limit is not None: + specs = specs[: args.limit] + print( + f"config: shape_source=runtime, cases={len(specs)}, " + f"warmup={args.warmup}, iters={args.iters}" + ) + results = [ + _benchmark_runtime_spec( + spec, + warmup=args.warmup, + iters=args.iters, + seed=args.seed, + ) + for spec in specs + ] + else: + args.shape_presets = args.shape_presets or args.model_shapes or ["all"] + args.shape_tier = args.shape_tier or ("ci" if ci else "smoke") + if args.decode_seq_len is not None: + args.decode_seq_lens = [args.decode_seq_len] + if args.prefill_extend_len is not None: + args.prefill_extend_lens = [args.prefill_extend_len] + + specs = list(_make_shape_specs(args)) + if args.limit is not None: + specs = specs[: args.limit] + print( + f"config: shape_source=preset, shape_presets={args.shape_presets}, " + f"shape_tier={args.shape_tier}, ratios={args.ratios}, " + f"modes={args.modes}, batch_sizes={args.batch_sizes}, " + f"decode_batch_sizes={args.decode_batch_sizes}, " + f"prefill_batch_sizes={args.prefill_batch_sizes}, " + f"head_dim={args.head_dim}, decode_seq_lens={args.decode_seq_lens}, " + f"prefill_extend_lens={args.prefill_extend_lens}, " + f"cases={len(specs)}, warmup={args.warmup}, iters={args.iters}" + ) + results = [ + _benchmark_spec( + spec, + warmup=args.warmup, + iters=args.iters, + seed=args.seed, + ) + for spec in specs + ] + + print(_format_table(results)) + if args.csv: + _write_csv(args.csv, results) + print(f"\nWrote CSV: {args.csv}") + return 0 + + +def _parse_args(argv: list[str]) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Run DeepSeek V4 compress state tests, runtime-shape replay benchmark, " + "or synthetic Flash/Pro shape benchmark." + ) + ) + parser.add_argument("--benchmark", action="store_true") + parser.add_argument( + "--shape-source", + choices=("runtime", "preset"), + default="runtime", + help=( + "runtime replays captured serving shapes; preset runs synthetic " + "Flash/Pro/custom shape grids." + ), + ) + parser.add_argument( + "--shape-presets", + nargs="+", + choices=("flash", "pro", "all", "custom"), + default=None, + help=( + "Preset shapes for --shape-source preset. 'all' covers Flash and Pro. " + "Use 'custom' with manual shape args." + ), + ) + parser.add_argument( + "--model-shapes", + nargs="+", + choices=("flash", "pro", "all", "custom"), + default=None, + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--shape-tier", + choices=("ci", "smoke", "full"), + default=None, + help=( + "Shape grid size for preset shapes. Defaults to ci in CI and smoke " + "otherwise." + ), + ) + parser.add_argument("--ratios", type=int, nargs="+", choices=(4, 128), default=None) + parser.add_argument( + "--modes", + nargs="+", + choices=("decode", "prefill"), + default=None, + ) + parser.add_argument("--batch-sizes", type=int, nargs="+", default=None) + parser.add_argument("--decode-batch-sizes", type=int, nargs="+", default=None) + parser.add_argument("--prefill-batch-sizes", type=int, nargs="+", default=None) + parser.add_argument("--head-dim", type=int, default=512) + parser.add_argument( + "--decode-seq-lens", + type=int, + nargs="+", + default=None, + help="Decode seq_len values. Overrides preset decode seq_len grid.", + ) + parser.add_argument( + "--prefill-extend-lens", + type=int, + nargs="+", + default=None, + help="Prefill extend_len values. Overrides preset prefill extend_len grid.", + ) + parser.add_argument( + "--decode-seq-len", + type=int, + default=None, + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--prefill-extend-len", + type=int, + default=None, + help=argparse.SUPPRESS, + ) + parser.add_argument("--warmup", type=int, default=None) + parser.add_argument("--iters", type=int, default=None) + parser.add_argument("--seed", type=int, default=20260603) + parser.add_argument( + "--limit", + type=int, + default=None, + help="Limit benchmark cases after shape expansion.", + ) + parser.add_argument("--csv", type=Path, default=None) + return parser.parse_args(argv) + + +if __name__ == "__main__": + if "--benchmark" in sys.argv: + sys.exit(_run_benchmark(_parse_args(sys.argv[1:]))) + sys.exit(pytest.main([__file__, "-v"]))