| 526 | } |
| 527 | |
| 528 | llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_context * ctx, int idx, bool grammar_first) { |
| 529 | llama_synchronize(ctx); |
| 530 | |
| 531 | // start measuring sampling time after the llama_context synchronization in order to not measure any ongoing async operations |
| 532 | const auto tm = gsmpl->tm(); |
| 533 | |
| 534 | llama_token id = LLAMA_TOKEN_NULL; |
| 535 | |
| 536 | auto & grmr = gsmpl->grmr; |
| 537 | auto & rbudget = gsmpl->rbudget; |
| 538 | auto & chain = gsmpl->chain; |
| 539 | auto & cur_p = gsmpl->cur_p; // initialized by set_logits |
| 540 | |
| 541 | // Check if a backend sampler has already sampled a token in which case we |
| 542 | // return that token id directly. |
| 543 | { |
| 544 | id = llama_get_sampled_token_ith(ctx, idx); |
| 545 | |
| 546 | if (id != LLAMA_TOKEN_NULL) { |
| 547 | LOG_DBG("%s: Backend sampler selected token: '%d'. Will not run any CPU samplers\n", __func__, id); |
| 548 | |
| 549 | GGML_ASSERT(!gsmpl->grmr && "using grammar in combination with backend sampling is not supported"); |
| 550 | GGML_ASSERT(!gsmpl->rbudget && "using reasoning budget in combination with backend sampling is not supported"); |
| 551 | |
| 552 | // TODO: simplify |
| 553 | gsmpl->cur.resize(1); |
| 554 | gsmpl->cur[0] = { id, 0.0f, 1.0f }; |
| 555 | cur_p = { gsmpl->cur.data(), gsmpl->cur.size(), 0, true }; |
| 556 | |
| 557 | return id; |
| 558 | } |
| 559 | } |
| 560 | |
| 561 | gsmpl->set_logits(ctx, idx); |
| 562 | |
| 563 | // apply reasoning budget first |
| 564 | llama_sampler_apply(rbudget, &cur_p); |
| 565 | |
| 566 | if (grammar_first && grammar_should_apply(gsmpl)) { |
| 567 | llama_sampler_apply(grmr, &cur_p); |
| 568 | } |
| 569 | |
| 570 | llama_sampler_apply(chain, &cur_p); |
| 571 | |
| 572 | id = cur_p.data[cur_p.selected].id; |
| 573 | |
| 574 | if (grammar_first || !grammar_should_apply(gsmpl)) { |
| 575 | return id; |
| 576 | } |
| 577 | |
| 578 | // check if it the sampled token fits the grammar (grammar-based rejection sampling) |
| 579 | { |
| 580 | llama_token_data single_token_data = { id, 1.0f, 0.0f }; |
| 581 | llama_token_data_array single_token_data_array = { &single_token_data, 1, -1, false }; |
| 582 | |
| 583 | llama_sampler_apply(grmr, &single_token_data_array); |
| 584 | |
| 585 | const bool is_valid = single_token_data_array.data[0].logit != -INFINITY; |