5080deploy: merge upstream #559, #531, #269, #468, #461, #484, #464, #266, #572

docs/cli.md

@@ -41,6 +41,7 @@ parsers all resolve automatically from the checkpoint and the GPU.
 |---|---|---|
 | `--host` | 127.0.0.1 | Bind address |
 | `--port` | 1919 | Bind port |
+| `--dist-port` | `--port` + 1 | Internal TP rendezvous port (loopback-only regardless of `--host`) |
 | `--gpu` | GPU 0 | GPU to run on: a UUID from `nvidia-smi -L` or an `nvidia-smi` index; see [below](#choosing-a-gpu) |
 | `--max-running-requests` | 4 | Max concurrently running requests |
 | `--max-output-tokens` | 32768 | Default output budget for requests that omit one |

python/freetoken/attention/__init__.py

@@ -26,6 +26,7 @@ class BackendInfo:
     supported_types: frozenset[AttnType]
     requires_flashinfer: bool = False
     requires_sgl_kernel: bool = False
+    requires_sm90: bool = False
     requires_sm100: bool = False
     # Allowed page sizes (None -> any). Config-time resolution coerces to the last
     # entry when the resolved page_size is not in the list.
@@ -77,6 +78,7 @@ def create_fi_backend(config: ModelConfig):
     BackendInfo(
         supported_types=frozenset({AttnType.FULL}),
         requires_sgl_kernel=True,
+        requires_sm90=True,
     ),
 )
 def create_fa_backend(config: ModelConfig):

python/freetoken/core.py

@@ -13,6 +13,7 @@
     from freetoken.kvcache.kv_host_offload import KVHostOffloader
     from freetoken.kvcache.linear_state_pool import LinearStatePool
     from freetoken.moe.offload_cache import OffloadMoeCache
+    from freetoken.tokenizer.detokenize import DecodeStatus
 
 
 @dataclass
@@ -103,6 +104,7 @@ 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
+    stop_decode_status: DecodeStatus | None = None
 
     def __post_init__(self) -> None:
         assert self.input_ids.is_cpu

python/freetoken/engine/config.py

@@ -88,11 +88,17 @@ class EngineConfig:
     # DSV4 window/full ratio directly. Used only when swa_num_pages_override is None (a runtime
     # rebuild can pin an absolute window instead).
     swa_full_tokens_ratio: float = 0.2
+    # Force one decode step after this many consecutive prefill steps, so a long chunked
+    # prefill cannot starve in-flight decodes. None/0 keeps the historical prefill-first
+    # order (the scheduler's own TODO names this: "support other policies: e.g. DECODE
+    # first"). At 8 the cost is ~1% of prefill wall time.
+    decode_interleave_every: int | None = None
     # Absolute window-pool size in the pool's own pages (usable, dummy excluded); None -> use the
     # ratio default above. A runtime cache rebuild sets this (num_swa_pages) to pin the window
     # regardless of the full anchor; the ratio is the startup default and the fallback.
     swa_num_pages_override: int | None = None
     distributed_timeout: float = 1800.0  # ranks reach the first collective minutes apart on a 100+ GiB offload load
+    distributed_port: int = 2333
     use_dummy_weight: bool = False
     use_pynccl: bool = True
     max_seq_len_override: int | None = None
@@ -166,4 +172,4 @@ def max_forward_len(self) -> int:
 
     @property
     def distributed_addr(self) -> str:
-        return "tcp://127.0.0.1:2333"
+        return f"tcp://127.0.0.1:{self.distributed_port}"

python/freetoken/engine/engine.py

@@ -5,6 +5,7 @@
 import gc
 import math
 import os
+import socket
 from datetime import timedelta
 from typing import Any, Dict, Iterable, NamedTuple, Tuple
 
@@ -128,6 +129,9 @@ def _backend_requirements_met(name: str) -> bool:
         return False
     if any(i.requires_sgl_kernel for i in infos) and not _sgl_flash_attn_available():
         return False
+    # requires_sm90 also admits sm_100 (FA4). Hopper-only would reject B200.
+    if any(i.requires_sm90 for i in infos) and not (is_sm90_family() or is_sm100_family()):
+        return False
     if any(i.requires_sm100 for i in infos) and not is_sm100_family():
         return False
     return True
@@ -284,6 +288,13 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att
                 "not installed. Install it with `pip install 'freetoken[sgl]'` (or "
                 "'freetoken[accel]'), or use --attention-backend triton."
             )
