[Bug] ngram verify: keep batch.seq_lens_sum in sync after accept (#26753)

2026-05-29 18:00:11 -07:00
parent 95cd2fd29f
commit c2ac37dbcc
1 changed files with 2 additions and 0 deletions
@@ -465,6 +465,8 @@ class NgramVerifyInput(SpecInput):

        batch.seq_lens.add_(self.num_accept_tokens)
        batch.seq_lens_cpu.add_(num_accept_tokens_cpu)
+        # Keep seq_lens_sum in sync; attn backends size kv_indices from it.
+        batch.seq_lens_sum += int(num_accept_tokens_cpu.sum())

        return logits_output, self.accept_tokens, num_correct_drafts