[Piggidragon]

Closes #42 (GPU sampling of the verify positions under -sm tensor; after the GPU sampler of Strata by @Niko1221, MIT).

Problem

After every verify step the target's logits are copied to the host whole: 248k floats per verified position. The CPU sampler chain then runs over the whole vocabulary. Under -sm tensor the 3060's share of each row goes over PCIe 3.0 x4.

nsys of a docker-compose edit with MTP + --spec-select (#47):

• after a 49-token verify (a long n-gram draft) both GPUs idle 7.7 ms, 5.5 ms of it the copy;
• after a 3-token MTP verify they idle 1.2 ms, 0.34 ms of it the copy.

That is ~126 ms of ~3 s of generation.

Change

• llama_set_logits_topk(ctx, k). After the decode's graph, each device runs a small graph (top_k + get_rows) on its share of the logits, i.e. the rows of the output head it holds. The graph is queued on that device's stream, and only the k candidates per share and row are copied.
  • The sampled-logits API (llama_get_sampled_logits_ith / _candidates_ith) reports them with their tokens. This also works for the reordered head of --spec-draft-vocab (spec : the MTP draft head as a view of the reordered output head, without memory of its own #46).
  • llama_get_logits_ith / llama_get_logits copy a whole row from the device on request, until a graph is computed again.
  • The mode is used only when all outputs of the decode come from one ubatch and the logits are in device memory. Otherwise all logits are copied as before.
  • CUDA TOP_K for 3 rows takes 44 us (4070, 155k columns) / 50 us (3060, 93k); for 49 rows 148 / 412 us.
• common_sampler_logits_topk(params, vocab). It returns the k that a sampler chain needs to sample the same token: its top_k (1 when greedy) plus the tokens it can lower before that point. Those are logit biases <= 0, the model's suppressed tokens, and penalties of the last n tokens. Production: 20 + 64 = 84.
  • It returns 0 (all logits) for chains that read the whole distribution: n_probs, mirostat, DRY, top-n-sigma, top-p/min-p/typical/XTC before the top-k, dynamic temperature, positive biases.
  • A grammar that rejects the sampled token, and a token forced by the reasoning budget, sample over a whole row (common_sampler_sample). So tool-call grammars still work, and they cost a row copy only when they intervene.
• Server: --logits-topk (opt-in). k is set per batch to the largest that any active slot needs, or 0 if one slot needs all logits.
• ggml-meta: ggml_backend_meta_n_simple_backends / ggml_backend_meta_get_simple_backend.

Results

Qwen3.8 27B UD-Q4_K_S, 4070 + 3060, -sm tensor, MTP 2 + --spec-select + draft list; 4 requests per task (median), temperature 1.0 / greedy:

task	without	--logits-topk (run 1)	(run 2)
docker-compose edit	115.2 / 120.0 t/s	133.2 / 132.3	131.3 / 136.8
code edit	70.7 / 76.0	70.5 / 76.0	72.2 / 76.6
German letter	48.0 / 55.7	48.9 / 56.7	48.8 / 56.6
explanation	57.3 / 63.1	58.2 / 64.4	58.2 / 64.2
mean	72.8 / 78.7	77.7 / 82.4	77.6 / 83.6

(The "without" column is run 1; the base of run 2 was 71.5 / 76.8 with some CPU load.)

• Exactness: greedy and seeded temperature-1.0 outputs are bitwise identical with and without (compose, explanation, German: 6 of 6).
• test-llama-archs -a '^(llama|qwen35|qwen35moe)$': 64 of 64 passed, with a new check on every device config, also with a draft-vocabulary list (reordered head).
  • The candidates hold the k largest logits of each row with their tokens and values.
  • A whole row read on request equals the row without top-k, bitwise.
• VRAM: a few hundred KiB of graph buffers per device; the TOP_K temporaries come from the CUDA pool, ~4 MB for 98 rows.

🤖 Generated with Claude Code

[Piggidragon]

Review: 3 findings, fixed in 1005c47

1. HIGH: whole rows gone before the sampling (server abort).

• Cause: with --phase-aware-workspace, the MTP draft context computes in the target's buffers before post_decode() samples.
• Effect: a grammar resample or a token forced by the reasoning budget hit the assert in topk_full_row. Reproduced with thinking_budget_tokens: 8 and with a grammar request.
• Fix:
  • The server takes all logits for a slot whose grammar constrains the next token (common_sampler_needs_all_logits). A lazy grammar that waits for its trigger (llama_sampler_grammar_waiting) constrains none, so tool-call requests keep the top-k until the call starts.
  • A draft verification stops at the token that starts the grammar. That token counts as sampled, and the next step samples from all logits.
  • A replay cannot stop early, so a replay of a slot with a grammar takes all logits.
  • The token the reasoning budget forces needs no row: it joins the candidates (common_reasoning_budget_get_forced).

2. HIGH: top-k with the outputs in the first of several ubatches. The later ubatches overwrote the rows. Top-k now runs only when the outputs are in the last ubatch of the decode. test-llama-archs checks outputs in the first and in the last of three ubatches.

3. MEDIUM: penalties that raise logits. A negative frequency or presence penalty, or a repeat penalty < 1, now takes all logits.

Verified (production model and args, --logits-topk)

case	before	after
reasoning budget 8 tokens	abort	ok, output equals the run without top-k
grammar 7[0-9]{3}	abort	ok, 7189 as without top-k
grammar-heavy slot generating while another slot prefills 44k tokens	abort	both ok
frequency_penalty -1.0, seed 42	different text	same text as without top-k
reviewer's multi-ubatch repro (outputs in the first ubatch)	abort	whole rows equal the reference

• Tool-call chats (two tools, three prompts, greedy and seeded), MTP drafts: the output is the same with and without --logits-topk, compared without the random call ids. Two runs without top-k were also equal to each other.
• Production drafts (ngram-simple,draft-mtp --spec-select, draft vocabulary, compact mask): grammar, budget, tool calls and the 44k-prefill case all without abort.
• test-llama-archs -a '^(llama|qwen35|qwen35moe)$': 64/64 passed.

🤖 Generated with Claude Code