File size: 3,871 Bytes
4a28d4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | #include "src/turbomind/generation/guided_decoding.h"
#include "src/turbomind/comm/host_comm.h"
#include "src/turbomind/core/allocator.h"
#include "src/turbomind/engine/batch.h"
#include "src/turbomind/kernels/apply_token_bitmask_inplace_cuda.h"
#include "xgrammar/matcher.h"
#include <dlpack/dlpack.h>
namespace turbomind {
struct GuidedDecoding::Data {
Tensor_<int32_t> bitmask;
bool active{};
std::vector<std::shared_ptr<xgrammar::GrammarMatcher>> matchers;
};
GuidedDecoding::GuidedDecoding(const BaseGenerationParam& base, const comm::HostComm& tp_group, int phases):
BaseGenerationParam{base}, //
tp_group_{tp_group->Split(0, 0)} // duplicate to avoid data race
{
const auto bitmask_size = xgrammar::GetBitmaskSize(vocab_size_padded_);
bitmask_buf_ = {{max_batch_size_, bitmask_size}, kCPUpinned};
output_ids_buf_ = {max_batch_size_, kCPUpinned};
for (int i = 0; i < phases; ++i) {
auto& d = data_.emplace_back(std::make_shared<Data>());
d->bitmask = empty_like(bitmask_buf_);
}
}
void GuidedDecoding::Setup(int phase, TensorMap& env)
{
auto& d = *data_.at(phase);
auto& b = *env.at("batch").data<BatchData*>()[0];
d.matchers.clear();
d.active = false;
for (const auto& r : b.rc) {
if (d.matchers.emplace_back(r->req->matcher)) {
d.active = true;
}
}
}
void GuidedDecoding::FillMask(int phase, TensorMap& env)
{
if (auto& d = *data_.at(phase); d.active) {
static_assert(sizeof(ssize_t) == sizeof(int64_t));
DLTensor dlbitmask{bitmask_buf_.data(),
DLDevice{kDLCPU, 0},
bitmask_buf_.ndim(),
xgrammar::GetBitmaskDLType(),
(int64_t*)bitmask_buf_.shape().data(),
nullptr,
0};
if (tp_group_->rank() == 0) {
for (size_t i = 0; i < d.matchers.size(); ++i) {
if (const auto& matcher = d.matchers[i]; matcher && !matcher->IsTerminated()) {
matcher->FillNextTokenBitmask(&dlbitmask, i);
}
else {
std::fill_n(bitmask_buf_.data() + i * bitmask_buf_.stride(0),
bitmask_buf_.stride(0),
static_cast<int32_t>(-1));
}
}
}
}
}
void GuidedDecoding::ApplyMask(int phase, TensorMap& env)
{
if (auto& d = *data_.at(phase); d.active) {
const ssize_t numel = d.matchers.size() * bitmask_buf_.stride(0);
if (tp_group_->n_ranks() > 1) {
// bcast the data instead of `bitmask_buf` instance (which may avoid copying the data)
comm::Broadcast(tp_group_, bitmask_buf_.data(), numel, 0);
}
Copy(bitmask_buf_.buffer(), numel, d.bitmask.buffer());
// Use logits shape(0) instead of d.matchers.size() to ensure dimension match.
// d.matchers.size() is the total number of requests in batch, but logits may be
// sliced to only include requests that are still generating (generation_size).
auto logits = env.at("logits");
ApplyTokenBitmaskInplace(logits, d.bitmask.slice(0, logits.shape(0)));
}
}
void GuidedDecoding::Update(int phase, TensorMap& env)
{
if (auto& d = *data_.at(phase); d.active) {
Copy(env.at("output_ids").buffer(), d.matchers.size(), output_ids_buf_);
core::Context::stream().Sync();
if (tp_group_->rank() == 0) {
for (size_t i = 0; i < d.matchers.size(); ++i) {
if (const auto& matcher = d.matchers[i]; matcher && !matcher->IsTerminated()) {
matcher->AcceptToken(output_ids_buf_[i]);
}
}
}
}
}
} // namespace turbomind
|