Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Tianlei Wu (tianleiwu) merged 2 commits into from
Oct 26, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a4f693d
[CUDA] Support batch_size >= 1 in SpeculativePrefixCommit
tianleiwu Oct 26, 2026
1dc5b37
lintrunner
tianleiwu Oct 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 33 additions & 13 deletions onnxruntime/contrib_ops/cuda/llm/speculative_fold.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,25 +36,39 @@ class SpeculativePrefixCommit final : public CudaKernel {
const auto* sine = context->Input<Tensor>(9);
ORT_RETURN_IF_NOT(tokens && ids && status && ctl && committed && prefixes && prefix2 && ptrs && cosine && sine,
"SpeculativePrefixCommit: missing required input");
const int64_t rows = tokens->Shape().Size();
ORT_RETURN_IF_NOT(rows > 0 && rows <= INT_MAX && ids->Shape().Size() >= rows &&
const auto& tok_shape = tokens->Shape();
const int64_t batch = tok_shape.NumDimensions() >= 2
? tok_shape.SizeToDimension(tok_shape.NumDimensions() - 1)
: int64_t{1};
const int64_t rows = tok_shape.NumDimensions() >= 2
? tok_shape[tok_shape.NumDimensions() - 1]
: tok_shape.Size();
ORT_RETURN_IF_NOT(batch > 0 && batch <= INT_MAX && rows > 0 && rows <= INT_MAX &&
ids->Shape().Size() >= batch * rows && ids->Shape().Size() % batch == 0 &&
ids->Shape().Size() / batch <= INT_MAX &&
status->Shape().Size() <= INT_MAX,
"SpeculativePrefixCommit: invalid token/status shape");
ORT_RETURN_IF_NOT(ctl->Shape().Size() >= (head_rows_ > 2 ? 16 : 8) &&
Count(committed, 2) && prefixes->Shape().Size() >= (head_rows_ > 2 ? 2 * head_rows_ : 2) &&
Count(prefix2, 2) && Count(ptrs, 2),
const int64_t min_ctl = (head_rows_ > 2 ? 16 : 8);
const int64_t min_pref = (head_rows_ > 2 ? 2 * head_rows_ : 2);
ORT_RETURN_IF_NOT(ctl->Shape().Size() >= batch * min_ctl && ctl->Shape().Size() % batch == 0 &&
ctl->Shape().Size() / batch <= INT_MAX &&
Count(committed, 2 * batch) &&
prefixes->Shape().Size() >= batch * min_pref && prefixes->Shape().Size() % batch == 0 &&
prefixes->Shape().Size() / batch <= INT_MAX &&
(Count(prefix2, 2 * batch) || Count(prefix2, 2)) &&
(Count(ptrs, 2 * batch) || Count(ptrs, 2)),
"SpeculativePrefixCommit: invalid control/context shape");
const auto& rope = cosine->Shape();
ORT_RETURN_IF_NOT(rope.NumDimensions() == 2 && rope[0] > 0 && rope[0] <= INT_MAX &&
rope[1] > 0 && rope[1] <= INT_MAX && sine->Shape() == rope,
"SpeculativePrefixCommit: rotary caches must have matching [positions, width] shapes");
auto* out_ids = context->Output(0, TensorShape({1, head_rows_}));
auto* head_c = context->Output(1, TensorShape({1}));
auto* slots = context->Output(2, TensorShape({1, head_rows_}));
auto* lengths = context->Output(3, TensorShape({1}));
auto* out_cos = context->Output(4, TensorShape({1, head_rows_, rope[1]}));
auto* out_sin = context->Output(5, TensorShape({1, head_rows_, rope[1]}));
auto* selection = context->Output(6, TensorShape({1}));
auto* out_ids = context->Output(0, TensorShape({batch, head_rows_}));
auto* head_c = context->Output(1, TensorShape({batch}));
auto* slots = context->Output(2, TensorShape({batch, head_rows_}));
auto* lengths = context->Output(3, TensorShape({batch}));
auto* out_cos = context->Output(4, TensorShape({batch, head_rows_, rope[1]}));
auto* out_sin = context->Output(5, TensorShape({batch, head_rows_, rope[1]}));
auto* selection = context->Output(6, TensorShape({batch}));
int64_t* ctl_ptr = const_cast<int64_t*>(ctl->Data<int64_t>());
int64_t* context_ptr = const_cast<int64_t*>(committed->Data<int64_t>());
if (context->OutputCount() > 7) {
Expand All @@ -69,6 +83,11 @@ class SpeculativePrefixCommit final : public CudaKernel {
context_ptr = updated_context->MutableData<int64_t>();
}
}
const int ids_stride = static_cast<int>(ids->Shape().Size() / batch);
const int ctl_stride = static_cast<int>(ctl->Shape().Size() / batch);
const int prefix_stride = static_cast<int>(prefixes->Shape().Size() / batch);
const int prefix2_stride = (Count(prefix2, 2 * batch) && batch > 1) ? 2 : 0;
const int ptrs_stride = (Count(ptrs, 2 * batch) && batch > 1) ? 2 : 0;
return speculative_fold::LaunchMid(
tokens->Data<int64_t>(), ids->Data<int64_t>(), status->Data<int32_t>(), static_cast<int>(status->Shape().Size()),
ctl_ptr, context_ptr, prefixes->Data<int64_t>(),
Expand All @@ -79,7 +98,8 @@ class SpeculativePrefixCommit final : public CudaKernel {
lengths->MutableData<int32_t>(), reinterpret_cast<uint16_t*>(out_cos->MutableData<BFloat16>()),
reinterpret_cast<uint16_t*>(out_sin->MutableData<BFloat16>()), selection->MutableData<int64_t>(),
static_cast<int>(rows), head_rows_, static_cast<int>(rope[0]), static_cast<int>(rope[1]), capacity_, Stream(context),
true, ctl->Data<int64_t>(), static_cast<int>(ctl->Shape().Size()), committed->Data<int64_t>());
true, ctl->Data<int64_t>(), ctl_stride, committed->Data<int64_t>(),
static_cast<int>(batch), ids_stride, prefix_stride, prefix2_stride, ptrs_stride);
}

private:
Expand Down
128 changes: 82 additions & 46 deletions onnxruntime/contrib_ops/cuda/llm/speculative_fold.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,87 +22,119 @@ __global__ void MidKernel(const int64_t* tokens, const int64_t* ids, const int32
int64_t* head_ids, int32_t* head_c, int64_t* slots, int32_t* lengths,
uint16_t* head_cos, uint16_t* head_sin, int64_t* selection,
int rows, int head_rows, int rope_rows, int rope_width, int capacity,
bool publish_ring) {
bool publish_ring, int batch_size, int ids_stride,
int prefix_stride, int prefix2_stride, int ptrs_stride) {
const int b = static_cast<int>(blockIdx.x);
if (b >= batch_size) return;
const int eff_ids_stride = ids_stride > 0 ? ids_stride : rows;
const int eff_ctl_stride = ctl_count > 0 ? ctl_count : (head_rows > 2 ? 16 : 8);
const int eff_prefix_stride = prefix_stride > 0 ? prefix_stride : (head_rows > 2 ? 2 * head_rows : 2);
const int64_t* tokens_b = tokens + static_cast<int64_t>(b) * rows;
const int64_t* ids_b = ids + static_cast<int64_t>(b) * eff_ids_stride;
const int64_t* input_ctl_b = input_ctl != nullptr ? input_ctl + static_cast<int64_t>(b) * eff_ctl_stride : nullptr;
int64_t* ctl_b = ctl + static_cast<int64_t>(b) * eff_ctl_stride;
const int64_t* input_context_b = input_context != nullptr ? input_context + static_cast<int64_t>(b) * 2 : nullptr;
int64_t* context_b = context + static_cast<int64_t>(b) * 2;
const int64_t* prefixes_b = prefixes + static_cast<int64_t>(b) * eff_prefix_stride;
const int64_t* prefix2_b = (prefix2_stride > 0 || b == 0)
? (prefix2 + static_cast<int64_t>(b) * prefix2_stride)
: (eff_prefix_stride >= 4 ? prefixes_b + 2 : prefix2);
int64_t* head_ids_b = head_ids + static_cast<int64_t>(b) * head_rows;
int64_t* slots_b = slots + static_cast<int64_t>(b) * head_rows;
uint16_t* head_cos_b = head_cos + static_cast<int64_t>(b) * head_rows * rope_width;
uint16_t* head_sin_b = head_sin + static_cast<int64_t>(b) * head_rows * rope_width;

__shared__ int64_t length;
__shared__ int64_t cache_length;
__shared__ int64_t errors;
__shared__ bool muted;
__shared__ int64_t local_ptrs[2];
if (threadIdx.x == 0) {
if (input_ctl != nullptr && input_ctl != ctl) {
for (int i = 0; i < ctl_count; ++i) ctl[i] = input_ctl[i];
if (input_ctl_b != nullptr && input_ctl_b != ctl_b) {
for (int i = 0; i < eff_ctl_stride; ++i) ctl_b[i] = input_ctl_b[i];
}
if (input_context_b != nullptr && input_context_b != context_b) {
context_b[0] = input_context_b[0];
context_b[1] = input_context_b[1];
}
if (input_context != nullptr && input_context != context) {
context[0] = input_context[0];
context[1] = input_context[1];
if (ptrs_stride > 0) {
local_ptrs[0] = ptrs[2 * b + 0];
local_ptrs[1] = ptrs[2 * b + 1];
} else {
local_ptrs[1] = ptrs[1];
local_ptrs[0] = (ptrs[0] != 0 && ptrs[1] > 0 &&
ptrs[1] <= LLONG_MAX / (8 * static_cast<int64_t>(sizeof(int64_t)) * batch_size))
? ptrs[0] + static_cast<int64_t>(b) * ptrs[1] * 8 * static_cast<int64_t>(sizeof(int64_t))
: ptrs[0];
}
muted = ctl[3] != 0;
length = muted ? 0 : ctl[1];
cache_length = muted ? 0 : ctl[7];
muted = ctl_b[3] != 0;
length = muted ? 0 : ctl_b[1];
cache_length = muted ? 0 : ctl_b[7];
errors = 0;
if (length < 0 || length > static_cast<int64_t>(rope_rows) - head_rows - 1) errors |= 8;
if (cache_length < 0 || cache_length > capacity - head_rows) errors |= 4;
if (!muted && !ValidRing(ptrs, ctl[4])) errors |= 16;
if (!muted && !ValidRing(local_ptrs, ctl_b[4])) errors |= 16;
if (rows > head_rows && !muted) errors |= 4;
if (errors) {
ctl[5] |= errors;
ctl[3] = 1;
ctl_b[5] |= errors;
ctl_b[3] = 1;
}
}
__syncthreads();
for (int64_t i = threadIdx.x; i < static_cast<int64_t>(head_rows) * rope_width; i += blockDim.x) {
const int64_t position = length + 1 + i / rope_width;
head_cos[i] = !errors ? cosine[position * rope_width + i % rope_width] : 0;
head_sin[i] = !errors ? sine[position * rope_width + i % rope_width] : 0;
head_cos_b[i] = !errors ? cosine[position * rope_width + i % rope_width] : 0;
head_sin_b[i] = !errors ? sine[position * rope_width + i % rope_width] : 0;
}
if (threadIdx.x != 0) return;
if (errors) {
for (int row = 0; row < head_rows; ++row) head_ids[row] = slots[row] = 0;
head_c[0] = lengths[0] = 0;
selection[0] = -1;
for (int row = 0; row < head_rows; ++row) head_ids_b[row] = slots_b[row] = 0;
head_c[b] = lengths[b] = 0;
selection[b] = -1;
return;
}
int accepted = 0;
if (!muted) {
while (accepted + 1 < rows && accepted + 1 < head_rows &&
tokens[accepted] == ids[accepted + 1]) ++accepted;
tokens_b[accepted] == ids_b[accepted + 1]) ++accepted;
}
for (int row = 0; row < head_rows; ++row) {
head_ids[row] = row < accepted ? ids[row + 1] : tokens[row < rows ? row : rows - 1];
slots[row] = cache_length + row;
head_ids_b[row] = row < accepted ? ids_b[row + 1] : tokens_b[row < rows ? row : rows - 1];
slots_b[row] = cache_length + row;
}
head_c[0] = static_cast<int32_t>(cache_length);
lengths[0] = static_cast<int32_t>(cache_length + head_rows - 1);
selection[0] = accepted;
head_c[b] = static_cast<int32_t>(cache_length);
lengths[b] = static_cast<int32_t>(cache_length + head_rows - 1);
selection[b] = accepted;
if (muted) return;
int64_t validation = 0;
for (int i = 0; i < status_count; ++i) validation |= status[i];
const int64_t* prefix = accepted == 1 ? prefix2 : prefixes + 2 * accepted;
context[0] = prefix[0];
context[1] = prefix[1];
const int64_t pending = tokens[accepted];
ctl[0] = pending;
ctl[1] = length + accepted + 1;
ctl[7] = cache_length + accepted + 1;
++ctl[4];
ctl[5] |= validation ? 2 : 0;
ctl[6] = accepted;
volatile int64_t* record = reinterpret_cast<int64_t*>(ptrs[0]) + 8 * ((ctl[4] - 1) % ptrs[1]);
const int64_t* prefix = accepted == 1 ? prefix2_b : prefixes_b + 2 * accepted;
context_b[0] = prefix[0];
context_b[1] = prefix[1];
const int64_t pending = tokens_b[accepted];
ctl_b[0] = pending;
ctl_b[1] = length + accepted + 1;
ctl_b[7] = cache_length + accepted + 1;
++ctl_b[4];
ctl_b[5] |= validation ? 2 : 0;
ctl_b[6] = accepted;
volatile int64_t* record = reinterpret_cast<int64_t*>(local_ptrs[0]) + 8 * ((ctl_b[4] - 1) % local_ptrs[1]);
auto pack_lo_hi = [](int64_t lo, int64_t hi, bool has_hi) -> int64_t {
if (!has_hi) return lo;
return static_cast<int64_t>(
static_cast<uint64_t>(static_cast<uint32_t>(lo)) |
(static_cast<uint64_t>(static_cast<uint32_t>(hi)) << 32));
};
record[1] = pack_lo_hi(tokens[0], rows > 4 ? tokens[4] : 0, rows > 4);
record[2] = pack_lo_hi(rows > 1 ? tokens[1] : tokens[0], rows > 5 ? tokens[5] : 0, rows > 5);
record[3] = pack_lo_hi(accepted, rows > 6 ? tokens[6] : 0, rows > 6);
record[1] = pack_lo_hi(tokens_b[0], rows > 4 ? tokens_b[4] : 0, rows > 4);
record[2] = pack_lo_hi(rows > 1 ? tokens_b[1] : tokens_b[0], rows > 5 ? tokens_b[5] : 0, rows > 5);
record[3] = pack_lo_hi(accepted, rows > 6 ? tokens_b[6] : 0, rows > 6);
record[4] = validation;
record[5] = pack_lo_hi(pending, rows > 7 ? tokens[7] : 0, rows > 7);
record[6] = ctl[1];
record[7] = rows > 2 ? pack_lo_hi(tokens[2], rows > 3 ? tokens[3] : 0, rows > 3) : 0;
record[5] = pack_lo_hi(pending, rows > 7 ? tokens_b[7] : 0, rows > 7);
record[6] = ctl_b[1];
record[7] = rows > 2 ? pack_lo_hi(tokens_b[2], rows > 3 ? tokens_b[3] : 0, rows > 3) : 0;
if (publish_ring) {
__threadfence_system();
record[0] = ctl[4];
record[0] = ctl_b[4];
}
}

Expand All @@ -118,11 +150,15 @@ Status LaunchMid(const int64_t* tokens, const int64_t* ids, const int32_t* statu
int64_t* head_ids, int32_t* head_c, int64_t* slots, int32_t* lengths,
uint16_t* head_cos, uint16_t* head_sin, int64_t* selection,
int rows, int head_rows, int rope_rows, int rope_width, int capacity, cudaStream_t stream,
bool publish_ring, const int64_t* input_ctl, int ctl_count, const int64_t* input_context) {
MidKernel<<<1, 256, 0, stream>>>(tokens, ids, status, status_count, input_ctl, ctl, ctl_count,
input_context, context, prefixes, prefix2, ptrs,
cosine, sine, head_ids, head_c, slots, lengths, head_cos, head_sin, selection,
rows, head_rows, rope_rows, rope_width, capacity, publish_ring);
bool publish_ring, const int64_t* input_ctl, int ctl_count, const int64_t* input_context,
int batch_size, int ids_stride, int prefix_stride, int prefix2_stride, int ptrs_stride) {
if (batch_size <= 0) return Status::OK();
MidKernel<<<static_cast<unsigned int>(batch_size), 256, 0, stream>>>(
tokens, ids, status, status_count, input_ctl, ctl, ctl_count,
input_context, context, prefixes, prefix2, ptrs,
cosine, sine, head_ids, head_c, slots, lengths, head_cos, head_sin, selection,
rows, head_rows, rope_rows, rope_width, capacity, publish_ring,
batch_size, ids_stride, prefix_stride, prefix2_stride, ptrs_stride);
return LaunchStatus(cudaGetLastError());
}

Expand Down
3 changes: 2 additions & 1 deletion onnxruntime/contrib_ops/cuda/llm/speculative_fold.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ Status LaunchMid(const int64_t* tokens, const int64_t* ids, const int32_t* statu
uint16_t* head_cos, uint16_t* head_sin, int64_t* selection,
int rows, int head_rows, int rope_rows, int rope_width, int capacity, cudaStream_t stream,
bool publish_ring = true, const int64_t* input_ctl = nullptr, int ctl_count = 0,
const int64_t* input_context = nullptr);
const int64_t* input_context = nullptr, int batch_size = 1, int ids_stride = 0,
int prefix_stride = 0, int prefix2_stride = 0, int ptrs_stride = 0);

} // namespace onnxruntime::contrib::cuda::speculative_fold
Loading
Loading