lablup / lablup/mlxcel

feat(sampling): add hyperparameter-free p-less sampling (`p_less`)

Open
#1,373 0 comments 0 reactions 0 assignees View on GitHub
area:cli area:core area:inference modeltype:text priority:low status:blocked type:enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.