Skip to content

Commit fbbbb4f

Browse files
committed
fix: fix initialized bitmask buffer when mixing guided requests and non-guided requests
1 parent 98a9117 commit fbbbb4f

File tree

1 file changed

+9
-10
lines changed

1 file changed

+9
-10
lines changed

src/turbomind/layers/sampling_layers/GuidedDecodeMaskLayer.cc

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -47,17 +47,16 @@ void GuidedDecodeMaskLayer<T>::Forward(TensorMap& args)
4747

4848
TM_CHECK(bsz == matchers_.size());
4949

50-
const auto bitmask_size = bitmask_buf_.shape(1);
51-
std::vector<int64_t> bitmask_shape = {bsz, bitmask_size};
50+
const auto bitmask_size = bitmask_buf_.shape(1);
51+
std::vector<int64_t> bitmask_shape = {bsz, bitmask_size};
52+
int32_t* data = bitmask_buf_.data();
53+
size_t total_elements = bitmask_shape[0] * bitmask_shape[1];
54+
bool need_apply = false;
55+
DLTensor bitmask_dltensor{
56+
data, DLDevice{kDLCPU, 0}, bitmask_buf_.ndim(), xgrammar::GetBitmaskDLType(), bitmask_shape.data(), nullptr, 0};
57+
58+
std::fill(data, data + total_elements, 0xffffffff);
5259

53-
DLTensor bitmask_dltensor{bitmask_buf_.data(),
54-
DLDevice{kDLCPU, 0},
55-
bitmask_buf_.ndim(),
56-
xgrammar::GetBitmaskDLType(),
57-
bitmask_shape.data(),
58-
nullptr,
59-
0};
60-
bool need_apply = false;
6160
for (size_t i = 0; i < bsz; ++i) {
6261
const auto& matcher = matchers_[i];
6362
if (matcher) {

0 commit comments

Comments
 (0)