feat(sampling): add locally typical sampling (typical_p) #1377

Summary

mlxcel has no locally typical sampler. Typical sampling keeps the tokens whose surprisal -log p is closest to the entropy of the row, accumulating probability mass in that order until it reaches typical_p, and masks the rest to -inf. It is a widely supported knob (typical_p, 1.0 disables) on the llama-server request schema that mlxcel advertises compatibility with. This issue adds the filter on the row-filter hook introduced by #1375, the typical_p request field on every endpoint, and a --typical-p CLI flag, and keeps uniform batches on the fused batch path.

Current behavior

grep -rn typical_p src/ returns nothing. SamplingConfig (src/lib/mlxcel-core/src/generate.rs:810-874) has no field for it, SamplingParams (src/server/types/request.rs:567-640) ignores the key (serde default struct, unknown fields dropped), and the C++ fused chain (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:5039-5125) implements only temperature, top_k, top_p, and min_p. The nearest code is the reference top_p_filter in src/lib/mlxcel-core/src/sampling.rs:1275-1326, which already does the argsort / cumsum / inverse-permutation dance this filter needs, and the row-filter hook apply_row_filters from #1375.

Expected behavior

New SamplingConfig field typical_p: f32, default 1.0 (disabled). Valid range (0.0, 1.0]; values <= 0.0 or > 1.0 (or non-finite) are rejected at the request layer with HTTP 400, message typical_p must be in (0.0, 1.0]. 1.0 disables the filter and adds no graph nodes.

Filter definition, per row of [B, V] logits, in float32:

f        = logits.astype(f32)
logp     = log_softmax(f, axis=-1)                         # -inf entries stay -inf
p        = exp(logp)                                       # masked entries are exactly 0
plogp    = where(p > 0, p * logp, 0)                       # avoids 0 * -inf = NaN
H        = -sum(plogp, axis=-1, keepdims)                  # entropy, [B, 1]
shifted  = where(isfinite(logp), abs((-logp) - H), +inf)  # masked entries sort last
order    = argsort(shifted, axis=-1)                       # ascending: most typical first
p_sorted = take_along_axis(p, order, -1)
cum_incl = cumsum(p_sorted, axis=-1)                       # inclusive
cum_excl = cum_incl - p_sorted                             # mass strictly before each token in typicality order
inv      = argsort(order, axis=-1)                         # inverse permutation
cum_excl_orig = take_along_axis(cum_excl, inv, -1)
keep     = (cum_excl_orig < typical_p) & isfinite(logp)
out      = where(keep, logits, -inf)

The most typical token always has cum_excl == 0 < typical_p, so at least one token survives. Unlike top-p, the argmax can be dropped. For p = [0.45, 0.275, 0.275] the entropy is H = 1.069 nats and the surprisals are [0.799, 1.291, 1.291], so shifted = [0.271, 0.222, 0.222]: the two 0.275 tokens are more typical than the argmax, their mass 0.55 accumulates first, and typical_p = 0.3 masks index 0. This is why the filter must be skipped on the greedy path (temperature == 0.0 || top_k == 1), which apply_row_filters already does.

