fix: DSV4 BCG compress-prefill plan OOB on underfilled (tiny) prefill replay (#27747)

2026-06-10 20:02:24 -07:00
parent f4f30d7d23
commit f8b0a120b8
1 changed files with 5 additions and 1 deletions
@@ -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];