From f8b0a120b8027a0c66679153b2bdf818936266b1 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Wed, 10 Jun 2026 20:02:24 -0700 Subject: [PATCH] fix: DSV4 BCG compress-prefill plan OOB on underfilled (tiny) prefill replay (#27747) --- python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh index 3e4aaaf5f..4e2f2ed28 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh +++ b/python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh @@ -203,7 +203,11 @@ __global__ __launch_bounds__(1024, 1) // if (is_mtp_extend) { // Path 1: token-driven. Each global token id maps to exactly one (batch_id, j). const uint32_t E = s_max_extend; - for (uint32_t k = tx; k < num_q; k += block_size) { + // num_q is the padded buffer size (graph bucket), not the work size: cap the + // loop at the real token count so batch_id = k / E stays < batch_size on an + // underfilled replay; Stage D pads [counter, num_q) with invalid. + const uint32_t num_real_q = params.batch_size * E; + for (uint32_t k = tx; k < num_real_q; k += block_size) { const uint32_t batch_id = k / E; const uint32_t j = k % E; const int32_t pl = s_prefix_len[batch_id];