Ordering within the chain (fixed by #1375): top_n_sigma -> p_less -> typical_p, all after the penalties and before XTC and the fused temperature / top_k / top_p / min_p chain. The entropy is therefore computed on the penalized, untempered distribution.

Fused batch path: FusedSampleParams gains typical_p, compared bitwise; uniform batches stay on the single-dispatch path through apply_row_filters inside batched_fused_sample.

Implementation plan

1. src/lib/mlxcel-core/src/generate.rs: pub typical_p: f32 on SamplingConfig, 1.0 in Default and greedy().
2. src/lib/mlxcel-core/src/sampling.rs:
   • pub(crate) fn typical_p_filter(logits: &MlxArray, typical_p: f32) -> UniquePtr<MlxArray> implementing the formula with ffi::astype, ffi::log_softmax, ffi::exp, ffi::greater, ffi::where_cond, ffi::multiply, ffi::sum_axis, ffi::negative, ffi::subtract, ffi::abs, ffi::isfinite, ffi::argsort, ffi::take_along_axis, ffi::cumsum(.., -1, false, true), ffi::less, ffi::logical_and. Reuse the inverse-permutation idiom from top_p_filter (argsort(order)).
   • Fill in the typical_p arm of apply_row_filters (enabled when 0.0 < typical_p < 1.0), positioned after p_less and before XTC.
   • FusedSampleParams: pub typical_p: f32, populated in from_config, compared in matches.
3. Request plumbing, same files and pattern as feat(sampling): add top-n-sigma logit filtering (top_n_sigma) with a batch-safe row-filter hook #1375: SamplingParams::typical_p: Option<f32>, RequestOptionOverrides::typical_p, resolve unwrap_or(1.0), ResolvedSamplingParams::typical_p threaded through both build_sampling_config branches, validate_typical_p(Option<f32>) in routes/chat.rs called from chat / completions / responses, NativeCompletionRequest::typical_p (llama-server field name), disaggregated wire struct, bench_decode.rs.
4. CLI: SamplingOptions in src/main.rs: #[arg(long = "typical-p", alias = "typical", default_value_t = 1.0, value_name = "FLOAT")] pub(crate) typical_p: f32, validated (0, 1]; 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:

• typical_p_filter_matches_host_reference: 40 random rows (V in 8..200, logits N(0, 3), typical_p in {0.2, 0.5, 0.9, 0.95}); compute the kept set on the host in f64 with the same formula and require equality for every token whose cum_excl is more than 1e-4 away from typical_p.
• typical_p_filter_one_disables: typical_p = 1.0 returns the input unchanged (pointer-equal / no new nodes) and every entry stays finite.
• typical_p_filter_smaller_keeps_fewer: kept counts for 0.2, 0.5, 0.9, 0.99 are non-decreasing.
• typical_p_filter_can_drop_argmax: row with p = [0.45, 0.275, 0.275] (logits ln p), typical_p = 0.3: index 0 is masked.
• typical_p_filter_ignores_neg_inf_entries: appending -inf entries changes neither entropy nor the kept set.
• typical_p_skipped_on_greedy: temperature = 0.0, typical_p = 0.3 on the row above still returns index 0.
• batched_fused_sample_honors_typical_p: 64 rows [10, 0, 0, 0, 0], typical_p = 0.3, temperature 1.0: every draw is index 0.
• fused_sample_params_matches_compares_typical_p.
• Request layer: chat_rejects_typical_p_zero, chat_rejects_typical_p_above_one, completions_accepts_typical_p_in_range, responses_rejects_typical_p_negative, native_completion_maps_typical_p, build_sampling_config_threads_typical_p.

(b) Real checkpoint: mlx-community/Qwen3-4B-4bit.

./target/release/mlxcel generate -m models/Qwen3-4B-4bit -p "Explain why the sky is blue in three sentences." -n 64 --temp 1.0 --typical-p 0.5 --seed 1

Acceptance: finite logits, fluent output; greedy (--temp 0) output is token-identical with and without --typical-p 0.5; a server batch of identical "typical_p": 0.5 requests logs the fused batch dispatch; "typical_p": 0 and "typical_p": 1.5 return 400.

Acceptance criteria

• SamplingConfig::typical_p exists, defaults to 1.0, and is honored by sample_token_optimized, batched_fused_sample, effective_token_distribution, and sample_token_with_distribution
• typical_p == 1.0 leaves every existing token stream byte-identical
• Field accepted on /v1/chat/completions, /v1/completions, /v1/responses, native /completion, and the disaggregated wire protocol; out-of-range values rejected with 400 on the OpenAI-shaped endpoints, while the native route follows b10621's limit-free schema and maps out-of-domain values to the explicit disabled form. Beyond the issue's scope, the server-wide --typical / --typical-p flag landed on both server binaries because the --typical manifest entry is owned by this issue and had to leave deferred before it closes
• --typical-p 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 pattern, and the request-plumbing template).