Repository navigation
feat(sampling): implement XTC sampling end to end - #802
Conversation
The server request schema already accepted xtc_probability and xtc_threshold (src/server/types/request.rs) but nothing read them, so clients setting these fields silently got default sampling. Core filter (src/lib/mlxcel-core/src/sampling.rs): apply_xtc_filter removes all but the least-probable token among those exceeding xtc_threshold (no-op below two candidates), always preserving ids in an allowlist. apply_xtc_step gates the filter with a draw from the same per-request seeded global MLX RNG the fused categorical sampler consumes (mlx::core::random's default key is a thread-local sequence split synchronously at graph-construction time, so drawing here before the categorical draw keeps the whole step reproducible for a fixed seed). Both are wired into sample_token_optimized_core alongside the existing penalty steps, gated on xtc_probability > 0.0 so a disabled request never advances the RNG stream or pays any array-op cost. Plumbing (mirrors the DRY parameter path from request to sampler): SamplingConfig gained xtc_probability, xtc_threshold, and xtc_special_token_ids (src/lib/mlxcel-core/src/generate.rs); ResolvedSamplingParams and build_sampling_config thread the first two through for both the greedy and sampled branches (src/execution/sampling.rs); RequestOptionOverrides and build_server_generate_options resolve per-request values against a disabled default, since XTC has no server-level CLI default (src/server/request_options.rs); chat.rs's build_generate_options maps SamplingParams into the overrides, which the /v1/completions and /v1/responses routes inherit for free since they already call it. The special-token allowlist (tokenizer newline id(s) plus the full merged end-of-sequence set) is resolved once per model load (resolve_xtc_newline_token_ids in src/server/model_worker.rs, wired into both scheduler-construction call sites) and combined with the per-request merged EOS set in BatchScheduler::enqueue_request, only when xtc_probability > 0.0. Request validation: chat_completions rejects xtc_threshold outside 0.0..=0.5 and xtc_probability outside 0.0..=1.0 with a 400 invalid_request_error, via a small pure validate_xtc_params helper alongside the existing top_logprobs range check. Tests: mlxcel-core sampling:: covers the filter (keeps least-probable, allowlist survives, no-op below two/zero candidates) and the step gate (0.0 never fires, 1.0 always fires, a mid value is reproducible for a fixed seed). Server-side: request_options:: and execution::sampling:: cover default/override resolution; chat:: covers the validation helper's boundaries; model_worker_tests:: covers newline-id resolution. All pre-existing tests in the touched modules continue to pass unchanged, and every non-XTC construction site was updated to keep the previous defaults, so requests that never set these fields behave byte-identically to before. Closes #772
XTC was silently skipped for the entire decode phase. config_supports_fused_batch only excluded rows that need token history or carry a token bias, so an XTC-only request (no penalties, no bias) was admitted to the batched fused sampler, which calls fused_sample directly and never runs apply_xtc_step. Both the batched decode dispatch and the async lookahead priming go through this gate via row_supports_fused_batch, so XTC applied only to the first token (finish_prefill uses the per-row sampler) and was dropped for every token after it, failing issue #772's acceptance criterion that xtc_probability=1.0 observably changes the distribution. XTC (xtc_probability > 0.0) is now treated as a per-row logit edit like a non-empty token bias, so it disqualifies the fused fast path and runs through the per-row sampler for every decode token. XTC range validation (the documented 400) was wired only into /v1/chat/completions. /v1/completions and /v1/responses both reach XTC through the shared build_generate_options and both flatten SamplingParams, so out-of-range xtc_probability / xtc_threshold bypassed the 400 and still mutated sampling. validate_xtc_params is now pub(crate) and called in the completions and create_response handlers, reading the same fields those handlers already forward into the sampling config. Adds a fused-batch gate unit test and out-of-range XTC rejection tests for both /v1/completions and /v1/responses. Refs #772.
Implementation Review SummaryIntent
Findings Addressed
Remaining Items
Verification
Fix commit: 7d09dd7 |
`docs/python-client.md`'s "Sampling parameters" section lists the server-specific knobs that are not part of the OpenAI schema (top_k, min_p, repetition_penalty, DRY settings, loop-detection fields) but never mentioned xtc_probability / xtc_threshold, which PR #802 wires end to end. Add both fields to that list along with their valid ranges (xtc_probability 0.0-1.0, xtc_threshold 0.0-0.5, both defaulting to disabled) and a one-line note that out-of-range values return a 400 on /v1/chat/completions, /v1/completions, and /v1/responses. No other doc under docs/ enumerates per-request sampling fields (responses-api.md only tracks the official OpenAI subset), and there is no docs/ko twin yet, so this is the only place XTC's request-facing surface needed a mention. Test coverage was already reviewed as sufficient (21 XTC tests: 8 in mlxcel-core, 13 in the server crate) and no source changed here. Validation: - cargo fmt --all -- --check: clean Refs #772
PR Finalization CompleteTest coverageReviewed the existing 21 XTC tests (8 in Documentation
Updated the PR body: it was stale after the review-cycle fix commit, still describing the pre-fix state (validation only in Lint/format
Pushed commit |
## Summary Implement XTC (Exclude Top Choices) sampling end to end: the server request schema already accepted `xtc_probability` and `xtc_threshold` but nothing read them, so clients setting these fields silently got default sampling. ## What changed - `src/lib/mlxcel-core/src/sampling.rs`: new `apply_xtc_filter` (removes all but the least-probable token among those exceeding `xtc_threshold`, no-op below two candidates, never removes allowlisted ids) and `apply_xtc_step` (gates the filter with a draw from the same per-request seeded global MLX RNG the fused categorical sampler consumes, keeping the whole decode step reproducible for a fixed seed). Both are wired into `sample_token_optimized_core`, gated on `xtc_probability > 0.0` so a disabled request never advances the RNG stream or pays any array-op cost. - `src/lib/mlxcel-core/src/generate.rs`: `SamplingConfig` gains `xtc_probability`, `xtc_threshold`, and `xtc_special_token_ids`. - `src/execution/sampling.rs`: `ResolvedSamplingParams` / `build_sampling_config` thread the new fields through for both the greedy and sampled branches. - `src/server/request_options.rs`: `RequestOptionOverrides` / `build_server_generate_options` resolve per-request values against a disabled default (XTC has no server-level CLI default). - `src/server/routes/chat.rs`: `build_generate_options` maps `SamplingParams` into the overrides, and the shared `validate_xtc_params` helper (`0.0..=0.5` for `xtc_threshold`, `0.0..=1.0` for `xtc_probability`) sits next to the existing `top_logprobs` range check. It now runs in `chat_completions`, `completions`, and `create_response`, each returning a 400 `invalid_request_error` before any generation work begins. - `src/server/model_worker.rs` / `src/server/batch/scheduler.rs`: the special-token allowlist (tokenizer newline id(s) plus the full merged end-of-sequence set) is resolved once per model load (`resolve_xtc_newline_token_ids`, wired into both scheduler-construction call sites) and combined with the per-request merged EOS set in `BatchScheduler::enqueue_request`, only when `xtc_probability > 0.0`. - `src/server/routes/native_completion.rs`, `src/distributed/disaggregated/serving_protocol.rs`, `src/bin/bench_decode.rs`, `src/commands/generate.rs`, `src/commands/chat_tests.rs`: updated every other `SamplingConfig` / `ResolvedSamplingParams` / `RequestOptionOverrides` construction site to keep the previous disabled defaults, so nothing outside the OpenAI-compatible chat/completions/responses routes changes behavior. ## Review-cycle fixes - **CRITICAL, fused-batch bypass**: `config_supports_fused_batch` only excluded rows that need token history or carry a token bias, so an XTC-only request (no penalties, no bias) was admitted to the batched fused sampler, which calls `fused_sample` directly and never runs `apply_xtc_step`. Both the batched decode dispatch and the async lookahead priming go through this gate via `row_supports_fused_batch`, so XTC applied only to the first token (`finish_prefill` uses the per-row sampler) and was silently dropped for every token after it, failing issue lablup#772's acceptance criterion that `xtc_probability=1.0` observably changes the distribution. XTC is now treated as a per-row logit edit like a non-empty token bias: `xtc_probability > 0.0` disqualifies the fused fast path, so every decode token runs through the per-row sampler. - **MEDIUM, route validation coverage**: `validate_xtc_params` was wired only into `/v1/chat/completions`. `/v1/completions` and `/v1/responses` both reach XTC through the shared `build_generate_options` and both flatten `SamplingParams`, so an out-of-range `xtc_probability` / `xtc_threshold` bypassed the documented 400 and still mutated sampling on those two routes. The helper is now `pub(crate)` and called from the `completions` and `create_response` handlers too. ## Test plan - [x] `cargo build --release --features cuda --workspace` (zero errors) - [x] `cargo test --release --features cuda -p mlxcel-core sampling::`: 8 XTC tests: filter keeps the least-probable token, allowlist survives removal, no-op below two/zero candidates, gate at 0.0 never fires, gate at 1.0 always fires, mid-probability gate reproducible for a fixed seed, and the fused-batch gate disqualifies `xtc_probability > 0.0` - [x] `cargo test --release --features cuda --lib execution::sampling::`: 2 passed - [x] `cargo test --release --features cuda --lib request_options::`: 12 passed - [x] `cargo test --release --features cuda --lib server::routes::chat::`: 27 passed (5 `validate_xtc_params` tests) - [x] `cargo test --release --features cuda --lib server::routes::completions::`: 3 new tests: out-of-range `xtc_probability`, out-of-range `xtc_threshold`, in-range accepted - [x] `cargo test --release --features cuda --lib server::routes::responses::`: 3 new tests: out-of-range `xtc_probability`, out-of-range `xtc_threshold`, in-range accepted - [x] `cargo test --release --features cuda --lib server::batch::scheduler::tests::`: 69 passed - [x] `cargo test --release --features cuda --lib model_worker`: 28 passed (2 `resolve_xtc_newline_token_ids` tests) - [x] `cargo test --release --features cuda --lib serving_protocol`: 10 passed - [x] `cargo test --release --features cuda --bin mlxcel chat_options_new_sets_conventional_defaults`: 1 passed - [x] `cargo test --release --features cuda --all-targets` (server crate, full run): exit 0 - [x] `cargo fmt --all -- --check`: clean Closes lablup#772
Summary
Implement XTC (Exclude Top Choices) sampling end to end: the server request schema already accepted
xtc_probabilityandxtc_thresholdbut nothing read them, so clients setting these fields silently got default sampling.What changed
src/lib/mlxcel-core/src/sampling.rs: newapply_xtc_filter(removes all but the least-probable token among those exceedingxtc_threshold, no-op below two candidates, never removes allowlisted ids) andapply_xtc_step(gates the filter with a draw from the same per-request seeded global MLX RNG the fused categorical sampler consumes, keeping the whole decode step reproducible for a fixed seed). Both are wired intosample_token_optimized_core, gated onxtc_probability > 0.0so a disabled request never advances the RNG stream or pays any array-op cost.src/lib/mlxcel-core/src/generate.rs:SamplingConfiggainsxtc_probability,xtc_threshold, andxtc_special_token_ids.src/execution/sampling.rs:ResolvedSamplingParams/build_sampling_configthread the new fields through for both the greedy and sampled branches.src/server/request_options.rs:RequestOptionOverrides/build_server_generate_optionsresolve per-request values against a disabled default (XTC has no server-level CLI default).src/server/routes/chat.rs:build_generate_optionsmapsSamplingParamsinto the overrides, and the sharedvalidate_xtc_paramshelper (0.0..=0.5forxtc_threshold,0.0..=1.0forxtc_probability) sits next to the existingtop_logprobsrange check. It now runs inchat_completions,completions, andcreate_response, each returning a 400invalid_request_errorbefore any generation work begins.src/server/model_worker.rs/src/server/batch/scheduler.rs: the special-token allowlist (tokenizer newline id(s) plus the full merged end-of-sequence set) is resolved once per model load (resolve_xtc_newline_token_ids, wired into both scheduler-construction call sites) and combined with the per-request merged EOS set inBatchScheduler::enqueue_request, only whenxtc_probability > 0.0.src/server/routes/native_completion.rs,src/distributed/disaggregated/serving_protocol.rs,src/bin/bench_decode.rs,src/commands/generate.rs,src/commands/chat_tests.rs: updated every otherSamplingConfig/ResolvedSamplingParams/RequestOptionOverridesconstruction site to keep the previous disabled defaults, so nothing outside the OpenAI-compatible chat/completions/responses routes changes behavior.Review-cycle fixes
config_supports_fused_batchonly excluded rows that need token history or carry a token bias, so an XTC-only request (no penalties, no bias) was admitted to the batched fused sampler, which callsfused_sampledirectly and never runsapply_xtc_step. Both the batched decode dispatch and the async lookahead priming go through this gate viarow_supports_fused_batch, so XTC applied only to the first token (finish_prefilluses the per-row sampler) and was silently dropped for every token after it, failing issue feat(sampling): implement XTC sampling (request params are accepted but ignored) #772's acceptance criterion thatxtc_probability=1.0observably changes the distribution. XTC is now treated as a per-row logit edit like a non-empty token bias:xtc_probability > 0.0disqualifies the fused fast path, so every decode token runs through the per-row sampler.validate_xtc_paramswas wired only into/v1/chat/completions./v1/completionsand/v1/responsesboth reach XTC through the sharedbuild_generate_optionsand both flattenSamplingParams, so an out-of-rangextc_probability/xtc_thresholdbypassed the documented 400 and still mutated sampling on those two routes. The helper is nowpub(crate)and called from thecompletionsandcreate_responsehandlers too.Test plan
cargo build --release --features cuda --workspace(zero errors)cargo test --release --features cuda -p mlxcel-core sampling::: 8 XTC tests: filter keeps the least-probable token, allowlist survives removal, no-op below two/zero candidates, gate at 0.0 never fires, gate at 1.0 always fires, mid-probability gate reproducible for a fixed seed, and the fused-batch gate disqualifiesxtc_probability > 0.0cargo test --release --features cuda --lib execution::sampling::: 2 passedcargo test --release --features cuda --lib request_options::: 12 passedcargo test --release --features cuda --lib server::routes::chat::: 27 passed (5validate_xtc_paramstests)cargo test --release --features cuda --lib server::routes::completions::: 3 new tests: out-of-rangextc_probability, out-of-rangextc_threshold, in-range acceptedcargo test --release --features cuda --lib server::routes::responses::: 3 new tests: out-of-rangextc_probability, out-of-rangextc_threshold, in-range acceptedcargo test --release --features cuda --lib server::batch::scheduler::tests::: 69 passedcargo test --release --features cuda --lib model_worker: 28 passed (2resolve_xtc_newline_token_idstests)cargo test --release --features cuda --lib serving_protocol: 10 passedcargo test --release --features cuda --bin mlxcel chat_options_new_sets_conventional_defaults: 1 passedcargo test --release --features cuda --all-targets(server crate, full run): exit 0cargo fmt --all -- --check: cleanCloses #772