+        if info.requires_sm90 and not (is_sm90_family() or is_sm100_family()):
+            raise RuntimeError(
+                f"Attention backend {config.attention_backend!r} requires a compute capability "
+                "9.x GPU (FA3, sm_90 cubins) or 10.x GPU (FA4). Ada (sm_89) and consumer "
+                "Blackwell (sm_12x) are not in the published sgl_kernel flash_ops set. "
+                "Use --attention-backend fi (or triton) instead."
+            )
         if info.requires_sm100 and not is_sm100_family():
             raise RuntimeError(
                 f"Attention backend {config.attention_backend!r} requires a compute capability "
@@ -611,29 +622,47 @@ def __init__(self, config: EngineConfig):
         if config.prefill_warmup:
             self._warmup_prefill()
 
-    def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup:
-        if config.tp_info.size == 1 or config.use_pynccl:
-            torch.distributed.init_process_group(
-                backend="gloo",
-                rank=config.tp_info.rank,
-                world_size=config.tp_info.size,
-                timeout=timedelta(seconds=config.distributed_timeout),
-                init_method=config.distributed_addr,
+    def _make_distributed_store(self, config: EngineConfig) -> torch.distributed.Store:
+        """The rendezvous store's C10d server ignores the host it's given and always listens
+        on every interface, so an unauthenticated TCPStore ends up reachable off-box even
+        when distributed_addr says 127.0.0.1. Rank 0 pre-binds the listening socket to
+        loopback itself and hands the fd to TCPStore (master_listen_fd) to force that."""
+        timeout = timedelta(seconds=config.distributed_timeout)
+        if not config.tp_info.is_primary():
+            return torch.distributed.TCPStore(
+                "127.0.0.1", config.distributed_port, config.tp_info.size,
+                is_master=False, timeout=timeout, multi_tenant=True,
             )
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        sock.bind(("127.0.0.1", config.distributed_port))
+        sock.listen(1024)
+        self._distributed_listen_addr = sock.getsockname()
+        # TCPStore takes ownership of the fd and closes it with the store; detach so the
+        # Python socket never closes that number again after it has been reused.
+        return torch.distributed.TCPStore(
+            "127.0.0.1", config.distributed_port, config.tp_info.size,
+            is_master=True, timeout=timeout, multi_tenant=True, master_listen_fd=sock.detach(),
+        )
+
+    def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup:
+        store = self._make_distributed_store(config)
+        use_gloo = config.tp_info.size == 1 or config.use_pynccl
+        torch.distributed.init_process_group(
+            backend="gloo" if use_gloo else "nccl",
+            rank=config.tp_info.rank,
+            world_size=config.tp_info.size,
+            timeout=timedelta(seconds=config.distributed_timeout),
+            store=store,
+        )
+        if use_gloo:
             tp_cpu_group = torch.distributed.group.WORLD
             assert tp_cpu_group is not None
             max_bytes = (
                 config.max_forward_len * config.model_config.hidden_size * self.dtype.itemsize
             )
             enable_pynccl_distributed(config.tp_info, tp_cpu_group, max_bytes)
         else:
-            torch.distributed.init_process_group(
-                backend="nccl",
-                rank=config.tp_info.rank,
-                world_size=config.tp_info.size,
-                timeout=timedelta(seconds=config.distributed_timeout),
-                init_method=config.distributed_addr,
-            )
             tp_cpu_group = torch.distributed.new_group(backend="gloo")
             assert tp_cpu_group is not None
         return tp_cpu_group

python/freetoken/kvcache/hybrid_radix_cache.py

@@ -116,7 +116,10 @@ def insert(self, input_ids: torch.Tensor, kv_indices: torch.Tensor,
         node, prefix_len = self._walk(input_ids)
         if prefix_len != insert_len:
             new_node = RadixTreeNode(self.key_fn)
-            new_node.set_key_value(input_ids[prefix_len:], kv_indices[prefix_len:].clone())
+            # A view would keep the request's full output-budget token buffer alive.
+            new_node.set_key_value(
+                input_ids[prefix_len:].clone(), kv_indices[prefix_len:].clone()
+            )
             new_node.set_parent(node)
             self.full_evictable += new_node.length
             node = new_node

python/freetoken/kvcache/radix_cache.py

@@ -167,7 +167,10 @@ def insert_prefix(self, input_ids: torch.Tensor, indices: torch.Tensor) -> Inser
         node, prefix_len = self._tree_walk(input_ids)
         if prefix_len != insert_len:  # NOTE: prefix_len < insert_len
             new_node = RadixTreeNode(self.key_fn)
-            new_node.set_key_value(input_ids[prefix_len:], indices[prefix_len:].clone())
+            # A view would keep the request's full output-budget token buffer alive.
+            new_node.set_key_value(
+                input_ids[prefix_len:].clone(), indices[prefix_len:].clone()
+            )
             new_node.set_parent(node)
             self.evictable_size += new_node.length
             node = new_node

python/freetoken/kvcache/swa_radix_cache.py

@@ -215,7 +215,8 @@ def insert(self, input_ids: torch.Tensor, kv_indices: torch.Tensor,
         # then a live (non-tombstone) leaf for the in-window remainder (never a tombstone leaf --
         # the clamp + the free_swa -page_size margin guarantee a live tail).
         if total < insert_len:
-            suffix_ids = input_ids[total:]
+            # Both suffix nodes can share this copy without retaining the request buffer.
+            suffix_ids = input_ids[total:].clone()
             suffix_kv = kv_indices[total:].clone()
             boundary = max(0, min(swa_evicted_seqlen, insert_len) - total)
             boundary = min(boundary, max(0, len(suffix_ids) - self.page_size))

python/freetoken/scheduler/interleave.py

@@ -0,0 +1,73 @@
+"""Decode-interleave scheduling policy.
+
+Background (measured on production, 2026-09-15)
+-----------------------------------------------
+``Scheduler._schedule_next_batch`` picks::
+
+    batch = (prefill_manager.schedule_next_batch(budget)
+             or decode_manager.schedule_next_batch())
+
+Prefill wins unconditionally.  A long prompt is chunked at the window pool's
+budget (384 tokens in production), so a 300k-token context is 797 consecutive
+prefill steps -- and while they run, *decoding requests already in flight are
+never scheduled*.  Measured on the live service:
+
+    four long cold prefills: 680/797/506/250 chunks, 267/317/190/93 s
+    prefill-burst share of a 15-minute window: 70%
+
+so an in-flight request waits up to ~5 minutes for its next token.  The
+scheduler already carries the marker for this::
+
+    # TODO: support other policies: e.g. DECODE first
+
+Reversing the order outright is wrong -- it starves prefill, and a request that
+never prefills never starts.  What is needed is *interleaving*: let prefill keep
+priority, but force a decode step every N prefill steps so in-flight requests
+make progress.  The cost is small by construction: measured chunk cost is
+0.388 s, and one decode step for a handful of requests is a fraction of that,
+so every 8th step costs ~1% of prefill wall time while bounding the stall at
+8 chunks (~3 s) instead of the whole burst (~300 s).
+
+This module holds the *policy* only -- a counter and a decision -- so it can be
+unit-tested without a Scheduler, an engine, or a GPU.
+"""
+
+from __future__ import annotations
+
+
+class DecodeInterleavePolicy:
+    """Decide whether to force a decode step between prefill steps.
+
+    ``every`` is the number of consecutive prefill steps allowed before a decode
+    step is forced.  ``None`` or <= 0 disables the policy, reproducing the
+    historical prefill-first behaviour bit-for-bit.
+
+    The counter only advances on prefill steps that were actually scheduled, so
+    a scheduler with nothing to prefill never accumulates credit, and the first
+    decode of a burst is not delayed.
+    """
+
+    def __init__(self, every: int | None = None) -> None:
+        self.every = int(every) if every else None
+        self._prefill_streak = 0
+
+    @property
+    def enabled(self) -> bool:
+        return self.every is not None and self.every > 0
+
+    def note_prefill(self) -> None:
+        """Record that a prefill step was scheduled in this slot."""
+        self._prefill_streak += 1
+
+    def note_decode(self) -> None:
+        """Record that a decode step was scheduled in this slot (streak resets)."""
+        self._prefill_streak = 0
+
+    def wants_decode(self) -> bool:
+        """True when the streak has reached the threshold and a decode is due."""
+        if not self.enabled:
+            return False
+        return self._prefill_streak >= self.every
+
+    def reset(self) -> None:
+        self._prefill_streak = 0

python/freetoken/scheduler/scheduler.py

@@ -32,6 +32,7 @@
 from .cache import CacheManager
 from .config import SchedulerConfig
 from .decode import DecodeManager
+from .interleave import DecodeInterleavePolicy
 from .io import SchedulerIOMixin
 from .mm import cut_image_spans, plan_mm_batch
 from .prefill import ChunkedReq, PrefillManager
@@ -92,6 +93,14 @@ def __init__(self, config: SchedulerConfig):
             ) or getattr(self.engine.kv_cache, "sliding_window_size", None),
         )
         self.decode_manager = DecodeManager(config.page_size)
+        # Bound how long a run of prefill steps may starve in-flight decodes. A long
+        # prompt chunked at the window budget is hundreds of consecutive prefill steps;
+        # measured on an 8xRTX4090 DSV4 deployment, that left an already-decoding
+        # request unscheduled for 267-317 s. Unset reproduces the historical
+        # prefill-first order exactly.
+        self._interleave = DecodeInterleavePolicy(
+            getattr(config, "decode_interleave_every", None)
+        )
         self._bidirectional_mm = any(getattr(g, "bidirectional_mm_blocks", False) for g in config.model_config.attention_groups)
         self.prefill_manager = PrefillManager(
             self.cache_manager,
@@ -462,16 +471,38 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
         self.send_result(reply)
 
     def _match_stop_str(self, req: Req) -> str | None:
-        """First stop string present in this request's generated tail, else None. Decodes
-        only a short suffix (bounded by the longest stop string's char length, so a stop of
-        N chars spans at most N tokens) to keep the per-step cost small."""
+        """Match stops against the same incrementally decoded text as the frontend."""
+        from freetoken.tokenizer.detokenize import DecodeStatus
+
         stop_strs = req.sampling_params.stop_strs
         prompt_len = req.max_device_len - req.output_len
-        if len(req.input_ids) <= prompt_len:
+        end = len(req.input_ids)
+        # The frontend omits a terminal EOS, including at the output limit with ignore_eos.
+        # The limit is the delivered length, as in hit_length: under overlap device_len (and
+        # so can_decode) runs a step ahead of the host.
+        if (
+            end > prompt_len
+            and end >= req.max_device_len
+            and int(req.input_ids[-1]) in self.eos_token_ids
+        ):
+            end -= 1
+        if end <= prompt_len:
             return None
+        if req.stop_decode_status is None:
+            req.stop_decode_status = DecodeStatus(
+                decoded_ids=[],
+                decoded_str="",
+                read_offset=0,
+                surr_offset=0,
+                sent_offset=0,
+                skip_special=req.sampling_params.skip_special_tokens,
+            )
+        state = req.stop_decode_status
+        state.decoded_ids.extend(req.input_ids[prompt_len + len(state.decoded_ids) : end].tolist())
+        tail = state.decode(self.tokenizer)
         max_chars = max(len(s) for s in stop_strs)
-        tail_start = max(prompt_len, len(req.input_ids) - (max_chars + 1))
-        tail = self.tokenizer.decode(req.input_ids[tail_start:].tolist())
+        # Bound the text used for matching, retaining token context for incomplete characters.
+        state.decoded_str = state.decoded_str[-max_chars:]
         for s in stop_strs:
             if s in tail:
                 return s
@@ -916,11 +947,28 @@ def _gather_multimodal(self, batch: Batch) -> None:
             )
 
     def _schedule_next_batch(self) -> ForwardInput | None:
-        # TODO: support other policies: e.g. DECODE first
-        batch = (
-            self.prefill_manager.schedule_next_batch(self.prefill_budget)
-            or self.decode_manager.schedule_next_batch()
-        )
+        # Prefill keeps priority -- a request that never prefills never starts -- but a
+        # long run of prefill chunks must not starve in-flight decodes. When the policy
+        # says a decode is due, take one if a decode is actually runnable; otherwise stay
+        # with prefill (never spend a slot idle just to keep a promise).
+        # getattr, not attribute access: the accounting tests drive this method on a
+        # partially built Scheduler, and interleaving must be an opt-in that a stripped
+        # object simply does not have rather than something that breaks it. A missing
+        # policy is the historical prefill-first order.
+        policy = getattr(self, "_interleave", None)
+        batch = None
+        if policy is not None and policy.wants_decode():
+            batch = self.decode_manager.schedule_next_batch()
+            if batch is not None:
+                policy.note_decode()
+        if batch is None:
+            batch = self.prefill_manager.schedule_next_batch(self.prefill_budget)
+            if batch is not None and policy is not None:
+                policy.note_prefill()
+        if batch is None:
+            batch = self.decode_manager.schedule_next_batch()
+            if batch is not None and policy is not None:
+                policy.note_decode()
         if batch is None:
             return None
         forward_input = self._prepare_batch(batch)