diff --git a/docs/cli.md b/docs/cli.md index d33c065bca..06077d8a27 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -114,6 +114,7 @@ See [models.md](models.md#moe-strategies) for what each strategy does. | `--moe-cpu-threads` | physical cores | CPU worker threads for the cpu/hybrid executor | | `--moe-cpu-layers` | all on GPU | With `offload`: which MoE layers decode on CPU (`3,7,11`, a count, a fraction, or `auto`). `auto` is for Windows/WSL only, where CUDA pinned memory is capped; every value needs an expert format the CPU executor serves (bf16, fp8_block, nvfp4, mxfp4) | | `--moe-hybrid-max-fetch` | auto | With `hybrid`: max experts fetched over PCIe per layer per step; rest computed on CPU | +| `--moe-prefill-demand-tokens` | 0 | Prefill: a chunk of at most this many tokens loads only the experts it routes to, through the expert cache as decode does (GPU decode only); 0 streams whole layers | | `--moe-prefill-hit-d2d` | off | Prefill: copy cache-hit experts device-side, stream only misses (CUDA >= 13) | | `--disable-moe-prefill-overlap` | overlap on | Disable the two-buffer prefill copy overlap | @@ -202,7 +203,7 @@ so a client can gate its attachment controls without reading the checkpoint conf | `--image-min-tokens`, `--image-max-tokens` | processor defaults | Per-image token budget: the image processor resizes every image to take between these many tokens, converted to the family's own units by its processor. A family with fixed budgets honors the maximum only and refuses one below its smallest budget at start-up | | `--mm-processor-kwargs` | none | JSON object of extra keyword arguments for the checkpoint's image processor call, for knobs the token budget does not cover; applied after the budget, so an explicit key wins | | `--mm-embed-cache-device` | cpu | Where encoded image embeddings live between prefill chunks. `cpu` keeps them out of the VRAM budget; `cuda` skips the copy back | -| `--allowed-media-domains` | any | Comma-separated hostname allowlist for image URLs; requests for other domains are rejected with a 400. Empty allows any domain | +| `--allowed-media-domains` | any | Comma-separated hostname allowlist for image URLs; requests for other domains are rejected with a 400. Empty allows any domain. Redirects are re-checked against the list and may not land on a private or loopback address | | `--allowed-local-media-path` | off | Directory `file://` image refs may be read from; unset rejects local files | ## ft shell diff --git a/python/freetoken/core.py b/python/freetoken/core.py index ab5b81b0af..a7811ed712 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -30,7 +30,8 @@ class SamplingParams: # min_p: drop tokens whose probability is below min_p x the top probability. min_p: float = 0.0 # OpenAI presence/frequency penalties over the generated tokens; repetition_penalty is the - # HF/vLLM multiplicative penalty over prompt + generated tokens (1.0 = off). + # HF/vLLM multiplicative penalty over prompt + generated tokens (1.0 = off). The sampler + # keeps their state on the device (Req.output_token_counts, Req.prompt_token_mask). presence_penalty: float = 0.0 frequency_penalty: float = 0.0 repetition_penalty: float = 1.0 @@ -60,11 +61,10 @@ def is_greedy(self) -> bool: @property def needs_logits_processing(self) -> bool: + # the penalties are not LogitsPlan work: the sampler applies them from its own + # device-side state (Sampler.prepare's penalties). return bool( self.min_p > 0.0 - or self.presence_penalty != 0.0 - or self.frequency_penalty != 0.0 - or self.repetition_penalty != 1.0 or self.logit_bias or self.min_tokens > 0 ) @@ -106,7 +106,10 @@ class Req: # handler must not free resources under an in-flight forward; it sets this flag and # _process_last_data frees the request when the batch drains (after copy_done.synchronize). aborted: bool = False + output_token_counts: torch.Tensor | None = field(default=None, init=False, repr=False) stop_decode_status: DecodeStatus | None = None + # [vocab] bool on the device: the prompt's ids, for repetition_penalty. Built once. + prompt_token_mask: torch.Tensor | None = field(default=None, init=False, repr=False) def __post_init__(self) -> None: assert self.input_ids.is_cpu diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 5567164573..98aaffb32b 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -50,6 +50,7 @@ class EngineConfig: moe_cache_auto: bool = False kv_reserve_tokens: int = 8192 # KV floor for --moe-cache-auto; small by design (MoE-priority) moe_cache_policy: str = "lru" + moe_prefill_demand_tokens: int = 0 moe_prefill_overlap: bool = True # Prefill hit/miss split: serve cache-resident experts D2D during prefill # prefetch instead of re-streaming the full layer over PCIe. Needs CUDA >= 12.8 diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 8dff9f3119..a8e0aaff8a 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -494,6 +494,9 @@ def __init__(self, config: EngineConfig): logger.info_rank0(f"Free memory before loading model: {mem_GB(init_free_memory)}") # ======================= Model initialization ======================== + from freetoken.kernel.triton.autotune_cache import bound_autotune_flush_buffer + + bound_autotune_flush_buffer(self.device.index) set_rope_device(self.device) with torch.device("meta"), torch_dtype(config.dtype): self.model = create_model(config.model_config) @@ -947,6 +950,7 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: cache_size=config.moe_cache_size, device=self.device, cache_policy=config.moe_cache_policy, + prefill_demand_tokens=config.moe_prefill_demand_tokens, prefill_overlap=config.moe_prefill_overlap, prefill_hit_d2d=config.moe_prefill_hit_d2d, quant_format=banks.quant_format, @@ -1868,6 +1872,7 @@ def _auto_cpu_layers(config: EngineConfig, num_moe_layers: int, *, reserved: int "moe_cpu_layers": None, "moe_cpu_threads": 0, "moe_hybrid_max_fetch": -1, + "moe_prefill_demand_tokens": 0, "moe_prefill_overlap": True, "moe_prefill_hit_d2d": False, "expert_load": "auto", diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index 89d2a6fb88..9d12b5f571 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -1,7 +1,7 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import TYPE_CHECKING, List +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, List, NamedTuple import torch from freetoken.utils import is_sm90_supported, nvtx_annotate @@ -13,41 +13,49 @@ @dataclass class LogitsPlan: """Inputs of the logits processors for the rows of one batch that asked for any - (presence / frequency / repetition penalties, logit_bias, min_tokens, min_p). - - Built on the host in ``Sampler.prepare`` from the requests' SamplingParams and - token histories, applied on the device by ``apply_logits_processors`` before the - sampling kernel. ``None`` when no request in the batch uses a processor, so the - default path stays exactly as it was. Row indices inside the plan are LOCAL (0..r-1, - in the order of ``rows``); ``rows`` maps them back to the batch. + (logit_bias, min_tokens, min_p). The penalties are not here: the sampler applies + them from per-request device state (``RowPenalty``). + + Built on the host in ``Sampler.prepare`` from the requests' SamplingParams, applied + on the device by ``apply_logits_processors`` before the sampling kernel. ``None`` + when no request in the batch uses a processor, so the default path stays exactly + as it was. Row indices inside the plan are LOCAL (0..r-1, in the order of ``rows``); + ``rows`` maps them back to the batch. """ rows: torch.Tensor # [r] int64 batch rows temps: torch.Tensor # [r, 1] float32 temperatures (min_p works on probabilities) - presence: torch.Tensor | None = None # [r, 1] float32 - frequency: torch.Tensor | None = None # [r, 1] float32 - repetition: torch.Tensor | None = None # [r, 1] float32 (1.0 = off) - # Generated ids per row, right-padded with vocab_size (a scratch column that is - # dropped), for the presence / frequency counts and the repetition mask. - hist_ids: torch.Tensor | None = None # [r, L] int64 - # Prompt ids per row, same padding: repetition_penalty also covers the prompt. - prompt_ids: torch.Tensor | None = None # [r, P] int64 bias_rows: torch.Tensor | None = None # [b] int64 local rows bias_ids: torch.Tensor | None = None # [b] int64 bias_vals: torch.Tensor | None = None # [b] float32 - # Rows still under min_tokens and the ids (EOS + stop_token_ids, padded) they must - # not sample yet. + # Rows still under min_tokens and the ids (EOS + stop_token_ids, padded with + # vocab_size, a scratch column that is dropped) they must not sample yet. min_rows: torch.Tensor | None = None # [m] int64 local rows min_ids: torch.Tensor | None = None # [m, S] int64 min_p: torch.Tensor | None = None # [r, 1] float32 (0 = off for that row) +class RowPenalty(NamedTuple): + """One batch row's presence / frequency / repetition penalty. ``counts`` are the + request's generated-token counts ([V] int32), updated on the device after each draw so + overlap scheduling never sees them late; ``prompt_mask`` ([V] bool) marks the prompt's + ids for the repetition penalty and is None when that is off.""" + + row: int + counts: torch.Tensor + presence: float + frequency: float + repetition: float = 1.0 + prompt_mask: torch.Tensor | None = None + + @dataclass class BatchSamplingArgs: temperatures: torch.Tensor | None top_k: torch.Tensor | None = None top_p: torch.Tensor | None = None greedy_mask: torch.Tensor | None = None + penalties: list[RowPenalty] = field(default_factory=list) plan: LogitsPlan | None = None logprob_rows: torch.Tensor | None = None max_top_logprobs: int = 0 @@ -60,47 +68,39 @@ def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> def _padded_rows(rows: list[list[int]], pad: int, device: torch.device) -> torch.Tensor: - """[len(rows), max_len] int64 with ``pad`` on the right. A synchronous copy: the - histories can be long (prompt + output) and a pinned buffer must not be freed under - an in-flight non_blocking copy.""" + """[len(rows), max_len] int64 with ``pad`` on the right.""" width = max(1, max((len(r) for r in rows), default=1)) - out = torch.full((len(rows), width), pad, dtype=torch.int64) - for i, r in enumerate(rows): - if r: - out[i, : len(r)] = torch.as_tensor(r, dtype=torch.int64) - return out.to(device) + return make_device_tensor([r + [pad] * (width - len(r)) for r in rows], torch.int64, device) + + +def apply_penalties(logits: torch.Tensor, penalties: list[RowPenalty]) -> torch.Tensor: + """A float32 copy of ``logits`` with each row's penalties applied in vLLM's order: + repetition (over prompt + generated ids), then frequency and presence (generated).""" + logits = logits.float().clone() + for row, counts, presence, frequency, repetition, prompt_mask in penalties: + if repetition != 1.0: + x = logits[row] + seen = counts > 0 + if prompt_mask is not None: + seen = seen | prompt_mask + # HF / vLLM semantics: divide positive logits, multiply negative ones. + penalized = torch.where(x > 0, x / repetition, x * repetition) + logits[row] = torch.where(seen, penalized, x) + if presence or frequency: + logits[row] -= frequency * counts + presence * (counts > 0) + return logits def apply_logits_processors( logits: torch.Tensor, plan: LogitsPlan, vocab_size: int ) -> torch.Tensor: - """A float32 copy of ``logits`` with the plan applied to its rows, in vLLM's order: - repetition -> presence -> frequency penalties, logit_bias, the min_tokens mask, then - min_p. Rows outside the plan are copied unchanged.""" + """A float32 copy of ``logits`` with the plan applied to its rows: logit_bias, the + min_tokens mask, then min_p. Rows outside the plan are copied unchanged.""" out = logits.to(torch.float32, copy=True) sub = out[plan.rows] # [r, V] (advanced indexing copies) - r = sub.shape[0] device = sub.device neg_inf = float("-inf") - if plan.hist_ids is not None: - counts = torch.zeros((r, vocab_size + 1), dtype=torch.float32, device=device) - counts.scatter_add_(1, plan.hist_ids, torch.ones_like(plan.hist_ids, dtype=torch.float32)) - counts = counts[:, :vocab_size] - if plan.repetition is not None: - seen = counts > 0 - if plan.prompt_ids is not None: - in_prompt = torch.zeros((r, vocab_size + 1), dtype=torch.bool, device=device) - in_prompt.scatter_(1, plan.prompt_ids, torch.ones_like(plan.prompt_ids, dtype=torch.bool)) - seen = seen | in_prompt[:, :vocab_size] - # HF / vLLM semantics: divide positive logits, multiply negative ones. - penalized = torch.where(sub > 0, sub / plan.repetition, sub * plan.repetition) - sub = torch.where(seen, penalized, sub) - if plan.presence is not None: - sub = sub - plan.presence * (counts > 0).to(sub.dtype) - if plan.frequency is not None: - sub = sub - plan.frequency * counts - if plan.bias_ids is not None: sub.index_put_((plan.bias_rows, plan.bias_ids), plan.bias_vals, accumulate=True) @@ -155,11 +155,7 @@ class Sampler: vocab_size: int def _plan(self, batch: Batch, params) -> LogitsPlan | None: - """The LogitsPlan for this batch, or None when no request asks for a processor. - Token histories come from the host-side ``req.input_ids``; under overlap - scheduling the previous step's token may not be appended yet, so a penalty sees - the history one token late. That is the accepted cost of keeping the counts off - the device.""" + """The LogitsPlan for this batch, or None when no request asks for a processor.""" MIN_T = 1e-6 rows = [i for i, p in enumerate(params) if p.needs_logits_processing] if not rows: @@ -172,36 +168,6 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None: ) pad = self.vocab_size - want_hist = any( - p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0 - for _, p in picked - ) - if want_hist: - hist: list[list[int]] = [] - prompts: list[list[int]] = [] - want_prompt = False - for req, p in picked: - prompt_len = req.max_device_len - req.output_len - ids = req.input_ids - penalized = p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0 - hist.append(ids[prompt_len:].tolist() if penalized else []) - if p.repetition_penalty != 1.0: - want_prompt = True - # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE); - # send them to the scratch column, not past the scatter target - prompts.append(ids[:prompt_len].clamp(max=pad).tolist()) - else: - prompts.append([]) - plan.hist_ids = _padded_rows(hist, pad, self.device) - if want_prompt: - plan.prompt_ids = _padded_rows(prompts, pad, self.device) - if any(p.presence_penalty != 0.0 for _, p in picked): - plan.presence = make_device_tensor([[p.presence_penalty] for _, p in picked], torch.float32, self.device) - if any(p.frequency_penalty != 0.0 for _, p in picked): - plan.frequency = make_device_tensor([[p.frequency_penalty] for _, p in picked], torch.float32, self.device) - if any(p.repetition_penalty != 1.0 for _, p in picked): - plan.repetition = make_device_tensor([[p.repetition_penalty] for _, p in picked], torch.float32, self.device) - bias_rows: list[int] = [] bias_ids: list[int] = [] bias_vals: list[float] = [] @@ -234,9 +200,48 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None: plan.min_p = make_device_tensor([[p.min_p] for _, p in picked], torch.float32, self.device) return plan + def _prompt_mask(self, req) -> torch.Tensor: + """[V] bool marking the request's prompt ids, uploaded once per request.""" + prompt_len = req.max_device_len - req.output_len + # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE); + # send them to the scratch column, not past the scatter target + ids = req.input_ids[:prompt_len].to(torch.int64).clamp(max=self.vocab_size) + if self.device.type == "cuda": + ids = ids.pin_memory().to(self.device, non_blocking=True) + mask = torch.zeros(self.vocab_size + 1, dtype=torch.bool, device=self.device) + mask[ids] = True + return mask[: self.vocab_size] + def prepare(self, batch: Batch) -> BatchSamplingArgs: params = [r.sampling_params for r in batch.reqs] is_greedy = [p.is_greedy for p in params] + penalties = [] + for row, req in enumerate(batch.reqs): + p = req.sampling_params + repetition = p.repetition_penalty + if not (p.presence_penalty or p.frequency_penalty or repetition != 1.0): + continue + if not req.can_decode: + continue + if req.output_token_counts is None: + req.output_token_counts = torch.zeros( + self.vocab_size, dtype=torch.int32, device=self.device + ) + prompt_mask = None + if repetition != 1.0: + if req.prompt_token_mask is None: + req.prompt_token_mask = self._prompt_mask(req) + prompt_mask = req.prompt_token_mask + penalties.append( + RowPenalty( + row, + req.output_token_counts, + p.presence_penalty, + p.frequency_penalty, + repetition, + prompt_mask, + ) + ) plan = self._plan(batch, params) want_logprobs = [p.logprobs for p in params] logprob_rows = ( @@ -250,6 +255,7 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: if all(is_greedy): return BatchSamplingArgs( temperatures=None, + penalties=penalties, plan=plan, logprob_rows=logprob_rows, max_top_logprobs=max_top_logprobs, @@ -281,6 +287,7 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: top_k=top_k, top_p=top_p, greedy_mask=greedy_mask, + penalties=penalties, plan=plan, logprob_rows=logprob_rows, max_top_logprobs=max_top_logprobs, @@ -289,16 +296,26 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: @nvtx_annotate("Sampler") def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: with torch.cuda.nvtx.range("Sampler"): + if args.penalties: + logits = apply_penalties(logits, args.penalties) + # after the penalties: min_p must see the final distribution if args.plan is not None: logits = apply_logits_processors(logits, args.plan, self.vocab_size) if args.temperatures is None: # greedy sampling - return torch.argmax(logits, dim=-1) - tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) - if args.greedy_mask is not None: - # Mixed batches still run probability sampling for all rows, but - # greedy rows must follow argmax's deterministic tie-breaking. - greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype) - tokens = torch.where(args.greedy_mask, greedy_tokens, tokens) + tokens = torch.argmax(logits, dim=-1) + else: + tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) + if args.greedy_mask is not None: + # Mixed batches still run probability sampling for all rows, but + # greedy rows must follow argmax's deterministic tie-breaking. + greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype) + tokens = torch.where(args.greedy_mask, greedy_tokens, tokens) + # Update on the sampling stream: overlapped scheduling can prepare the next + # batch before the previous token reaches Req.input_ids on the CPU. + for row, counts, *_ in args.penalties: + counts.scatter_add_( + 0, tokens[row : row + 1].long(), counts.new_ones(1) + ) return tokens def compute_logprobs( diff --git a/python/freetoken/kernel/triton/autotune_cache.py b/python/freetoken/kernel/triton/autotune_cache.py index ea95768368..2e57fadde9 100644 --- a/python/freetoken/kernel/triton/autotune_cache.py +++ b/python/freetoken/kernel/triton/autotune_cache.py @@ -34,4 +34,28 @@ {"cache_results": _CACHE_RESULTS} if _SUPPORTS_AUTOTUNE_CACHE else {} ) -__all__ = ["autotune_cache_kwargs"] + + +def bound_autotune_flush_buffer(device_index: int) -> None: + """Size triton's L2-flush buffer for autotune benchmarks to the device instead of a flat 256 MB. + + A kernel autotunes on its first launch of each key, which can be the first request, after the + cache planner has handed out nearly all VRAM; a 256 MB transient there is an OOM on a small + card. Four times the L2 still evicts everything between timed runs. + """ + import torch + + try: + from triton.backends.nvidia.driver import CudaDriver + except ImportError: + return + l2 = int(getattr(torch.cuda.get_device_properties(device_index), "L2_cache_size", 0) or 0) + nbytes = min(256 << 20, max(8 << 20, 4 * l2)) + + def get_empty_cache_for_benchmark(self): + return torch.empty(nbytes // 4, dtype=torch.int, device="cuda") + + CudaDriver.get_empty_cache_for_benchmark = get_empty_cache_for_benchmark + + +__all__ = ["autotune_cache_kwargs", "bound_autotune_flush_buffer"] diff --git a/python/freetoken/kernel/triton/hc.py b/python/freetoken/kernel/triton/hc.py index 45d71bd731..acae534c45 100644 --- a/python/freetoken/kernel/triton/hc.py +++ b/python/freetoken/kernel/triton/hc.py @@ -326,9 +326,11 @@ def _hc_combine_norm_kernel( tl.store(out_ptr + row * stride_out + offs, out, mask=mask_inner) out = out.to(tl.float32) - # Keep the two-axis reduction: flattening the padded tile is ~40% slower - # at decode sizes. - sum_sq = tl.sum(tl.sum(out * out, axis=1), axis=0) + # The standalone grouped RMSNorm reduces one padded [next_pow2(HC_DIM)] + # vector; reshape keeps the identical tree (same element + # order, same masked zeros), so the fused statistic is bit-exact vs the split. + flat = tl.reshape(out * out, [NUM_TILES_PAD * BLOCK_SIZE]) + sum_sq = tl.sum(flat) rrms = tl.rsqrt(sum_sq / HC_DIM + EPS) if launch_pdl: diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 71858b54d0..b19a362cf0 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -278,8 +278,9 @@ def prefill_forward( # ------------------------------------------------------------------ # Data movement -- one decision tree for every quant format (the banks # registry makes the cache machinery bank-count agnostic). Decode loads - # on demand; prefill streams whole layers, double-buffered when overlap - # is enabled. The kernels only ever see bank views plus row indices; + # on demand, and so does a short enough prefill chunk; a longer one + # streams whole layers, double-buffered when overlap is enabled. The + # kernels only ever see bank views plus row indices; # which kernel runs is decided afterwards, in ``_expert_gemm``. # ------------------------------------------------------------------ @@ -420,21 +421,23 @@ def _prefill_routed( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor: - """Prefill movement: stream whole layers -- double-buffered behind the - previous layer's GEMMs when ``prefill_overlap`` is on, else a synchronous - ``materialize_layer``. In both, position == expert id, so the routing ids - pass through unmapped.""" + """Prefill movement. A chunk within ``prefill_demand_tokens`` loads its routed experts + into the slot cache as decode does and runs the same grouped GEMM over the slots. A + longer one streams whole layers -- double-buffered behind the previous layer's GEMMs + when ``prefill_overlap`` is on, else a synchronous ``materialize_layer`` -- where + position == expert id, so the routing ids pass through unmapped.""" cache = self.offload_cache assert cache is not None if ( - 0 < hidden_states.shape[0] <= _SMALL_PREFILL_TOKENS + 0 < hidden_states.shape[0] <= max(cache.prefill_demand_tokens, _SMALL_PREFILL_TOKENS) and cache.decode_target in ("gpu", "hybrid") and not cache.is_unpinned_layer(self.layer_id) ): # A short extension (an agent turn over a cached prefix) touches a fraction of each - # layer's experts; fetching those beats streaming all of them (~68 GB here). An - # unpinned layer's copy_missing presumes position == expert id, so it streams below; - # unpinned layers only exist with the prefill overlap off, so no layer can enter + # layer's experts; fetching those beats streaming all of them (~68 GB here). It runs + # the prefill GEMM over the slot cache (#676), so it matches the streamed layer bit for + # bit. An unpinned layer's copy_missing presumes position == expert id, so it streams + # below; unpinned layers only exist with the prefill overlap off, so no layer can enter # the overlap double buffer after an earlier layer skipped begin_prefill. return self._expert_gemm( cache, @@ -442,9 +445,9 @@ def _prefill_routed( topk_weights, self._ensure_unique(cache, topk_ids), views=cache.bank_views(), - n=None, + n=cache.cache_size, alphas=cache.alphas_for_slots(self.layer_id), - is_prefill=False, + is_prefill=True, ) if cache.prefill_overlap: views = self._wait_prefill_overlap(cache) @@ -490,7 +493,7 @@ def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torc pin_since = cache.step.clone() if cache.cache_policy_id else None for start in range(0, n, _ENSURE_CHUNK_IDS): part = uniq[start : start + _ENSURE_CHUNK_IDS] # a contiguous view: ensure rewrites it in place - cache.ensure_experts(self.layer_id, part, update_state=False, pin_since=pin_since) # not a decode step + cache.ensure_experts(self.layer_id, part, update_state=False, pin_since=pin_since, counts_as_decode=False) # not a decode step cache.copy_missing() return uniq[inverse].view_as(topk_ids) diff --git a/python/freetoken/mm/media.py b/python/freetoken/mm/media.py index 5c3c4ef789..a51e5bc98b 100644 --- a/python/freetoken/mm/media.py +++ b/python/freetoken/mm/media.py @@ -1,6 +1,8 @@ """Client-supplied image handling: ref collection, the accept gate, and byte fetching. Raises plain ValueError on bad input; the server layer maps it to its wire error type. +The message names the class of failure only: no local path, allowlist root, errno or +library text reaches the client (that detail goes to the server log). """ from __future__ import annotations @@ -8,10 +10,15 @@ import base64 from typing import TYPE_CHECKING, Any +from freetoken.utils import init_logger + if TYPE_CHECKING: from freetoken.server.args import ServerArgs +logger = init_logger(__name__) + _MAX_IMAGE_BYTES = 32 << 20 +_MAX_REDIRECTS = 5 def collect_image_refs(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: @@ -56,6 +63,57 @@ def _check_media_domain(url: str, config: ServerArgs) -> None: ) +def _check_redirect_target(url: str, config: ServerArgs) -> None: + """A redirect hop is the remote server's choice, not the client's URL: it must pass the + allowlist like the original, and may never land on a loopback / private / link-local + address (an allowed public host 302-ing into the operator's network is the SSRF shape).""" + import ipaddress + from urllib.parse import urlparse + + if not url.startswith(("http://", "https://")): + raise ValueError("the URL redirects to a non-http(s) location") + try: + _check_media_domain(url, config) + except ValueError: + raise ValueError("the URL redirects outside the allowed domains") from None + host = (urlparse(url).hostname or "").rstrip(".").lower() + if host == "localhost" or host.endswith(".localhost"): + raise ValueError("the URL redirects to a private or loopback address") + try: + ip = ipaddress.ip_address(host) + except ValueError: + return # a name; the allowlist above is the gate for names (no DNS resolution here) + if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped is not None: + ip = ip.ipv4_mapped + if not ip.is_global or ip.is_multicast: # covers loopback, private, link-local, CGNAT, unspecified + raise ValueError("the URL redirects to a private or loopback address") + + +async def _fetch_remote_media(url: str, config: ServerArgs) -> bytes: + """GET an http(s) ref, following redirects one hop at a time so every target is re-checked.""" + import httpx + + _check_media_domain(url, config) + async with httpx.AsyncClient(timeout=30.0, follow_redirects=False) as client: + for _ in range(_MAX_REDIRECTS + 1): + async with client.stream("GET", url) as resp: + if resp.next_request is not None: # a 3xx with a Location: httpx built the hop + url = str(resp.next_request.url) + _check_redirect_target(url, config) + continue + resp.raise_for_status() + declared = resp.headers.get("content-length", "") + if declared.isdigit() and int(declared) > _MAX_IMAGE_BYTES: + raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes") + buf = bytearray() + async for chunk in resp.aiter_bytes(): + buf += chunk + if len(buf) > _MAX_IMAGE_BYTES: + raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes") + return bytes(buf) + raise ValueError(f"the URL redirects more than {_MAX_REDIRECTS} times") + + def _load_local_media(url: str, config: ServerArgs) -> bytes: """Read a file:// ref; the resolved path must be a strict subpath of --allowed-local-media-path.""" from pathlib import Path @@ -70,41 +128,53 @@ def _load_local_media(url: str, config: ServerArgs) -> bytes: # resolve() follows symlinks, so a link inside the root escaping it is rejected too resolved = filepath.resolve() if Path(root).resolve() not in resolved.parents: - raise ValueError( - f"the file path {filepath} must be a subpath of " - f"--allowed-local-media-path {root}" - ) + raise ValueError("the file path must be a subpath of --allowed-local-media-path") return resolved.read_bytes() +async def _fetch_one(ref: dict[str, Any], config: ServerArgs) -> bytes: + """One ref to bytes. Our own ValueErrors are already client-safe and pass through; anything + a library raises is reduced to its failure class and logged with the detail.""" + data = ref.get("data") or "" + try: + if ref.get("kind") == "b64": + return base64.b64decode(data) + if data.startswith("data:"): + return base64.b64decode(data.split(",", 1)[1]) + except Exception as exc: # noqa: BLE001 -- binascii.Error / IndexError, input-driven + raise ValueError("image data is not valid base64") from exc + if data.startswith(("http://", "https://")): + import httpx + + try: + return await _fetch_remote_media(data, config) + except ValueError: + raise + except httpx.HTTPStatusError as exc: + raise ValueError(f"the image URL returned HTTP {exc.response.status_code}") from exc + except httpx.TimeoutException as exc: + raise ValueError("timed out fetching the image URL") from exc + except Exception as exc: # noqa: BLE001 -- transport / protocol / bad URL + logger.warning("image URL fetch failed: %r", exc) + raise ValueError("the image URL is unreachable") from exc + if data.startswith("file://"): + try: + return _load_local_media(data, config) + except ValueError: + raise + except OSError as exc: + logger.warning("local image read failed: %r", exc) + raise ValueError("the local file is missing or unreadable") from exc + raise ValueError("unsupported image source (expect http(s)/file url or base64)") + + async def fetch_image_bytes(refs: list[dict[str, Any]], config: ServerArgs) -> list[bytes]: """Resolve collected image refs (URLs / base64) to raw bytes, in order.""" out: list[bytes] = [] for ref in refs: - data = ref.get("data") or "" try: - if ref.get("kind") == "b64": - out.append(base64.b64decode(data)) - elif data.startswith("data:"): - out.append(base64.b64decode(data.split(",", 1)[1])) - elif data.startswith(("http://", "https://")): - import httpx - - _check_media_domain(data, config) - # redirect targets are not re-checked against the allowlist - async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client: - resp = await client.get(data) - resp.raise_for_status() - if len(resp.content) > _MAX_IMAGE_BYTES: - raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes") - out.append(resp.content) - elif data.startswith("file://"): - out.append(_load_local_media(data, config)) - else: - raise ValueError( - "unsupported image source (expect http(s)/file url or base64)" - ) - except Exception as exc: # noqa: BLE001 -- input-driven, client-classifiable + out.append(await _fetch_one(ref, config)) + except ValueError as exc: raise ValueError(f"could not load image: {exc}") from exc return out diff --git a/python/freetoken/models/qwen4_exp/hc.py b/python/freetoken/models/qwen4_exp/hc.py index f13808dccd..1574cd473f 100644 --- a/python/freetoken/models/qwen4_exp/hc.py +++ b/python/freetoken/models/qwen4_exp/hc.py @@ -17,6 +17,7 @@ from freetoken.kernel.triton.hc import ( grouped_gemma_rmsnorm, hc_combine, + hc_combine_norm, hc_gate_mix, hc_silu, ) @@ -119,14 +120,16 @@ def _down(self, rn: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: # both slices keep unit inner stride, so the kernels read them without a copy return down[:, : self.lowrank], down[:, self.lowrank : self.lowrank + self.hc_count] - def _mix_kernel(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: - rn = grouped_gemma_rmsnorm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) + def _mix_kernel(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]: + if rn is None: + rn = grouped_gemma_rmsnorm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) lora, s = self._down(rn) gate = self.input_mix_weight_up.forward(hc_silu(lora, self.hc_count)) return hc_gate_mix(rn, gate, self.hc_count), s - def _mix_torch(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: - rn = grouped_plus_one_rms_norm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) + def _mix_torch(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]: + if rn is None: + rn = grouped_plus_one_rms_norm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) lora, s = self._down(rn) lora = F.silu(lora.float() / self.hc_count) gate = self.input_mix_weight_up.forward(lora.to(R.dtype)) @@ -140,9 +143,12 @@ def _combine_torch(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> t out = out + y.float().unsqueeze(-2) * inject.unsqueeze(-1) return out.flatten(-2).to(R.dtype) - def mix(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: - """Return the block input ``x [T, hidden]`` and the inject logits ``s [T, hc_count]`` (None if no combine).""" - return self._mix_kernel(R) if R.is_cuda else self._mix_torch(R) + def mix(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]: + """Return the block input ``x [T, hidden]`` and the inject logits ``s [T, hc_count]`` (None if no combine). + + ``rn`` is the pre-normalized R from a preceding :meth:`combine_norm` (same norm weight as + this block's ``hc_norm``), letting the caller fuse the two kernels; None recomputes it.""" + return self._mix_kernel(R, rn) if R.is_cuda else self._mix_torch(R, rn) def combine(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> torch.Tensor: """Inject the block output ``y [T, hidden]`` back into every stream of ``R``.""" @@ -150,5 +156,15 @@ def combine(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> torch.Te return hc_combine(R, y, s, self.hc_count) return self._combine_torch(R, y, s) + def combine_norm( + self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor, next_block: "GatedResidual" + ) -> Tuple[torch.Tensor, torch.Tensor | None]: + """combine() fused with ``next_block``'s hc_norm; returns (R', Rn'). CPU keeps the split path (Rn' = None).""" + if not R.is_cuda: + return self.combine(R, y, s), None + return hc_combine_norm( + R, y, s, next_block.hc_norm.weight, next_block.hc_norm.eps, self.hc_count + ) + __all__ = ["GatedResidual", "GroupedPlusOneRMSNorm", "grouped_plus_one_rms_norm"] diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 7f4f826807..6ce2da985b 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -75,17 +75,19 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N ) @nvtx_annotate("Layer_{}", layer_id_field="_layer_id") - def forward(self, hidden: torch.Tensor, batch: Batch) -> torch.Tensor: + def forward(self, hidden: torch.Tensor, batch: Batch, rn: torch.Tensor | None = None): + """Returns (R, mlp_out, inject) un-combined: the caller fuses the mlp combine with the + next consumer's norm (next layer's attn mix, or the final mixer).""" if self.ple is not None: hidden = hidden + self.ple.forward(hidden, batch) - block_input, inject = self.attn_hyper_connection.mix(hidden) + block_input, inject = self.attn_hyper_connection.mix(hidden, rn=rn) if self._is_linear: block_output = self.linear_attn.forward(block_input) else: block_output = self.self_attn.forward(block_input, batch) - hidden = self.attn_hyper_connection.combine(hidden, block_output, inject) - block_input, inject = self.mlp_hyper_connection.mix(hidden) - return self.mlp_hyper_connection.combine(hidden, self.mlp.forward(block_input), inject) + hidden, rn_mlp = self.attn_hyper_connection.combine_norm(hidden, block_output, inject, self.mlp_hyper_connection) + block_input, inject = self.mlp_hyper_connection.mix(hidden, rn=rn_mlp) + return hidden, self.mlp.forward(block_input), inject class Qwen4ExpModel(BaseOP): @@ -121,13 +123,25 @@ def forward(self, input_ids: torch.Tensor, batch: Batch) -> torch.Tensor: meta = build_ple_metadata(batch, self._ple[0].args, input_ids.device) for ple in self._ple: # gather the pinned-host PLE rows while the early layers run ple.start_prefetch(batch, meta) - for layer in self.layers.op_list: - hidden = layer.forward(hidden, batch) + # The last block of each layer combines through the layer boundary, fusing the + # combine with the next consumer's hc_norm -- except before a PLE layer, whose + # `hidden + ple(hidden)` happens after the boundary and changes what the norm sees. + layers = self.layers.op_list + rn = None + for i, layer in enumerate(layers): + hidden, mlp_out, inject = layer.forward(hidden, batch, rn=rn) + nxt = layers[i + 1] if i + 1 < len(layers) else None + if nxt is not None and nxt.ple is not None: + hidden = layer.mlp_hyper_connection.combine(hidden, mlp_out, inject) + rn = None + else: + consumer = nxt.attn_hyper_connection if nxt is not None else self.hyper_connection_mixer + hidden, rn = layer.mlp_hyper_connection.combine_norm(hidden, mlp_out, inject, consumer) if meta is not None: # single writer: the layers only read the context, so a second PLE layer's # prefetch sees the un-rolled window commit_ngram_context(meta, getattr(batch, "fla_metadata", None)) - return self.hyper_connection_mixer.mix(hidden)[0] + return self.hyper_connection_mixer.mix(hidden, rn=rn)[0] class Qwen4ExpForCausalLM(BaseLLMModel): diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py index a90aa4c792..b0b9c01c71 100644 --- a/python/freetoken/moe/fused.py +++ b/python/freetoken/moe/fused.py @@ -44,6 +44,9 @@ def fused_topk( return fused_topk_softmax(gating_output, topk, renormalize, num_token_non_padded) +_SGL_ALIGN_MAX_EXPERTS = 1024 + + def moe_align_block_size( topk_ids: torch.Tensor, block_size: int, num_experts: int ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -86,7 +89,10 @@ def moe_align_block_size( """ from freetoken.kernel.backend import is_sgl_kernel_installed - if not is_sgl_kernel_installed(): + # sgl_kernel's align kernel takes num_experts + 1 counters in one block and reads out of + # bounds past 1024 of them (0.4.5 and 0.4.9). A prefill over an offload slot cache passes the + # slot count here (#676), which is thousands; the Triton kernel has no such limit. + if not is_sgl_kernel_installed() or num_experts + 1 > _SGL_ALIGN_MAX_EXPERTS: from freetoken.kernel.triton.moe_align import ( moe_align_block_size as triton_moe_align_block_size, ) diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 246b6b248f..8f0d8635dc 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -114,6 +114,9 @@ class OffloadMoeCache: cache_size: int device: torch.device cache_policy: str = "lru" + # A prefill chunk of at most this many tokens admits the experts it routes to into the slot + # cache, as decode does, instead of streaming every expert of each layer; 0 streams them all. + prefill_demand_tokens: int = 0 prefill_overlap: bool = False # Prefill hit/miss split: experts already resident in the slot cache (slots # >= 2 * num_experts) are gathered device-side into the double buffer instead @@ -172,6 +175,10 @@ def __post_init__(self) -> None: raise ValueError( f"--moe-cache-policy {self.cache_policy} needs GPU decode, not {self.decode_target!r}" ) + if self.prefill_demand_tokens and self.decode_target != "gpu": + raise ValueError( + f"--moe-prefill-demand-tokens needs GPU decode; this cache decodes on {self.decode_target!r}" + ) if self.layout is None: assert self.quant_format in _BANK_SCHEMAS, f"unknown quant_format {self.quant_format!r}" # Attached by the engine for decode_target == "cpu" (CpuMoeExecutor); None @@ -1051,6 +1058,9 @@ def release_prefill_layer(self, layer_id: int) -> None: self._prefill_buffer_has_release_event[buffer_id] = True self._prefill_buffer_released[buffer_id] = True + def loads_prefill_on_demand(self, num_tokens: int) -> bool: + return 0 < num_tokens <= self.prefill_demand_tokens + def ensure_experts( self, layer_id: int, @@ -1059,6 +1069,7 @@ def ensure_experts( update_state: bool = True, router_logits: torch.Tensor | None = None, pin_since: torch.Tensor | None = None, + counts_as_decode: bool = True, ) -> None: """``update_state=False`` (small-prefill ensures) keeps a scored policy's state and decode-step clock still; ``router_logits`` (``[rows, num_experts]``) feed the ``rule`` @@ -1067,7 +1078,7 @@ def ensure_experts( already keeps those slots.""" from freetoken.moe.offload_kernels import ensure_experts, ensure_experts_scored - if self.collect_decode_freq: + if self.collect_decode_freq and counts_as_decode and update_state: # ``expert_ids`` still holds raw expert ids here (the kernel rewrites them to # slot ids in place), so snapshot the routing histogram before that happens. ids = expert_ids.reshape(-1).long() @@ -1080,7 +1091,7 @@ def ensure_experts( lowpri = self.prefetch_on width = expert_ids.numel() if self.cache_policy_id == 0 and not lowpri: - ensure_experts(self, layer_id, expert_ids) + ensure_experts(self, layer_id, expert_ids, record_stats=counts_as_decode and update_state) else: # prefetch on: LRU runs the vendored kernel (flashlib's cannot key low-priority slots) ensure_experts_scored( diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index 5993738d10..8abe5e9c63 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -16,7 +16,7 @@ ) -def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None: +def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor, *, record_stats: bool = True) -> None: """Make this layer's routed experts resident; rewrite ``expert_ids`` to slot ids. Delegates to flashlib's slot cache. ``id_base`` maps this layer's expert ids into the @@ -35,7 +35,7 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None: cache.src_indices, cache.evict_slots, cache.num_indices, - stats=cache.lru_stats[layer_id] if cache.collect_stats else None, + stats=cache.lru_stats[layer_id] if cache.collect_stats and record_stats else None, id_base=layer_id * cache.num_experts, ) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index f9e00f9f84..343c4c2c30 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -923,6 +923,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--moe-prefill-demand-tokens", + type=int, + default=ServerArgs.moe_prefill_demand_tokens, + help=( + "A prefill chunk of at most this many tokens loads only the experts it routes to, " + "through the expert cache as decode does, instead of every expert of each layer. " + "0 streams whole layers for every chunk." + ), + ) + parser.add_argument( "--enable-special-token-ckpt", action="store_true", diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 7f94efc250..0fc8113bc8 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -68,9 +68,6 @@ def _thinking_type(req: Any) -> str | None: if value in ("enabled", "disabled"): return value return None - - - def chat_request_to_genspec( req: ChatCompletionRequest, model_sampling: dict[str, Any], @@ -305,8 +302,10 @@ async def handle_chat_completion( if logprobs_conflict is not None: return create_error_response(logprobs_conflict, param="logprobs") if _response_format_unsupported(req.response_format): + rtype = (req.response_format or {}).get("type") return create_error_response( - "response_format json_object/json_schema is not supported (no constrained decoding)", + f"response_format.type {rtype!r} is not supported " + "(supported: text, json_object, json_schema)", param="response_format", ) if not 1 <= req.n <= MAX_N: @@ -1061,8 +1060,12 @@ def _usage( def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool: - # We have no constrained/guided decoding; only plain text ('text' or unset) is honored. - return response_format is not None and response_format.get("type") not in (None, "text") + # Soft-accept OpenAI json_object/json_schema. FreeToken has no constrained + # decoding; clients that already ask for JSON in the prompt still need the + # wire field accepted (HTTP 200) instead of a hard 400. Unknown types 400. + if response_format is None: + return False + return response_format.get("type") not in (None, "text", "json_object", "json_schema") def _completion_unsupported_reason(req: CompletionRequest) -> str | None: @@ -1072,7 +1075,11 @@ def _completion_unsupported_reason(req: CompletionRequest) -> str | None: if logprobs_error is not None: return logprobs_error if _response_format_unsupported(req.response_format): - return "response_format json_object/json_schema is not supported (no constrained decoding)" + rtype = (req.response_format or {}).get("type") + return ( + f"response_format.type {rtype!r} is not supported " + "(supported: text, json_object, json_schema)" + ) return None diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py index abd2e1f939..1829f35f32 100644 --- a/tests/engine/test_logits_processors.py +++ b/tests/engine/test_logits_processors.py @@ -8,7 +8,13 @@ import pytest import torch from freetoken.core import SamplingParams -from freetoken.engine.sample import LogitsPlan, Sampler, apply_logits_processors +from freetoken.engine.sample import ( + LogitsPlan, + RowPenalty, + Sampler, + apply_logits_processors, + apply_penalties, +) V = 8 CPU = torch.device("cpu") @@ -32,33 +38,17 @@ def _ids(rows: list[list[int]]) -> torch.Tensor: return out -def test_presence_and_frequency_penalties_count_generated_tokens(): - logits = torch.zeros((2, V)) - plan = _plan( - [1], - hist_ids=_ids([[3, 3, 5]]), - presence=torch.tensor([[0.5]]), - frequency=torch.tensor([[0.25]]), - ) - out = apply_logits_processors(logits, plan, V) - assert torch.equal(out[0], logits[0]) # row outside the plan untouched - assert out[1, 3].item() == -(0.5 + 2 * 0.25) - assert out[1, 5].item() == -(0.5 + 0.25) - assert out[1, 0].item() == 0.0 - - def test_repetition_penalty_covers_prompt_and_output_with_hf_semantics(): logits = torch.tensor([[2.0, -2.0, 2.0, -2.0, 1.0, 1.0, 1.0, 1.0]]) - plan = _plan( - [0], - hist_ids=_ids([[0, 1]]), # generated: 0 (positive), 1 (negative) - prompt_ids=_ids([[2, 3]]), # prompt: 2 (positive), 3 (negative) - repetition=torch.tensor([[2.0]]), - ) - out = apply_logits_processors(logits, plan, V) + counts = torch.zeros(V, dtype=torch.int32) + counts[[0, 1]] = 1 # generated: 0 (positive), 1 (negative) + prompt = torch.zeros(V, dtype=torch.bool) + prompt[[2, 3]] = True # prompt: 2 (positive), 3 (negative) + out = apply_penalties(logits, [RowPenalty(0, counts, 0.0, 0.0, 2.0, prompt)]) assert out[0, 0].item() == 1.0 and out[0, 2].item() == 1.0 # divided assert out[0, 1].item() == -4.0 and out[0, 3].item() == -4.0 # multiplied assert out[0, 4].item() == 1.0 # unseen + assert logits[0, 0].item() == 2.0 # the caller's logits stay untouched def test_logit_bias_adds_per_row(): @@ -134,6 +124,9 @@ def _req( output_len=max_tokens, max_device_len=len(prompt) + max_tokens, sampling_params=SamplingParams(**sp), + can_decode=True, + output_token_counts=None, + prompt_token_mask=None, ) @@ -164,17 +157,15 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them(): plan = args.plan assert plan is not None assert plan.rows.tolist() == [1, 2, 3] - assert plan.hist_ids[0].tolist()[:2] == [3, 3] # row 1's generated tokens - assert plan.prompt_ids[0].tolist()[:2] == [1, 2] + assert [pen.row for pen in args.penalties] == [1] + assert batch.reqs[1].prompt_token_mask.nonzero().flatten().tolist() == [1, 2] assert plan.bias_rows.tolist() == [0] and plan.bias_ids.tolist() == [5] assert plan.min_rows.tolist() == [1] # local row of the third request only assert plan.min_ids[0].tolist() == [0, 7] logits = torch.zeros((4, V)) out = apply_logits_processors(logits, plan, V) - assert ( - out[1, 3].item() == -1.0 - ) # presence on the repeated 3 (repetition on 0 is a no-op) + assert out[1, 3].item() == 0.0 # the penalties are not plan work assert out[1, 5].item() == 2.0 assert out[2, 0].item() == float("-inf") and out[2, 7].item() == float("-inf") assert torch.isfinite(out[3]).all() @@ -182,11 +173,62 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them(): assert sampler.sample(logits, args)[1].item() == 5 +def test_penalties_ride_the_device_state_not_the_plan(): + # the plan must not apply them a second time on top of the sampler's own counts + sampler = Sampler(CPU, V) + only = SimpleNamespace(reqs=[_req([1], [], 4, presence_penalty=1.0, frequency_penalty=0.5)]) + args = sampler.prepare(only) + assert args.plan is None and [row for row, *_ in args.penalties] == [0] + + req = _req([1], [], 4, presence_penalty=1.0, logit_bias=[[5, 0.25]]) + args = sampler.prepare(SimpleNamespace(reqs=[req])) + assert args.plan is not None and [row for row, *_ in args.penalties] == [0] + req.output_token_counts[3] = 1 + logits = torch.zeros((1, V)) + logits[0, 3] = 1.5 + # once: 3 -> 0.5 beats 5 -> 0.25; twice would drop 3 to -0.5 and pick 5 + assert sampler.sample(logits, args).item() == 3 + assert req.output_token_counts[3].item() == 2 + + def test_repetition_penalty_ignores_multimodal_placeholder_ids_in_the_prompt(): # image tokens carry pseudo-ids >= MM_PAD_SHIFT_VALUE; scattering them into the # [rows, V+1] prompt mask indexed out of bounds (a device-side assert on CUDA) sampler = Sampler(CPU, V) batch = SimpleNamespace(reqs=[_req([1, 1_000_000, 2], [3], 4, repetition_penalty=2.0)]) args = sampler.prepare(batch) - out = apply_logits_processors(torch.ones(1, V), args.plan, V) + out = apply_penalties(torch.ones(1, V), args.penalties) assert out[0, 1] == 0.5 and out[0, 2] == 0.5 and out[0, 4] == 1.0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_min_p_with_temperature_draws_only_the_kept_tokens(): + # min_p leaves -inf over nearly the whole row; the sampling softmax must not turn that + # into NaN (which the Pascal fallback then samples as a uniform row) + V = 32768 + sampler = Sampler(torch.device("cuda"), V) + req = _req([1, 2], [], 4, temperature=0.8, min_p=0.1) + args = sampler.prepare(SimpleNamespace(reqs=[req])) + logits = torch.full((1, V), -10.0, device="cuda") + kept = [100, 20000, 31000] + logits[0, kept] = torch.tensor([5.0, 4.5, 4.0], device="cuda") + for _ in range(32): + assert sampler.sample(logits, args).item() in kept + + +def test_repetition_penalty_reads_the_prompt_once_and_counts_output_on_the_device(): + # the prompt (up to the whole context) is uploaded once per request, not every step, + # and the generated side comes from the counts sample() keeps, not the host history + sampler = Sampler(CPU, V) + req = _req([1, 2], [], 4, repetition_penalty=2.0) + batch = SimpleNamespace(reqs=[req]) + args = sampler.prepare(batch) + mask = req.prompt_token_mask + logits = torch.tensor([[0.0, 1.0, 1.0, 4.0, 0.5, 0.0, 0.0, 0.0]]) + assert sampler.sample(logits, args).item() == 3 + assert req.output_token_counts.tolist() == [0, 0, 0, 1, 0, 0, 0, 0] + req.input_ids = torch.tensor([9, 9, 3], dtype=torch.int32) # a stale host view is not read + args = sampler.prepare(batch) + assert req.prompt_token_mask is mask + out = apply_penalties(logits, args.penalties) + assert out[0].tolist() == [0.0, 0.5, 0.5, 2.0, 0.5, 0.0, 0.0, 0.0] diff --git a/tests/engine/test_sample.py b/tests/engine/test_sample.py new file mode 100644 index 0000000000..28e7195b58 --- /dev/null +++ b/tests/engine/test_sample.py @@ -0,0 +1,69 @@ +"""Presence/frequency penalties in engine.sample.Sampler (greedy and sampled rows).""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.core import SamplingParams +from freetoken.engine import sample as S + + +def _batch(*params: SamplingParams, can_decode: bool = True): + reqs = [ + SimpleNamespace(sampling_params=p, can_decode=can_decode, output_token_counts=None) + for p in params + ] + return SimpleNamespace(reqs=reqs) + + +def test_no_penalty_allocates_nothing(): + sampler = S.Sampler(torch.device("cpu"), 5) + batch = _batch(SamplingParams(temperature=0.0)) + args = sampler.prepare(batch) + assert args.penalties == [] and batch.reqs[0].output_token_counts is None + + +def test_frequency_penalty_counts_generated_tokens_on_the_greedy_path(): + sampler = S.Sampler(torch.device("cpu"), 5) + batch = _batch(SamplingParams(temperature=0.0, frequency_penalty=10.0)) + logits = torch.tensor([[0.0, 5.0, 1.0, 0.0, 0.0]]) + + assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [1] + # Token 1 now carries -10, so the runner-up wins; the caller's logits stay untouched. + assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [2] + assert batch.reqs[0].output_token_counts.tolist() == [0, 1, 1, 0, 0] + assert logits[0, 1] == 5.0 + + +def test_intermediate_prefill_chunks_are_not_counted(): + sampler = S.Sampler(torch.device("cpu"), 5) + batch = _batch(SamplingParams(temperature=0.0, presence_penalty=1.0), can_decode=False) + assert sampler.prepare(batch).penalties == [] + + +def test_penalized_logits_reach_the_probability_sampler(monkeypatch): + """Sampled rows (the torch top-k/top-p fallback on Pascal) see the penalized logits.""" + seen = {} + + def fake_sample_impl(logits, temperatures, top_k, top_p): + seen["logits"] = logits.clone() + return torch.argmax(logits, dim=-1) + + monkeypatch.setattr(S, "sample_impl", fake_sample_impl) + monkeypatch.setattr(S, "make_device_tensor", lambda data, dtype, device: torch.tensor(data, dtype=dtype)) + sampler = S.Sampler(torch.device("cpu"), 4) + batch = _batch( + SamplingParams(temperature=1.0, top_k=2, presence_penalty=3.0), + SamplingParams(temperature=1.0, top_k=2), + ) + batch.reqs[0].output_token_counts = torch.tensor([0, 2, 0, 0], dtype=torch.int32) + logits = torch.tensor([[0.0, 4.0, 2.0, 0.0], [0.0, 4.0, 2.0, 0.0]]) + + tokens = sampler.sample(logits, sampler.prepare(batch)) + + assert seen["logits"][0].tolist() == [0.0, 1.0, 2.0, 0.0] # presence: -3 once, not per count + assert seen["logits"][1].tolist() == [0.0, 4.0, 2.0, 0.0] + assert tokens.tolist() == [2, 1] diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py index 6bd204604d..88fffcc2ed 100644 --- a/tests/engine/test_sample_logprobs.py +++ b/tests/engine/test_sample_logprobs.py @@ -2,6 +2,7 @@ import torch +from freetoken.core import SamplingParams from freetoken.engine.sample import BatchSamplingArgs, Sampler @@ -137,8 +138,8 @@ def test_compute_logprobs_with_prepare_keeps_max_top_logprobs_clamped(monkeypatc batch = SimpleNamespace( reqs=[ - SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=True, top_logprobs=17, is_greedy=True)), - SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=False, top_logprobs=0, is_greedy=True)), + SimpleNamespace(sampling_params=SamplingParams(logprobs=True, top_logprobs=17)), + SimpleNamespace(sampling_params=SamplingParams(logprobs=False, top_logprobs=0)), ] ) args = sampler.prepare(batch) diff --git a/tests/kernels/test_hc_combine_norm.py b/tests/kernels/test_hc_combine_norm.py new file mode 100644 index 0000000000..7f46755479 --- /dev/null +++ b/tests/kernels/test_hc_combine_norm.py @@ -0,0 +1,48 @@ +"""Pins for the fused hc combine + grouped RMSNorm path (hc_combine_norm). + +hc_combine_norm replaces the standalone pair (hc_combine -> grouped_gemma_rmsnorm) +in the Qwen4-exp decoder loop. The swap is only sound if it is BITWISE identical: +the hyper-connection streams feed every later layer, and a bf16 flip here compounds +into different greedy tokens by the end of the stack. We assert torch.equal, not +assert_close: an approx-tolerance test would not have caught the regression this +path shipped with at first attempt -- the fused kernel reduced squares over the +[8, block] tile view (a two-axis tree) while the standalone kernel reduces the +flat row tree; fp32 ulps that differ between trees flip bf16 rounding. Both paths +on the same data must return the same bytes: fails before the reduction-tree fix +in _hc_combine_norm_kernel, passes after. +""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kernel.triton.hc import hc_combine, grouped_gemma_rmsnorm, hc_combine_norm + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +HC, H = 4, 2560 # model geometry: the reduction tree specializes on the block layout +DIM = HC * H +EPS = 1e-6 + + +def _inputs(seed: int, T: int): + g = torch.Generator(device="cuda").manual_seed(seed) + bf = lambda shape: torch.randn(shape, generator=g, device="cuda", dtype=torch.float32).to(torch.bfloat16) + return ( + bf((T, DIM)), # residual streams R + bf((T, H)), # block output y + bf((T, HC)), # inject scores s + (torch.randn(DIM, generator=g, device="cuda", dtype=torch.float32) * 0.1).to(torch.bfloat16), # norm weight + ) + + +@pytest.mark.parametrize("T", [1, 8]) +@pytest.mark.parametrize("seed", range(6)) +def test_combine_norm_matches_split_bitwise(seed: int, T: int) -> None: + R, y, s, w = _inputs(seed, T) + r2_split = hc_combine(R, y, s, HC) + rn_split = grouped_gemma_rmsnorm(r2_split, w, EPS, HC) + r2_fused, rn_fused = hc_combine_norm(R, y, s, w, EPS, HC) + assert torch.equal(r2_fused, r2_split) + assert torch.equal(rn_fused, rn_split) diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py index 7bdbc95681..a31d6ff763 100644 --- a/tests/moe/test_nvfp4_backends.py +++ b/tests/moe/test_nvfp4_backends.py @@ -397,6 +397,41 @@ def test_triton_prefill_large_batch_matches_dequant_reference(hidden, inter, tok _assert_close(out, ref) +@cuda +def test_triton_demand_prefill_matches_the_streamed_layer_bit_for_bit(): + """A short prefill chunk loads only its routed experts, through the slot cache, and runs the + same grouped GEMM over the slots: its output must equal the whole-layer stream's exactly, and + its admissions must not count as decode misses.""" + from freetoken.moe.fused_nvfp4 import fused_experts_nvfp4 + + device = torch.device("cuda") + torch.manual_seed(15) + M = 12 + hidden = torch.randn(M, H, dtype=torch.bfloat16, device=device) / 4 + topk_ids = torch.rand(M, E, device=device).argsort(dim=1)[:, :TOPK].to(torch.int32).contiguous() + topk_weights = torch.rand(M, TOPK, dtype=torch.float32, device=device) + + streamed_cache, _ = _triton_cache(device) + streamed_cache.materialize_layer(1) + streamed_cache.copy_missing() + streamed = fused_experts_nvfp4( + hidden, *streamed_cache.bank_views(E), topk_weights, topk_ids.clone(), E, "silu", False + ) + + demand_cache, _ = _triton_cache(device, cache_size=2 * S) + demand_cache.collect_stats = True + slots = topk_ids.clone() + demand_cache.ensure_experts(1, slots, counts_as_decode=False) + demand_cache.copy_missing() + demanded = fused_experts_nvfp4( + hidden, *demand_cache.bank_views(), topk_weights, slots, demand_cache.cache_size, "silu", False + ) + + assert not torch.equal(slots, topk_ids), "the admission must have rewritten expert ids to slots" + assert torch.equal(demanded, streamed) + assert int(demand_cache.lru_stats.sum()) == 0 + + @cuda def test_triton_swigluoai_matches_dequant_reference(): """MiniMax-M3's swigluoai routed experts through the Triton prefill grouped GEMM @@ -604,3 +639,22 @@ def test_b12x_pack_keeps_per_layer_banks_and_flat_alphas(): assert len(cache.bank_sources["gate_up"]) == L assert sum(t.shape[0] for t in cache.bank_sources["gate_up"]) == total assert cache.gate_up_alpha.shape == (total,) + + +@cuda +@pytest.mark.parametrize("num_experts", [1023, 2100]) +def test_moe_align_handles_a_slot_cache_wider_than_1024(num_experts): + """A prefill over the offload slot cache aligns against its slot count (#676): sgl_kernel's + align reads out of bounds past 1024 counters, so a 2100-slot cache must still align.""" + from freetoken.moe.fused import moe_align_block_size + + torch.manual_seed(0) + ids = torch.randint(0, num_experts, (300, 10), device="cuda", dtype=torch.int32) + sorted_ids, expert_ids, padded = moe_align_block_size(ids, 16, num_experts) + n = int(padded) + flat = ids.flatten() + real = sorted_ids[:n][sorted_ids[:n] < flat.numel()] + assert sorted(real.tolist()) == list(range(flat.numel())) + block_of = torch.arange(n, device="cuda") // 16 + rows = sorted_ids[:n] < flat.numel() + assert torch.equal(flat[sorted_ids[:n][rows].long()], expert_ids[block_of[rows]]) diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 19e25eafe3..086b1f3d83 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -125,6 +125,50 @@ def fake_fused( assert calls["topk_ids"].tolist() == [[2, 1]] +def test_offload_moe_layer_short_prefill_loads_its_experts_through_the_slot_cache(monkeypatch): + layer, cache = _make_layer_and_cache() + cache.prefill_demand_tokens = 1 + topk_weights = torch.tensor([[0.7, 0.3]], dtype=torch.float32) + hidden_states = torch.randn(1, 8) + calls = {} + + monkeypatch.setattr( + "freetoken.layers.moe.fused_topk", + lambda *, hidden_states, gating_output, topk, renormalize: (topk_weights, torch.tensor([[2, 1]], dtype=torch.int32)), + ) + monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.setdefault("materialized", layer_id)) + + def ensure(layer_id, expert_ids, *, counts_as_decode=True, **kwargs): + # the layer hands the unique routed ids (decode-width chunks); rewrite each to a slot + calls["ensured"] = (layer_id, expert_ids.tolist(), counts_as_decode) + expert_ids.add_(3) + + monkeypatch.setattr(cache, "ensure_experts", ensure) + monkeypatch.setattr(cache, "copy_missing", lambda: calls.setdefault("copied", True)) + + def fake_fused(hidden_states, w1, w2, got_topk_weights, got_topk_ids, activation, apply_router_weight_on_input, + act_alpha=1.0, act_limit=float("inf")): + calls["w1"] = w1 + calls["topk_ids"] = got_topk_ids.clone() + return hidden_states + + monkeypatch.setattr("freetoken.moe.fused.fused_experts_impl", fake_fused) + + out = layer.prefill_forward(hidden_states, torch.randn(1, 4)) + + assert out is hidden_states + assert "materialized" not in calls + assert calls["ensured"] == (0, [1, 2], False) + assert calls["copied"] is True + # the grouped GEMM reads the whole slot cache at the slots the admission wrote + assert calls["w1"].shape[0] == cache.cache_size + assert calls["topk_ids"].tolist() == [[5, 4]] + + calls.clear() + layer.prefill_forward(torch.randn(2, 8), torch.randn(2, 4)) + assert calls["materialized"] == 0 and "ensured" not in calls + + def test_offload_moe_layer_prefill_overlap_prefetches_layers_into_two_buffers(monkeypatch): from freetoken.moe.offload_cache import OffloadMoeCache @@ -240,6 +284,22 @@ def test_offload_moe_cache_marlin_rejects_slot_count_beyond_kernel_limit(): ) +@pytest.mark.parametrize("decode_target", ["cpu", "hybrid"]) +def test_offload_moe_cache_rejects_demand_prefill_without_gpu_decode(decode_target): + # a demand-loaded chunk admits its experts into GPU decode's slot cache, which cpu and hybrid decode do not read + from freetoken.moe.offload_cache import OffloadMoeCache + + with pytest.raises(ValueError, match="needs GPU decode"): + OffloadMoeCache( + num_layers=1, + num_experts=4, + cache_size=6, + device=torch.device("cpu"), + decode_target=decode_target, + prefill_demand_tokens=64, + ) + + def test_prefill_overlap_prefetch_invalidates_borrowed_unified_cache_slots(): from freetoken.moe.offload_cache import OffloadMoeCache @@ -614,6 +674,37 @@ def test_adjust_config_converts_moe_cache_rate_to_cache_size(monkeypatch): assert is_offload_moe_strategy(config.moe_strategy) +def test_adjust_config_drops_the_demand_prefill_threshold_for_a_dense_model(): + from types import SimpleNamespace + + from freetoken.distributed import DistributedInfo + from freetoken.engine.config import EngineConfig + from freetoken.engine.engine import _adjust_config + + config = EngineConfig( + model_path="/tmp/freetoken-test-model", + tp_info=DistributedInfo(rank=0, size=1), + dtype=torch.float16, + attention_backend="triton", + moe_prefill_demand_tokens=512, + ) + object.__setattr__( + config, + "model_config", + SimpleNamespace( + has_swa_attention=False, + has_linear_attention=False, + is_moe=False, + num_layers=10, + expert_quant="none", + ), + ) + + _adjust_config(config) + + assert config.moe_prefill_demand_tokens == 0 + + def test_graph_capture_reuses_warm_offload_cache_before_capture(monkeypatch): import freetoken.core as core from freetoken.core import Context, Req, get_global_ctx diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py index 8b7b9d802c..f7971c67d5 100644 --- a/tests/moe/test_small_prefill_on_demand.py +++ b/tests/moe/test_small_prefill_on_demand.py @@ -51,7 +51,7 @@ def fake_gemm(cache_, hs, w, ids, *, views, n, alphas, is_prefill): hs = torch.randn(tokens, 16) layer._prefill_routed(hs, torch.full((tokens, 2), 0.5), torch.zeros(tokens, 2, dtype=torch.int32)) if on_demand: - assert calls == ["ensure", "copy"] and got == {"n": None, "is_prefill": False} + assert calls == ["ensure", "copy"] and got == {"n": cache.cache_size, "is_prefill": True} else: assert "ensure" not in calls and "begin_prefill" not in calls[1:] and got["n"] == layer.num_experts @@ -97,7 +97,7 @@ def test_small_prefill_path_does_not_depend_on_how_many_experts_a_layer_routes_t ids = torch.tensor([[5, 0], [1, 7], [2, 5], [4, 0]], dtype=torch.int32) layer._prefill_routed(torch.randn(4, 16), torch.full((4, 2), 0.5), ids.clone()) assert calls == [("ensure", [0, 1, 2, 4]), ("copy",), ("ensure", [5, 7]), ("copy",)] - assert got == {"ids": (ids + 100).tolist(), "n": None, "is_prefill": False} + assert got == {"ids": (ids + 100).tolist(), "n": cache.cache_size, "is_prefill": True} def test_small_prefill_refuses_more_experts_than_the_cache_holds(monkeypatch): diff --git a/tests/server/test_mm_input.py b/tests/server/test_mm_input.py index a694a50a84..33561acdcd 100644 --- a/tests/server/test_mm_input.py +++ b/tests/server/test_mm_input.py @@ -10,6 +10,7 @@ import pytest +from freetoken.mm.media import _MAX_IMAGE_BYTES as _MAX from freetoken.mm.media import collect_image_refs, fetch_image_bytes, image_reject_reason from freetoken.server.generation import GenerationError, render_messages from freetoken.server.stats import derive_model_card @@ -205,3 +206,123 @@ def test_render_messages_rejects_an_image_in_a_merged_system_message(): ]}, {"role": "user", "content": "hi"}, ]) +# ---- error bodies and redirect policy (no network: httpx.MockTransport stands in for the remote) ---- + + +def _mock_remote(monkeypatch, routes): + """Serve ``routes`` ({url: (status, headers, body) | Exception}) through the client media.py builds; + returns the list of URLs the transport actually saw.""" + import httpx + + seen = [] + + def handler(request): + seen.append(str(request.url)) + route = routes[str(request.url)] + if isinstance(route, Exception): + raise route + status, headers, body = route + return httpx.Response(status, headers=headers, content=body, request=request) + + real = httpx.AsyncClient + monkeypatch.setattr(httpx, "AsyncClient", lambda **kw: real(transport=httpx.MockTransport(handler), **kw)) + return seen + + +def _fetch(url, config): + return asyncio.run(fetch_image_bytes([{"kind": "url", "data": url}], config)) + + +def test_local_media_error_bodies_carry_no_path_root_or_errno(tmp_path): + root = tmp_path / "media" + root.mkdir() + config = _config(allowed_local_media_path=str(root)) + outside = tmp_path / "secret.png" + outside.write_bytes(b"x") + (root / "dir").mkdir() + + cases = { + f"file://{outside}": "subpath", # escapes the root + f"file://{root / 'missing.png'}": "missing or unreadable", # ENOENT + f"file://{root / 'dir'}": "missing or unreadable", # EISDIR + } + for url, expect in cases.items(): + with pytest.raises(ValueError, match=expect) as ei: + _fetch(url, config) + body = str(ei.value) + assert str(tmp_path) not in body and "Errno" not in body, body + assert "secret.png" not in body and "missing.png" not in body, body + + +def test_remote_fetch_failures_are_classified_without_library_text(monkeypatch): + import httpx + + big = _MAX + 1 + seen = _mock_remote(monkeypatch, { + "https://cdn.example.com/404.png": (404, {}, b""), + "https://cdn.example.com/down.png": httpx.ConnectError("[Errno -2] Name or service not known"), + "https://cdn.example.com/slow.png": httpx.ReadTimeout("timed out"), + "https://cdn.example.com/declared.png": (200, {"content-length": str(big)}, b""), + "https://cdn.example.com/streamed.png": (200, {}, b"\0" * big), + }) + config = _config(allowed_media_domains="cdn.example.com") + for name, expect in [("404", "HTTP 404"), ("down", "unreachable"), ("slow", "timed out"), + ("declared", "exceeds"), ("streamed", "exceeds")]: + with pytest.raises(ValueError, match=expect) as ei: + _fetch(f"https://cdn.example.com/{name}.png", config) + assert "Errno" not in str(ei.value) + assert len(seen) == 5 + + +def test_redirects_are_rechecked_against_the_allowlist(monkeypatch): + seen = _mock_remote(monkeypatch, { + "https://cdn.example.com/a.png": (302, {"location": "https://evil.com/a.png"}, b""), + "https://cdn.example.com/b.png": (302, {"location": "/real.png"}, b""), + "https://cdn.example.com/real.png": (200, {}, b"fakepng"), + "https://evil.com/a.png": (200, {}, b"pwned"), + }) + config = _config(allowed_media_domains="cdn.example.com") + with pytest.raises(ValueError, match="redirects outside the allowed domains") as ei: + _fetch("https://cdn.example.com/a.png", config) + assert "evil.com" not in str(ei.value) + # a same-host relative redirect is followed and served + assert _fetch("https://cdn.example.com/b.png", config) == [b"fakepng"] + # the disallowed hop was never requested + assert seen == ["https://cdn.example.com/a.png", "https://cdn.example.com/b.png", "https://cdn.example.com/real.png"] + + +@pytest.mark.parametrize("target", [ + "http://127.0.0.1/x", "http://10.0.0.1/x", "http://169.254.169.254/latest/meta-data", + "http://[::1]/x", "http://[::ffff:10.0.0.1]/x", "http://localhost/x", "http://0.0.0.0/x", + "http://100.64.0.1/x", +]) +def test_redirects_to_private_or_loopback_addresses_are_refused(monkeypatch, target): + seen = _mock_remote(monkeypatch, { + "https://public.example.com/a.png": (302, {"location": target}, b""), + target: (200, {}, b"internal"), + }) + # even with an empty allowlist (any domain admitted) a redirect may not reach inside + with pytest.raises(ValueError, match="private or loopback"): + _fetch("https://public.example.com/a.png", _config()) + assert seen == ["https://public.example.com/a.png"] + + +def test_redirect_hop_count_and_scheme_are_bounded(monkeypatch): + chain = {f"https://h.example.com/{i}.png": (302, {"location": f"https://h.example.com/{i + 1}.png"}, b"") + for i in range(7)} + chain["https://h.example.com/7.png"] = (200, {}, b"fakepng") + chain["https://h.example.com/file.png"] = (302, {"location": "file:///etc/passwd"}, b"") + seen = _mock_remote(monkeypatch, chain) + with pytest.raises(ValueError, match="redirects more than"): + _fetch("https://h.example.com/0.png", _config()) + assert len(seen) == 6 # the original plus five hops; the sixth hop is never requested + with pytest.raises(ValueError, match="non-http"): + _fetch("https://h.example.com/file.png", _config()) + # a short chain that stays public is fine + assert _fetch("https://h.example.com/5.png", _config()) == [b"fakepng"] + + +def test_bad_base64_is_named_without_decoder_text(): + for ref in ({"kind": "b64", "data": "not base64!!"}, {"kind": "url", "data": "data:image/png;base64"}): + with pytest.raises(ValueError, match="not valid base64"): + asyncio.run(fetch_image_bytes([ref], _config())) diff --git a/tests/utils/test_mp_coalesced.py b/tests/utils/test_mp_coalesced.py index e599db7265..c64ee1c5c3 100644 --- a/tests/utils/test_mp_coalesced.py +++ b/tests/utils/test_mp_coalesced.py @@ -92,7 +92,7 @@ def test_sub_queue_coalesced(self): sub.stop() -async def test_async_pull_coalesced(): +def test_async_pull_coalesced(): import sys from freetoken.utils.mp import ZmqAsyncPullQueue as _AsyncPull @@ -101,16 +101,19 @@ async def test_async_pull_coalesced(): # zmq.asyncio needs a Selector loop on Windows; the project runs on Linux CI. if sys.platform == "win32": pytest.skip("zmq.asyncio requires SelectorEventLoop on Windows (prod runs on Linux)") - push = ZmqAsyncPushQueue(_addr(5604), create=True, encoder=lambda o: o) - pull = ZmqAsyncPullQueue(_addr(5604), create=False, decoder=_identity_decoder) - try: - await push.put({"n": 1}) - await push.put({"n": 2}) - await asyncio.sleep(0.1) - - # Even if both messages arrive in a single frame, two gets must succeed. - assert await pull.get() == {"n": 1} - assert await pull.get() == {"n": 2} - finally: - push.stop() - pull.stop() \ No newline at end of file + async def scenario(): + push = ZmqAsyncPushQueue(_addr(5604), create=True, encoder=lambda o: o) + pull = ZmqAsyncPullQueue(_addr(5604), create=False, decoder=_identity_decoder) + try: + await push.put({"n": 1}) + await push.put({"n": 2}) + await asyncio.sleep(0.1) + + # Even if both messages arrive in a single frame, two gets must succeed. + assert await pull.get() == {"n": 1} + assert await pull.get() == {"n": 2} + finally: + push.stop() + pull.stop() + + asyncio.run(scenario())