feat(sampling): add hyperparameter-free p-less sampling (`p_less`)
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Summary
mlxcel has no p-less sampler. P-less sampling is a parameter-free truncation: it keeps every token whose probability is at least the collision probability `L = sum_v p(v)^2` of the temperature-scaled distribution (the exponential of the order-2 Renyi entropy) and masks the rest to `-inf`. Since `L <= max_v p(v)`, the most likely token always survives, and a sharply peaked row collapses to its argmax while a flat row keeps most of the vocabulary. This issue adds the filter as a boolean `p_less` request field and `--p-less` CLI flag on the row-filter hook from #1375.
## Current behavior
`grep -rn p_less src/` returns nothing. `SamplingConfig` (`src/lib/mlxcel-core/src/generate.rs:810-874`) has no such field and the fused chain (`src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:5039-5125`) implements temperature, top_k, top_p, and min_p only. The nearest code is `min_p_filter` at `src/lib/mlxcel-core/src/sampling.rs:1219-1228` (softmax, per-row max, threshold, `where`), which has the same shape with a different threshold.
## Expected behavior
New `SamplingConfig` field `p_less: bool`, default `false`. Request field `p_less: bool` (absent resolves to `false`); no range validation needed.
Filter definition, per row of `[B, V]` logits, in float32, using the request temperature `T` (the filter is defined on the tempered distribution, so it takes the temperature explicitly even though the fused chain applies `T` again later; the fused chain's own scaling is unchanged):
```
f = logits.astype(f32)
probs = softmax(f / T, axis=-1) # -inf entries become exactly 0
L = sum(probs * probs, axis=-1, keepdims) # collision probability, [B, 1]
keep = probs >= L # masked entries (0 >= L) stay masked, L > 0 always
out = where(keep, logits, -inf)
```
Properties the tests pin: the argmax always survives (`L <= p_max`); `[10, 0, 0, 0, 0]` at `T = 1` keeps only index 0; raising `T` never shrinks the kept set for a fixed row; rows are independent.
Ordering within the chain (fixed by #1375): `top_n_sigma -> p_less -> typical_p`, after the penalties, before XTC and the fused chain; skipped on the greedy path (`temperature == 0.0 || top_k == 1`), where `T` would be 0 anyway.
Fused batch path: `FusedSampleParams` gains `p_less: bool`, compared by equality; uniform batches stay single-dispatch.
## Implementation plan
1. `src/lib/mlxcel-core/src/generate.rs`: `pub p_less: bool` on `SamplingConfig`, `false` in `Default` and `greedy()`.
2. `src/lib/mlxcel-core/src/sampling.rs`:
- `pub(crate) fn p_less_filter(logits: &MlxArray, temperature: f32) -> UniquePtr` with `ffi::astype`, `ops::divide_scalar(.., temperature)` (skip when `temperature == 1.0`), `ffi::softmax`, `ffi::square` or `ffi::multiply`, `ffi::sum_axis(.., -1, true)`, `ffi::greater_equal`, `ffi::where_cond` against `ffi::full_f32(&[1], f32::NEG_INFINITY, FLOAT32)`.
- Fill in the `p_less` arm of `apply_row_filters` (between `top_n_sigma` and `typical_p`), passing `params.temperature`.
- `FusedSampleParams`: `pub p_less: bool`, populated in `from_config`, compared in `matches`.
3. Request plumbing, same files and pattern as #1375: `SamplingParams::p_less: Option`, `RequestOptionOverrides::p_less`, resolve `unwrap_or(false)`, `ResolvedSamplingParams::p_less` through both `build_sampling_config` branches, `NativeCompletionRequest::p_less`, disaggregated wire struct, `bench_decode.rs`. No validator (a bool has no range).
4. CLI: `SamplingOptions` in `src/main.rs`: `#[arg(long = "p-less")] pub(crate) p_less: bool`; thread through `src/commands/generate.rs` and `src/commands/chat.rs`.
5. Docs: sampling-parameters table and `--help`.
## Validation
(a) Unit tests in the `sampling.rs` `mod tests`, `src/server/request_options_tests.rs`, `src/execution/sampling_tests.rs`:
- `p_less_filter_matches_host_reference`: 40 random rows (V in 8..200, logits N(0, 3), `T` in {0.5, 0.7, 1.0, 1.3, 2.0}); host f64 reference `keep = p >= sum(p^2)`; require equality for every token with `|p - L| > 1e-5`.
- `p_less_filter_argmax_always_survives`: 50 random rows and temperatures; the argmax is never `-inf`.
- `p_less_filter_peaked_keeps_top_only`: `[10, 0, 0, 0, 0]`, `T = 1` keeps exactly `{0}`.
- `p_less_filter_higher_temp_keeps_more`: kept counts for `T = 0.5, 1, 2, 5` on `[4, 2, 1, 0.5, 0, -1]` are non-decreasing.
- `p_less_filter_rows_independent`: `[[10,0,0,0,0],[1,0.9,0.8,0.7,0.6]]` keeps `{0}` in row 0 and more than one index in row 1.
- `p_less_filter_ignores_neg_inf_entries`: appending `-inf` entries does not change `L` or the kept set.
- `batched_fused_sample_honors_p_less`: 128 rows `[10,0,0,0,0]`, `p_less = true`, `T = 1`: every draw is index 0.
- `fused_sample_params_matches_compares_p_less`.
- Request layer: `chat_accepts_p_less_bool`, `native_completion_maps_p_less`, `build_sampling_config_threads_p_less`.
(b) Real checkpoint: `mlx-community/Qwen3-4B-4bit`.
```
./target/release/mlxcel generate -m models/Qwen3-4B-4bit -p "Write a haiku about rain." -n 64 --temp 1.5 --p-less --seed 1
```
Acceptance: fluent output at `--temp 1.5` with `--p-less` (and visibly less coherent without it); greedy output token-identical with and without the flag; `"p_less": true` accepted on all three OpenAI-style endpoints and native `/completion`.
## Acceptance criteria
- [ ] `SamplingConfig::p_less` exists, defaults to `false`, and is honored by `sample_token_optimized`, `batched_fused_sample`, `effective_token_distribution`, and `sample_token_with_distribution`
- [ ] `p_less == false` leaves every existing token stream byte-identical
- [ ] Field accepted on `/v1/chat/completions`, `/v1/completions`, `/v1/responses`, native `/completion`, and the disaggregated wire protocol
- [ ] `--p-less` on `mlxcel generate` and `mlxcel chat`
- [ ] Unit tests listed above pass; real-checkpoint check above passes
- [ ] cargo test --workspace --profile test-fast --features metal,accelerate passes
- [ ] cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass
## Dependencies
Blocked by #1375 (provides `apply_row_filters`, the `FusedSampleParams` extension, and the request-plumbing template).
Contributor guide
Research direction
Start with the sampling flow in src/lib/mlxcel-core/src/sampling.rs and the SamplingConfig definition in src/lib/mlxcel-core/src/generate.rs, after reviewing the row-filter and fused-sampling changes from blocked issue #1375. Trace request plumbing through src/server/request_options_tests.rs and src/execution/sampling_tests.rs, then inspect the CLI paths in src/main.rs and src/commands/. Done means the listed unit, integration, formatting, clippy, workspace, and checkpoint checks pass.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- ai
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 48/100