diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py index eda3b61f55..8e52845003 100644 --- a/python/freetoken/attention/__init__.py +++ b/python/freetoken/attention/__init__.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Protocol +from typing import TYPE_CHECKING, Literal, Protocol from freetoken.utils import Registry, init_logger @@ -33,6 +33,8 @@ class BackendInfo: # Whether forward() honors a per-call AttentionSpec (window/sm_scale/sinks). # Non-consumers raise on a non-None spec instead of silently dropping it. consumes_attn_spec: bool = False + # How the backend captures a graph of `rows_per_req` query rows per request: "device" keeps all metadata on device, "host_plan" plans on the host before each replay, "none" cannot. + uniform_extend_graph: Literal["none", "host_plan", "device"] = "none" SUPPORTED_ATTENTION_BACKENDS = Registry[BackendCreator]("Attention Backend") @@ -58,6 +60,7 @@ def create_trtllm_backend(config: ModelConfig): BackendInfo( supported_types=frozenset({AttnType.FULL}), requires_flashinfer=True, + uniform_extend_graph="host_plan", ), ) def create_fi_backend(config: ModelConfig): @@ -71,6 +74,7 @@ def create_fi_backend(config: ModelConfig): BackendInfo( supported_types=frozenset({AttnType.FULL}), requires_sgl_kernel=True, + uniform_extend_graph="device", ), ) def create_fa_backend(config: ModelConfig): @@ -84,6 +88,7 @@ def create_fa_backend(config: ModelConfig): BackendInfo( supported_types=frozenset({AttnType.FULL, AttnType.SWA}), consumes_attn_spec=True, + uniform_extend_graph="device", ), ) def create_triton_backend(config: ModelConfig): diff --git a/python/freetoken/attention/base.py b/python/freetoken/attention/base.py index 846e99447a..0217299292 100644 --- a/python/freetoken/attention/base.py +++ b/python/freetoken/attention/base.py @@ -70,7 +70,8 @@ def forward( def prepare_metadata(self, batch: Batch) -> None: ... @abstractmethod - def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: ... + def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None: + """Static buffers for graphs of ``rows_per_req`` query rows per request: 1 for decode, more for a uniform extend.""" @abstractmethod def prepare_for_capture(self, batch: Batch) -> None: ... @@ -116,8 +117,8 @@ def prepare_metadata(self, batch: Batch) -> None: backend = self.prefill_backend if batch.is_prefill else self.decode_backend return backend.prepare_metadata(batch) - def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: - self.decode_backend.init_capture_graph(max_seq_len, bs_list) + def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None: + self.decode_backend.init_capture_graph(max_seq_len, bs_list, rows_per_req) def prepare_for_capture(self, batch: Batch) -> None: self.decode_backend.prepare_for_capture(batch) diff --git a/python/freetoken/attention/dsa.py b/python/freetoken/attention/dsa.py index e8c7b3c2a5..b565b01f51 100644 --- a/python/freetoken/attention/dsa.py +++ b/python/freetoken/attention/dsa.py @@ -341,7 +341,9 @@ def _prefill(self, md, layer_id, q_nope, q_pe, batch, inputs) -> torch.Tensor: return o # ----- CUDA graph (decode) ---------------------------------------------------------- - def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: + def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None: + if rows_per_req != 1: + raise NotImplementedError("dsa captures decode graphs only") self.max_seq_len = max_seq_len self.capture_bs = sorted(bs_list) max_bs = max(bs_list) diff --git a/python/freetoken/attention/dsa_indexer_kpool.py b/python/freetoken/attention/dsa_indexer_kpool.py index 2b2492aa46..2db80ea886 100644 --- a/python/freetoken/attention/dsa_indexer_kpool.py +++ b/python/freetoken/attention/dsa_indexer_kpool.py @@ -56,8 +56,8 @@ def __init__(self, config: "ModelConfig") -> None: # ----- CUDA-graph decode staging ------------------------------------------------------ # The tail rings are keyed by Req.table_idx; a captured decode step must read it # from a static buffer restaged per replay (rows/kvlen precedent in the parent). - def init_capture_graph(self, max_seq_len: int, bs_list) -> None: - super().init_capture_graph(max_seq_len, bs_list) + def init_capture_graph(self, max_seq_len: int, bs_list, rows_per_req: int = 1) -> None: + super().init_capture_graph(max_seq_len, bs_list, rows_per_req) self._slots_buf = torch.zeros(max(bs_list), dtype=torch.int64, device=self.device) def _stage_decode(self, batch: "Batch", bs: int, table_idx: torch.Tensor) -> None: diff --git a/python/freetoken/attention/dsv41_sparse.py b/python/freetoken/attention/dsv41_sparse.py index baaa2cf019..0d827ffb3e 100644 --- a/python/freetoken/attention/dsv41_sparse.py +++ b/python/freetoken/attention/dsv41_sparse.py @@ -212,7 +212,9 @@ def _prefill_metadata(self, batch: Batch) -> DSV41AttnMetadata: decoder_rows = torch.cat(rows).to(self.device, non_blocking=True) if rows else torch.empty(0, dtype=torch.int64, device=self.device) return DSV41AttnMetadata(last_indices=last_indices, segments=segments, decoder_segments=dec, decoder_rows=decoder_rows, decoder_pad=pad) - def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: + def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None: + if rows_per_req != 1: + raise NotImplementedError("dsv41_sparse captures decode graphs only") assert self.capture is None, "Capture already initialized." self.max_graph_bs = max(bs_list) self.capture = DSV41CaptureData.create(self.max_graph_bs, max_seq_len, self.device) diff --git a/python/freetoken/attention/dsv4_sparse.py b/python/freetoken/attention/dsv4_sparse.py index 2c740e2a48..bb68411af0 100644 --- a/python/freetoken/attention/dsv4_sparse.py +++ b/python/freetoken/attention/dsv4_sparse.py @@ -190,7 +190,9 @@ def prepare_metadata(self, batch: Batch) -> None: last_indices=last, table_rows=self._table_rows(batch), window_ar=self._window_ar ) - def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: + def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None: + if rows_per_req != 1: + raise NotImplementedError("dsv4_sparse captures decode graphs only") assert self.capture is None, "Capture already initialized." self.max_graph_bs = max(bs_list) self.capture = DSV4CaptureData.create(self.max_graph_bs, max_seq_len, self.device) diff --git a/python/freetoken/attention/fa.py b/python/freetoken/attention/fa.py index 782ef040e2..07212e1373 100644 --- a/python/freetoken/attention/fa.py +++ b/python/freetoken/attention/fa.py @@ -46,6 +46,7 @@ def __init__(self, config: ModelConfig): self.kvcache = ctx.kv_cache self.page_size = ctx.page_size self.capture: FACaptureData | None = None + self.capture_rows_per_req = 1 self.max_graph_bs = 0 self.capture_bs: List[int] = [] self.scale = config.head_dim**-0.5 @@ -120,23 +121,26 @@ def prepare_metadata(self, batch: Batch) -> None: page_table=new_page_table, ) - def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: + def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None: assert self.capture is None, "Capture already initialized." max_bs = max(bs_list) - capture = FACaptureData.create(max_bs, max_seq_len // self.page_size, self.kvcache.device) + capture = FACaptureData.create(max_bs, max_seq_len // self.page_size, self.kvcache.device, rows_per_req=rows_per_req) self.max_graph_bs = max_bs self.capture = capture self.capture_bs = sorted(bs_list) + self.capture_rows_per_req = rows_per_req def prepare_for_capture(self, batch: Batch) -> None: assert (bs := batch.size) in self.capture_bs and self.capture capture = self.capture + rows_per_req = self.capture_rows_per_req + capture.seq_lens[:bs].fill_(rows_per_req) metadata = FAMetadata( cu_seqlens_k=capture.cu_seqlens_k[: bs + 1], cu_seqlens_q=capture.cu_seqlens_q[: bs + 1], cache_seqlens=capture.seq_lens[:bs], max_seqlen_k=capture.page_table.size(1) * self.page_size, - max_seqlen_q=1, # decode only + max_seqlen_q=rows_per_req, page_table=capture.page_table[:bs, :], ) batch.attn_metadata = metadata @@ -145,7 +149,7 @@ def prepare_for_replay(self, batch: Batch) -> None: metadata, bs = batch.attn_metadata, batch.padded_size assert isinstance(metadata, FAMetadata) assert self.capture is not None and bs in self.capture_bs - # cu_seqlens_q is always [0, 1, 2, ..., bs] for decode (i.e. no-op) + # cu_seqlens_q is the constant arange * rows_per_req under a graph (i.e. no-op) table_len = metadata.page_table.size(1) self.capture.cu_seqlens_k[: bs + 1].copy_(metadata.cu_seqlens_k) self.capture.seq_lens[:bs].copy_(metadata.cache_seqlens) diff --git a/python/freetoken/attention/fi.py b/python/freetoken/attention/fi.py index c9e58538f7..175c21ffaf 100644 --- a/python/freetoken/attention/fi.py +++ b/python/freetoken/attention/fi.py @@ -139,8 +139,11 @@ def __init__(self, config: ModelConfig) -> None: # for cuda graph self.capture_bs: List[int] = [] self.max_graph_bs = 0 - self.graph_wrappers: Dict[int, CUDAGraphBatchDecodeWithPagedKVCacheWrapper] = {} + self.graph_wrappers: Dict[ + int, CUDAGraphBatchDecodeWithPagedKVCacheWrapper | BatchPrefillWithPagedKVCacheWrapper + ] = {} self.capture: FICaptureData | None = None + self.capture_rows_per_req = 1 self.last_event = torch.cuda.Event() self.last_event.record() @@ -266,14 +269,15 @@ def reset_capture(self) -> None: super().reset_capture() self.graph_wrappers = {} - def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: + def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None: assert self.capture is None, "Capture already initialized." max_bs = max(bs_list) - capture = FICaptureData.create(max_bs, max_seq_len, self.kvcache.device) + capture = FICaptureData.create(max_bs, max_seq_len, self.kvcache.device, rows_per_req=rows_per_req) capture.page_table = capture.page_table.view(-1) # use 1D as ragged indices self.max_graph_bs = max_bs self.capture = capture self.capture_bs = sorted(bs_list) + self.capture_rows_per_req = rows_per_req @cached_property def use_tensor_cores(self) -> bool: @@ -283,22 +287,41 @@ def use_tensor_cores(self) -> bool: GQA = self.config.num_qo_heads // self.config.num_kv_heads return GQA >= 4 - def prepare_for_capture(self, batch: Batch) -> None: - from flashinfer import CUDAGraphBatchDecodeWithPagedKVCacheWrapper + def _make_graph_wrapper( + self, bs: int + ) -> CUDAGraphBatchDecodeWithPagedKVCacheWrapper | BatchPrefillWithPagedKVCacheWrapper: + from flashinfer import BatchPrefillWithPagedKVCacheWrapper, CUDAGraphBatchDecodeWithPagedKVCacheWrapper + + capture = self.capture + assert capture is not None + if self.capture_rows_per_req > 1: + wrapper = BatchPrefillWithPagedKVCacheWrapper( + self.float_workspace_buffer, + kv_layout="NHD", + backend="fa2", + use_cuda_graph=True, + qo_indptr_buf=capture.cu_seqlens_q[: bs + 1], + paged_kv_indptr_buf=capture.cu_seqlens_k[: bs + 1], + paged_kv_indices_buf=capture.indices, + paged_kv_last_page_len_buf=capture.one_tensor[:bs], + ) + else: + wrapper = CUDAGraphBatchDecodeWithPagedKVCacheWrapper( + self.float_workspace_buffer, + kv_layout="NHD", + use_tensor_cores=self.use_tensor_cores, + indptr_buffer=capture.cu_seqlens_k[: bs + 1], + indices_buffer=capture.indices, + last_page_len_buffer=capture.one_tensor[:bs], + ) + wrapper._backend = "fa2" + wrapper._int_workspace_buffer = self.int_workspace_buffer + return wrapper + def prepare_for_capture(self, batch: Batch) -> None: bs = batch.size assert bs in self.capture_bs and bs not in self.graph_wrappers and self.capture - capture = self.capture - self.graph_wrappers[bs] = CUDAGraphBatchDecodeWithPagedKVCacheWrapper( - self.float_workspace_buffer, - kv_layout="NHD", - use_tensor_cores=self.use_tensor_cores, - indptr_buffer=capture.cu_seqlens_k[: bs + 1], - indices_buffer=capture.indices, - last_page_len_buffer=capture.one_tensor[:bs], - ) - self.graph_wrappers[bs]._backend = "fa2" - self.graph_wrappers[bs]._int_workspace_buffer = self.int_workspace_buffer + self.graph_wrappers[bs] = self._make_graph_wrapper(bs) self.prepare_metadata(batch) metadata = batch.attn_metadata assert isinstance(metadata, FIMetadata) diff --git a/python/freetoken/attention/m3_sparse.py b/python/freetoken/attention/m3_sparse.py index 82a98a5a78..dc1869cb5c 100644 --- a/python/freetoken/attention/m3_sparse.py +++ b/python/freetoken/attention/m3_sparse.py @@ -368,7 +368,9 @@ def _prefill(self, md, layer_id: int, slot: int, q, index_q, batch: Batch) -> to return out # ----- CUDA graph (decode) ------------------------------------------------------------ - def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: + def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None: + if rows_per_req != 1: + raise NotImplementedError("m3_sparse captures decode graphs only") self.inner.init_capture_graph(max_seq_len, bs_list) self.capture_bs = sorted(bs_list) max_bs = max(bs_list) diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py index 8ce63c4242..3e37d5f888 100644 --- a/python/freetoken/attention/qsa_sparse.py +++ b/python/freetoken/attention/qsa_sparse.py @@ -489,7 +489,9 @@ def _scratch(self, name: str, rows: int, *shape: int, dtype: torch.dtype) -> tor return torch.empty((rows, *shape), dtype=dtype, device=self.device) # ----- CUDA graph (decode) -------------------------------------------------------------- - def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: + def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None: + if rows_per_req != 1: + raise NotImplementedError("qsa_sparse captures decode graphs only") self.capture_bs = sorted(bs_list) max_bs = max(bs_list) width = get_global_ctx().page_table.shape[1] diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py index 7f29c481bb..071ee0d454 100644 --- a/python/freetoken/attention/triton.py +++ b/python/freetoken/attention/triton.py @@ -32,15 +32,16 @@ def create( num_q_heads: int, max_head_dim: int, scratch_rows: int, + rows_per_req: int = 1, **kwargs, ): return cls( seq_lens=torch.ones((max_bs,), dtype=torch.int32, device=device), - positions=torch.zeros((max_bs,), dtype=torch.int32, device=device), + positions=torch.zeros((max_bs * rows_per_req,), dtype=torch.int32, device=device), cu_seqlens_k=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device), - cu_seqlens_q=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device), + cu_seqlens_q=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device) * rows_per_req, page_table=torch.zeros((max_bs, max_seq_len), dtype=torch.int32, device=device), - q_to_req=torch.arange(max_bs, dtype=torch.int32, device=device), + q_to_req=torch.arange(max_bs, dtype=torch.int32, device=device).repeat_interleave(rows_per_req), attn_logits=torch.empty((scratch_rows * num_q_heads * max_head_dim,), dtype=torch.float32, device=device), attn_lse=torch.empty((scratch_rows * num_q_heads,), dtype=torch.float32, device=device), **kwargs, @@ -71,6 +72,7 @@ def __init__(self, config: ModelConfig): self.kvcache = get_global_ctx().kv_cache self.device = self.kvcache.device self.capture: TritonCaptureData | None = None + self.capture_rows_per_req = 1 self.capture_bs: List[int] = [] self.max_graph_bs = 0 self.prefill_tile_min_q = 128 @@ -253,7 +255,7 @@ def prepare_metadata(self, batch: Batch) -> None: swa_indices=swa_indices, ) - def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: + def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None: assert self.capture is None, "Capture already initialized." max_bs = max(bs_list) self.capture = TritonCaptureData.create( @@ -263,7 +265,9 @@ def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: num_q_heads=max(1, self.num_q_heads), max_head_dim=max(1, self.max_head_dim), scratch_rows=max(bs * self._split_cap(bs) for bs in bs_list), + rows_per_req=rows_per_req, ) + self.capture_rows_per_req = rows_per_req if self._swa_capture_enabled(): self.capture.swa_page_table = torch.zeros( (max_bs, max_seq_len), @@ -290,6 +294,10 @@ def _point_to_capture(self, metadata: TritonMetadata, bs: int) -> None: indices = self.capture.page_table.view(-1) self.capture.cu_seqlens_q[: bs + 1].copy_(metadata.cu_seqlens_q_gpu) self.capture.cu_seqlens_k[: bs + 1].copy_(metadata.indptr) + if not metadata.is_decode: + # the extend kernel reads prefix_lens; seq_lens is otherwise unused by this backend, so it holds them + self.capture.seq_lens[:bs].copy_(metadata.prefix_lens) + metadata.prefix_lens = self.capture.seq_lens[:bs] total = metadata.indices.numel() indices[:total].copy_(metadata.indices) if metadata.swa_indices is not None: @@ -319,6 +327,24 @@ def prepare_for_capture(self, batch: Batch) -> None: bs = batch.size assert bs in self.capture_bs and self.capture is not None capture = self.capture + rows_per_req = self.capture_rows_per_req + if rows_per_req > 1: + n = bs * rows_per_req + # placeholders for a fresh request (kv rows == query rows, no prefix); replay overwrites both + capture.cu_seqlens_k[: bs + 1].copy_(capture.cu_seqlens_q[: bs + 1]) + capture.seq_lens[:bs].fill_(0) + batch.attn_metadata = TritonMetadata( + cu_seqlens_q_gpu=capture.cu_seqlens_q[: bs + 1], + indptr=capture.cu_seqlens_k[: bs + 1], + indices=capture.page_table.view(-1), + q_to_req=capture.q_to_req[:n], + q_positions=capture.positions[:n], + is_decode=False, + prefix_lens=capture.seq_lens[:bs], + max_q_len=rows_per_req, + swa_indices=self._capture_swa_indices(), + ) + return attn_logits, attn_lse = self._capture_scratch(bs) batch.attn_metadata = TritonMetadata( cu_seqlens_q_gpu=capture.cu_seqlens_q[: bs + 1], diff --git a/python/freetoken/attention/trtllm.py b/python/freetoken/attention/trtllm.py index 2984415628..20eb2fb538 100644 --- a/python/freetoken/attention/trtllm.py +++ b/python/freetoken/attention/trtllm.py @@ -138,7 +138,9 @@ def prepare_metadata(self, batch: Batch) -> None: page_table=new_page_table, ) - def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: + def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None: + if rows_per_req != 1: + raise NotImplementedError("trtllm captures decode graphs only") assert self.capture is None, "Capture already initialized." max_bs = max(bs_list) capture = TRTLLMCaptureData.create( diff --git a/python/freetoken/attention/utils.py b/python/freetoken/attention/utils.py index c45c394148..1de53e70d3 100644 --- a/python/freetoken/attention/utils.py +++ b/python/freetoken/attention/utils.py @@ -12,12 +12,13 @@ class BaseCaptureData: page_table: torch.Tensor @classmethod - def create(cls, max_bs: int, max_seq_len: int, device: torch.device, **kwargs): + def create(cls, max_bs: int, max_seq_len: int, device: torch.device, rows_per_req: int = 1, **kwargs): return cls( seq_lens=torch.ones((max_bs,), dtype=torch.int32, device=device), - positions=torch.zeros((max_bs,), dtype=torch.int32, device=device), + positions=torch.zeros((max_bs * rows_per_req,), dtype=torch.int32, device=device), cu_seqlens_k=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device), - cu_seqlens_q=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device), + # constant under a graph: every request has exactly `rows_per_req` query rows + cu_seqlens_q=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device) * rows_per_req, page_table=torch.zeros((max_bs, max_seq_len), dtype=torch.int32, device=device), **kwargs, ) diff --git a/python/freetoken/core.py b/python/freetoken/core.py index 812f9a4349..f37ef84f89 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -11,7 +11,7 @@ from freetoken.attention.linear import FLAMetadata from freetoken.kvcache import BaseCacheHandle, BaseKVCachePool from freetoken.kvcache.linear_state_pool import LinearStatePool - from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.models.blocks import FeatureSink @dataclass @@ -95,6 +95,15 @@ def complete_one(self) -> None: self.cached_len = self.device_len self.device_len += 1 + def widen_pending(self, k: int) -> None: + """Extend the pending token at cached_len by k provisional draft tokens.""" + self.device_len += k + + def commit_pending(self, n: int) -> None: + """Make n tokens final and leave the last committed one pending at cached_len.""" + self.cached_len += n + self.device_len = self.cached_len + 1 + def append_host(self, next_token: torch.Tensor) -> None: n = self.input_ids.numel() m = n + next_token.numel() @@ -118,7 +127,9 @@ def __repr__(self) -> str: @dataclass class Batch: reqs: List[Req] - phase: Literal["prefill", "decode"] + phase: Literal["prefill", "decode", "verify"] + # query rows per request: 1 on decode, more on verify, None on prefill (ragged) + rows_per_req: int | None = None # these fields should be set by scheduler input_ids: torch.Tensor = field(init=False) positions: torch.Tensor = field(init=False) @@ -160,6 +171,14 @@ class Batch: # _prepare_batch succeeds. Continuation chunks leave this empty, so accounting is # exactly-once. prompt_admissions: List[Tuple[int, int, int]] = field(default_factory=list, init=False) + # per-row hidden states a draft model reads from this forward; None when nothing reads them + features: "FeatureSink | None" = field(default=None, init=False) + + def __post_init__(self) -> None: + if self.rows_per_req is None and self.phase == "decode": + self.rows_per_req = 1 + # without it, input_mapping[0][::None] would read every row as its own request + assert self.phase != "verify" or self.rows_per_req is not None, "a verify batch needs rows_per_req" @property def is_prefill(self) -> bool: @@ -187,7 +206,6 @@ class Context: # NOTE: this table always treat page_size = 1 page_table: torch.Tensor = field(init=False) attn_backend: BaseAttnBackend = field(init=False) - moe_offload_cache: OffloadMoeCache | None = None kv_cache: BaseKVCachePool = field(init=False) # Per-request recurrent state for GatedDeltaNet layers; set by the engine for # hybrid linear-attention models, otherwise None. @@ -200,11 +218,37 @@ def batch(self) -> Batch: return self._batch @contextmanager - def forward_batch(self, batch: Batch): + def override( + self, + *, + kv_cache: BaseKVCachePool | None = None, + attn_backend: BaseAttnBackend | None = None, + hide_target_state: bool = False, + ): + """Swap the KV pool and attention backend for a draft forward; hide_target_state also hides the target's GDN pool.""" + swaps = {} + if kv_cache is not None: + swaps["kv_cache"] = kv_cache + if attn_backend is not None: + swaps["attn_backend"] = attn_backend + if hide_target_state: + swaps["linear_state_pool"] = None + saved = {name: getattr(self, name) for name in swaps} + try: + for name, value in swaps.items(): + setattr(self, name, value) + yield + finally: + for name, value in saved.items(): + setattr(self, name, value) + + @contextmanager + def forward_batch(self, batch: Batch, **override): assert self._batch is None, "Nested forward_batch is not allowed" try: self._batch = batch - yield + with self.override(**override): + yield finally: self._batch = None diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 77aaa3eefe..39379c4278 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -176,6 +176,11 @@ def max_seq_len(self) -> int: def max_forward_len(self) -> int: return self.max_seq_len + @property + def decode_rows_per_req(self) -> int: + """Query rows each running request takes per decode step; every per-request decode allowance derives from it.""" + return 1 + @property def distributed_addr(self) -> str: return "tcp://127.0.0.1:2333" diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 2de7f47a3e..eca7bddd69 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -47,7 +47,8 @@ def _clamp_decode_batch_to_pools(config: EngineConfig, pool_sizes: list[int]) -> small = [s for s in pool_sizes if s < num_experts] if not small: return - limit = min(small) // config.model_config.num_experts_per_tok + # a decode step routes every row of every request, decode_rows_per_req rows each + limit = min(small) // (config.model_config.num_experts_per_tok * config.decode_rows_per_req) if config.max_running_req > limit: logger.warning_rank0( f"the smallest MoE slot pool has {min(small)} slots; lowering max_running_req " @@ -111,6 +112,12 @@ def _page_table_width(max_seq_len: int, page_size: int) -> int: return align_ceil(align_ceil(max_seq_len, page_size), 32) +def _seq_limits(config, num_tokens: int) -> tuple[int, int]: + """(longest request, page-table width). A decode step may widen a request by decode_rows_per_req - 1 provisional rows past its length, so requests get that much less and the table still covers the widened tail.""" + limit = min(config.max_seq_len, num_tokens) + return limit - (config.decode_rows_per_req - 1), _page_table_width(limit, config.page_size) + + def _required_attn_types(model_config) -> frozenset[AttnType]: """Backend-driving attention types of this model, from the group-spec walk (single source shared with the pool factory and the KV cost model). getattr @@ -524,8 +531,7 @@ def __init__(self, config: EngineConfig): # ======================= Page table initialization ======================== # NOTE: 1. aligned to 128 bytes; 2. store raw locations instead of pages - self.max_seq_len = min(config.max_seq_len, num_tokens) - aligned_max_seq_len = _page_table_width(self.max_seq_len, config.page_size) + self.max_seq_len, aligned_max_seq_len = _seq_limits(config, num_tokens) self.ctx.page_table = self.page_table = torch.zeros( # + 1 for dummy request (config.max_running_req + 1, aligned_max_seq_len), dtype=torch.int32, @@ -564,16 +570,16 @@ def __init__(self, config: EngineConfig): self.graph_runner = GraphRunner( stream=self.stream, device=self.device, - model=self.model, + forward=self.model.forward, attn_backend=self.attn_backend, cuda_graph_bs=config.cuda_graph_bs, cuda_graph_max_bs=config.cuda_graph_max_bs, free_memory=init_free_memory, max_seq_len=aligned_max_seq_len, - vocab_size=config.model_config.vocab_size, dummy_req=self.dummy_req, moe_offload_cache=self.moe_offload_cache, mrope=config.model_config.model_is_mrope, + out_dtype=torch.float32, ) if config.attention_backend.split(",")[0] == "triton": # Prefill runs on the first comma part; warm its autotune cache. @@ -805,8 +811,8 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: f"(--moe-cache-scope per-layer gives each layer its own)" ) self._moe_slot_limit = _slot_limit(layer_methods) if layer_methods is not None else (method.slot_limit() if method is not None else None) - # with GPU decode, a pool below num_experts slots bounds the batch; the startup split keeps room for one token per step - pools = self._moe_pools(1) + # with GPU decode, a pool below num_experts slots bounds the batch; the startup split keeps room for one request's decode step + pools = self._moe_pools(config.decode_rows_per_req if decode_target == "gpu" else 1) size_cap = sum(p.max_num_slots for p in pools) if config.moe_cache_auto: size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method, max_slots=size_cap) @@ -902,7 +908,6 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: self._resolve_hybrid_fetch(config, cache) if locked_fetch: attach_locked_fetch(cache, config, staging, self._profile_layer_rows) - self.ctx.moe_offload_cache = cache self.moe_offload_cache = cache return cache @@ -1034,6 +1039,7 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: # Decode batches never exceed max_running_req, but CUDA-graph padding can # round a batch up to the largest captured size; cover both. max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, *(config.cuda_graph_bs or []), 1) + max_tokens *= config.decode_rows_per_req by_id = sorted(layers, key=lambda layer: layer.layer_id) layer_formats = [layer.quant_method.cpu_spec() for layer in by_id] executor = CpuMoeExecutor( @@ -1098,8 +1104,7 @@ def _resize_kv_pool(self, config, num_pages: int, num_swa_pages: int | None) -> def _refresh_seq_state(self, config) -> None: num_tokens = self.num_pages * config.page_size - self.max_seq_len = min(config.max_seq_len, num_tokens) - aligned_max_seq_len = _page_table_width(self.max_seq_len, config.page_size) + self.max_seq_len, aligned_max_seq_len = _seq_limits(config, num_tokens) if aligned_max_seq_len != self.page_table.shape[1]: # max_seq_len changed (e.g. KV grew past the startup token budget); the page table # columns must track it or new requests would index out of bounds. The scheduler @@ -1143,7 +1148,7 @@ def rebuild_runtime_cache( ) try: # a rebuild does not lower max_running_req again, so with GPU decode every pool must hold a full batch; hybrid caps its fetches instead - batch = config.max_running_req if self.moe_offload_cache.decode_target == "gpu" else 1 + batch = config.max_running_req * config.decode_rows_per_req if self.moe_offload_cache.decode_target == "gpu" else 1 moe_pool_sizes = split_slots(moe_cache_size, self._moe_pools(batch), self._moe_expert_counts) self.moe_offload_cache.validate_rebuild(moe_cache_size, moe_pool_sizes) except ValueError as e: @@ -1244,7 +1249,7 @@ def rebuild_runtime_cache( self.linear_state_pool.rebuild(num_mamba_slots + 1) # 3. Refresh max_seq_len (+ generic page table) for the new token budget. self._refresh_seq_state(config) - aligned_max_seq_len = _page_table_width(self.max_seq_len, config.page_size) + aligned_max_seq_len = self.page_table.shape[1] # 4. Re-capture CUDA graphs against the new tensors (reset_capture above re-armed # the backend; _sync_get_memory empties the cache so freed memory is reclaimed). gc.collect() @@ -1252,22 +1257,20 @@ def rebuild_runtime_cache( self.graph_runner = GraphRunner( stream=self.stream, device=self.device, - model=self.model, + forward=self.model.forward, attn_backend=self.attn_backend, cuda_graph_bs=prior_graph_bs, # reuse the startup-resolved set (see above) cuda_graph_max_bs=config.cuda_graph_max_bs, free_memory=free_min, max_seq_len=aligned_max_seq_len, - vocab_size=config.model_config.vocab_size, dummy_req=self.dummy_req, moe_offload_cache=self.moe_offload_cache, mrope=config.model_config.model_is_mrope, + out_dtype=torch.float32, ) - def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: - assert torch.cuda.current_stream() == self.stream - if batch.mm_gather_plan: - self._run_mm_encoder(batch) + def forward_logits(self, batch: Batch) -> torch.Tensor: + """The batch's logits from its captured graph, or eagerly when none fits.""" use_graph = self.graph_runner.can_use_cuda_graph(batch) with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, use_graph): logits = self.graph_runner.replay(batch) if use_graph else self.model.forward() @@ -1275,6 +1278,13 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: # One pinned read: surfaces a fired flag-handshake watchdog (dead coordinator # -> stale expert outputs) as a loud error instead of silent corruption. self.cpu_moe_executor.raise_if_unhealthy() + return logits + + def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: + assert torch.cuda.current_stream() == self.stream + if batch.mm_gather_plan: + self._run_mm_encoder(batch) + logits = self.forward_logits(batch) for req in batch.reqs: req.complete_one() diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index a995316ddc..42306918be 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -5,7 +5,7 @@ import gc import sys from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal import torch from freetoken.core import Batch, Req, get_global_ctx @@ -16,7 +16,6 @@ if TYPE_CHECKING: from freetoken.attention import BaseAttnBackend - from freetoken.models import BaseLLMModel from freetoken.moe.offload_cache import OffloadMoeCache logger = init_logger(__name__) @@ -27,56 +26,54 @@ class GraphCaptureBuffer: input_ids: torch.Tensor out_loc: torch.Tensor positions: torch.Tensor - # [3, bs] t/h/w rope positions; allocated only for mrope models (else None). + # [3, bs * rows_per_req] t/h/w rope positions; allocated only for mrope models (else None). mrope_positions: torch.Tensor | None - logits: torch.Tensor table_idx: torch.Tensor # per-request slot id for GatedDeltaNet state gather/scatter - # Decode GDN query indptr = arange(bs+1); a constant per captured bs, filled once. + # GDN query indptr = arange(bs+1) * rows_per_req; a constant per captured bs, filled once. fla_cu_seqlens: torch.Tensor + rows_per_req: int = 1 @classmethod - def init( - cls, bs: int, vocab_size: int, device: torch.device, mrope: bool = False - ) -> GraphCaptureBuffer: + def init(cls, bs: int, device: torch.device, mrope: bool = False, rows_per_req: int = 1) -> GraphCaptureBuffer: + n = bs * rows_per_req return GraphCaptureBuffer( - input_ids=torch.zeros(bs, dtype=torch.int32, device=device), - out_loc=torch.zeros(bs, dtype=torch.int32, device=device), - positions=torch.zeros(bs, dtype=torch.int32, device=device), + input_ids=torch.zeros(n, dtype=torch.int32, device=device), + out_loc=torch.zeros(n, dtype=torch.int32, device=device), + positions=torch.zeros(n, dtype=torch.int32, device=device), mrope_positions=( - torch.zeros(3, bs, dtype=torch.int32, device=device) if mrope else None + torch.zeros(3, n, dtype=torch.int32, device=device) if mrope else None ), - logits=torch.empty(bs, vocab_size, dtype=torch.float32, device=device), table_idx=torch.zeros(bs, dtype=torch.int32, device=device), - fla_cu_seqlens=torch.arange(bs + 1, dtype=torch.int32, device=device), + fla_cu_seqlens=torch.arange(bs + 1, dtype=torch.int32, device=device) * rows_per_req, + rows_per_req=rows_per_req, ) def set_batch(self, batch: Batch) -> None: from freetoken.attention.linear import FLAMetadata - _slice = slice(batch.padded_size) bs = batch.padded_size - batch.input_ids = self.input_ids[_slice] - batch.out_loc = self.out_loc[_slice] - batch.positions = self.positions[_slice] + _rows = slice(bs * self.rows_per_req) + batch.input_ids = self.input_ids[_rows] + batch.out_loc = self.out_loc[_rows] + batch.positions = self.positions[_rows] if self.mrope_positions is not None: - batch.mrope_positions = self.mrope_positions[:, _slice] - batch.linear_table_idx = self.table_idx[_slice] - # Decode GDN metadata reads the persistent cu_seqlens (constant arange) and the - # persistent table_idx slot map, so the captured kernels see stable addresses. + batch.mrope_positions = self.mrope_positions[:, _rows] + batch.linear_table_idx = self.table_idx[:bs] + # GDN metadata reads the persistent cu_seqlens (constant arange * rows_per_req) and table_idx slot map, so the captured kernels see stable addresses. batch.fla_metadata = FLAMetadata( - cu_seqlens=self.fla_cu_seqlens[: bs + 1], cache_indices=self.table_idx[_slice] + cu_seqlens=self.fla_cu_seqlens[: bs + 1], cache_indices=self.table_idx[:bs] ) def copy_from(self, batch: Batch) -> None: - _slice = slice(batch.padded_size) - self.input_ids[_slice] = batch.input_ids + _rows = slice(batch.padded_size * self.rows_per_req) + self.input_ids[_rows] = batch.input_ids if batch.out_loc is not None: - self.out_loc[_slice] = batch.out_loc - self.positions[_slice] = batch.positions + self.out_loc[_rows] = batch.out_loc + self.positions[_rows] = batch.positions if self.mrope_positions is not None: - self.mrope_positions[:, _slice] = batch.mrope_positions + self.mrope_positions[:, _rows] = batch.mrope_positions if batch.linear_table_idx is not None: - self.table_idx[_slice] = batch.linear_table_idx + self.table_idx[: batch.padded_size] = batch.linear_table_idx def _determine_cuda_graph_bs( @@ -130,26 +127,37 @@ def _upload_graph(graph: torch.cuda.CUDAGraph, stream: torch.cuda.Stream) -> Non class GraphRunner: + """CUDA graphs of ``forward`` over padded batches of ``rows_per_req`` rows per request, one graph per batch size; ``replay`` returns the real requests' rows of each output.""" + def __init__( self, stream: torch.cuda.Stream, device: torch.device, - model: BaseLLMModel, + forward: Callable[[], Any], attn_backend: BaseAttnBackend, cuda_graph_bs: List[int] | None, cuda_graph_max_bs: int | None, free_memory: int, max_seq_len: int, - vocab_size: int, dummy_req: Req, moe_offload_cache: OffloadMoeCache | None = None, mrope: bool = False, + *, + rows_per_req: int = 1, + phase: Literal["decode", "verify"] = "decode", + ctx_override: Dict[str, Any] | None = None, + batch_hook: Callable[[Batch], None] | None = None, + pool: Any = None, + out_dtype: torch.dtype | None = None, ) -> None: cuda_graph_bs = _determine_cuda_graph_bs( cuda_graph_bs=cuda_graph_bs, cuda_graph_max_bs=cuda_graph_max_bs, free_memory=free_memory, ) + # the padded rows replay the dummy request, so it must carry rows_per_req rows + assert rows_per_req == 1 or dummy_req.extend_len == rows_per_req, (rows_per_req, dummy_req.extend_len) + self.forward = forward self.attn_backend = attn_backend self.max_graph_bs = max(cuda_graph_bs) if cuda_graph_bs else 0 self.graph_bs_list = sorted(cuda_graph_bs) @@ -158,13 +166,33 @@ def __init__( self.mrope = mrope self.stream = stream self.device = device - self._capture_graphs(max_seq_len, vocab_size, model) + self.rows_per_req = rows_per_req + self.phase = phase + self.ctx_override = ctx_override or {} + self.batch_hook = batch_hook + self.graph_pool = pool + self.out_dtype = out_dtype + self._out: List[torch.Tensor] | None = None + self._out_is_tuple = False + self._out_rows: Dict[int, tuple[int, ...]] = {} + self._capture_graphs(max_seq_len) def _reset_moe_offload_cache(self) -> None: if self.moe_offload_cache is not None: self.moe_offload_cache.reset() - def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel): + def _store_output(self, bs: int, out: Any) -> None: + self._out_is_tuple = isinstance(out, tuple) + outs = out if self._out_is_tuple else (out,) + if self._out is None: + # the largest batch is captured first, so its outputs size the buffers + assert self.out_dtype is None or len(outs) == 1, "out_dtype applies to a single output" + self._out = [torch.empty_like(o, dtype=self.out_dtype or o.dtype) for o in outs] + for buf, o in zip(self._out, outs, strict=True): + buf[: o.shape[0]] = o + self._out_rows[bs] = tuple(o.shape[0] for o in outs) + + def _capture_graphs(self, max_seq_len: int): # Mark the post-weights "warmup" phase for /health: this stretch (graph capture — or the # remaining readiness work when graphs are disabled) moves no bytes, so without this the # loader would sit at 100% (last byte bar) until the ready ack. total=0 ⇒ the desktop @@ -175,7 +203,7 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel if self.max_graph_bs == 0: return logger.info_rank0("CUDA graph is disabled.") - self.attn_backend.init_capture_graph(max_seq_len=max_seq_len, bs_list=self.graph_bs_list) + self.attn_backend.init_capture_graph(max_seq_len=max_seq_len, bs_list=self.graph_bs_list, rows_per_req=self.rows_per_req) torch.cuda.synchronize(self.device) torch.cuda.empty_cache() @@ -185,9 +213,7 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel free_memory = get_free_memory(self.device) logger.info_rank0(f"Free GPU memory before capturing CUDA graphs: {mem_GB(free_memory)}") - self.buffer = GraphCaptureBuffer.init( - self.max_graph_bs, vocab_size, self.device, mrope=self.mrope - ) + self.buffer = GraphCaptureBuffer.init(self.max_graph_bs, self.device, mrope=self.mrope, rows_per_req=self.rows_per_req) self._reset_moe_offload_cache() pbar = tqdm( @@ -196,14 +222,16 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel unit="batch", disable=not get_tp_info().is_primary(), # disable for non-primary ranks ) - pool = None + pool = self.graph_pool for bs in pbar: free_memory = get_free_memory(self.device) pbar.desc = f"Capturing graphs: bs = {bs:<3} | avail_mem = {mem_GB(free_memory)}" pbar.refresh() graph = torch.cuda.CUDAGraph() - batch = Batch(reqs=[self.dummy_req] * bs, phase="decode") + batch = Batch(reqs=[self.dummy_req] * bs, phase=self.phase, rows_per_req=self.rows_per_req) batch.padded_reqs = batch.reqs + if self.batch_hook is not None: + self.batch_hook(batch) self.attn_backend.prepare_for_capture(batch) self.buffer.set_batch(batch) # capture on the dummy linear-state slot so GatedDeltaNet gather/scatter @@ -213,32 +241,41 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel if self.dummy_req.linear_slot_idx is not None else self.dummy_req.table_idx) self.buffer.table_idx[:bs].fill_(dummy_slot) - with get_global_ctx().forward_batch(batch): - self.buffer.logits[:bs] = model.forward() + with get_global_ctx().forward_batch(batch, **self.ctx_override): + self._store_output(bs, self.forward()) # Keep the offload cache warmed for capture. Resetting here forces # CUDA graph capture to replay cold-cache expert copies. with torch.cuda.graph(graph, pool=pool, stream=self.stream): - self.buffer.logits[:bs] = model.forward() + self._store_output(bs, self.forward()) _upload_graph(graph, self.stream) self._reset_moe_offload_cache() if pool is None: pool = graph.pool() # reuse cuda graph handle to reduce memory self.graph_map[bs] = graph + self.graph_pool = pool self._reset_moe_offload_cache() free_memory = get_free_memory(self.device) logger.info_rank0(f"Free GPU memory after capturing CUDA graphs: {mem_GB(free_memory)}") def can_use_cuda_graph(self, batch: Batch) -> bool: - return batch.is_decode and batch.size <= self.max_graph_bs + return ( + batch.phase == self.phase + and batch.rows_per_req == self.rows_per_req + and batch.size <= self.max_graph_bs + ) - def replay(self, batch: Batch) -> torch.Tensor: - assert self.can_use_cuda_graph(batch) + def replay(self, batch: Batch) -> Any: + assert self.can_use_cuda_graph(batch) and self._out is not None self.buffer.copy_from(batch) - g = self.graph_map[batch.padded_size] + padded = batch.padded_size + g = self.graph_map[padded] self.attn_backend.prepare_for_replay(batch) g.replay() - return self.buffer.logits[: batch.size] + outs = tuple( + buf[: rows // padded * batch.size] for buf, rows in zip(self._out, self._out_rows[padded], strict=True) + ) + return outs if self._out_is_tuple else outs[0] def pad_batch(self, batch: Batch) -> None: padded_size = ( # choose the first available batch size @@ -251,10 +288,11 @@ def pad_batch(self, batch: Batch) -> None: # NOTE: This must be called before freeing NCCL resources to prevent program hang def destroy_cuda_graphs(self) -> None: # Drop the CUDAGraph objects (and the shared mempool they hold) AND the static - # GraphCaptureBuffer tensors ([max_bs, vocab] logits + input/out_loc/positions/...). + # GraphCaptureBuffer tensors (input/out_loc/positions/...) and the output buffers (_out). # Dropping the references is the load-bearing step; without it a runtime rebuild's # free-before-alloc cannot reclaim this GPU memory. empty_cache() is left to the # caller / next capture (GraphRunner._capture_graphs already runs it). self.graph_map = {} self.buffer = None + self._out = None gc.collect() diff --git a/python/freetoken/layers/embedding.py b/python/freetoken/layers/embedding.py index f86670b14c..469b752da0 100644 --- a/python/freetoken/layers/embedding.py +++ b/python/freetoken/layers/embedding.py @@ -129,11 +129,16 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ctx = get_global_ctx() batch = ctx.batch bs = batch.size + if (features := getattr(batch, "features", None)) is not None: + features.final_hidden(x) if batch.is_prefill: indices = batch.attn_metadata.get_last_indices(bs) x = x[indices].contiguous() del indices + return self.logits(x) + def logits(self, x: torch.Tensor) -> torch.Tensor: + """Full-vocab logits of the rows of ``x``; reads no batch state.""" if self.tied_embedding is not None: logits = self.tied_embedding.quant_method.logits(self.tied_embedding, x, self.bias) else: @@ -143,7 +148,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: input_shape = logits.shape output_tensor = self._comm.all_gather(logits) - if bs == 1: + # keyed on rows, not batch size: a verify batch has several rows per request + if input_shape[0] == 1: return output_tensor.view(1, -1)[:, : self.num_embeddings] output_tensor = output_tensor.view((self.tp_size,) + input_shape) diff --git a/python/freetoken/llm/llm.py b/python/freetoken/llm/llm.py index 20ae087405..8b6a17a15f 100644 --- a/python/freetoken/llm/llm.py +++ b/python/freetoken/llm/llm.py @@ -94,8 +94,10 @@ def offline_send_result(self, reply: List[BaseTokenizerMsg]) -> None: continue assert isinstance(msg, DetokenizeMsg) status = self.status_map[msg.uid] - if not (msg.finished and msg.next_token in self.eos_token_ids): - status.output_ids.append(msg.next_token) + tokens = msg.next_tokens + if msg.finished and tokens and tokens[-1] in self.eos_token_ids: + tokens = tokens[:-1] + status.output_ids.extend(tokens) def generate( self, diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py index f091d0d369..56308d6ece 100644 --- a/python/freetoken/message/tokenizer.py +++ b/python/freetoken/message/tokenizer.py @@ -27,7 +27,8 @@ class BatchTokenizerMsg(BaseTokenizerMsg): @dataclass class DetokenizeMsg(BaseTokenizerMsg): uid: int - next_token: int + # tokens committed this step, in order: one per decode step, up to k+1 per speculative verify + next_tokens: List[int] finished: bool finish_reason: str | None = None # The stop string that ended generation (if any), so the detokenizer can trim it diff --git a/python/freetoken/models/blocks.py b/python/freetoken/models/blocks.py index 4f039fd3ba..a7a095a291 100644 --- a/python/freetoken/models/blocks.py +++ b/python/freetoken/models/blocks.py @@ -48,6 +48,21 @@ def place_encoder_weights(self, mode: str) -> None: ... +@runtime_checkable +class FeatureSink(Protocol): + """Where a forward hands per-row hidden states to a draft model; set on ``batch.features``.""" + + wanted_layers: frozenset[int] + + def final_hidden(self, x: torch.Tensor) -> None: + """Every row's hidden state before the LM head; the LM head calls it once per forward.""" + ... + + def layer_output(self, layer_id: int, x: torch.Tensor, num_streams: int = 1) -> None: + """A layer's output rows; called only for layers in ``wanted_layers``.""" + ... + + def embed_input_ids(embed_tokens, input_ids: torch.Tensor, batch: Batch) -> torch.Tensor: """Token embeddings of the batch; on a chunk with multimodal rows, batch.mm_embeds' leading columns replace the rows batch.mm_rows.""" if batch.mm_embeds is None: diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 4a6a297df9..6610b128cd 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -690,7 +690,7 @@ def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None: reloads banks. Tears down prefill-overlap buffers first (their views alias the old ``bank_caches``), frees the old GPU tensors, then reallocates. Slots cold-start after rebuild. Object identity is preserved so attached layers and - ``ctx.moe_offload_cache`` stay valid. + ``Engine.moe_offload_cache`` stay valid. """ assert self.bank_sources, "set_bank_sources must run before rebuild" sizes = [cache_size] if pool_sizes is None else list(pool_sizes) diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index d4982f9e5d..5a3b46015b 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -635,6 +635,19 @@ def _keys(self, ids: torch.Tensor, n: int) -> torch.Tensor: # ids stay below 2^31 (multimodal pseudo ids included), so a pair packs into one int64 return (ids[:m].to(torch.int64) << 31) | ids[1 : m + 1].to(torch.int64) + def rollback_tail(self, req: Req, keep_len: int, drop_len: int) -> None: + """Free the slots allocated for positions [keep_len, drop_len), rounded out to whole pages as allocate_paged allocates them. + + Call inside lazy_free_region, after the forward that wrote them has drained. + """ + start = div_ceil(keep_len, self.page_size) * self.page_size + end = div_ceil(drop_len, self.page_size) * self.page_size + if end > start: + tail = self.page_table[req.table_idx, start:end] + if self.swa_paged: + self._free_swa(tail) + self._free(tail) + def _page_to_token(self, pages: torch.Tensor) -> torch.Tensor: if self.page_size == 1: return pages diff --git a/python/freetoken/scheduler/decode.py b/python/freetoken/scheduler/decode.py index 8f74c12c78..c5c75ef2e4 100644 --- a/python/freetoken/scheduler/decode.py +++ b/python/freetoken/scheduler/decode.py @@ -9,6 +9,8 @@ @dataclass class DecodeManager: page_size: int + # slots a decode step may hold past each request's committed length (provisional rows) + token_headroom: int = 0 running_reqs: Set[Req] = field(default_factory=set) def filter_reqs(self, reqs: Iterable[Req]) -> None: @@ -26,7 +28,8 @@ def abort_req(self, uid: int) -> Req | None: @property def inflight_tokens(self) -> int: - tokens_reserved = (self.page_size - 1) * len(self.running_reqs) # 1 page reserved + # each request's partly filled page, plus the provisional rows a decode step may add + tokens_reserved = (self.page_size - 1 + self.token_headroom) * len(self.running_reqs) return sum(req.remain_len for req in self.running_reqs) + tokens_reserved def schedule_next_batch(self) -> Batch | None: diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py index de4bddfb4a..94e708fb1d 100644 --- a/python/freetoken/scheduler/prefill.py +++ b/python/freetoken/scheduler/prefill.py @@ -55,6 +55,8 @@ class PrefillAdder: # allocated only in allocate_paged (after the pass), so swa_available_size does not decrement # across the admission loop -- without this, successive admits all see the full pool. reserved_swa: int = 0 + # slots a decode step may hold past a request's own length; reserved with it + token_headroom: int = 0 def __post_init__(self) -> None: if not self.pass_budget: @@ -79,7 +81,7 @@ def _try_allocate_one(self, req: PendingReq): # TODO: better estimate policy extend_len = req.input_len - cached_len estimated_size = self._kv_reservation_size( - req.input_len + req.output_len, cached_len + req.input_len + req.output_len + self.token_headroom, cached_len ) if estimated_size + self.reserved_size > self.cache_manager.available_size: @@ -201,7 +203,7 @@ def _add_one_req( # page span here; charging raw tokens can admit several short requests against one page # even though allocate_paged() needs a separate page for each request. self.reserved_size += self._kv_reservation_size( - pending_req.input_len + pending_req.output_len, cached_len + pending_req.input_len + pending_req.output_len + self.token_headroom, cached_len ) # NOTE: update the tokens ids only; new pages will be allocated in the scheduler _slice = slice(cached_len, cached_len + chunk_size) @@ -277,6 +279,7 @@ class PrefillManager: decode_manager: DecodeManager encoder_cache: EncoderCache | None = None keep_images_whole: bool = False + token_headroom: int = 0 pending_list: List[PendingReq] = field(default_factory=list) def add_one_req(self, req: UserMsg) -> None: @@ -303,6 +306,7 @@ def schedule_next_batch(self, prefill_budget: int) -> Batch | None: table_manager=self.table_manager, encoder_cache=self.encoder_cache, keep_images_whole=self.keep_images_whole, + token_headroom=self.token_headroom, ) reqs: List[Req] = [] chunked_list: List[PendingReq] = [] diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index f16608cc6a..21ef19d9f8 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -93,7 +93,8 @@ def __init__(self, config: SchedulerConfig): None, ) or getattr(self.engine.kv_cache, "sliding_window_size", None), ) - self.decode_manager = DecodeManager(config.page_size) + token_headroom = config.decode_rows_per_req - 1 + self.decode_manager = DecodeManager(config.page_size, token_headroom=token_headroom) 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, @@ -101,6 +102,7 @@ def __init__(self, config: SchedulerConfig): self.decode_manager, encoder_cache=self.engine.encoder_cache, keep_images_whole=self._bidirectional_mm, + token_headroom=token_headroom, ) # some alias for easy access @@ -360,43 +362,9 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: # are freed below/already; shipping this token would append past the # client's terminal reply. continue - next_token = next_tokens_cpu[i] - req.append_host(next_token.unsqueeze(0)) - next_token = int(next_token.item()) - # EOS / stop-string -> "stop", output budget exhausted -> "length"; - # EOS and stop strings win over length. - # Overlap can advance device_len one token ahead of this host reply. - hit_length = req.input_ids.numel() >= req.max_device_len - hit_eos = ( - not req.sampling_params.ignore_eos and next_token in self.eos_token_ids - ) - matched_stop = ( - self._match_stop_str(req) - if not hit_eos and req.sampling_params.stop_strs - else None - ) - finished = hit_length or hit_eos or matched_stop is not None - finish_reason = ( - ("stop" if (hit_eos or matched_stop is not None) else "length") - if finished - else None - ) - if ( - next_token == self.toolcall_anchor_id - and req.toolcall_anchor_len is None - and not finished - ): - req.toolcall_anchor_len = req.input_ids.numel() - reply.append( - DetokenizeMsg( - uid=req.uid, - next_token=next_token, - finished=finished, - finish_reason=finish_reason, - matched_stop=matched_stop, - stop_strs=req.sampling_params.stop_strs or None, - ) - ) + msg, _ = self._commit_tokens(req, next_tokens_cpu[i : i + 1]) + reply.append(msg) + finished = msg.finished # NOTE: overlap scheduling may make the request freed twice, skip second free if finished and req not in self.finished_reqs: @@ -444,6 +412,37 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: ) self.send_result(reply) + def _commit_tokens(self, req: Req, tokens: torch.Tensor) -> Tuple[DetokenizeMsg, int]: + """Append sampled tokens in order until one finishes the request; returns the reply and how many tokens were committed.""" + params = req.sampling_params + committed: List[int] = [] + finished, finish_reason, matched_stop = False, None, None + for j in range(tokens.numel()): + req.append_host(tokens[j : j + 1]) + token = int(tokens[j].item()) + committed.append(token) + # EOS / stop-string -> "stop", output budget exhausted -> "length"; + # EOS and stop strings win over length. + # Overlap can advance device_len one token ahead of this host reply. + hit_length = req.input_ids.numel() >= req.max_device_len + hit_eos = not params.ignore_eos and token in self.eos_token_ids + matched_stop = self._match_stop_str(req) if not hit_eos and params.stop_strs else None + if hit_length or hit_eos or matched_stop is not None: + finished = True + finish_reason = "stop" if (hit_eos or matched_stop is not None) else "length" + break + if token == self.toolcall_anchor_id and req.toolcall_anchor_len is None: + req.toolcall_anchor_len = req.input_ids.numel() + msg = DetokenizeMsg( + uid=req.uid, + next_tokens=committed, + finished=finished, + finish_reason=finish_reason, + matched_stop=matched_stop, + stop_strs=params.stop_strs or None, + ) + return msg, len(committed) + 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 @@ -841,7 +840,7 @@ def _log_cache_geometry(self, event: str) -> None: def _prepare_batch(self, batch: Batch) -> ForwardInput: self.engine.graph_runner.pad_batch(batch) self._forward_iter += 1 - if batch.is_decode: + if not batch.is_prefill: # Free each decoding request's now-out-of-window SWA slots BEFORE the alloc below, # so they can back the new token -- this is what bounds the per-request swa # footprint during decode. (no-op unless the model is SWA / paged swa pool.) @@ -866,7 +865,7 @@ def _prepare_batch(self, batch: Batch) -> ForwardInput: write_mapping = _make_write_tuple(batch, self.device) batch.out_loc = self.engine.page_table[input_mapping] if self.engine.linear_state_pool is not None: - if batch.is_decode: + if not batch.is_prefill: # GPU GDN-state slot (one per padded request) for the decode gather/scatter; # lands in the CUDA-graph input buffer via copy_from. Gate on the cache mode, # NOT on whether any padded req has a linear_slot_idx -- the persistent dummy @@ -882,15 +881,15 @@ def _prepare_batch(self, batch: Batch) -> ForwardInput: slots, dtype=torch.int32, device="cpu", pin_memory=True ).to(self.device, non_blocking=True) else: - batch.linear_table_idx = input_mapping[0].to(torch.int32) + batch.linear_table_idx = input_mapping[0][:: batch.rows_per_req].to(torch.int32) # Per-forward GDN metadata (cu_seqlens / cache_indices / continuation flags), # built once here instead of rebuilt in each of the 30 GDN layers. For decode # under CUDA graph the persistent cu_seqlens buffer is supplied by set_batch. batch.fla_metadata = build_fla_metadata(batch, self.device) - if batch.is_decode: + if not batch.is_prefill: # This batch's padded per-row page-table rows. Backends that snapshot the table for # a captured replay (DSV4) read them in prepare_metadata / prepare_for_replay. - batch.active_table_idx = input_mapping[0].view(-1) + batch.active_table_idx = input_mapping[0][:: batch.rows_per_req].view(-1) self.engine.attn_backend.prepare_metadata(batch) return ForwardInput( batch=batch, diff --git a/python/freetoken/scheduler/status.py b/python/freetoken/scheduler/status.py index 76c20cd405..5e9ca0e0da 100644 --- a/python/freetoken/scheduler/status.py +++ b/python/freetoken/scheduler/status.py @@ -36,7 +36,9 @@ def report_batch( page_size: int, mamba_slots: tuple[int, int] | None = None, swa_tokens: tuple[int, int] | None = None, + generated: int | None = None, ) -> None: + """``generated``: tokens a decode step committed; None means one per request.""" if batch.is_prefill: self._report_prefill( batch, @@ -47,7 +49,7 @@ def report_batch( mamba_slots=mamba_slots, swa_tokens=swa_tokens, ) - elif batch.is_decode: + else: self._report_decode( batch, running_reqs=running_reqs, @@ -57,6 +59,7 @@ def report_batch( page_size=page_size, mamba_slots=mamba_slots, swa_tokens=swa_tokens, + generated=generated, ) def _report_prefill( @@ -103,9 +106,10 @@ def _report_decode( page_size: int, mamba_slots: tuple[int, int] | None = None, swa_tokens: tuple[int, int] | None = None, + generated: int | None = None, ) -> None: self._decode_forward_count += 1 - self._decode_generated_tokens += len(batch.reqs) + self._decode_generated_tokens += len(batch.reqs) if generated is None else generated if self._decode_forward_count % self.decode_log_interval != 0: return diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py index d138ac1a68..045e338918 100644 --- a/python/freetoken/tokenizer/detokenize.py +++ b/python/freetoken/tokenizer/detokenize.py @@ -103,8 +103,10 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: sent_offset=0, ) s = self.decode_map[msg.uid] - if not (msg.finished and msg.next_token in self.eos_token_ids): - s.decoded_ids.append(msg.next_token) + tokens = msg.next_tokens + if msg.finished and tokens and tokens[-1] in self.eos_token_ids: + tokens = tokens[:-1] + s.decoded_ids.extend(tokens) read_ids.append(s.decoded_ids[s.surr_offset :]) surr_ids.append(s.decoded_ids[s.surr_offset : s.read_offset]) diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py index fad66d746d..9888716050 100644 --- a/python/freetoken/tokenizer/server.py +++ b/python/freetoken/tokenizer/server.py @@ -236,7 +236,7 @@ def tokenize_worker( finished=msg.finished, finish_reason=msg.finish_reason, matched_stop=msg.matched_stop, - completion_tokens_delta=1, + completion_tokens_delta=len(msg.next_tokens), kv_used_pages=msg.kv_used_pages, kv_total_pages=msg.kv_total_pages, mamba_used_slots=msg.mamba_used_slots, diff --git a/tests/attention/test_capture_rows_per_req.py b/tests/attention/test_capture_rows_per_req.py new file mode 100644 index 0000000000..a7469f1ba5 --- /dev/null +++ b/tests/attention/test_capture_rows_per_req.py @@ -0,0 +1,105 @@ +"""Attention graphs captured for `rows_per_req` query rows per request (1 = decode, more = a uniform extend such as a speculative verify) replay what the eager path computes.""" + +from types import SimpleNamespace + +import pytest +import torch + +import freetoken.core as core +from freetoken.attention import attention_backend_info +from freetoken.core import Batch, Context, Req +from freetoken.distributed import set_tp_info, try_get_tp_info + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +HEADS, KV_HEADS, HEAD_DIM = 8, 2, 128 +BS, WIDTH = 2, 64 +CACHED = (5, 17) + + +def _backend(name: str): + config = SimpleNamespace(num_qo_heads=HEADS, num_kv_heads=KV_HEADS, head_dim=HEAD_DIM) + if name == "fa": + pytest.importorskip("sgl_kernel") + from freetoken.attention.fa import FlashAttentionBackend + + return FlashAttentionBackend(config) + if name == "fi": + pytest.importorskip("flashinfer") + from freetoken.attention.fi import FlashInferBackend + + return FlashInferBackend(config) + from freetoken.attention.triton import TritonAttentionBackend + + return TritonAttentionBackend(config) + + +def _context(monkeypatch) -> Context: + from freetoken.kvcache.mha_pool import MHAKVCache + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + device = torch.device("cuda") + ctx = Context(page_size=1) + # rows 0..BS-1 belong to the requests, row BS is the dummy request's (padded rows write there) + ctx.page_table = torch.arange((BS + 1) * WIDTH, dtype=torch.int32, device=device).view(BS + 1, WIDTH) + ctx.kv_cache = MHAKVCache(KV_HEADS, 1, HEAD_DIM, num_pages=(BS + 1) * WIDTH, page_size=1, dtype=torch.bfloat16, device=device) + n = (BS + 1) * WIDTH + ctx.kv_cache.store_kv( + torch.randn(n, KV_HEADS * HEAD_DIM, dtype=torch.bfloat16, device=device), + torch.randn(n, KV_HEADS * HEAD_DIM, dtype=torch.bfloat16, device=device), + torch.arange(n, dtype=torch.int32, device=device), + 0, + ) + monkeypatch.setattr(core, "_GLOBAL_CTX", ctx) + return ctx + + +def _batch(reqs, rows_per_req: int, ctx: Context) -> Batch: + batch = Batch(reqs=reqs, phase="decode" if rows_per_req == 1 else "verify", rows_per_req=rows_per_req) + batch.padded_reqs = reqs + batch.positions = torch.cat([torch.arange(r.cached_len, r.device_len, dtype=torch.int32) for r in reqs]).cuda() + batch.out_loc = torch.cat([ctx.page_table[r.table_idx, r.cached_len : r.device_len] for r in reqs]) + return batch + + +def _req(table_idx: int, cached: int, rows_per_req: int) -> Req: + return Req(input_ids=torch.zeros(cached + rows_per_req, dtype=torch.int32), table_idx=table_idx, cached_len=cached, + output_len=1, uid=table_idx, sampling_params=None, cache_handle=None) + + +@pytest.mark.parametrize("rows_per_req", [1, 4]) +@pytest.mark.parametrize("name", ["fa", "triton", "fi"]) +def test_graph_replay_matches_eager_attention(monkeypatch, name, rows_per_req): + assert attention_backend_info(name).uniform_extend_graph != "none" + ctx = _context(monkeypatch) + backend = _backend(name) + rows = BS * rows_per_req + q = torch.randn(rows, HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + k = torch.randn(rows, KV_HEADS * HEAD_DIM, dtype=torch.bfloat16, device="cuda") + v = torch.randn(rows, KV_HEADS * HEAD_DIM, dtype=torch.bfloat16, device="cuda") + + eager = _batch([_req(i, CACHED[i], rows_per_req) for i in range(BS)], rows_per_req, ctx) + backend.prepare_metadata(eager) + expected = backend.forward(q, k, v, 0, eager) + + backend.init_capture_graph(max_seq_len=WIDTH, bs_list=[BS], rows_per_req=rows_per_req) + capture = _batch([_req(BS, 0, rows_per_req)] * BS, rows_per_req, ctx) + backend.prepare_for_capture(capture) + static_q, static_k, static_v = q.clone(), k.clone(), v.clone() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + backend.forward(static_q, static_k, static_v, 0, capture) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out = backend.forward(static_q, static_k, static_v, 0, capture) + torch.cuda.current_stream().wait_stream(stream) + + replay = _batch([_req(i, CACHED[i], rows_per_req) for i in range(BS)], rows_per_req, ctx) + backend.prepare_metadata(replay) + backend.prepare_for_replay(replay) + capture.out_loc.copy_(replay.out_loc) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(out, expected, atol=1e-2, rtol=1e-2) diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index 44130d3549..5c1f93b348 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -569,7 +569,14 @@ def test_moe_pool_flags_and_decode_batch_clamp(): # GPU decode needs batch * top_k distinct experts to fit the smallest pool below num_experts config = SimpleNamespace( model_config=SimpleNamespace(num_experts=16, num_experts_per_tok=4), - max_running_req=4, cuda_graph_max_bs=None, cuda_graph_bs=[1, 2, 4], + max_running_req=4, cuda_graph_max_bs=None, cuda_graph_bs=[1, 2, 4], decode_rows_per_req=1, ) _clamp_decode_batch_to_pools(config, [12, 16, 9]) assert (config.max_running_req, config.cuda_graph_max_bs, config.cuda_graph_bs) == (2, 2, [1, 2]) + # two rows per request route twice the experts: 9 // (4 * 2) leaves one request + config = SimpleNamespace( + model_config=SimpleNamespace(num_experts=16, num_experts_per_tok=4), + max_running_req=4, cuda_graph_max_bs=None, cuda_graph_bs=[1, 2, 4], decode_rows_per_req=2, + ) + _clamp_decode_batch_to_pools(config, [12, 16, 9]) + assert (config.max_running_req, config.cuda_graph_max_bs, config.cuda_graph_bs) == (1, 1, [1]) diff --git a/tests/kvcache/test_kv_cache_rebuild.py b/tests/kvcache/test_kv_cache_rebuild.py index 0e3b7050e8..be1e87c3a4 100644 --- a/tests/kvcache/test_kv_cache_rebuild.py +++ b/tests/kvcache/test_kv_cache_rebuild.py @@ -227,7 +227,7 @@ def test_dsv4_refresh_seq_state_tracks_page_table_width(): P, mrr = 128, 4 config = SimpleNamespace( - max_seq_len=65536, page_size=P, max_running_req=mrr, + max_seq_len=65536, page_size=P, max_running_req=mrr, decode_rows_per_req=1, model_config=SimpleNamespace(dsv4_args=SimpleNamespace(window_size=P)), ) # attach_page_table only re-points the pool's full_loc_map; the decode snapshot belongs to diff --git a/tests/models/qwen4_exp/test_ple_disk.py b/tests/models/qwen4_exp/test_ple_disk.py index f71c9e6a0f..85a98dfdae 100644 --- a/tests/models/qwen4_exp/test_ple_disk.py +++ b/tests/models/qwen4_exp/test_ple_disk.py @@ -398,10 +398,9 @@ def forward(): stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): # Exercise the capture-time upload from #697 as well as the PLE dispatch hook. - runner = GraphRunner(stream=stream, device=torch.device("cuda"), model=SimpleNamespace(forward=forward), + runner = GraphRunner(stream=stream, device=torch.device("cuda"), forward=forward, attn_backend=attn, cuda_graph_bs=[8], cuda_graph_max_bs=None, - free_memory=1024**3, max_seq_len=16, - vocab_size=args.num_ngram_heads * disk.head_dim, dummy_req=req) + free_memory=1024**3, max_seq_len=16, dummy_req=req, out_dtype=torch.float32) batch = Batch(reqs=[req] * 8, phase="decode") runner.pad_batch(batch) batch.positions = torch.full((8,), 2, dtype=torch.int32, device="cuda") diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index a29d94e9af..3baddeb523 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -475,7 +475,7 @@ def fake_cuda_graph(graph, pool=None, stream=None): events.append("graph_exit") class FakeAttnBackend: - def init_capture_graph(self, max_seq_len, bs_list): + def init_capture_graph(self, max_seq_len, bs_list, rows_per_req=1): pass def prepare_for_capture(self, batch): @@ -511,13 +511,12 @@ def reset(self): GraphRunner( stream=None, device=torch.device("cpu"), - model=FakeModel(), + forward=FakeModel().forward, attn_backend=FakeAttnBackend(), cuda_graph_bs=[1], cuda_graph_max_bs=None, free_memory=1024, max_seq_len=1, - vocab_size=3, dummy_req=dummy_req, moe_offload_cache=FakeOffloadCache(), ) diff --git a/tests/scheduler/test_abort_inflight_prefill.py b/tests/scheduler/test_abort_inflight_prefill.py index 85c217859f..2ec5508a1b 100644 --- a/tests/scheduler/test_abort_inflight_prefill.py +++ b/tests/scheduler/test_abort_inflight_prefill.py @@ -72,6 +72,7 @@ def _setup(): _last_data=None, ) stub._free_req_resources = lambda req: Scheduler._free_req_resources(stub, req) + stub._commit_tokens = lambda req, tokens: Scheduler._commit_tokens(stub, req, tokens) return pool, cm, tm, dm, pm, sent, stub diff --git a/tests/scheduler/test_cost_accounting_core.py b/tests/scheduler/test_cost_accounting_core.py index ed741ef8cf..67ac53d19a 100644 --- a/tests/scheduler/test_cost_accounting_core.py +++ b/tests/scheduler/test_cost_accounting_core.py @@ -209,7 +209,7 @@ def test_normal_loop_sends_prior_sample_before_abort_terminal(): scheduler._schedule_next_batch = lambda: None sent = [] scheduler.send_result = lambda messages: sent.append(messages) - late_sample = DetokenizeMsg(uid=5, next_token=99, finished=False) + late_sample = DetokenizeMsg(uid=5, next_tokens=[99], finished=False) scheduler._process_last_data = lambda _data: scheduler.send_result([late_sample]) Scheduler.normal_loop(scheduler) diff --git a/tests/scheduler/test_rollback_tail.py b/tests/scheduler/test_rollback_tail.py new file mode 100644 index 0000000000..da74b188c2 --- /dev/null +++ b/tests/scheduler/test_rollback_tail.py @@ -0,0 +1,36 @@ +"""rollback_tail after a widened decode step: the request keeps exactly the pages of its committed tokens, the rest go back, and the next step allocates no page twice.""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.core import Req +from freetoken.scheduler.cache import CacheManager +from freetoken.utils import div_ceil + +K = 3 +NUM_PAGES = 64 + + +@pytest.mark.parametrize("page_size", [1, 4]) +@pytest.mark.parametrize("cached", [8, 9, 11]) +@pytest.mark.parametrize("n", range(1, K + 2)) +def test_rollback_keeps_committed_pages_only(page_size, cached, n): + page_table = torch.zeros(2, 64, dtype=torch.int32) + cm = CacheManager(NUM_PAGES, page_size, page_table, "naive") + req = Req(input_ids=torch.zeros(cached, dtype=torch.int32), table_idx=0, cached_len=0, output_len=32, + uid=0, sampling_params=SimpleNamespace(), cache_handle=None) + cm.allocate_paged([req]) + req.complete_one() # cached_len = cached, the pending token sits at position `cached` + + for _ in range(2): + entry = req.cached_len + req.widen_pending(K) + cm.allocate_paged([req]) + req.commit_pending(n) + with cm.lazy_free_region(): + cm.rollback_tail(req, req.cached_len, entry + K + 1) + kept = div_ceil(req.cached_len, page_size) + assert len(cm.free_slots) == NUM_PAGES - kept + assert not set(cm.free_slots.tolist()) & set(cm.page_table[0, : kept * page_size : page_size].tolist()) diff --git a/tests/server/test_message_wire.py b/tests/server/test_message_wire.py index fdc095921d..76d8ecc397 100644 --- a/tests/server/test_message_wire.py +++ b/tests/server/test_message_wire.py @@ -106,7 +106,7 @@ def test_user_reply_token_deltas_round_trip(): def test_detokenize_msg_carries_kv_usage_round_trip(): msg = DetokenizeMsg( - uid=3, next_token=42, finished=True, + uid=3, next_tokens=[42], finished=True, kv_used_pages=10, kv_total_pages=256, gpu_mem_bytes=1 << 30, mamba_used_slots=7, mamba_total_slots=64, swa_used_tokens=8448, swa_total_tokens=76800,