From 323f907ca6c1508c399616f38e348134679b9331 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 2 Sep 2026 13:17:41 +0800 Subject: [PATCH] feat: expose paged cache serving settings --- astrai/inference/cache/pool.py | 14 ++++++++ astrai/inference/engine.py | 4 +++ astrai/inference/network/app.py | 18 ++++++++++- astrai/inference/scheduler.py | 15 +++++++++ docs/developer/docker-serving.md | 2 ++ docs/guides/inference.md | 10 ++++++ docs/guides/params.md | 6 ++++ scripts/tools/server.py | 53 +++++++++++++++++++++++++++++++ tests/inference/test_cache.py | 25 +++++++++++++++ tests/inference/test_engine.py | 15 +++++++++ tests/inference/test_scheduler.py | 22 +++++++++++++ tests/inference/test_server.py | 25 ++++++++++++++- tests/test_serve_cli.py | 44 ++++++++++++++++++++++++- 13 files changed, 250 insertions(+), 3 deletions(-) diff --git a/astrai/inference/cache/pool.py b/astrai/inference/cache/pool.py index 05509020..2a12c0ba 100644 --- a/astrai/inference/cache/pool.py +++ b/astrai/inference/cache/pool.py @@ -104,6 +104,20 @@ def __init__( page_size: int = 1, n_tokens: Optional[int] = None, ): + if isinstance(page_size, bool) or not isinstance(page_size, int): + raise TypeError("page_size must be an integer") + if page_size <= 0: + raise ValueError("page_size must be positive") + if n_tokens is not None: + if isinstance(n_tokens, bool) or not isinstance(n_tokens, int): + raise TypeError("n_tokens must be an integer") + if n_tokens <= 0: + raise ValueError("n_tokens must be positive") + if n_tokens % page_size: + raise ValueError("n_tokens must be divisible by page_size") + elif page_size != 1: + raise ValueError("page_size requires n_tokens in paged mode") + self.page_size = page_size self.max_batch_size = max_batch_size self.max_seq_len = max_seq_len diff --git a/astrai/inference/engine.py b/astrai/inference/engine.py index e440cceb..8ab545e9 100644 --- a/astrai/inference/engine.py +++ b/astrai/inference/engine.py @@ -76,6 +76,8 @@ def __init__( cache: Optional[PagePool] = None, enable_cuda_graph: bool = True, backend: Optional[Union[str, ATTN_BACKEND, AttentionBackend, type]] = None, + kv_cache_page_size: int = 1, + kv_cache_tokens: Optional[int] = None, ): self.model = model self.tokenizer = tokenizer @@ -87,6 +89,8 @@ def __init__( cache=cache, enable_cuda_graph=enable_cuda_graph, backend=backend, + kv_cache_page_size=kv_cache_page_size, + kv_cache_tokens=kv_cache_tokens, ) self.scheduler.start() diff --git a/astrai/inference/network/app.py b/astrai/inference/network/app.py index 4292b93b..236c9a1a 100644 --- a/astrai/inference/network/app.py +++ b/astrai/inference/network/app.py @@ -111,6 +111,8 @@ def _create_engine( dtype: torch.dtype = torch.bfloat16, max_batch_size: int = 16, max_seq_len: Optional[int] = None, + kv_cache_tokens: Optional[int] = None, + kv_cache_page_size: int = 1, ) -> InferenceEngine: if not param_path.exists(): raise FileNotFoundError(f"Parameter directory not found: {param_path}") @@ -125,8 +127,18 @@ def _create_engine( tokenizer=tokenizer, max_batch_size=max_batch_size, max_seq_len=max_seq_len, + kv_cache_tokens=kv_cache_tokens, + kv_cache_page_size=kv_cache_page_size, + ) + cache_mode = "paged" if kv_cache_tokens is not None else "contiguous" + logger.info( + "Inference engine initialized with max_batch_size=%s, " + "kv_cache_mode=%s, kv_cache_tokens=%s, kv_cache_page_size=%s", + max_batch_size, + cache_mode, + kv_cache_tokens, + kv_cache_page_size, ) - logger.info(f"Inference engine initialized with max_batch_size={max_batch_size}") return engine @@ -189,6 +201,8 @@ def run_server( dtype: torch.dtype = torch.bfloat16, max_batch_size: int = 16, max_seq_len: Optional[int] = None, + kv_cache_tokens: Optional[int] = None, + kv_cache_page_size: int = 1, ): app = get_app() app.state.server_config = { @@ -197,6 +211,8 @@ def run_server( "param_path": param_path, "max_batch_size": max_batch_size, "max_seq_len": max_seq_len, + "kv_cache_tokens": kv_cache_tokens, + "kv_cache_page_size": kv_cache_page_size, } uvicorn.run( app, diff --git a/astrai/inference/scheduler.py b/astrai/inference/scheduler.py index f568fe29..35b8413b 100644 --- a/astrai/inference/scheduler.py +++ b/astrai/inference/scheduler.py @@ -42,6 +42,8 @@ def __init__( cache: Optional[PagePool] = None, enable_cuda_graph: bool = True, backend: Optional[Union[str, ATTN_BACKEND, AttentionBackend, type]] = None, + kv_cache_page_size: int = 1, + kv_cache_tokens: Optional[int] = None, ): config = model.config @@ -60,6 +62,11 @@ def __init__( head_dim = config.hidden_size // config.num_attention_heads if cache is not None: + if kv_cache_tokens is not None or kv_cache_page_size != 1: + raise ValueError( + "cache cannot be combined with kv_cache_tokens or " + "kv_cache_page_size" + ) self._cache = cache else: self._cache = PagePool( @@ -70,6 +77,8 @@ def __init__( max_seq_len=self.max_seq_len, device=self.device, dtype=self.dtype, + page_size=kv_cache_page_size, + n_tokens=kv_cache_tokens, ) self._metrics = MetricsCollector() @@ -125,6 +134,12 @@ def remove_task(self, task_id: str) -> bool: def get_stats(self) -> Dict[str, Any]: stats = self._task_mgr.get_stats() stats["kv_cache_tasks"] = self._task_cache.task_count + stats["kv_cache_mode"] = "contiguous" if self._cache.contiguous else "paged" + stats["kv_cache_tokens"] = self._cache.n_tokens + stats["kv_cache_page_size"] = self._cache.page_size + stats["kv_cache_prefix_caching"] = ( + not self._cache.contiguous and self._cache.page_size > 1 + ) return stats @property diff --git a/docs/developer/docker-serving.md b/docs/developer/docker-serving.md index 384457ea..b7815d7a 100644 --- a/docs/developer/docker-serving.md +++ b/docs/developer/docker-serving.md @@ -45,6 +45,8 @@ server: dtype: bfloat16 # bfloat16 | float16 | float32 max_batch_size: 16 max_seq_len: null # falls back to model config + kv_cache_tokens: null # set to enable paged allocation + kv_cache_page_size: 1 # values above 1 enable prefix caching ``` - Relative paths resolve from the YAML file's directory, not the current shell. diff --git a/docs/guides/inference.md b/docs/guides/inference.md index 1add4999..c0c9cb14 100644 --- a/docs/guides/inference.md +++ b/docs/guides/inference.md @@ -39,6 +39,11 @@ PagePool (top-level manager, orchestrates all layers) - **Contiguous (default)**: pre-allocates `max_batch_size * max_seq_len` token slots. `req_to_token` is a trivial linear mapping (`slot = req_idx * max_seq_len + pos`). No dynamic allocation. - **Paged** (`page_size=1` or `>1` with `n_tokens` set): shared token pool with on-demand allocation. `Allocator` provides ref-counted allocation and LRU eviction. When `page_size > 1`, `RadixCache` also enables prefix sharing. +The server exposes the same choice as `kv_cache_tokens` and +`kv_cache_page_size`. Leaving `kv_cache_tokens` unset preserves contiguous +allocation. Setting it enables paged allocation; a page size above 1 also +enables prefix caching. The token capacity must be divisible by the page size. + `RadixCache` indexes complete token pages as parent-linked radix edges. Lookup walks from the root and compares each page's exact token tuple, so an identical page can only be reused under the same parent prefix. Hash values are retained for introspection, but never determine a match. Only fully materialized KV pages enter the radix. A partial final page remains private to its request and is released when the request ends. On completion, the scheduler records the prompt plus generated tokens already decoded into KV; it excludes the final sampled token because that token has not yet passed through the model. A later request resumes prefill immediately after the longest complete-page hit. @@ -192,6 +197,8 @@ server: dtype: bfloat16 max_batch_size: 16 max_seq_len: null + kv_cache_tokens: 131072 # enables paged allocation + kv_cache_page_size: 64 # values above 1 enable prefix caching ``` ```bash @@ -199,6 +206,9 @@ python scripts/tools/server.py --config serve.yaml python scripts/tools/server.py --config serve.yaml --port 9000 # CLI wins ``` +`GET /stats` reports the effective `kv_cache_mode`, token capacity, page size, +and whether prefix caching is active. + In Docker, `scripts/serve.sh` drives the same YAML (a `runtime:` section controls ports/GPU/mounts); see [Docker Serving](../developer/docker-serving.md). diff --git a/docs/guides/params.md b/docs/guides/params.md index 9b23ebd8..bfdcb9d6 100644 --- a/docs/guides/params.md +++ b/docs/guides/params.md @@ -211,6 +211,8 @@ nohup python scripts/tools/train.py \ | `--dtype` | str | `bfloat16` | Model weights dtype (`bfloat16`, `float16`, `float32`) | | `--max_batch_size` | int | `16` | Maximum batch size for continuous batching | | `--max_seq_len` | int | model config `max_position_embeddings` | Maximum sequence length (KV cache size + prompt truncation) | +| `--kv_cache_tokens` | int | `None` | Shared KV token capacity; setting it enables paged allocation | +| `--kv_cache_page_size` | int | `1` | Paged allocation size; values above 1 enable prefix caching | | `--reload` | flag | `False` | Enable auto-reload for development | Usage: @@ -230,7 +232,11 @@ server: dtype: bfloat16 max_batch_size: 16 max_seq_len: null + kv_cache_tokens: 131072 + kv_cache_page_size: 64 ``` +`kv_cache_tokens` must be positive and divisible by `kv_cache_page_size`. +Leave it unset to retain the contiguous cache default. `serve.yaml` also carries a `runtime:` section for the Docker wrapper; see [Docker Serving](../developer/docker-serving.md). diff --git a/scripts/tools/server.py b/scripts/tools/server.py index 82b537a2..61a5e014 100644 --- a/scripts/tools/server.py +++ b/scripts/tools/server.py @@ -17,6 +17,8 @@ "dtype", "max_batch_size", "max_seq_len", + "kv_cache_tokens", + "kv_cache_page_size", ) @@ -62,6 +64,25 @@ def _as_int(value, name: str) -> int | None: raise click.UsageError(f"{name} must be an integer, got {value!r}") from None +def _validate_kv_cache_config( + kv_cache_tokens: int | None, kv_cache_page_size: int +) -> None: + if kv_cache_page_size <= 0: + raise click.UsageError("server.kv_cache_page_size must be positive") + if kv_cache_tokens is None: + if kv_cache_page_size != 1: + raise click.UsageError( + "server.kv_cache_page_size requires server.kv_cache_tokens" + ) + return + if kv_cache_tokens <= 0: + raise click.UsageError("server.kv_cache_tokens must be positive") + if kv_cache_tokens % kv_cache_page_size: + raise click.UsageError( + "server.kv_cache_tokens must be divisible by server.kv_cache_page_size" + ) + + def _resolve_server_config( config_path: str, passed_kwargs: dict, @@ -79,11 +100,21 @@ def _resolve_server_config( _as_int(resolved["max_batch_size"], "server.max_batch_size") or 16 ) resolved["max_seq_len"] = _as_int(resolved["max_seq_len"], "server.max_seq_len") + resolved["kv_cache_tokens"] = _as_int( + resolved.get("kv_cache_tokens"), "server.kv_cache_tokens" + ) + page_size = _as_int( + resolved.get("kv_cache_page_size", 1), "server.kv_cache_page_size" + ) + resolved["kv_cache_page_size"] = 1 if page_size is None else page_size resolved["reload"] = bool(resolved["reload"]) if resolved["dtype"] not in _DTYPES: raise click.UsageError( f"server.dtype must be one of {', '.join(_DTYPES)}, got {resolved['dtype']!r}" ) + _validate_kv_cache_config( + resolved["kv_cache_tokens"], resolved["kv_cache_page_size"] + ) return resolved @@ -126,6 +157,18 @@ def _resolve_server_config( default=None, help="Maximum sequence length (KV cache size + prompt truncation). Uses model config if not set.", ) +@click.option( + "--kv_cache_tokens", + type=int, + default=None, + help="Shared KV token capacity. Setting it enables paged allocation.", +) +@click.option( + "--kv_cache_page_size", + type=int, + default=1, + help="Paged KV allocation size. Values above 1 enable prefix caching.", +) @click.pass_context def server_command( ctx, @@ -138,6 +181,8 @@ def server_command( dtype, max_batch_size, max_seq_len, + kv_cache_tokens, + kv_cache_page_size, ): """Launch inference server (OpenAI-compatible API).""" if config_path: @@ -150,6 +195,8 @@ def server_command( "dtype": dtype, "max_batch_size": max_batch_size, "max_seq_len": max_seq_len, + "kv_cache_tokens": kv_cache_tokens, + "kv_cache_page_size": kv_cache_page_size, } explicit_keys = { key @@ -165,8 +212,12 @@ def server_command( dtype = resolved["dtype"] max_batch_size = resolved["max_batch_size"] max_seq_len = resolved["max_seq_len"] + kv_cache_tokens = resolved["kv_cache_tokens"] + kv_cache_page_size = resolved["kv_cache_page_size"] click.echo(f"Config: {config_path}") + _validate_kv_cache_config(kv_cache_tokens, kv_cache_page_size) + dtype_map = { "bfloat16": torch.bfloat16, "float16": torch.float16, @@ -186,6 +237,8 @@ def server_command( param_path=Path(param_path), max_batch_size=max_batch_size, max_seq_len=max_seq_len, + kv_cache_tokens=kv_cache_tokens, + kv_cache_page_size=kv_cache_page_size, ) diff --git a/tests/inference/test_cache.py b/tests/inference/test_cache.py index 2145c770..d9098324 100644 --- a/tests/inference/test_cache.py +++ b/tests/inference/test_cache.py @@ -1,5 +1,6 @@ """Unit tests for inference cache components.""" +import pytest import torch from astrai.inference.cache import ( @@ -231,6 +232,30 @@ def _make_contiguous_pool(**kwargs): return PagePool(**defaults) +@pytest.mark.parametrize( + ("kwargs", "error", "message"), + [ + ({"page_size": 0}, ValueError, "page_size must be positive"), + ({"page_size": True}, TypeError, "page_size must be an integer"), + ({"n_tokens": 0}, ValueError, "n_tokens must be positive"), + ({"n_tokens": True}, TypeError, "n_tokens must be an integer"), + ( + {"page_size": 8, "n_tokens": 10}, + ValueError, + "n_tokens must be divisible by page_size", + ), + ( + {"page_size": 8}, + ValueError, + "page_size requires n_tokens in paged mode", + ), + ], +) +def test_page_pool_rejects_invalid_capacity_settings(kwargs, error, message): + with pytest.raises(error, match=message): + _make_contiguous_pool(**kwargs) + + def test_page_pool_contiguous_task_alloc_free(): pool = _make_contiguous_pool() task_cache = _make_task_cache(pool) diff --git a/tests/inference/test_engine.py b/tests/inference/test_engine.py index 049b728f..b6175848 100644 --- a/tests/inference/test_engine.py +++ b/tests/inference/test_engine.py @@ -288,6 +288,21 @@ def test_engine_passes_backend_to_scheduler(): assert MockSched.call_args.kwargs["backend"] == "torch_native" +def test_engine_passes_paged_cache_settings_to_scheduler(): + mock_model, mock_tokenizer = _make_engine_mocks() + + with patch("astrai.inference.engine.InferenceScheduler") as MockSched: + InferenceEngine( + mock_model, + mock_tokenizer, + kv_cache_tokens=4096, + kv_cache_page_size=64, + ) + + assert MockSched.call_args.kwargs["kv_cache_tokens"] == 4096 + assert MockSched.call_args.kwargs["kv_cache_page_size"] == 64 + + def test_generate_captures_calling_backend_context(): mock_model, mock_tokenizer = _make_engine_mocks() captured = [] diff --git a/tests/inference/test_scheduler.py b/tests/inference/test_scheduler.py index 4e65cbd1..d68454a2 100644 --- a/tests/inference/test_scheduler.py +++ b/tests/inference/test_scheduler.py @@ -246,6 +246,28 @@ def get_stats(): assert stats["total_tasks"] >= 0 +def test_scheduler_exposes_paged_cache_settings(mock_model_and_tokenizer): + mock_model, mock_tokenizer = mock_model_and_tokenizer + + scheduler = InferenceScheduler( + model=mock_model, + tokenizer=mock_tokenizer, + max_batch_size=4, + max_seq_len=64, + device="cpu", + kv_cache_tokens=512, + kv_cache_page_size=8, + ) + try: + stats = scheduler.get_stats() + assert stats["kv_cache_mode"] == "paged" + assert stats["kv_cache_tokens"] == 512 + assert stats["kv_cache_page_size"] == 8 + assert stats["kv_cache_prefix_caching"] is True + finally: + scheduler.stop() + + def _make_real_scheduler(device): """Build a scheduler backed by a tiny real model for run_batch tests.""" cfg = make_rollout_config(max_position_embeddings=64) diff --git a/tests/inference/test_server.py b/tests/inference/test_server.py index 926e4526..120e8058 100644 --- a/tests/inference/test_server.py +++ b/tests/inference/test_server.py @@ -6,7 +6,7 @@ import torch from astrai.inference import get_app -from astrai.inference.network.app import _create_engine +from astrai.inference.network.app import _create_engine, run_server from astrai.model.transformer import AutoRegressiveLM from astrai.serialization import save_model from tests.helpers import CHAT_TEMPLATE, build_test_tokenizer, make_tiny_config @@ -31,6 +31,24 @@ def test_health_with_model(client, loaded_model): assert data["model_loaded"] is True +def test_run_server_records_paged_cache_settings(tmp_path, monkeypatch): + captured = {} + + def fake_uvicorn_run(app, **kwargs): + captured.update(app.state.server_config) + + monkeypatch.setattr("astrai.inference.network.app.uvicorn.run", fake_uvicorn_run) + run_server( + tmp_path, + device="cpu", + kv_cache_tokens=4096, + kv_cache_page_size=64, + ) + + assert captured["kv_cache_tokens"] == 4096 + assert captured["kv_cache_page_size"] == 64 + + def test_chat_completions_non_stream(client, loaded_model): """POST /v1/chat/completions with stream=false returns OpenAI-style JSON.""" @@ -244,9 +262,14 @@ def test_chat_completions_real_engine(tmp_path, client): dtype=torch.float32, max_batch_size=1, max_seq_len=64, + kv_cache_tokens=128, + kv_cache_page_size=8, ) try: get_app().state.engine = engine + stats = engine.get_stats() + assert stats["kv_cache_mode"] == "paged" + assert stats["kv_cache_prefix_caching"] is True response = client.post( "/v1/chat/completions", json={ diff --git a/tests/test_serve_cli.py b/tests/test_serve_cli.py index 46a84436..6002e0f5 100644 --- a/tests/test_serve_cli.py +++ b/tests/test_serve_cli.py @@ -22,6 +22,8 @@ def _passed() -> dict: "dtype": "bfloat16", "max_batch_size": 16, "max_seq_len": None, + "kv_cache_tokens": None, + "kv_cache_page_size": 1, } @@ -62,6 +64,29 @@ def test_resolve_config_rejects_bad_dtype(tmp_path): _resolve_server_config(str(config_path), _passed()) +@pytest.mark.parametrize( + ("yaml_body", "message"), + [ + ("kv_cache_tokens: 0", "server.kv_cache_tokens must be positive"), + ("kv_cache_page_size: 0", "server.kv_cache_page_size must be positive"), + ( + "kv_cache_page_size: 64", + "server.kv_cache_page_size requires server.kv_cache_tokens", + ), + ( + "kv_cache_tokens: 100\n kv_cache_page_size: 64", + "server.kv_cache_tokens must be divisible", + ), + ], +) +def test_resolve_config_rejects_invalid_kv_cache_settings(tmp_path, yaml_body, message): + config_path = tmp_path / "serve.yaml" + config_path.write_text(f"server:\n {yaml_body}\n", encoding="utf-8") + + with pytest.raises(click.UsageError, match=message): + _resolve_server_config(str(config_path), _passed()) + + def test_server_command_rejects_bad_yaml_dtype(tmp_path): config_path = tmp_path / "serve.yaml" config_path.write_text("server:\n dtype: fp8\n", encoding="utf-8") @@ -76,7 +101,12 @@ def test_server_command_merges_yaml_and_cli(tmp_path, monkeypatch): """Full CLI path: YAML values apply, explicit CLI flags override, args reach run_server.""" config_path = tmp_path / "serve.yaml" config_path.write_text( - "server:\n device: cpu\n dtype: float16\n max_batch_size: 8\n", + "server:\n" + " device: cpu\n" + " dtype: float16\n" + " max_batch_size: 8\n" + " kv_cache_tokens: 4096\n" + " kv_cache_page_size: 64\n", encoding="utf-8", ) captured = {} @@ -95,6 +125,18 @@ def fake_run_server(**kwargs): assert captured["dtype"] == torch.float16 assert captured["max_batch_size"] == 32 assert captured["port"] == 8000 + assert captured["kv_cache_tokens"] == 4096 + assert captured["kv_cache_page_size"] == 64 + + +def test_server_command_rejects_invalid_cli_cache_settings(): + result = CliRunner().invoke( + server_command, + ["--kv_cache_tokens", "100", "--kv_cache_page_size", "64"], + ) + + assert result.exit_code == 2 + assert "server.kv_cache_tokens must be divisible" in result.output def test_config_option_rejects_missing_file(tmp_path):