diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 4d96f8694..69f0d2782 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -197,7 +197,10 @@ request, and the engine passes each config to its resource when the request is i There are two ``ResourceReqConfig`` subclasses: - ``SamplingReqConfig`` holds ``temperature``, ``top_k``, ``top_p``, - ``repetition_penalty`` and ``ignore_eos``. The conductor fills in the per-request seed. + ``repetition_penalty``, ``min_p`` and ``ignore_eos``. The conductor fills in the + per-request seed. ``min_p`` follows the HF processor order (after the penalty and + temperature, before top-k/top-p) and needs ``enable_min_p=True`` on the node's + ``SamplerSpec``, which adds the filter to that node's captured sampler only. - ``KVReqConfig`` holds ``needed_labels``, ``needed_labels_per_node`` and ``needed_labels_per_node_walk``. These name the cache streams that the request will actually read. In a PD-disaggregated deployment, a KV transfer then copies only those diff --git a/mstar/engine/resources/sampler/config.py b/mstar/engine/resources/sampler/config.py index 8f8c3ab66..c1543916b 100644 --- a/mstar/engine/resources/sampler/config.py +++ b/mstar/engine/resources/sampler/config.py @@ -26,6 +26,11 @@ class SamplerSpec(NodeResourceSpec): # actually runs on a given step is settled per step from the resident # requests' `repetition_penalty` — see `SamplerResource.admit`. enable_repetion_penalty: bool = True + # Same kind of capability for min-p: whether the node's sampler (eager and + # graph-captured) carries the filter. It costs two passes over ``[B, V]`` + # per step, so nodes that never ask for it pay nothing; a request that + # sets ``min_p`` on a node without it is refused at ingest. + enable_min_p: bool = False @property def resource_class(self) -> "type[Resource]": @@ -41,6 +46,10 @@ class SamplingReqConfig(ResourceReqConfig): top_p: float = 1 ignore_eos: bool = False # used for benchmark parity repetition_penalty: float = 1 + # Min-p (HF ``MinPLogitsWarper`` / vLLM ``min_p``): drop every token whose + # probability is below ``min_p`` times the most likely token's, after the + # penalty and temperature and before top-k/top-p. 0 disables. + min_p: float = 0.0 _seed: int = 0 # set by the conductor def apply_conductor_config( diff --git a/mstar/engine/resources/sampler/resource.py b/mstar/engine/resources/sampler/resource.py index c82e35802..9ef009668 100644 --- a/mstar/engine/resources/sampler/resource.py +++ b/mstar/engine/resources/sampler/resource.py @@ -25,9 +25,11 @@ def __init__( vocab_size: int | None, enable_repetion_penalty: bool, device: torch.device, - comm_group: JointGroups | None=None + comm_group: JointGroups | None=None, + enable_min_p: bool = False, ): self._track_seen_tokens = enable_repetion_penalty + self._enable_min_p = enable_min_p self._vocab_size = vocab_size if self._track_seen_tokens else None self._sampler = Sampler( device=device, @@ -91,6 +93,7 @@ def build(cls, spec: SamplerSpec, info: EngineResourceInfo): enable_repetion_penalty=spec.enable_repetion_penalty, device=info.device, comm_group=info.joint_comm_group, + enable_min_p=spec.enable_min_p, ) def build_cuda_graph_buffers( @@ -113,6 +116,7 @@ def build_cuda_graph_buffers( tp_group=self._comm_group, vocab_size=self._vocab_size, cg_slots=self._cg_slots, + enable_min_p=self._enable_min_p, ) def ingest_request(self, rid: str, overrides: SamplingReqConfig | None=None): @@ -125,8 +129,17 @@ def ingest_request(self, rid: str, overrides: SamplingReqConfig | None=None): ) # Read off the resolved config rather than `overrides`, so a request # that leaves the penalty unset takes the same default the sampler will. - if self._sampler._sampling_config[rid].repetition_penalty != 1.0: + resolved = self._sampler._sampling_config[rid] + if resolved.repetition_penalty != 1.0: self._penalty_rids.add(rid) + if resolved.min_p > 0 and not self._enable_min_p: + # the graph-captured sampler has no min-p buffer, so honouring it + # eagerly but not in graph would make sampling depend on the path + self._sampler.remove_request(rid) + raise ValueError( + f"request {rid!r} asks for min_p={resolved.min_p} but the node's " + "SamplerSpec has enable_min_p=False" + ) if self._cg_buffers is not None: self._cg_buffers.register_request( rid, sampling_config=self._sampler._sampling_config[rid] diff --git a/mstar/engine/resources/sampler/utils.py b/mstar/engine/resources/sampler/utils.py index cf9bb2807..752a6619f 100644 --- a/mstar/engine/resources/sampler/utils.py +++ b/mstar/engine/resources/sampler/utils.py @@ -226,6 +226,20 @@ def fused_temperature_softmax( return probs +def apply_min_p(probs: torch.Tensor, min_p: torch.Tensor) -> torch.Tensor: + """Min-p filter on a ``[B, V]`` distribution: zero every probability below + ``min_p`` times the row's largest, then renormalise. + + The same tokens HF's ``MinPLogitsWarper`` keeps (``probs >= min_p * max``, + so the argmax always survives). Rows with ``min_p == 0`` and one-hot + (greedy) rows come back unchanged. No CPU branches or data-dependent + shapes, so it can sit inside a captured graph. + """ + threshold = probs.amax(dim=-1, keepdim=True) * min_p[:, None] + kept = torch.where(probs >= threshold, probs, torch.zeros_like(probs)) + return kept / kept.sum(dim=-1, keepdim=True) + + @dataclass class SamplingConfig: # Sizes the per-request seen-token mask for the repetition penalty. When set, @@ -238,6 +252,7 @@ class SamplingConfig: top_p: float = 1 ignore_eos: bool = False # used for benchmark parity repetition_penalty: float = 1 + min_p: float = 0.0 # 0 = disabled; see ``SamplingReqConfig.min_p`` _seed: int = 0 # set by the conductor def set_seed(self, seed: int): @@ -364,6 +379,10 @@ def sample( top_k = torch.tensor([c.top_k for c in configs], device=logits.device, dtype=torch.int32) top_p = torch.tensor([c.top_p for c in configs], device=logits.device) r_pen = torch.tensor([c.repetition_penalty for c in configs], device=logits.device) + min_p = ( + torch.tensor([c.min_p for c in configs], device=logits.device) + if any(c.min_p > 0 for c in configs) else None + ) seed = torch.tensor([c.seed for c in configs], device=logits.device, dtype=torch.long) rand_offset = torch.tensor( [self._step_offset.get(rid, 0) for rid in request_ids], @@ -400,6 +419,7 @@ def sample( all_top_k_zero=all_top_k_zero, seed=seed, rand_offset=rand_offset, + min_p=min_p, ) # TODO: make this scatter async. Currently runs 2 kernels per rid @@ -444,6 +464,7 @@ def _sample_cuda( all_top_k_zero: bool | None, seed: torch.Tensor | None, rand_offset: torch.Tensor | None, + min_p: torch.Tensor | None = None, ) -> torch.Tensor: """Sample normalized CUDA inputs with FlashInfer.""" import flashinfer @@ -465,6 +486,8 @@ def _sample_cuda( seen_mask=seen_token_mask, include_greedy=run_greedy, ) + if min_p is not None: + probs = apply_min_p(probs, min_p) result = flashinfer.sampling.top_p_sampling_from_probs( probs, top_p, deterministic=True, @@ -478,6 +501,8 @@ def _sample_cuda( seen_mask=seen_token_mask, include_greedy=run_greedy, ) + if min_p is not None: + probs = apply_min_p(probs, min_p) result = flashinfer.sampling.top_k_top_p_sampling_from_probs( probs, top_k, top_p, deterministic=True, @@ -497,6 +522,7 @@ def _sample_xpu( any_top_k_zero: bool | None, seed: torch.Tensor | None, rand_offset: torch.Tensor | None, + min_p: torch.Tensor | None = None, ) -> torch.Tensor: """Sample normalized XPU inputs with vllm-xpu-kernels.""" import vllm_xpu_kernels._xpu_C # noqa: F401 @@ -516,6 +542,11 @@ def _sample_xpu( ) scores = (scores / safe_temperature[:, None]).contiguous() greedy_tokens = scores.argmax(dim=-1) if run_greedy else None + if min_p is not None: + # the kernel samples from raw logits, so the filter masks those + probs = scores.softmax(dim=-1) + threshold = probs.amax(dim=-1, keepdim=True) * min_p[:, None] + scores = scores.masked_fill(probs < threshold, float("-inf")).contiguous() # The XPU kernel accepts one CPU [seed, offset] pair per invocation. # Invoke it per row to preserve independent request RNG streams. @@ -603,6 +634,7 @@ def sample_tokens( all_top_k_zero: bool | None = None, seed: torch.Tensor | None = None, rand_offset: torch.Tensor | None = None, + min_p: float | torch.Tensor | None = None, ) -> torch.Tensor: """Sample tokens from logits with temperature, top-k, top-p, and repetition penalty. @@ -618,6 +650,9 @@ def sample_tokens( branch entirely. None = unknown → run the full path. any_top_k_zero: CPU-side hint. When False, skips the `top_k == 0 → vocab` masked_fill. None = unknown → run the full path. + min_p: Scalar or per-request tensor [batch_size]; None/0 = disabled. + Applied to the temperature-scaled, penalised distribution before + top-k/top-p (the HF processor order). Returns: tokens: [batch_size] sampled token IDs. @@ -630,6 +665,8 @@ def sample_tokens( top_p = _to_tensor(top_p, batch_size, logits.device) if seen_token_mask is not None: repetition_penalty = _to_tensor(repetition_penalty, batch_size, logits.device) + if min_p is not None: + min_p = _to_tensor(min_p, batch_size, logits.device) # Default to the conservative "unknown → do the work" path. run_greedy = True if any_greedy is None else any_greedy @@ -646,6 +683,7 @@ def sample_tokens( all_top_k_zero, seed, rand_offset, + min_p=min_p, ) elif logits.device.type == "xpu": return _sample_xpu( @@ -659,6 +697,7 @@ def sample_tokens( any_top_k_zero, seed, rand_offset, + min_p=min_p, ) else: raise ValueError( @@ -701,6 +740,7 @@ def sample_cuda_graphable_gpu( apply_penalty: bool = False, rep_penalty: torch.Tensor | None = None, seen_tokens: torch.Tensor | None = None, + min_p: torch.Tensor | None = None, ) -> torch.Tensor: """Deterministic per-batch top-k/top-p sampling for graph-captured code. @@ -729,6 +769,8 @@ def sample_cuda_graphable_gpu( apply_penalty: when True, ``rep_penalty`` + ``seen_tokens`` are applied. rep_penalty: ``[batch_size]`` float tensor (1.0 = disabled per row). seen_tokens: ``[batch_size, vocab_size]`` bool mask of seen tokens. + min_p: ``[batch_size]`` float tensor (0.0 = disabled per row); None + leaves the filter out of the captured graph entirely. Returns: ``[batch_size]`` int64 sampled token IDs. FlashInfer's default @@ -744,6 +786,8 @@ def sample_cuda_graphable_gpu( seen_mask=seen_tokens if apply_penalty else None, include_greedy=True, ) + if min_p is not None: + probs = apply_min_p(probs, min_p) top_k = torch.where(top_k > 0, top_k, logits.shape[1]) # NOTE: this is NOT batch-invariant — flashinfer's deterministic RNG # folds the batch row index into philox, so identical (probs, seed, @@ -769,6 +813,8 @@ class CudaGraphableSampler(BaseSampler): # that don't opt into seen-token tracking (then ``apply_penalty`` is a no-op). rep_penalty_buf: torch.Tensor | None = None seen_tokens_buf: torch.Tensor | None = None # [bs, V] bool + # ``None`` for submodules whose ``SamplerSpec`` leaves ``enable_min_p`` off. + min_p_buf: torch.Tensor | None = None tp_group: "CommGroup | None" = None # noqa: F821 # Set during graph capture, and used by the cuda graph runner to determine @@ -787,6 +833,7 @@ def sample( apply_penalty=apply_penalty, rep_penalty=self.rep_penalty_buf, seen_tokens=self.seen_tokens_buf, + min_p=self.min_p_buf, ) self.offset_buf += 1 codes = self._broadcast_tokens(codes) @@ -973,6 +1020,9 @@ class SamplerBuffers: # only for submodules that opt in by declaring a vocab size (e.g. the # Qwen3-Omni Talker). ``None`` => the CUDA-graph path applies no penalty. seen_tokens: "MaskBuffer | None" = None + # Per-request min-p; allocated only for submodules whose spec enables it, + # so every other node's captured sampler is unchanged. + min_p: "Buffer | None" = None # Master cache capacity (grown by doubling when more requests are # concurrently registered than the per-step buffer holds). _master_capacity: int = field(default=0, repr=False) @@ -1010,7 +1060,10 @@ def tracks_seen_tokens(self) -> bool: return self.seen_tokens is not None def _scalar_buffers(self) -> list[Buffer]: - return [self.temperature, self.top_k, self.top_p, self.seed, self.rep_penalty] + bufs = [self.temperature, self.top_k, self.top_p, self.seed, self.rep_penalty] + if self.min_p is not None: + bufs.append(self.min_p) + return bufs @classmethod def allocate( @@ -1020,6 +1073,7 @@ def allocate( tp_group: "CommGroup | None" = None, # noqa: F821 vocab_size: int | None = None, cg_slots: int = 1, + enable_min_p: bool = False, ) -> "SamplerBuffers": """Allocate sampling buffers for ``max_batch_size``. @@ -1053,6 +1107,7 @@ def mk(dtype: torch.dtype, default: float, slots=cg_slots) -> Buffer: offset=mk(torch.long, 0, slots=1), tp_group=tp_group, seen_tokens=seen_tokens, + min_p=mk(torch.float32, 0.0) if enable_min_p else None, _master_capacity=cap, cg_slots=cg_slots, _slot_idx_cpu=torch.zeros(cg_slots, max_batch_size, dtype=torch.long, pin_memory=pinned), @@ -1074,6 +1129,7 @@ def slice_for_bs(self, bs: int, cg_slot: int = 0) -> dict[str, Any]: "offset_buf": self.offset.slot_view(cg_slot, bs), "rep_penalty_buf": self.rep_penalty.slot_view(cg_slot, bs), "seen_tokens_buf": self.seen_tokens.slot_view(cg_slot, bs) if self.seen_tokens is not None else None, + "min_p_buf": self.min_p.slot_view(cg_slot, bs) if self.min_p is not None else None, "tp_group": self.tp_group, } @@ -1104,6 +1160,8 @@ def _write_master_row(self, slot: int, cfg: SamplingConfig) -> None: self.top_p.write_master_row(slot, p) self.seed.write_master_row(slot, cfg.seed) self.rep_penalty.write_master_row(slot, float(cfg.repetition_penalty)) + if self.min_p is not None: + self.min_p.write_master_row(slot, float(cfg.min_p) if cfg.temperature > 0 else 0.0) def _grow_master(self, new_capacity: int) -> None: """Double-and-copy the master buffers up to at least ``new_capacity``. diff --git a/test/modular/test_sampler_min_p.py b/test/modular/test_sampler_min_p.py new file mode 100644 index 000000000..a961aa754 --- /dev/null +++ b/test/modular/test_sampler_min_p.py @@ -0,0 +1,143 @@ +"""Min-p sampling: the filter's semantics, and that the knob reaches the eager +sampler's per-request config, the graph buffers, and is refused on nodes whose +``SamplerSpec`` does not enable it. + +The filter itself is checked against the HF ``MinPLogitsWarper`` definition +(``probs < min_p * max_prob`` is dropped) re-derived here on logits, so the +probability-space implementation cannot drift from the reference order: +penalty -> temperature -> min-p -> top-k/top-p. +""" + +from __future__ import annotations + +import sys +from dataclasses import asdict +from types import SimpleNamespace + +sys.path.insert(0, ".") + +import pytest +import torch + +from mstar.engine.resources.sampler.config import SamplerSpec, SamplingReqConfig +from mstar.engine.resources.sampler.resource import SamplerResource +from mstar.engine.resources.sampler.utils import ( + Sampler, + SamplerBuffers, + SamplingConfig, + apply_min_p, +) + +CPU = torch.device("cpu") + + +def _hf_min_p(logits: torch.Tensor, min_p: torch.Tensor) -> torch.Tensor: + """transformers.MinPLogitsWarper on logits, then softmax.""" + probs = logits.softmax(dim=-1) + threshold = min_p[:, None] * probs.amax(dim=-1, keepdim=True) + return logits.masked_fill(probs < threshold, float("-inf")).softmax(dim=-1) + + +def test_apply_min_p_matches_the_hf_warper(): + torch.manual_seed(0) + logits = torch.randn(4, 64) * 3 + min_p = torch.tensor([0.0, 0.05, 0.3, 1.0]) + + ours = apply_min_p(logits.softmax(dim=-1), min_p) + + torch.testing.assert_close(ours, _hf_min_p(logits, min_p), atol=1e-6, rtol=1e-5) + torch.testing.assert_close(ours.sum(dim=-1), torch.ones(4)) + + +def test_min_p_zero_is_the_identity_and_one_keeps_only_the_argmax(): + torch.manual_seed(1) + probs = torch.randn(3, 16).softmax(dim=-1) + + torch.testing.assert_close(apply_min_p(probs, torch.zeros(3)), probs) + one_hot = apply_min_p(probs, torch.ones(3)) + assert torch.equal(one_hot.argmax(dim=-1), probs.argmax(dim=-1)) + torch.testing.assert_close(one_hot.amax(dim=-1), torch.ones(3)) + + +def test_greedy_one_hot_rows_pass_through(): + """The prep kernel turns temperature-0 rows into a one-hot; min-p must not + disturb them (the threshold is the single kept probability itself).""" + probs = torch.zeros(2, 8) + probs[0, 5] = 1.0 + probs[1, 2] = 1.0 + + torch.testing.assert_close(apply_min_p(probs, torch.tensor([0.1, 0.9])), probs) + + +def test_request_config_reaches_the_eager_sampler(): + sampler = Sampler(device=CPU) + sampler.add_request("r") + sampler.set_config("r", **asdict(SamplingReqConfig(min_p=0.05, temperature=0.8))) + + assert sampler._sampling_config["r"].min_p == 0.05 + assert sampler._sampling_config["r"].temperature == 0.8 + + +def test_graph_buffers_carry_min_p_only_when_enabled(): + on = SamplerBuffers.allocate(max_batch_size=4, device=CPU, enable_min_p=True) + off = SamplerBuffers.allocate(max_batch_size=4, device=CPU) + + assert off.slice_for_bs(2)["min_p_buf"] is None + assert on.slice_for_bs(2)["min_p_buf"].shape == (2,) + + on._write_master_row(1, SamplingConfig(min_p=0.1, temperature=0.7)) + off._write_master_row(1, SamplingConfig(min_p=0.1, temperature=0.7)) # no buffer, no error + assert on.min_p.master[1].item() == pytest.approx(0.1) + # greedy rows sample a one-hot; the filter is written inert for them + on._write_master_row(2, SamplingConfig(min_p=0.1, temperature=0.0)) + assert on.min_p.master[2].item() == 0.0 + + +def test_resource_refuses_min_p_without_the_capability(): + plain = SamplerResource(vocab_size=None, enable_repetion_penalty=False, device=CPU) + with pytest.raises(ValueError, match="enable_min_p=False"): + plain.ingest_request("r", SamplingReqConfig(min_p=0.05)) + assert "r" not in plain._sampler._sampling_config + + plain.ingest_request("ok", SamplingReqConfig(min_p=0.0)) # default stays fine + assert plain._sampler._sampling_config["ok"].min_p == 0.0 + + capable = SamplerResource( + vocab_size=None, enable_repetion_penalty=False, device=CPU, enable_min_p=True, + ) + capable.ingest_request("r", SamplingReqConfig(min_p=0.05)) + assert capable._sampler._sampling_config["r"].min_p == 0.05 + + +def test_spec_capability_defaults_off(): + spec = SamplerSpec(resource_key="sampler", nodes={"lm"}, vocab_size=None) + assert spec.enable_min_p is False + # and build() forwards it, so a spec that opts in gets a capable resource + info = SimpleNamespace(device=CPU, joint_comm_group=None) + assert SamplerResource.build(spec, info)._enable_min_p is False + on = SamplerSpec(resource_key="sampler", nodes={"lm"}, vocab_size=None, enable_min_p=True) + assert SamplerResource.build(on, info)._enable_min_p is True + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="FlashInfer sampler requires CUDA") +def test_min_p_one_is_greedy_for_every_seed_on_cuda(): + pytest.importorskip("flashinfer") + from mstar.engine.resources.sampler.utils import sample_cuda_graphable_gpu, sample_tokens + + dev = torch.device("cuda") + torch.manual_seed(0) + logits = torch.randn(4, 128, device=dev) + expected = logits.argmax(dim=-1) + for seed in range(8): + seeds = torch.full((4,), seed, device=dev, dtype=torch.long) + offsets = torch.zeros(4, device=dev, dtype=torch.long) + eager = sample_tokens( + logits, temperature=1.0, min_p=1.0, seed=seeds, rand_offset=offsets, + any_greedy=False, any_top_k_zero=True, all_top_k_zero=True, + ) + graph = sample_cuda_graphable_gpu( + logits, torch.ones(4, device=dev), torch.zeros(4, device=dev, dtype=torch.int32), + torch.ones(4, device=dev), seeds, offsets, min_p=torch.ones(4, device=dev), + ) + assert torch.equal(eager.to(expected.dtype), expected) + assert torch.equal(graph, expected)