diff --git a/TECHNICAL_REPORTS/1482-locally-typical-sampling-20260828.en.md b/TECHNICAL_REPORTS/1482-locally-typical-sampling-20260828.en.md new file mode 100644 index 000000000..8b15bca95 --- /dev/null +++ b/TECHNICAL_REPORTS/1482-locally-typical-sampling-20260828.en.md @@ -0,0 +1,46 @@ +# Technical Report: PR #1482 - Locally typical sampling (typical_p) + +**Date**: 2026-08-28 + +**Status**: Completed + +**Languages**: Rust + +**Risk Level**: Medium + +## Executive Summary + +Adds locally typical sampling (`typical_p`, default `1.0` disabled) on the #1375 row-filter hook: tokens whose surprisal is closest to the row entropy are kept in typicality order until `typical_p` probability mass accumulates. Unlike the two earlier filters this one flips its two llama-compat manifest entries (`--typical`, `field:typical_p`) to `supported`, which pulled the server-wide flag on both binaries and the native `/completion` field into the same PR. + +## 1. Problem Statement + +mlxcel had no typical sampler, and the `--typical` / `field:typical_p` manifest entries are owned by issue #1377 itself: the epic's closure rule requires every entry owned by an issue to leave `deferred` before that issue closes, and the llama-compat gates tie a supported option claim to both server binaries and a supported field claim to `NativeCompletionRequest`. The scope therefore had to cover the engine semantics, both flag surfaces, the native field, and the manifest flip in one change. + +## 2. Technical Decisions + +### 2.1 Chain position: an idempotent top-k pre-mask + +The review caught the load-bearing subtlety: b10621's chain is `top_k -> typ_p` and its typical sampler re-softmaxes the surviving candidates, so entropy comes from the renormalized top-k distribution. mlxcel's hook runs before the C++ chain's top_k, so the filter initially computed entropy over the full vocabulary, diverging for essentially every enabled request (server default top_k 40). The fix masks the row to top-k ahead of the typical filter whenever `1 < top_k < vocab`; masking to top-k is idempotent, so the C++ chain's own top_k re-selects the same set and no stage runs out of order. A regression test uses a row where the two orders provably disagree (full-vocab typicality keeps only token 1, top-k-renormalized typicality keeps only token 0). + +### 2.2 Value domains split by surface + +OpenAI-shaped endpoints validate `(0.0, 1.0]` and return 400 otherwise. The native `/completion` field mirrors b10621's limit-free schema: a present out-of-domain value resolves to the explicit disabled `1.0`, still overriding the server default the way an upstream request value replaces it. The server-wide `--typical` flag accepts what b10621 accepts but folds out-of-domain values to disabled at startup with a warning, and `/props` reports the resolved value so the operator can read it back. The CLI generate/chat flag rejects out-of-domain values at parse time. + +### 2.3 Effective-value normalization and serde default + +`FusedSampleParams::from_config` and the speculative-window `sampling_config_eq` compare `effective_typical_p()` (greedy and out-of-domain fold to `1.0`), so necessarily-identical rows stay on the fused batch, pipelined lookahead, and shared speculative windows. The disaggregated wire struct defaults the field to `1.0` through a serde default fn, because the f32 zero default would turn an older peer's frame into an invalid always-on cutoff. + +## 3. Change Summary + +| Item | Value | +|------|-------| +| Files changed | 35 (+8 in the review follow-up) | +| Lines | +884 / -33 across both commits | +| Tests added | 25 | + +Validation: 91 core sampling tests green including a 40-case f64 host-reference test; compat integration test validates the flipped `--typical` claim against both freshly built binaries; manifest checker green with zero remaining #1377-owned deferred entries; real-checkpoint check on Qwen3-4B-4bit (fluent at `typical_p 0.5` with and without `top_k 40`, greedy byte-identical control). + +## 4. Follow-up Actions + +- #1373 (`p_less`) fills the remaining slot in `apply_row_filters`. +- #1436 completes the remaining b10621 sampling semantics; the `+inf` handling difference between `typical_p_filter` (masked) and `top_n_sigma_filter` (kept) is documented in the filter docs. diff --git a/TECHNICAL_REPORTS/1482-locally-typical-sampling-20260828.ko.md b/TECHNICAL_REPORTS/1482-locally-typical-sampling-20260828.ko.md new file mode 100644 index 000000000..409a7090a --- /dev/null +++ b/TECHNICAL_REPORTS/1482-locally-typical-sampling-20260828.ko.md @@ -0,0 +1,46 @@ +# 기술 보고서: PR #1482 - Locally typical sampling (typical_p) + +**작성일**: 2026-08-28 + +**상태**: 완료 + +**언어**: Rust + +**위험도**: Medium + +## 요약 + +#1375 row-filter 훅 위에 locally typical sampling(`typical_p`, 기본값 `1.0` 비활성)을 추가했다. surprisal이 행 엔트로피에 가장 가까운 토큰부터 typicality 순서로 `typical_p` 확률 질량이 쌓일 때까지 남긴다. 앞의 두 필터와 달리 이 이슈가 소유한 llama-compat 매니페스트 엔트리 2건(`--typical`, `field:typical_p`)을 `supported`로 플립했고, 그 때문에 두 서버 바이너리의 서버 전역 플래그와 네이티브 `/completion` 필드가 같은 PR에 들어왔다. + +## 1. 문제 정의 + +mlxcel에는 typical 샘플러가 없었고, `--typical` / `field:typical_p` 엔트리는 이슈 #1377 소유다. 에픽의 종료 규칙상 이슈가 닫히기 전에 소유 엔트리 전부가 `deferred`를 벗어나야 하고, llama-compat 게이트는 supported 옵션 주장을 두 서버 바이너리에, supported 필드 주장을 `NativeCompletionRequest`에 묶는다. 따라서 엔진 의미론, 두 플래그 표면, 네이티브 필드, 매니페스트 플립이 한 변경에 함께 들어가야 했다. + +## 2. 기술적 선택과 그 이유 + +### 2.1 체인 위치: 멱등 top-k 프리마스크 + +리뷰가 핵심 미묘함을 잡았다. b10621의 체인은 `top_k -> typ_p`이고 typical 샘플러는 생존 후보를 다시 softmax하므로 엔트로피가 재정규화된 top-k 분포에서 나온다. mlxcel 훅은 C++ 체인의 top_k보다 앞서 실행되어 처음에는 전체 어휘로 엔트로피를 계산했고, 서버 기본 top_k가 40이라 사실상 모든 활성 요청이 관찰 가능하게 갈라졌다. 수정은 `1 < top_k < vocab`일 때 typical 필터 앞에 top-k 마스크를 먼저 적용한다. top-k 마스킹은 멱등이라 C++ 체인의 top_k가 같은 집합을 다시 고르며 어떤 단계도 순서가 어긋나지 않는다. 두 순서가 증명 가능하게 갈라지는 행(전체 어휘 기준은 토큰 1만, top-k 재정규화 기준은 토큰 0만 남는다)으로 회귀 테스트를 추가했다. + +### 2.2 표면별 값 도메인 분리 + +OpenAI 형태 엔드포인트는 `(0.0, 1.0]`을 검증하고 벗어나면 400을 반환한다. 네이티브 `/completion` 필드는 b10621의 무제한 스키마를 따라, 도메인 밖 값이 오면 명시적 비활성 `1.0`으로 해석하되 서버 기본값은 여전히 덮어쓴다. 서버 전역 `--typical` 플래그는 b10621처럼 받아들이되 시작 시 도메인 밖 값을 경고와 함께 비활성으로 접고, `/props`가 해석된 값을 보고해 운영자가 읽어볼 수 있다. CLI generate/chat 플래그는 파싱 시점에 거부한다. + +### 2.3 유효값 정규화와 serde 기본값 + +`FusedSampleParams::from_config`와 speculative 윈도우의 `sampling_config_eq`는 `effective_typical_p()`(greedy와 도메인 밖 값은 `1.0`으로 접힘)를 비교하므로, 필연적으로 동일한 행들이 fused batch, pipelined lookahead, 공유 speculative 윈도우에 남는다. 분산 와이어 구조체는 serde 기본값 함수로 `1.0`을 쓴다. f32의 0 기본값이면 구버전 피어의 프레임이 항상 켜진 잘못된 컷오프가 되기 때문이다. + +## 3. 변경 요약 + +| 항목 | 값 | +|-----|---| +| 변경 파일 | 35개 (리뷰 후속 +8) | +| 라인 | 두 커밋 합산 +884 / -33 | +| 테스트 추가 | 25개 | + +검증: 40케이스 f64 호스트 레퍼런스 테스트를 포함해 코어 sampling 테스트 91개 green, compat 통합 테스트가 새로 빌드한 두 바이너리에서 플립된 `--typical` 주장을 검증, 매니페스트 체커 green(#1377 소유 deferred 엔트리 0건), Qwen3-4B-4bit 실체크포인트 검증(`typical_p 0.5`에서 `top_k 40` 유무와 무관하게 유창, greedy 대조군 토큰 동일). + +## 4. 후속 조치 + +- #1373(`p_less`)이 `apply_row_filters`의 남은 슬롯을 채운다. +- #1436이 나머지 b10621 샘플링 의미론을 마무리한다. `typical_p_filter`(마스킹)와 `top_n_sigma_filter`(유지)의 `+inf` 처리 차이는 필터 문서에 기록해 두었다. diff --git a/compat/llama-server/b10621/sampling-and-grammar.json b/compat/llama-server/b10621/sampling-and-grammar.json index 63974ac95..dfb51d886 100644 --- a/compat/llama-server/b10621/sampling-and-grammar.json +++ b/compat/llama-server/b10621/sampling-and-grammar.json @@ -820,12 +820,20 @@ "default": "1.00, 1.0 = disabled", "description": "locally typical sampling, parameter p (default: 1.00, 1.0 = disabled)", "discovery": "help-entry-head", - "state": "deferred", + "state": "supported", "issue": 1377, - "test": null, - "notes": null, + "test": "tests/llama_compat_manifest.rs", + "notes": "Server-wide default via --typical / --typical-p on both binaries (mlxcel serve and mlxcel-server), default 1.0 = disabled, matching b10621. Chain position matches b10621 (top_k -> typ_p -> top_p -> min_p): when top_k is active, an idempotent Rust-side top-k mask runs ahead of the typical filter so the typicality statistics come from the renormalized top-k survivors, and the fused chain's own top_k then re-selects the same set. An out-of-range startup value folds to disabled with a warning, mirroring the upstream sampler treating such values as inert.", "divergence": [], - "mlxcel": null + "mlxcel": { + "accepted_spellings": [ + "--typical", + "--typical-p" + ], + "defaults": [ + "1" + ] + } }, { "kind": "option", @@ -1462,12 +1470,14 @@ "field_type": "num", "description": "Enable locally typical sampling with parameter p (1.0 = disabled)", "discovery": "source-schema-declaration", - "state": "deferred", + "state": "supported", "issue": 1377, - "test": null, - "notes": null, + "test": "src/server/llama_compat_tests.rs", + "notes": "b10621 declares the field with no schema limits; a present value outside the enabled range (0.0, 1.0) resolves to the explicit disabled form 1.0 and still overrides the server-wide default. Chain position matches b10621: typicality is computed over the renormalized top-k survivors (see the --typical entry).", "divergence": [], - "mlxcel": null + "mlxcel": { + "field": "typical_p" + } }, { "kind": "native_request_field", diff --git a/docs/python-client.md b/docs/python-client.md index 05090ab07..9711ef8ae 100644 --- a/docs/python-client.md +++ b/docs/python-client.md @@ -109,7 +109,7 @@ print(llm.chat(messages, max_tokens=128, temperature=0.3)) ## Sampling parameters -Generation methods accept OpenAI sampling fields directly: `max_tokens`, `temperature`, `top_p`, `stop`, `seed`, `presence_penalty`, `frequency_penalty`, `logit_bias`, and `response_format`. Server-specific knobs that are not part of the OpenAI schema (`top_k`, `min_p`, `top_n_sigma`, `repetition_penalty`, DRY settings, `xtc_probability` / `xtc_threshold`, and the vLLM-compatible loop-detection fields `max_pattern_size` / `min_pattern_size` / `min_count`) are forwarded in the request body. `xtc_probability` (`0.0`-`1.0`, default `0.0`, disabled) is the per-step chance that XTC (Exclude Top Choices) removes all but the single least-probable token among those whose probability exceeds `xtc_threshold` (`0.0`-`0.5`, default `0.1`). The threshold is tested against the renormalised distribution that remains after `top_k`, `top_p` and `min_p` have run, not against the raw model distribution, matching the llama-server chain; both are validated on `/v1/chat/completions`, `/v1/completions`, and `/v1/responses`, and an out-of-range value returns a 400 before generation starts. `top_n_sigma` (default `0.0`, disabled) keeps only the tokens whose raw logit lies within N standard deviations of the row maximum before the temperature / top-k / top-p / min-p chain runs, which keeps sampling stable at high temperature; it must be finite and `>= 0.0` on the OpenAI-shaped endpoints (a negative value returns a 400). The native `/completion` spelling of the field lands together with its llama-server compatibility-manifest entry (#1436). You can also pass an explicit `extra_body={...}` for arbitrary server fields; values you set there win on conflict. +Generation methods accept OpenAI sampling fields directly: `max_tokens`, `temperature`, `top_p`, `stop`, `seed`, `presence_penalty`, `frequency_penalty`, `logit_bias`, and `response_format`. Server-specific knobs that are not part of the OpenAI schema (`top_k`, `min_p`, `top_n_sigma`, `typical_p`, `repetition_penalty`, DRY settings, `xtc_probability` / `xtc_threshold`, and the vLLM-compatible loop-detection fields `max_pattern_size` / `min_pattern_size` / `min_count`) are forwarded in the request body. `xtc_probability` (`0.0`-`1.0`, default `0.0`, disabled) is the per-step chance that XTC (Exclude Top Choices) removes all but the single least-probable token among those whose probability exceeds `xtc_threshold` (`0.0`-`0.5`, default `0.1`). The threshold is tested against the renormalised distribution that remains after `top_k`, `top_p` and `min_p` have run, not against the raw model distribution, matching the llama-server chain; both are validated on `/v1/chat/completions`, `/v1/completions`, and `/v1/responses`, and an out-of-range value returns a 400 before generation starts. `top_n_sigma` (default `0.0`, disabled) keeps only the tokens whose raw logit lies within N standard deviations of the row maximum before the temperature / top-k / top-p / min-p chain runs, which keeps sampling stable at high temperature; it must be finite and `>= 0.0` on the OpenAI-shaped endpoints (a negative value returns a 400). The native `/completion` spelling of the field lands together with its llama-server compatibility-manifest entry (#1436). `typical_p` (default `1.0`, disabled) enables locally typical sampling: tokens whose surprisal is closest to the row entropy are kept, in typicality order, until `typical_p` probability mass accumulates; it must be finite and in `(0.0, 1.0]` on the OpenAI-shaped endpoints (400 otherwise), the native `/completion` field mirrors llama-server's limit-free schema by resolving a present out-of-range value to the explicit disabled form, and the server-wide `--typical` / `--typical-p` flag sets the default applied when a request omits the field. You can also pass an explicit `extra_body={...}` for arbitrary server fields; values you set there win on conflict. ```python llm.generate("Once upon a time", max_tokens=200, top_p=0.9, top_k=40, min_p=0.05) diff --git a/src/bin/bench_decode.rs b/src/bin/bench_decode.rs index 91622b53a..662a68dc3 100644 --- a/src/bin/bench_decode.rs +++ b/src/bin/bench_decode.rs @@ -297,6 +297,7 @@ fn sampling_config(model_path: &Path) -> SamplingConfig { xtc_probability: 0.0, xtc_threshold: 0.1, top_n_sigma: 0.0, + typical_p: 1.0, stop_token_ids: mlxcel::read_eos_token_ids(model_path), }) } diff --git a/src/bin/mlx_server.rs b/src/bin/mlx_server.rs index c476b0a09..e1b2ac9c6 100644 --- a/src/bin/mlx_server.rs +++ b/src/bin/mlx_server.rs @@ -966,6 +966,15 @@ struct ServerArgs { #[arg(long = "min-p", default_value_t = 0.05)] min_p: f32, + /// Default locally typical sampling, parameter p (1.0 = disabled) + #[arg( + long = "typical", + alias = "typical-p", + value_name = "N", + default_value_t = 1.0 + )] + typical_p: f32, + /// Random seed (-1 = random) #[arg(short = 's', long = "seed", default_value_t = -1)] seed: i64, @@ -2095,6 +2104,7 @@ fn build_startup_input(mut args: ServerArgs) -> anyhow::Result crate::GenerateArgs { top_k: 0, min_p: 0.0, top_n_sigma: 0.0, + typical_p: 1.0, repetition_penalty: 1.0, dry_multiplier: 0.0, dry_base: 1.75, diff --git a/src/commands/serve.rs b/src/commands/serve.rs index 1ab4f64c8..ceedc4768 100644 --- a/src/commands/serve.rs +++ b/src/commands/serve.rs @@ -514,6 +514,7 @@ fn build_startup_input(mut args: crate::ServeArgs) -> anyhow::Result crate::ServeArgs { top_k: 50, top_p: 0.95, min_p: 0.2, + typical_p: 1.0, seed: -1, repeat_last_n: 32, repeat_penalty: 1.2, diff --git a/src/distributed/disaggregated/decode_scheduler_tests.rs b/src/distributed/disaggregated/decode_scheduler_tests.rs index b0278d141..b1700018c 100644 --- a/src/distributed/disaggregated/decode_scheduler_tests.rs +++ b/src/distributed/disaggregated/decode_scheduler_tests.rs @@ -21,6 +21,7 @@ fn default_sampling() -> SerializableSamplingState { frequency_penalty: 0.0, presence_penalty: 0.0, top_n_sigma: 0.0, + typical_p: 1.0, stop_token_ids: Vec::new(), } } diff --git a/src/distributed/disaggregated/prefill_scheduler_tests.rs b/src/distributed/disaggregated/prefill_scheduler_tests.rs index a28cc6a1a..a0f48cf3b 100644 --- a/src/distributed/disaggregated/prefill_scheduler_tests.rs +++ b/src/distributed/disaggregated/prefill_scheduler_tests.rs @@ -23,6 +23,7 @@ fn default_sampling() -> SerializableSamplingState { frequency_penalty: 0.0, presence_penalty: 0.0, top_n_sigma: 0.0, + typical_p: 1.0, stop_token_ids: Vec::new(), } } diff --git a/src/distributed/disaggregated/serving_protocol.rs b/src/distributed/disaggregated/serving_protocol.rs index 0410ceadb..1b3450751 100644 --- a/src/distributed/disaggregated/serving_protocol.rs +++ b/src/distributed/disaggregated/serving_protocol.rs @@ -246,6 +246,7 @@ pub fn sampling_to_serializable(config: &SamplingConfig) -> SerializableSampling frequency_penalty: config.frequency_penalty, presence_penalty: config.presence_penalty, top_n_sigma: config.top_n_sigma, + typical_p: config.typical_p, stop_token_ids: config.stop_token_ids.clone(), } } @@ -270,6 +271,7 @@ pub fn sampling_from_serializable(state: &SerializableSamplingState) -> Sampling frequency_penalty: state.frequency_penalty, presence_penalty: state.presence_penalty, top_n_sigma: state.top_n_sigma, + typical_p: state.typical_p, stop_token_ids: state.stop_token_ids.clone(), token_bias: Default::default(), // Loop detection is not yet serialized across the disaggregated diff --git a/src/distributed/disaggregated/serving_protocol_tests.rs b/src/distributed/disaggregated/serving_protocol_tests.rs index 275dbff3f..cf2a9d294 100644 --- a/src/distributed/disaggregated/serving_protocol_tests.rs +++ b/src/distributed/disaggregated/serving_protocol_tests.rs @@ -219,6 +219,7 @@ fn sampling_round_trips_through_the_serializable_mirror() { config.stop_token_ids = vec![13, 14]; config.dry_sequence_breakers = vec![198]; config.top_n_sigma = 1.5; + config.typical_p = 0.5; let restored = sampling_from_serializable(&sampling_to_serializable(&config)); @@ -233,6 +234,7 @@ fn sampling_round_trips_through_the_serializable_mirror() { assert_eq!(restored.stop_token_ids, vec![13, 14]); assert_eq!(restored.dry_sequence_breakers, vec![198]); assert_eq!(restored.top_n_sigma, 1.5); + assert_eq!(restored.typical_p, 0.5); } #[test] @@ -251,3 +253,19 @@ fn sampling_state_from_an_older_peer_without_top_n_sigma_defaults_to_disabled() assert_eq!(state.top_n_sigma, 0.0); assert_eq!(sampling_from_serializable(&state).top_n_sigma, 0.0); } + +#[test] +fn sampling_state_from_an_older_peer_without_typical_p_defaults_to_disabled() { + // The serde default must be the DISABLED value 1.0, not the f32 zero + // default, which would be an invalid always-on cutoff. + let mut value = serde_json::to_value(sampling_to_serializable(&SamplingConfig::greedy())) + .expect("serialize sampling state"); + value + .as_object_mut() + .expect("sampling state serializes as an object") + .remove("typical_p"); + let state: SerializableSamplingState = + serde_json::from_value(value).expect("deserialize legacy sampling state"); + assert_eq!(state.typical_p, 1.0); + assert_eq!(sampling_from_serializable(&state).typical_p, 1.0); +} diff --git a/src/distributed/kv_cache_serde/tests.rs b/src/distributed/kv_cache_serde/tests.rs index e3a2b3380..7d36cea9f 100644 --- a/src/distributed/kv_cache_serde/tests.rs +++ b/src/distributed/kv_cache_serde/tests.rs @@ -202,6 +202,7 @@ fn full_state_serialize_deserialize_round_trip() { frequency_penalty: 0.0, presence_penalty: 0.0, top_n_sigma: 0.0, + typical_p: 1.0, stop_token_ids: vec![2], }), token_history: vec![100, 200, 300], diff --git a/src/distributed/kv_cache_serde/types.rs b/src/distributed/kv_cache_serde/types.rs index 43aac6bdf..7f18fcc4a 100644 --- a/src/distributed/kv_cache_serde/types.rs +++ b/src/distributed/kv_cache_serde/types.rs @@ -280,9 +280,20 @@ pub struct SerializableSamplingState { /// to the disabled baseline instead of failing the handoff. #[serde(default)] pub top_n_sigma: f32, + /// Locally typical sampling (#1377). The default fn (not the f32 zero + /// default, which would be an INVALID enabled value) resolves an older + /// peer's frame to the disabled baseline `1.0`. + #[serde(default = "default_typical_p_disabled")] + pub typical_p: f32, pub stop_token_ids: Vec, } +/// The `1.0` "disabled" default for [`SerializableSamplingState::typical_p`] +/// on frames from peers that predate the field. +fn default_typical_p_disabled() -> f32 { + 1.0 +} + /// Complete serializable cache state for one sequence. /// /// This is the top-level type that gets serialized to the wire format diff --git a/src/execution/sampling.rs b/src/execution/sampling.rs index 54dd822ac..2f8275a54 100644 --- a/src/execution/sampling.rs +++ b/src/execution/sampling.rs @@ -49,6 +49,9 @@ pub struct ResolvedSamplingParams { /// Top-n-sigma logit filter (`0.0` = disabled). Non-negative and finite, /// enforced at the request layer. pub top_n_sigma: f32, + /// Locally typical sampling cutoff (`1.0` = disabled). In `(0.0, 1.0]` + /// and finite, enforced at the request layer. + pub typical_p: f32, pub stop_token_ids: Vec, } @@ -80,6 +83,7 @@ pub fn build_sampling_config(params: ResolvedSamplingParams) -> SamplingConfig { // row-filter hook gates on `temperature == 0.0 || top_k == 1`), // so the config faithfully mirrors what the request resolved. top_n_sigma: params.top_n_sigma, + typical_p: params.typical_p, stop_token_ids: params.stop_token_ids, ..SamplingConfig::greedy() } @@ -101,6 +105,7 @@ pub fn build_sampling_config(params: ResolvedSamplingParams) -> SamplingConfig { xtc_probability: params.xtc_probability, xtc_threshold: params.xtc_threshold, top_n_sigma: params.top_n_sigma, + typical_p: params.typical_p, stop_token_ids: params.stop_token_ids, token_bias: TokenBiasMap::default(), // Loop detection defaults to disabled here. The server control diff --git a/src/execution/sampling_tests.rs b/src/execution/sampling_tests.rs index ac9488a33..6452dd4dd 100644 --- a/src/execution/sampling_tests.rs +++ b/src/execution/sampling_tests.rs @@ -32,6 +32,7 @@ fn sample_params() -> ResolvedSamplingParams { xtc_probability: 0.4, xtc_threshold: 0.15, top_n_sigma: 1.5, + typical_p: 0.4, stop_token_ids: vec![1, 2], } } @@ -110,6 +111,20 @@ fn build_sampling_config_keeps_dry_sequence_breakers_at_zero_temperature() { assert_eq!(config.top_p, 1.0); } +#[test] +fn build_sampling_config_threads_typical_p_in_both_branches() { + let params = sample_params(); + assert_eq!(params.typical_p, 0.4); + let config = build_sampling_config(params); + assert_eq!(config.typical_p, 0.4); + + let mut greedy_params = sample_params(); + greedy_params.temperature = 0.0; + let config = build_sampling_config(greedy_params); + assert_eq!(config.temperature, 0.0); + assert_eq!(config.typical_p, 0.4); +} + #[test] fn build_sampling_config_threads_top_n_sigma_in_both_branches() { // Stochastic branch. diff --git a/src/lib/mlxcel-core/src/generate.rs b/src/lib/mlxcel-core/src/generate.rs index 82885d80d..11ff4ceda 100644 --- a/src/lib/mlxcel-core/src/generate.rs +++ b/src/lib/mlxcel-core/src/generate.rs @@ -887,6 +887,16 @@ pub struct SamplingConfig { /// path (`temperature == 0.0 || top_k == 1`) skips the filter entirely. /// See [`crate::sampling::apply_row_filters`]. pub top_n_sigma: f32, + /// Locally typical sampling (`1.0` = disabled, the default). Keeps the + /// tokens whose surprisal `-log p` is closest to the row entropy, + /// accumulating probability mass in that typicality order until it + /// reaches `typical_p`, and masks the rest to `-inf`. Unlike top-p the + /// argmax can be dropped, which is why the greedy path + /// (`temperature == 0.0 || top_k == 1`) skips the filter entirely. + /// Valid enabled range `(0.0, 1.0)`; `1.0` is a zero-overhead no-op that + /// preserves the bit-exact baseline. See + /// [`crate::sampling::apply_row_filters`]. + pub typical_p: f32, } impl Default for SamplingConfig { @@ -912,6 +922,7 @@ impl Default for SamplingConfig { xtc_threshold: 0.1, xtc_special_token_ids: Vec::new(), top_n_sigma: 0.0, + typical_p: 1.0, } } } @@ -940,6 +951,7 @@ impl SamplingConfig { xtc_threshold: 0.1, xtc_special_token_ids: Vec::new(), top_n_sigma: 0.0, + typical_p: 1.0, } } @@ -972,6 +984,23 @@ impl SamplingConfig { } } + /// The typical-p value the sampler will actually apply: `1.0` (disabled) + /// when the config is greedy (`temperature == 0.0 || top_k == 1`), or + /// when the raw field is outside the enabled range `(0.0, 1.0)` or + /// non-finite. Batch-uniformity gates compare THIS value, mirroring + /// [`SamplingConfig::effective_top_n_sigma`], so rows whose sampled + /// outputs are necessarily identical stay on the shared fused paths. + pub fn effective_typical_p(&self) -> f32 { + if self.temperature == 0.0 || self.top_k == 1 { + return 1.0; + } + if self.typical_p.is_finite() && self.typical_p > 0.0 && self.typical_p < 1.0 { + self.typical_p + } else { + 1.0 + } + } + /// Check if any penalty-based sampling is enabled pub fn needs_token_history(&self) -> bool { self.repetition_penalty != 1.0 diff --git a/src/lib/mlxcel-core/src/sampling.rs b/src/lib/mlxcel-core/src/sampling.rs index 4e1fadc38..f77c445b1 100644 --- a/src/lib/mlxcel-core/src/sampling.rs +++ b/src/lib/mlxcel-core/src/sampling.rs @@ -501,11 +501,11 @@ fn preprocess_logits_for_sampling( last_logits }; - // Row filters (top-n-sigma; #1377 / #1373 add typical_p and p_less here) - // run on the penalized, untempered logits: after every history-based - // penalty and before the fused C++ chain (temperature / top_k / top_p / - // min_p / XTC). Disabled filters return `last_logits` unchanged with no - // new graph nodes, preserving the bit-exact baseline. + // Row filters (top-n-sigma and typical_p; #1373 adds p_less) run on the + // penalized, untempered logits: after every history-based penalty and + // before the fused C++ chain (temperature / top_k / top_p / min_p / + // XTC). Disabled filters return `last_logits` unchanged with no new + // graph nodes, preserving the bit-exact baseline. apply_row_filters(last_logits, &FusedSampleParams::from_config(config)) } @@ -664,6 +664,10 @@ pub struct FusedSampleParams { /// RNG-free, so it stays fused-eligible: [`apply_row_filters`] applies it /// to the whole `[B, V]` batch before the single fused dispatch. pub top_n_sigma: f32, + /// Locally typical sampling cutoff (`1.0` disables). Row-wise, + /// history-free and RNG-free like `top_n_sigma`, and applied by the same + /// [`apply_row_filters`] hook before the single fused dispatch. + pub typical_p: f32, } impl FusedSampleParams { @@ -681,6 +685,7 @@ impl FusedSampleParams { top_p: config.top_p, min_p: config.min_p, top_n_sigma: config.effective_top_n_sigma(), + typical_p: config.effective_typical_p(), } } @@ -695,6 +700,7 @@ impl FusedSampleParams { && self.top_p.to_bits() == other.top_p.to_bits() && self.min_p.to_bits() == other.min_p.to_bits() && self.top_n_sigma.to_bits() == other.top_n_sigma.to_bits() + && self.typical_p.to_bits() == other.typical_p.to_bits() } } @@ -811,15 +817,17 @@ fn token_ids_to_host(tokens: &MlxArray) -> Vec { /// top_n_sigma -> p_less (#1373) -> typical_p (#1377) /// ``` /// -/// (the latter two land in their own issues; their slots in this function fix -/// the order here so neither can accidentally reorder the chain). +/// (`p_less` lands in its own issue; its slot in this function fixes the +/// order here so it cannot accidentally reorder the chain). /// /// Returns the input pointer unchanged, adding NO graph nodes, when every /// filter is disabled or when the config is greedy (`temperature == 0.0 || /// top_k == 1`), so the disabled baseline and the greedy path both stay /// byte-identical. A non-finite or non-positive `top_n_sigma` counts as /// disabled, which also makes the b10621 `-1.0` "disabled" sentinel inert -/// rather than a mask-everything threshold. +/// rather than a mask-everything threshold; a `typical_p` outside the open +/// interval `(0.0, 1.0)` (b10621 documents `1.0` as disabled and leaves the +/// low end of the range undeclared) counts as disabled the same way. /// /// Used by: [`preprocess_logits_for_sampling`] (per-row sampler, and through /// it [`effective_token_distribution`] / [`sample_token_with_distribution`]), @@ -839,7 +847,24 @@ pub fn apply_row_filters( logits = top_n_sigma_filter(&logits, params.top_n_sigma); } // p_less (#1373) slot: runs here, after top_n_sigma, when it lands. - // typical_p (#1377) slot: runs here, after p_less, when it lands. + if params.typical_p.is_finite() && params.typical_p > 0.0 && params.typical_p < 1.0 { + // b10621's default sampler chain is `... top_n_sigma; top_k; typ_p; + // top_p; min_p ...`: typical sampling runs on the RENORMALIZED top-k + // survivors, so its entropy and typicality ranking are computed over + // the truncated distribution, not the full vocabulary. mlxcel's + // top_k lives in the C++ chain AFTER this hook, so to reproduce the + // b10621 chain position the top-k mask is applied here first, ahead + // of the typical filter. The C++ chain's own top_k then re-selects + // the same set (masking to top-k is idempotent), so no stage runs + // out of order and the fused dispatch is unchanged. + if params.top_k > 1 { + let vocab = ffi::array_shape(&logits).last().copied().unwrap_or(0); + if params.top_k < vocab { + logits = top_k_filter(&logits, params.top_k); + } + } + logits = typical_p_filter(&logits, params.typical_p); + } logits } @@ -893,6 +918,73 @@ pub(crate) fn top_n_sigma_filter(logits: &MlxArray, n_sigma: f32) -> UniquePtr typ_p` chain order does. Entries that are not finite in the +/// input (`-inf` masking from token bias / penalties / top-k, or a NaN from +/// a broken forward) are mapped to `-inf` BEFORE the softmax, so they carry +/// zero probability, contribute nothing to the entropy, sort to the end of +/// the typicality order, and stay masked in the output. Note this includes +/// `+inf`: a softmax cannot represent an infinite logit, so unlike +/// [`top_n_sigma_filter`] (which keeps `+inf` entries), this filter treats +/// them as masked. Both inputs are model failures; the divergence is +/// documented rather than papered over. +/// +/// The most typical token always has an exclusive cumulative mass of +/// `0 < typical_p`, so at least one token survives. Unlike top-p, the +/// argmax CAN be dropped when it is less typical than the mid-probability +/// tokens, which is why [`apply_row_filters`] skips this filter on the +/// greedy path. +/// +/// Rows are independent (`axis = -1` throughout), and the output carries the +/// original logits dtype (the mask fill is cast to it, mirroring +/// [`top_n_sigma_filter`]). +/// +/// Used by: [`apply_row_filters`], unit tests +pub(crate) fn typical_p_filter(logits: &MlxArray, typical_p: f32) -> UniquePtr { + let neg_inf_f32 = ffi::full_f32(&[1], f32::NEG_INFINITY, dtype::FLOAT32); + // Sanitize in f32: every non-finite entry (already-masked -inf, NaN, + // +inf) becomes -inf so it cannot poison the softmax statistics. + let raw = ffi::astype(logits, dtype::FLOAT32); + let f = ffi::where_cond(&ffi::isfinite(&raw), &raw, &neg_inf_f32); + let logp = ffi::log_softmax(&f, -1); // -inf entries stay -inf + let p = ffi::exp(&logp); // masked entries are exactly 0 + let zero = ffi::full_f32(&[1], 0.0, dtype::FLOAT32); + // p * log p with the 0 * -inf = NaN case forced to 0. + let plogp = ffi::where_cond(&ffi::greater(&p, &zero), &ffi::multiply(&p, &logp), &zero); + // Entropy per row, [B, 1]. + let entropy = ffi::negative(&ffi::sum_axis(&plogp, -1, true)); + let finite_lp = ffi::isfinite(&logp); + // |(-log p) - H|: distance from typicality. Masked entries sort last. + let dev = ffi::abs(&ffi::subtract(&ffi::negative(&logp), &entropy)); + let pos_inf = ffi::full_f32(&[1], f32::INFINITY, dtype::FLOAT32); + let shifted = ffi::where_cond(&finite_lp, &dev, &pos_inf); + // Ascending sort: most typical first. + let order = ffi::argsort(&shifted, -1); + let p_sorted = ffi::take_along_axis(&p, &order, -1); + // Exclusive cumulative mass strictly before each token in typicality + // order: inclusive cumsum minus the entry itself. + let cum_incl = ffi::cumsum(&p_sorted, -1, false, true); + let cum_excl = ffi::subtract(&cum_incl, &p_sorted); + // Inverse permutation back to vocabulary order. + let inv = ffi::argsort(&order, -1); + let cum_excl_orig = ffi::take_along_axis(&cum_excl, &inv, -1); + let tp = ffi::full_f32(&[1], typical_p, dtype::FLOAT32); + let keep = ffi::logical_and(&ffi::less(&cum_excl_orig, &tp), &finite_lp); + // Mask fill in the ORIGINAL dtype (see `top_n_sigma_filter` on why an + // f32 fill would silently promote f16/bf16 logits). + let neg_inf = ffi::full_f32(&[1], f32::NEG_INFINITY, ffi::array_dtype(logits)); + ffi::where_cond(&keep, logits, &neg_inf) +} + /// Apply repetition penalty to logits. /// /// For tokens in history: @@ -1424,9 +1516,10 @@ pub(crate) fn min_p_filter(logits: &MlxArray, min_p: f32) -> UniquePtr /// Production runs top-k inside the C++ chain or inside the rejection kernel; /// this is the reference copy the tests hold those against. /// -/// Used by: unit tests (`fused_sample_probs_equals_softmax_of_filtered_over_t` -/// in `sampling::tests`) -#[allow(dead_code)] +/// Used by: [`apply_row_filters`] (the idempotent pre-`typical_p` top-k +/// mask that pins b10621's `top_k -> typ_p` chain position), unit tests +/// (`fused_sample_probs_equals_softmax_of_filtered_over_t` in +/// `sampling::tests`) pub(crate) fn top_k_filter(logits: &MlxArray, k: i32) -> UniquePtr { let neg_logits = ffi::negative(logits); let indices = ffi::argpartition(&neg_logits, k - 1, -1); @@ -3481,4 +3574,372 @@ mod tests { "distribution not normalized: {total}" ); } + + // -- typical_p row filter (#1377) -- + + /// Host-side f64 reference of `typical_p_filter`'s kept set. + fn typical_p_reference_keep(logits: &[f32], typical_p: f64) -> Vec { + let finite: Vec = logits.iter().map(|x| x.is_finite()).collect(); + let max = logits + .iter() + .filter(|x| x.is_finite()) + .fold(f64::NEG_INFINITY, |m, &x| m.max(x as f64)); + let z: f64 = logits + .iter() + .filter(|x| x.is_finite()) + .map(|&x| ((x as f64) - max).exp()) + .sum(); + let p: Vec = logits + .iter() + .map(|&x| { + if x.is_finite() { + ((x as f64) - max).exp() / z + } else { + 0.0 + } + }) + .collect(); + let entropy: f64 = -p + .iter() + .filter(|&&pi| pi > 0.0) + .map(|&pi| pi * pi.ln()) + .sum::(); + let mut order: Vec = (0..logits.len()).collect(); + let dev: Vec = p + .iter() + .zip(&finite) + .map(|(&pi, &fin)| { + if fin && pi > 0.0 { + ((-pi.ln()) - entropy).abs() + } else { + f64::INFINITY + } + }) + .collect(); + order.sort_by(|&a, &b| dev[a].partial_cmp(&dev[b]).unwrap()); + let mut keep = vec![false; logits.len()]; + let mut cum = 0.0f64; + for &i in &order { + if !finite[i] { + break; + } + if cum < typical_p { + keep[i] = true; + } + cum += p[i]; + } + keep + } + + /// Deterministic LCG so the reference test needs no rand dependency. + fn lcg_next(state: &mut u64) -> f32 { + *state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + // top 24 bits -> [0, 1) + ((*state >> 40) as f32) / (1u64 << 24) as f32 + } + + #[test] + fn typical_p_filter_matches_host_reference() { + let mut rng = 0x1377_u64; + for case in 0..40 { + let v = 8 + (case * 5) % 193; + let logits_host: Vec = (0..v).map(|_| (lcg_next(&mut rng) - 0.5) * 6.0).collect(); + let tp = [0.2f32, 0.5, 0.9, 0.95][case % 4]; + let logits = ffi::from_slice_f32(&logits_host, &[1, v as i32]); + let out = to_vec_f32(&typical_p_filter(&logits, tp)); + let reference = typical_p_reference_keep(&logits_host, tp as f64); + // Compare only tokens whose exclusive cumulative mass is not + // within 1e-4 of the boundary, where f32-vs-f64 rounding can + // legitimately flip the strict comparison. + let max = logits_host + .iter() + .cloned() + .fold(f32::NEG_INFINITY, f32::max); + let z: f64 = logits_host.iter().map(|&x| ((x - max) as f64).exp()).sum(); + let p: Vec = logits_host + .iter() + .map(|&x| ((x - max) as f64).exp() / z) + .collect(); + let entropy: f64 = -p.iter().map(|&pi| pi * pi.ln()).sum::(); + let mut order: Vec = (0..v).collect(); + order.sort_by(|&a, &b| { + ((-p[a].ln()) - entropy) + .abs() + .partial_cmp(&((-p[b].ln()) - entropy).abs()) + .unwrap() + }); + let mut cum_excl = vec![0.0f64; v]; + let mut cum = 0.0f64; + for &i in &order { + cum_excl[i] = cum; + cum += p[i]; + } + for i in 0..v { + if (cum_excl[i] - tp as f64).abs() <= 1e-4 { + continue; + } + assert_eq!( + out[i].is_finite(), + reference[i], + "case {case} (V={v}, typical_p={tp}): token {i} diverges from the host reference" + ); + } + } + } + + #[test] + fn typical_p_one_disables_via_hook() { + // typical_p = 1.0 must add zero graph nodes: the exact input pointer + // comes back from the hook. + let logits = ffi::from_slice_f32(&[0.0, 1.0, 2.0], &[1, 3]); + let ptr_before = &*logits as *const MlxArray; + let params = FusedSampleParams::from_config(&SamplingConfig::default()); + assert_eq!(params.typical_p, 1.0); + let out = apply_row_filters(logits, ¶ms); + assert_eq!(ptr_before, &*out as *const MlxArray); + + // And with a DIFFERENT filter active, typical_p = 1.0 must add + // nothing on top of it: the hook output equals the direct + // top-n-sigma filter byte for byte. + let logits = ffi::from_slice_f32(&[0.0, 1.0, 2.0, 3.0, 4.0], &[1, 5]); + let cfg = SamplingConfig { + top_n_sigma: 1.0, + ..Default::default() + }; + let via_hook = apply_row_filters( + ffi::from_slice_f32(&[0.0, 1.0, 2.0, 3.0, 4.0], &[1, 5]), + &FusedSampleParams::from_config(&cfg), + ); + let direct = top_n_sigma_filter(&logits, 1.0); + assert_eq!(to_vec_f32(&via_hook), to_vec_f32(&direct)); + } + + #[test] + fn typical_p_runs_on_the_renormalized_top_k_survivors() { + // b10621's chain order is `top_k -> typ_p`, so the typicality + // statistics come from the RENORMALIZED top-k distribution. This row + // is built so the two orders disagree: logits [5.0, 4.9, 0 x 50]. + // + // Full-vocabulary entropy (~1.60 nats) is inflated by the 50 tail + // tokens, making token 1 more typical than token 0, so typ_p = 0.3 + // over the full row would keep ONLY token 1. Over the renormalized + // top-2 survivors (p = [0.525, 0.475], H = 0.692) token 0 is the + // more typical one and typ_p = 0.3 keeps ONLY token 0. + let mut row = vec![0.0f32; 52]; + row[0] = 5.0; + row[1] = 4.9; + let logits = ffi::from_slice_f32(&row, &[1, 52]); + let cfg = SamplingConfig { + top_k: 2, + typical_p: 0.3, + ..Default::default() + }; + let out = apply_row_filters(logits, &FusedSampleParams::from_config(&cfg)); + let v = to_vec_f32(&out); + assert!( + v[0].is_finite(), + "token 0 must survive: typicality must be computed over the \ + renormalized top-k set (b10621 order), not the full vocabulary" + ); + assert_eq!( + v[1], + f32::NEG_INFINITY, + "token 1 is less typical within the top-2 set" + ); + for (i, x) in v.iter().enumerate().skip(2) { + assert_eq!( + *x, + f32::NEG_INFINITY, + "tail token {i} escaped the top-k mask" + ); + } + } + + #[test] + fn typical_p_top_k_premask_skips_oversized_k() { + // top_k >= vocab is a no-op upstream; the pre-mask must skip it + // rather than hand argpartition an out-of-range kth index. + let logits = ffi::from_slice_f32(&[0.45f32.ln(), 0.275f32.ln(), 0.275f32.ln()], &[1, 3]); + let cfg = SamplingConfig { + top_k: 64, + typical_p: 0.3, + ..Default::default() + }; + let out = apply_row_filters(logits, &FusedSampleParams::from_config(&cfg)); + let v = to_vec_f32(&out); + // Same result as the no-top_k case: the less-typical argmax drops. + assert_eq!(v[0], f32::NEG_INFINITY); + assert!(v[1].is_finite() || v[2].is_finite()); + } + + #[test] + fn typical_p_filter_smaller_keeps_fewer() { + let logits_host: Vec = (0..50).map(|i| ((i * 7) % 13) as f32 * 0.3).collect(); + let logits = ffi::from_slice_f32(&logits_host, &[1, 50]); + let mut prev = 0usize; + for tp in [0.2f32, 0.5, 0.9, 0.99] { + let kept = to_vec_f32(&typical_p_filter(&logits, tp)) + .iter() + .filter(|x| x.is_finite()) + .count(); + assert!( + kept >= prev, + "kept count decreased from {prev} to {kept} at typical_p={tp}" + ); + assert!( + kept > 0, + "at least one token must survive at typical_p={tp}" + ); + prev = kept; + } + } + + #[test] + fn typical_p_filter_can_drop_argmax() { + // p = [0.45, 0.275, 0.275]: H = 1.069 nats, surprisals + // [0.799, 1.291, 1.291], deviations [0.271, 0.222, 0.222]. The two + // 0.275 tokens are MORE typical than the argmax and their mass 0.55 + // accumulates first, so typical_p = 0.3 masks index 0. + let logits = ffi::from_slice_f32(&[0.45f32.ln(), 0.275f32.ln(), 0.275f32.ln()], &[1, 3]); + let v = to_vec_f32(&typical_p_filter(&logits, 0.3)); + assert_eq!( + v[0], + f32::NEG_INFINITY, + "the less-typical argmax must be dropped" + ); + assert!(v[1].is_finite() || v[2].is_finite()); + } + + #[test] + fn typical_p_filter_ignores_neg_inf_entries() { + let base = [0.45f32.ln(), 0.275f32.ln(), 0.275f32.ln()]; + let with_masked = [ + 0.45f32.ln(), + 0.275f32.ln(), + 0.275f32.ln(), + f32::NEG_INFINITY, + f32::NEG_INFINITY, + ]; + let a = to_vec_f32(&typical_p_filter(&ffi::from_slice_f32(&base, &[1, 3]), 0.3)); + let b = to_vec_f32(&typical_p_filter( + &ffi::from_slice_f32(&with_masked, &[1, 5]), + 0.3, + )); + for i in 0..3 { + assert_eq!( + a[i].is_finite(), + b[i].is_finite(), + "appending -inf entries changed the kept set at index {i}" + ); + } + assert_eq!(b[3], f32::NEG_INFINITY); + assert_eq!(b[4], f32::NEG_INFINITY); + } + + #[test] + fn typical_p_filter_nan_entry_does_not_mask_the_row() { + // A NaN logit is sanitized to -inf before the softmax, so it cannot + // poison the entropy or the kept set. + let base = [0.45f32.ln(), 0.275f32.ln(), 0.275f32.ln()]; + let with_nan = [0.45f32.ln(), 0.275f32.ln(), 0.275f32.ln(), f32::NAN]; + let a = to_vec_f32(&typical_p_filter(&ffi::from_slice_f32(&base, &[1, 3]), 0.3)); + let b = to_vec_f32(&typical_p_filter( + &ffi::from_slice_f32(&with_nan, &[1, 4]), + 0.3, + )); + for i in 0..3 { + assert_eq!(a[i].is_finite(), b[i].is_finite()); + } + assert_eq!(b[3], f32::NEG_INFINITY); + } + + #[test] + fn typical_p_filter_preserves_dtype() { + let logits_f32 = ffi::from_slice_f32(&[0.0, 1.0, 2.0, 3.0], &[1, 4]); + let logits_f16 = ffi::astype(&logits_f32, dtype::FLOAT16); + let filtered = typical_p_filter(&logits_f16, 0.5); + assert_eq!(ffi::array_dtype(&filtered), dtype::FLOAT16); + } + + #[test] + fn typical_p_skipped_on_greedy() { + // Greedy must still return the argmax even when typical_p would + // drop it (the can_drop_argmax row above). + let logits = ffi::from_slice_f32(&[0.45f32.ln(), 0.275f32.ln(), 0.275f32.ln()], &[1, 1, 3]); + let mut config = SamplingConfig::greedy(); + config.typical_p = 0.3; + let (token, _) = sample_token_optimized(&logits, &config, &[]); + ffi::eval(&token); + assert_eq!(ffi::item_i32(&token), 0); + } + + #[test] + fn batched_fused_sample_honors_typical_p() { + // 64 rows [10, 0, 0, 0, 0]: the argmax holds ~99.98% of the mass, so + // it is by far the most typical token and typical_p = 0.3 keeps only + // it; every draw must be index 0. + const B: usize = 64; + let mut flat = Vec::with_capacity(B * 5); + for _ in 0..B { + flat.extend_from_slice(&[10.0f32, 0.0, 0.0, 0.0, 0.0]); + } + let logits = ffi::from_slice_f32(&flat, &[B as i32, 1, 5]); + let cfg = SamplingConfig { + typical_p: 0.3, + ..Default::default() + }; + let params = FusedSampleParams::from_config(&cfg); + assert_eq!(params.typical_p, 0.3); + let tokens = batched_fused_sample(&logits, ¶ms); + assert_eq!(tokens, vec![0; B]); + } + + #[test] + fn fused_sample_params_matches_compares_typical_p() { + let base = FusedSampleParams::from_config(&SamplingConfig::with_temperature(0.7)); + assert!(base.matches(&base)); + let diff = FusedSampleParams { + typical_p: 0.5, + ..base + }; + assert!(!base.matches(&diff)); + assert!(diff.matches(&diff)); + } + + #[test] + fn fused_sample_params_normalizes_inert_typical_p() { + // Greedy rows and out-of-domain values normalize to the disabled 1.0 + // so necessarily-identical rows stay on the shared fused paths. + let mut greedy = SamplingConfig::greedy(); + greedy.typical_p = 0.5; + assert_eq!(FusedSampleParams::from_config(&greedy).typical_p, 1.0); + for bad in [0.0f32, -0.5, 1.5, f32::NAN, f32::INFINITY] { + let cfg = SamplingConfig { + typical_p: bad, + ..Default::default() + }; + assert_eq!( + FusedSampleParams::from_config(&cfg).typical_p, + 1.0, + "typical_p={bad} must normalize to the disabled form" + ); + } + } + + #[test] + fn effective_token_distribution_zeroes_typical_p_masked_tokens() { + // The #902 distribution must carry zero mass on the dropped argmax. + let logits = ffi::from_slice_f32(&[0.45f32.ln(), 0.275f32.ln(), 0.275f32.ln()], &[1, 1, 3]); + let cfg = SamplingConfig { + typical_p: 0.3, + ..Default::default() + }; + let probs = effective_token_distribution(&logits, &cfg, &[]); + let v = to_vec_f32(&probs); + assert_eq!(v[0], 0.0, "dropped argmax kept probability {}", v[0]); + let total: f32 = v.iter().sum(); + assert!((total - 1.0).abs() < 1e-5); + } } diff --git a/src/main.rs b/src/main.rs index 2626b8a6b..26f9a80e2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -670,6 +670,17 @@ pub(crate) struct SamplingOptions { )] pub(crate) top_n_sigma: f32, + /// Locally typical sampling: keep tokens whose surprisal is closest to + /// the row entropy until P probability mass accumulates (1.0 = disabled) + #[arg( + long = "typical-p", + alias = "typical", + default_value_t = 1.0, + value_name = "FLOAT", + value_parser = parse_typical_p + )] + pub(crate) typical_p: f32, + /// Repetition penalty multiplier #[arg(long, default_value_t = 1.0, value_name = "FLOAT")] pub(crate) repetition_penalty: f32, @@ -712,6 +723,19 @@ fn parse_non_negative_f32(s: &str) -> Result { } } +/// Clap value parser: a finite f32 in the half-open interval (0, 1]. +/// +/// Used by: `--typical-p` (reject zero, negative, above-one, and non-finite +/// values at parse time; `1.0` is the explicit disabled form). +fn parse_typical_p(s: &str) -> Result { + let v: f32 = s.parse().map_err(|e| format!("not a number: {e}"))?; + if v.is_finite() && v > 0.0 && v <= 1.0 { + Ok(v) + } else { + Err(format!("must be in (0.0, 1.0], got {v}")) + } +} + /// Clap value parser: an f32 in the closed interval [0, 1]. /// /// Used by: `--diffusion-threshold` (fail fast at startup instead of @@ -1644,6 +1668,15 @@ pub(crate) struct ServeArgs { #[arg(long, default_value_t = 0.05)] min_p: f32, + /// Default locally typical sampling, parameter p (1.0 = disabled) + #[arg( + long = "typical", + alias = "typical-p", + value_name = "N", + default_value_t = 1.0 + )] + typical_p: f32, + /// Random seed (-1 = random) #[arg(short = 's', long, default_value_t = -1)] seed: i64, diff --git a/src/server/batch/scheduler.rs b/src/server/batch/scheduler.rs index 2a07e36f9..fa8e15bbc 100644 --- a/src/server/batch/scheduler.rs +++ b/src/server/batch/scheduler.rs @@ -7324,7 +7324,7 @@ impl BatchScheduler { ) -> Option { let logits = self.lookahead_forward(seq_ids, input)?; let last_logits = mlxcel_core::slice_last_logits(&logits); - // Same pre-fused row filters (top-n-sigma) as `batched_fused_sample`, + // Same pre-fused row filters (top-n-sigma, typical_p) as `batched_fused_sample`, // so the pipelined lookahead samples from the identical distribution // as the synchronous fused path it accelerates. A no-op adding no // graph nodes while every filter is disabled. diff --git a/src/server/batch/speculative_burst.rs b/src/server/batch/speculative_burst.rs index 571aa89b6..2b5ba771b 100644 --- a/src/server/batch/speculative_burst.rs +++ b/src/server/batch/speculative_burst.rs @@ -2848,6 +2848,7 @@ pub(crate) fn sampling_config_eq(a: &SamplingConfig, b: &SamplingConfig) -> bool // Compare the value the sampler will actually apply, so greedy rows // differing only in an inert top_n_sigma still share a window. && a.effective_top_n_sigma().to_bits() == b.effective_top_n_sigma().to_bits() + && a.effective_typical_p().to_bits() == b.effective_typical_p().to_bits() && a.seed == b.seed && a.stop_token_ids == b.stop_token_ids && a.token_bias.is_empty() diff --git a/src/server/batch/speculative_burst_tests.rs b/src/server/batch/speculative_burst_tests.rs index cd416a058..44af0882f 100644 --- a/src/server/batch/speculative_burst_tests.rs +++ b/src/server/batch/speculative_burst_tests.rs @@ -1704,6 +1704,29 @@ fn sampling_config_eq_ignores_inert_greedy_top_n_sigma() { ); } +#[test] +fn sampling_config_eq_distinguishes_typical_p() { + let a = SamplingConfig::default(); + let b = SamplingConfig { + typical_p: 0.5, + ..SamplingConfig::default() + }; + assert!(!super::speculative_burst::sampling_config_eq(&a, &b)); +} + +#[test] +fn sampling_config_eq_ignores_inert_greedy_typical_p() { + let a = SamplingConfig::greedy(); + let b = SamplingConfig { + typical_p: 0.5, + ..SamplingConfig::greedy() + }; + assert!( + super::speculative_burst::sampling_config_eq(&a, &b), + "greedy rows differing only in an inert typical_p must share a batched window" + ); +} + #[test] fn sampling_config_eq_distinguishes_stop_token_ids() { let a = SamplingConfig::default(); diff --git a/src/server/cli_input.rs b/src/server/cli_input.rs index 1ea1a0a18..4ffbde16b 100644 --- a/src/server/cli_input.rs +++ b/src/server/cli_input.rs @@ -165,6 +165,7 @@ pub struct ServerStartupInput { pub top_p: f32, pub top_p_was_set: bool, pub min_p: f32, + pub typical_p: f32, pub seed: i64, pub repeat_last_n: usize, pub repeat_penalty: f32, @@ -884,6 +885,7 @@ impl ServerStartupInput { top_p: self.top_p, top_p_was_set: self.top_p_was_set, min_p: self.min_p, + typical_p: self.typical_p, seed: resolve_seed(self.seed), repeat_last_n: self.repeat_last_n, repeat_penalty: self.repeat_penalty, diff --git a/src/server/cli_input_tests.rs b/src/server/cli_input_tests.rs index 5cd199321..2e0370931 100644 --- a/src/server/cli_input_tests.rs +++ b/src/server/cli_input_tests.rs @@ -98,6 +98,7 @@ fn sample_input() -> ServerStartupInput { top_p: 0.9, top_p_was_set: false, min_p: 0.1, + typical_p: 1.0, seed: 42, repeat_last_n: 64, repeat_penalty: 1.1, diff --git a/src/server/config.rs b/src/server/config.rs index 992d4553d..533f3e90c 100644 --- a/src/server/config.rs +++ b/src/server/config.rs @@ -415,6 +415,9 @@ pub struct ServerConfig { pub default_top_p: f32, pub default_top_k: i32, pub default_min_p: f32, + /// Server-wide default for locally typical sampling (`1.0` = disabled), + /// set by `--typical` / `--typical-p` (#1377). + pub default_typical_p: f32, pub default_repetition_penalty: f32, pub default_repetition_context_size: usize, pub default_max_tokens: usize, @@ -742,6 +745,7 @@ impl Default for ServerConfig { default_top_p: 0.95, default_top_k: 40, default_min_p: 0.05, + default_typical_p: 1.0, default_repetition_penalty: 1.0, default_repetition_context_size: 64, default_max_tokens: 512, diff --git a/src/server/request_options.rs b/src/server/request_options.rs index b3fb413f0..ebc8e1b76 100644 --- a/src/server/request_options.rs +++ b/src/server/request_options.rs @@ -83,6 +83,10 @@ pub(crate) struct RequestOptionOverrides { /// Top-n-sigma logit-filter override. Validated at the request layer /// (finite, `>= 0.0`) before it reaches here. pub top_n_sigma: Option, + /// Locally typical sampling override. Validated at the request layer + /// (finite, in `(0.0, 1.0]`) on the OpenAI-shaped endpoints; the native + /// `/completion` route sanitizes out-of-domain values itself. + pub typical_p: Option, pub stop_sequences: Option>, pub priority: RequestPriority, /// per-request thinking-token budget override. @@ -259,6 +263,11 @@ pub(crate) fn build_server_generate_options( // Like XTC, top-n-sigma is request-only with no server-level default: // an absent field always resolves to the disabled baseline (`0.0`). top_n_sigma: overrides.top_n_sigma.unwrap_or(0.0), + // Unlike XTC and top-n-sigma, typical_p follows the classic sampler + // knobs (--temp / --top-p / --min-p): the server-wide --typical / + // --typical-p default applies when the request omits the field, + // matching llama-server's sampling-params surface. + typical_p: overrides.typical_p.unwrap_or(config.default_typical_p), stop_token_ids: Vec::new(), }); diff --git a/src/server/request_options_tests.rs b/src/server/request_options_tests.rs index 01bc77a46..ca4f8c561 100644 --- a/src/server/request_options_tests.rs +++ b/src/server/request_options_tests.rs @@ -124,6 +124,10 @@ fn build_server_generate_options_uses_server_defaults() { assert_eq!(options.sampling.xtc_threshold, 0.1); // Like XTC, top-n-sigma is request-only: absent resolves to disabled. assert_eq!(options.sampling.top_n_sigma, 0.0); + // typical_p follows the classic sampler knobs: the server default + // applies (and defaults to the disabled 1.0). + assert_eq!(options.sampling.typical_p, config.default_typical_p); + assert_eq!(options.sampling.typical_p, 1.0); assert_eq!(options.stop_sequences, None); } @@ -150,6 +154,7 @@ fn build_server_generate_options_applies_request_overrides() { xtc_probability: Some(0.7), xtc_threshold: Some(0.2), top_n_sigma: Some(2.5), + typical_p: Some(0.35), stop_sequences: Some(vec!["stop".to_string()]), priority: crate::server::batch::RequestPriority::High, reasoning_budget: crate::server::config::ReasoningBudgetOverride::default(), @@ -179,6 +184,7 @@ fn build_server_generate_options_applies_request_overrides() { assert_eq!(options.sampling.xtc_probability, 0.7); assert_eq!(options.sampling.xtc_threshold, 0.2); assert_eq!(options.sampling.top_n_sigma, 2.5); + assert_eq!(options.sampling.typical_p, 0.35); assert_eq!(options.stop_sequences, Some(vec!["stop".to_string()])); } diff --git a/src/server/routes/chat.rs b/src/server/routes/chat.rs index 08a705403..47b87f8dd 100644 --- a/src/server/routes/chat.rs +++ b/src/server/routes/chat.rs @@ -273,6 +273,9 @@ pub async fn chat_completions( if let Err(message) = validate_top_n_sigma(request.params.top_n_sigma) { return ErrorResponse::new(message, "invalid_request_error").into_response(); } + if let Err(message) = validate_typical_p(request.params.typical_p) { + return ErrorResponse::new(message, "invalid_request_error").into_response(); + } // Reject image, audio and video content blocks the loaded checkpoint // cannot consume, before any referenced URL or file is read. Capability is @@ -1280,6 +1283,22 @@ pub(crate) fn validate_top_n_sigma(top_n_sigma: Option) -> Result<(), &'sta Ok(()) } +/// Validate the locally-typical-sampling parameter. +/// +/// `typical_p` must be finite and in `(0.0, 1.0]` when set (`1.0` disables +/// the filter); an absent field (`None`) is always valid and lets the +/// server-wide `--typical` default apply. `0.0` is rejected rather than +/// treated as "keep one token": b10621's schema leaves the low end of the +/// range undeclared, and a zero cutoff would otherwise mask every token. +pub(crate) fn validate_typical_p(typical_p: Option) -> Result<(), &'static str> { + if let Some(value) = typical_p + && !(value.is_finite() && value > 0.0 && value <= 1.0) + { + return Err("typical_p must be in (0.0, 1.0]"); + } + Ok(()) +} + /// Maximum number of tools allowed in a single request. pub(crate) const MAX_TOOLS: usize = 128; @@ -1473,6 +1492,7 @@ pub(crate) fn build_generate_options( xtc_probability: params.xtc_probability, xtc_threshold: params.xtc_threshold, top_n_sigma: params.top_n_sigma, + typical_p: params.typical_p, stop_sequences: params.stop.clone(), priority: RequestPriority::default(), // the caller (non_stream_chat_completion / @@ -1591,6 +1611,25 @@ mod tests { ); } + #[test] + fn validate_typical_p_accepts_unset_disabled_and_in_range() { + assert!(validate_typical_p(None).is_ok()); + assert!(validate_typical_p(Some(1.0)).is_ok()); + assert!(validate_typical_p(Some(0.5)).is_ok()); + assert!(validate_typical_p(Some(0.001)).is_ok()); + } + + #[test] + fn chat_rejects_out_of_domain_typical_p() { + for bad in [0.0f32, -0.5, 1.5, f32::NAN, f32::INFINITY] { + assert_eq!( + validate_typical_p(Some(bad)), + Err("typical_p must be in (0.0, 1.0]"), + "typical_p={bad} must be rejected" + ); + } + } + #[test] fn validate_xtc_params_checks_threshold_before_probability() { // Both fields invalid: the threshold check runs first. diff --git a/src/server/routes/completions.rs b/src/server/routes/completions.rs index 3fdab15a7..aa230f47d 100644 --- a/src/server/routes/completions.rs +++ b/src/server/routes/completions.rs @@ -42,7 +42,7 @@ use crate::tokenizer::MlxcelTokenizer; use super::chat::{ build_generate_options, decode_token, parse_priority_header, structured_error_to_response, - validate_top_n_sigma, validate_xtc_params, + validate_top_n_sigma, validate_typical_p, validate_xtc_params, }; fn generation_error_to_response(err: anyhow::Error) -> ErrorResponse { @@ -172,6 +172,9 @@ pub async fn completions( if let Err(message) = validate_top_n_sigma(request.params.top_n_sigma) { return ErrorResponse::new(message, "invalid_request_error").into_response(); } + if let Err(message) = validate_typical_p(request.params.typical_p) { + return ErrorResponse::new(message, "invalid_request_error").into_response(); + } // validate thinking_budget_tokens early. let effective_max_tokens = resolve_server_max_tokens(&state.config, request.params.max_tokens); @@ -493,6 +496,24 @@ mod tests { assert!(validate_top_n_sigma(request.params.top_n_sigma).is_ok()); } + #[test] + fn completions_rejects_zero_typical_p() { + let request: CompletionRequest = + serde_json::from_str(r#"{"model":"m","prompt":"hi","typical_p":0.0}"#).unwrap(); + assert_eq!( + validate_typical_p(request.params.typical_p), + Err("typical_p must be in (0.0, 1.0]") + ); + } + + #[test] + fn completions_accepts_typical_p_in_range() { + let request: CompletionRequest = + serde_json::from_str(r#"{"model":"m","prompt":"hi","typical_p":0.5}"#).unwrap(); + assert_eq!(request.params.typical_p, Some(0.5)); + assert!(validate_typical_p(request.params.typical_p).is_ok()); + } + #[test] fn completions_accepts_in_range_xtc() { let request: CompletionRequest = serde_json::from_str( diff --git a/src/server/routes/native_completion.rs b/src/server/routes/native_completion.rs index 302604f5c..08e99474f 100644 --- a/src/server/routes/native_completion.rs +++ b/src/server/routes/native_completion.rs @@ -347,6 +347,7 @@ fn native_generation_settings( "min_p": sampling.min_p, "xtc_probability": sampling.xtc_probability, "xtc_threshold": sampling.xtc_threshold, + "typical_p": sampling.typical_p, "repeat_penalty": sampling.repetition_penalty, "presence_penalty": sampling.presence_penalty, "frequency_penalty": sampling.frequency_penalty, @@ -621,6 +622,19 @@ fn build_native_generate_options( // manifest flip to land in one change, so until then the filter // stays at the disabled baseline on this endpoint. top_n_sigma: None, + // b10621 declares `typical_p` with no schema limits and treats + // any value at or above 1.0 as disabled, so a present value + // outside the enabled range (0.0, 1.0) maps to the explicit + // disabled form 1.0. It still overrides a server-wide --typical + // default, exactly like an upstream request value replaces the + // server default; only an ABSENT field lets the default apply. + typical_p: request.typical_p.map(|v| { + if v.is_finite() && v > 0.0 && v < 1.0 { + v + } else { + 1.0 + } + }), stop_sequences: request.stop.clone(), priority: RequestPriority::default(), // the caller fills this from the validated request @@ -646,3 +660,61 @@ fn build_native_generate_options( }, ) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::server::ServerConfig; + + fn native_request(json: &str) -> NativeCompletionRequest { + serde_json::from_str(json).unwrap() + } + + #[test] + fn native_completion_maps_typical_p() { + let request = native_request(r#"{"prompt":"hi","typical_p":0.5}"#); + let options = build_native_generate_options(&ServerConfig::default(), &request, None); + assert_eq!(options.sampling.typical_p, 0.5); + } + + #[test] + fn native_completion_out_of_domain_typical_p_is_the_explicit_disable() { + // b10621 declares the field with no schema limits and treats values + // at or above 1.0 as disabled. A present out-of-domain value must + // resolve to the explicit disabled form (1.0) and override the + // server-wide default, not fall back to it. + let config = ServerConfig { + default_typical_p: 0.4, + ..Default::default() + }; + for body in [ + r#"{"prompt":"hi","typical_p":1.0}"#, + r#"{"prompt":"hi","typical_p":0.0}"#, + r#"{"prompt":"hi","typical_p":-1.0}"#, + r#"{"prompt":"hi","typical_p":2.5}"#, + ] { + let request = native_request(body); + let options = build_native_generate_options(&config, &request, None); + assert_eq!(options.sampling.typical_p, 1.0, "body: {body}"); + } + } + + #[test] + fn native_completion_absent_typical_p_takes_the_server_default() { + let config = ServerConfig { + default_typical_p: 0.4, + ..Default::default() + }; + let request = native_request(r#"{"prompt":"hi"}"#); + let options = build_native_generate_options(&config, &request, None); + assert_eq!(options.sampling.typical_p, 0.4); + } + + #[test] + fn native_generation_settings_echo_typical_p() { + let request = native_request(r#"{"prompt":"hi","typical_p":0.5}"#); + let options = build_native_generate_options(&ServerConfig::default(), &request, None); + let settings = native_generation_settings(&request, &options); + assert_eq!(settings["typical_p"], 0.5); + } +} diff --git a/src/server/routes/props.rs b/src/server/routes/props.rs index 8645ac55c..356562183 100644 --- a/src/server/routes/props.rs +++ b/src/server/routes/props.rs @@ -39,6 +39,7 @@ pub(crate) fn default_generation_settings(config: &ServerConfig) -> serde_json:: "top_k": config.default_top_k, "top_p": config.default_top_p, "min_p": config.default_min_p, + "typical_p": config.default_typical_p, "repeat_penalty": config.default_repetition_penalty, "repeat_last_n": config.default_repetition_context_size, "seed": config.default_seed.unwrap_or(u64::MAX), diff --git a/src/server/routes/props_tests.rs b/src/server/routes/props_tests.rs index e519b6701..f1298bc0f 100644 --- a/src/server/routes/props_tests.rs +++ b/src/server/routes/props_tests.rs @@ -83,6 +83,7 @@ fn props_reports_exactly_the_documented_key_set() { "temperature", "top_k", "top_p", + "typical_p", ], "the /props payload key set changed. Adding a key is fine, update this list; \ losing one is a silent break for anyone reading it back." diff --git a/src/server/routes/responses.rs b/src/server/routes/responses.rs index 928d0505f..bd6e5e9c4 100644 --- a/src/server/routes/responses.rs +++ b/src/server/routes/responses.rs @@ -66,7 +66,7 @@ fn generation_error_to_response(err: anyhow::Error) -> ErrorResponse { use super::chat::{ build_generate_options, build_prompt_cache_request_context, parse_priority_header, - validate_top_n_sigma, validate_xtc_params, + validate_top_n_sigma, validate_typical_p, validate_xtc_params, }; use crate::server::request_options::{chat_carries_loop_amplifier, resolve_server_max_tokens}; @@ -126,6 +126,9 @@ pub async fn create_response( if let Err(message) = validate_top_n_sigma(translated.chat_request.params.top_n_sigma) { return ErrorResponse::new(message, "invalid_request_error").into_response(); } + if let Err(message) = validate_typical_p(translated.chat_request.params.typical_p) { + return ErrorResponse::new(message, "invalid_request_error").into_response(); + } // Build the structured-output constraint (json_schema) up front. let structured = { @@ -1228,6 +1231,18 @@ mod tests { ); } + #[test] + fn responses_rejects_above_one_typical_p() { + let request: CreateResponseRequest = + serde_json::from_str(r#"{"model":"m","input":"hi","typical_p":1.5}"#).unwrap(); + let translated = responses_request_to_chat(&request, None, None).unwrap(); + assert_eq!(translated.chat_request.params.typical_p, Some(1.5)); + assert_eq!( + validate_typical_p(translated.chat_request.params.typical_p), + Err("typical_p must be in (0.0, 1.0]") + ); + } + #[test] fn responses_rejects_out_of_range_xtc_threshold() { let request: CreateResponseRequest = diff --git a/src/server/startup.rs b/src/server/startup.rs index aad792b36..d08545f3f 100644 --- a/src/server/startup.rs +++ b/src/server/startup.rs @@ -224,6 +224,7 @@ pub struct ServerStartupConfig { pub top_p: f32, pub top_p_was_set: bool, pub min_p: f32, + pub typical_p: f32, pub seed: Option, pub repeat_last_n: usize, pub repeat_penalty: f32, @@ -558,6 +559,7 @@ impl Default for ServerStartupConfig { top_p: 0.95, top_p_was_set: false, min_p: 0.05, + typical_p: 1.0, seed: None, repeat_last_n: 64, repeat_penalty: 1.0, @@ -736,6 +738,25 @@ pub(super) fn resolve_elastic_pp_config(startup: &ServerStartupConfig) -> Option /// model's `max_position_embeddings` read from `config.json`, with the /// historical 4096 used only when neither is available. A non-negative /// `--n-predict N` is taken verbatim. +/// Sanitize the server-wide `--typical` / `--typical-p` default: values +/// outside the enabled range `(0.0, 1.0)` other than the explicit disabled +/// form `1.0` (zero, negative, above one, non-finite) fold to `1.0` with a +/// startup warning, so the operator learns the flag is inert instead of the +/// sampler silently ignoring it and `/props` echoing a value that never +/// applies. b10621 starts with such values too (its sampler treats them as +/// disabled), so startup does not reject them. +fn sanitize_typical_p_default(value: f32) -> f32 { + if value.is_finite() && value > 0.0 && value <= 1.0 { + value + } else { + tracing::warn!( + "--typical {value} is outside the enabled range (0.0, 1.0]; \ + locally typical sampling stays disabled (1.0)" + ); + 1.0 + } +} + pub(super) fn resolve_default_max_tokens( n_predict: i32, per_slot_context_size: usize, @@ -1173,6 +1194,7 @@ pub(super) fn build_server_config( default_top_p: sampling_defaults.top_p, default_top_k: sampling_defaults.top_k, default_min_p: startup.min_p, + default_typical_p: sanitize_typical_p_default(startup.typical_p), default_repetition_penalty: startup.repeat_penalty, default_repetition_context_size: startup.repeat_last_n, default_max_tokens: resolve_default_max_tokens( diff --git a/src/server/startup_tests.rs b/src/server/startup_tests.rs index ca27506ce..a611441aa 100644 --- a/src/server/startup_tests.rs +++ b/src/server/startup_tests.rs @@ -124,6 +124,7 @@ fn build_server_config_applies_normalized_startup_values() { top_p: 0.95, top_k: 32, min_p: 0.05, + typical_p: 1.0, repeat_penalty: 1.2, repeat_last_n: 96, n_predict: -1, @@ -1424,3 +1425,17 @@ fn an_unrestricted_server_never_fails_the_mode_gate() { .expect("no mode flag, nothing to check"); } } + +#[test] +fn sanitize_typical_p_default_folds_out_of_domain_to_disabled() { + use super::sanitize_typical_p_default; + assert_eq!(sanitize_typical_p_default(0.5), 0.5); + assert_eq!(sanitize_typical_p_default(1.0), 1.0); + for bad in [0.0f32, -1.0, 1.5, f32::NAN, f32::INFINITY] { + assert_eq!( + sanitize_typical_p_default(bad), + 1.0, + "--typical {bad} must fold to the disabled default" + ); + } +} diff --git a/src/server/types/request.rs b/src/server/types/request.rs index d6c56de7e..d1b8b9fa6 100644 --- a/src/server/types/request.rs +++ b/src/server/types/request.rs @@ -657,6 +657,12 @@ pub struct SamplingParams { /// maximum. Must be finite and `>= 0.0`; validated at the request layer. pub top_n_sigma: Option, + /// Locally typical sampling cutoff (`1.0` = disabled). Keeps the tokens + /// whose surprisal is closest to the row entropy until `typical_p` + /// probability mass accumulates. Must be finite and in `(0.0, 1.0]`; + /// validated at the request layer. + pub typical_p: Option, + // OpenAI-compatible frequency/presence penalties /// Frequency penalty (0.0 = disabled) - penalizes based on frequency pub frequency_penalty: Option, @@ -1163,6 +1169,15 @@ pub struct NativeCompletionRequest { pub top_p: Option, /// Min-p sampling pub min_p: Option, + /// Locally typical sampling, parameter p (`1.0` = disabled). b10621 + /// declares the field with no schema limits + /// ( + /// leaves the range commented out), so the route sanitizes rather than + /// rejects: a present value outside the enabled range `(0.0, 1.0)` + /// resolves to the explicit disabled form `1.0`, overriding any + /// server-wide `--typical` default exactly like an upstream request + /// value replaces the server default. + pub typical_p: Option, /// Repetition penalty pub repeat_penalty: Option, /// Repetition penalty last N tokens diff --git a/tests/disaggregated_e2e.rs b/tests/disaggregated_e2e.rs index e4a16775f..ce9c99fd2 100644 --- a/tests/disaggregated_e2e.rs +++ b/tests/disaggregated_e2e.rs @@ -75,6 +75,7 @@ fn test_sampling_state() -> SerializableSamplingState { frequency_penalty: 0.0, presence_penalty: 0.0, top_n_sigma: 0.0, + typical_p: 1.0, stop_token_ids: vec![2], } }