[danielhanchen]

Fixes speculative drafts with backend sampling on multi-slot servers: with -np N, every slot except slot 0 gets wrong draft token ids, so DSpark / DFlash acceptance collapses for all concurrent requests but one.

Problem

Draft contexts sample on the backend by default (common_params_speculative::backend_sampling = true), with a top-k chain per sequence and one output per block position. llm_graph_context::build_sampling builds those chains for every sequence in every ubatch so the graph topology stays fixed, but only marks the chains of the sequence that is actually in the ubatch as outputs.

ggml_gallocr decides whether to re-plan a graph from node count and tensor sizes only, not from the output flags. A later graph with the same topology therefore reuses the address plan of the first one. Draft ubatches hold one sequence each, so the plan is made for the sequence-0 ubatch. There, the chains of sequences 1..N-1 are intermediates whose memory is freed and reused, which puts all candidate-id buffers of each such sequence at one address. When sequence k runs, every block position writes its top-k ids to that same buffer and the last one wins. The logits rows live elsewhere, so the probabilities look right while the ids are wrong.

Symptoms on a server with N slots and N concurrent requests:

• slot 0 drafts normally
• slots 1..N-1 draft the same token at every block position and accept almost nothing
• --no-spec-draft-backend-sampling hides the problem

Fix

In build_sampling, flag the outputs of every chain (sampled, probs, logits, candidates) whether its sequence is active or not, so the allocation plan no longer depends on which sequence the first ubatch held. The plain logits view is skipped as before. This costs a few small extra output tensors per chain.

A broader alternative would be to make ggml_gallocr re-plan when a node's output flag differs from the stored plan. That touches every backend, so this change stays local to the sampling graph.

Evidence

Same upstream base for both arms, CUDA build for sm_100 (B200), Qwen3-4B-Q8_0.gguf target, deepseek-ai/dspark_qwen3_4b_block7 draft converted with --target-model-dir, --spec-type draft-dspark -np 8 -fa on, 8 different greedy chat requests sent at once, then the same 8 one at a time:

build	requests	overall acceptance	per request
before	8 at once	0.203	0.01 0.59 0.67 0.00 0.21 0.00 0.17 0.02
before	one at a time	0.667	0.94 0.56 0.80 0.31 0.97 0.36 0.65 0.37
after	8 at once	0.695	0.94 0.59 0.82 0.42 0.97 0.36 0.71 0.37
after	one at a time	0.677	0.94 0.59 0.80 0.42 0.97 0.36 0.65 0.37

The one-at-a-time numbers are the per-request ceiling. Before the fix most concurrent requests stay far below it; after the fix all of them reach it. Even one at a time the old build loses some acceptance, because the server can hand a sequential request to a slot other than 0.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-29T13:28:29.672342Z	b240c56	PR opened
🔒 Security Review	✅ Completed 2026-10-29T13:29:10.847462Z	b240c56	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b240c56069

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Remove the hard wrapping from the new comment

The new comment splits one sentence across three fixed-width lines, contrary to the repository's explicit rule in AGENTS.md: "Do NOT split a line into multiple lines mid-sentence". Keep the allocator-lifetime explanation, but shorten it or leave the sentence unwrapped to follow the documented readability convention.

AGENTS.md reference: AGENTS.md:L84-L84

Useful? React with 👍 / 👎.