Skip to content

[CUDA] Support batch_size >= 1 in SpeculativePrefixCommit - #33626

Merged
Tianlei Wu (tianleiwu) merged 2 commits into
asonawane/main-ops-mtp-fold-statefrom
tlwu/20261026/speculative-prefix-commit-batched
Oct 26, 2026
Merged

Tianlei Wu (tianleiwu) merged 2 commits into
asonawane/main-ops-mtp-fold-statefrom
tlwu/20261026/speculative-prefix-commit-batched

Conversation

@tianleiwu

Copy link
Copy Markdown
Contributor

Summary

Extends com.microsoft::SpeculativePrefixCommit (introduced in #33474) to support batch_size >= 1 for batched speculative decoding:

  • Schema and shape inference (bert_defs.cc): Propagates the batch dimension from next_tokens ([batch, rows] or [rows]) to outputs 0..6 (head_ids, head_c, slot_ids, seqlens, cos_out, sin_out, sel).
  • CUDA kernel (speculative_fold.h/.cc/.cu): Launches MidKernel<<<batch_size, 256, 0, stream>>> with per-slot strides for ctl, ctx_commit, prefixes, prefix2, per-slot result ring buffers, and per-slot muted flags.
  • Unit test (speculative_fold_op_test.cc): Adds SupportsBatchedSpeculativePrefixCommitS8 covering batch = 3, head_rows = 8 with mixed acceptance lengths and a muted slot.

- Propagate batch dimension from next_tokens[0] in SpeculativePrefixCommit type and shape inference.
- Launch MidKernel across batch_size blocks with per-slot control blocks, contexts, ring buffers, and muted slot flags.
- Add SupportsBatchedSpeculativePrefixCommitS8 unit test.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@tianleiwu
Tianlei Wu (tianleiwu) force-pushed the tlwu/20261026/speculative-prefix-commit-batched branch from 954b3c1 to a4f693d Compare October 26, 2026 20:09
@tianleiwu
Tianlei Wu (tianleiwu) merged commit 7ac2cd9 into asonawane/main-ops-mtp-fold-state Oct 26, 2026
5 checks passed
@tianleiwu
Tianlei Wu (tianleiwu) deleted the tlwu/20261026/speculative-prefix-commit-batched branch October 26, 2026 20:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant