From fa81a2218cfcbcc6427657cc7d6742d97527e965 Mon Sep 17 00:00:00 2001 From: Jackson57279 Date: Sun, 6 Sep 2026 15:23:23 -0500 Subject: [PATCH 1/2] chore(xpu): snapshot in-flight Intel Arc / accelerator support Working-tree changes that predate the Modal + fp8-KV work and are unrelated to it, committed separately so the following commit reviews cleanly. Introduces the accelerator abstraction (freetoken/accelerator.py, hardware.py), the eager PyTorch XPU path for dense Llama, the torch-native attention backend and the XPU benchmark, plus their tests. Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 3 +- docs/intel-arc-pro-b60.md | 61 +++++++ python/freetoken/accelerator.py | 146 +++++++++++++++ python/freetoken/attention/torch_native.py | 44 +++++ python/freetoken/benchmark/xpu.py | 178 +++++++++++++++++++ python/freetoken/cli.py | 13 ++ python/freetoken/engine/graph.py | 16 +- python/freetoken/engine/sample.py | 13 +- python/freetoken/hardware.py | 72 ++++++++ python/freetoken/layers/activation.py | 4 + python/freetoken/layers/embedding.py | 24 ++- python/freetoken/layers/norm.py | 24 ++- python/freetoken/layers/rotary.py | 18 +- python/freetoken/llm/llm.py | 4 +- python/freetoken/models/config.py | 1 + python/freetoken/models/llama/config.py | 14 ++ python/freetoken/moe/offload_cache.py | 2 +- python/freetoken/moe/offload_stats.py | 12 ++ python/freetoken/scheduler/scheduler.py | 17 +- python/freetoken/utils/torch_utils.py | 6 + requirements-xpu.txt | 21 +++ setup.py | 9 +- tests/models/test_llama_config.py | 69 ++++++++ tests/moe/test_offload.py | 27 ++- tests/scheduler/test_cache_rebuild.py | 3 +- tests/test_accelerator.py | 195 +++++++++++++++++++++ tests/test_xpu_benchmark.py | 71 ++++++++ tests/test_xpu_llama_eager.py | 86 +++++++++ tests/test_xpu_torch_ops.py | 103 +++++++++++ 29 files changed, 1211 insertions(+), 45 deletions(-) create mode 100644 docs/intel-arc-pro-b60.md create mode 100644 python/freetoken/accelerator.py create mode 100644 python/freetoken/attention/torch_native.py create mode 100644 python/freetoken/benchmark/xpu.py create mode 100644 python/freetoken/hardware.py create mode 100644 python/freetoken/moe/offload_stats.py create mode 100644 requirements-xpu.txt create mode 100644 tests/models/test_llama_config.py create mode 100644 tests/test_accelerator.py create mode 100644 tests/test_xpu_benchmark.py create mode 100644 tests/test_xpu_llama_eager.py create mode 100644 tests/test_xpu_torch_ops.py diff --git a/README.md b/README.md index 2a56a0865..907bd55d2 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ FreeToken is an edge-native Mixture-of-Experts (MoE) serving engine designed for - **Semantic-Aware Caching**: Features semantic anchor checkpoints for recurrent state and KV caches, allowing agentic context edits (e.g., tool calls, thinking blocks) to avoid redundant context recomputation. - **Elastic Memory Management**: Supports dynamic, runtime VRAM re-allocation between expert caches and KV memory without engine restarts or weight reloading. - **Broad MoE & Ecosystem Support**: Supports frontier open-weight MoE models (e.g., DeepSeek-V4-Flash, Qwen3.6-35B-A3B, GLM-5.2) across various parameter scales and quantization formats (e.g., MXFP4, NVFP4, FP8, BF16), with Anthropic/OpenAI-compatible APIs for seamless integration with real-world coding and tool-calling agents (e.g., Codex, Claude Code, OpenCode, OpenClaw, DeepSeek Harness). -- **Diverse Consumer Hardware**: Scales across consumer laptops, gaming desktops, and workstation GPUs, with native support for NVIDIA RTX 30, RTX 40, and RTX 50 series GPUs. +- **Diverse Consumer Hardware**: Scales across consumer laptops, gaming desktops, and workstation GPUs, with native support for NVIDIA RTX 30, RTX 40, and RTX 50 series GPUs. Intel Arc Pro B60 has an experimental eager PyTorch XPU path for unquantized dense Llama models. ## Getting Started @@ -55,6 +55,7 @@ For More details: - [Quick start](https://github.com/FlashML-org/FreeToken/blob/main/docs/quickstart.md) - [Supported models](https://github.com/FlashML-org/FreeToken/blob/main/docs/models.md) - [CLI reference](https://github.com/FlashML-org/FreeToken/blob/main/docs/cli.md) +- [Intel Arc Pro B60](https://github.com/FlashML-org/FreeToken/blob/main/docs/intel-arc-pro-b60.md) ## Citation diff --git a/docs/intel-arc-pro-b60.md b/docs/intel-arc-pro-b60.md new file mode 100644 index 000000000..99a87b333 --- /dev/null +++ b/docs/intel-arc-pro-b60.md @@ -0,0 +1,61 @@ +# Intel Arc Pro B60 + +FreeToken has an experimental eager PyTorch XPU path for the 24 GB Intel Arc Pro B60. +The supported slice is intentionally narrow: unquantized, dense checkpoints whose +architecture is `LlamaForCausalLM`. Other architectures, quantized weights, MoE, +attention sinks, and tensor parallelism have not been validated and are rejected before +model allocation. + +## Platform + +- Ubuntu 24.04 or newer with a current Intel `xe` compute driver +- Resizable BAR enabled in firmware +- The PyTorch XPU wheel matching FreeToken's supported PyTorch version +- BF16 or FP16 weights that fit in the card's 24 GB VRAM + +Install the Intel driver and create a dedicated XPU environment. Do not install it over +a working FreeToken CUDA environment: + +```bash +uv venv --python 3.12 +source .venv/bin/activate +uv pip install -r requirements-xpu.txt +FREETOKEN_ACCELERATOR=xpu uv pip install -e . --no-deps +ft hardware --accelerator xpu +``` + +Serve an unquantized FTW Llama checkpoint with the portable attention backend: + +```bash +ft serve --model --accelerator xpu \ + --attention-backend torch --dtype bfloat16 +``` + +This path uses PyTorch scaled-dot-product attention and eager tensor operations. CUDA +graphs are disabled, and the naive prefix cache is selected automatically. Keep the model +resident, prefer BF16, and leave KV-cache headroom. Arc B-series defaults to the Level Zero +V2 adapter, which supports immediate command lists only. Start with that default: + +```bash +ft serve +``` + +Warm the model before measuring. Keep the queue mode that improves both time-to-first-token +and decode tokens/second. The eager path prioritizes correctness and compatibility; a +B60-specific fused-kernel path remains future optimization work. + +Use the built-in benchmark with identical workload geometry. Compare the default V2 path +against the legacy adapter only when short-kernel submission overhead is suspected: + +```bash +ft bench xpu --json > xpu-v2-default.json +SYCL_UR_USE_LEVEL_ZERO_V2=0 SYCL_PI_LEVEL_ZERO_USE_IMMEDIATE_COMMANDLISTS=0 \ + ft bench xpu --json > xpu-legacy-regular.json +SYCL_UR_USE_LEVEL_ZERO_V2=0 SYCL_PI_LEVEL_ZERO_USE_IMMEDIATE_COMMANDLISTS=1 \ + ft bench xpu --json > xpu-legacy-immediate.json +``` + +It records the device, PyTorch version, queue mode, dtype, geometry, memory-copy bandwidth, +prefill SDPA tokens per second, and decode SDPA steps per second. Run it while the GPU is idle +and keep the faster adapter/mode for the representative context size. Intel documents the +Arc B-series V2 behavior in its [Level Zero immediate command-list guide](https://www.intel.com/content/www/us/en/developer/articles/guide/level-zero-immediate-command-lists.html). diff --git a/python/freetoken/accelerator.py b/python/freetoken/accelerator.py new file mode 100644 index 000000000..0a0c6e301 --- /dev/null +++ b/python/freetoken/accelerator.py @@ -0,0 +1,146 @@ +from __future__ import annotations + +from enum import Enum +from types import ModuleType +from typing import TYPE_CHECKING, Final, Literal + +if TYPE_CHECKING: + from freetoken.engine.config import EngineConfig + from freetoken.models.config import ModelConfig + + +class AcceleratorKind(str, Enum): + CUDA = "cuda" + XPU = "xpu" + + +SUPPORTED_ACCELERATORS: Final = tuple(AcceleratorKind) + + +def select_accelerator( + requested: str, + cuda_available: bool, + xpu_available: bool, +) -> AcceleratorKind: + try: + choice = AcceleratorKind(requested) + except ValueError as exc: + if requested != "auto": + supported = ", ".join(("auto", *(kind.value for kind in SUPPORTED_ACCELERATORS))) + raise RuntimeError( + f"Unsupported accelerator {requested!r}; choose one of: {supported}." + ) from exc + else: + match choice: + case AcceleratorKind.CUDA: + if not cuda_available: + raise RuntimeError( + "CUDA was requested, but PyTorch cannot see a CUDA device." + ) + case AcceleratorKind.XPU: + if not xpu_available: + raise RuntimeError( + "Intel XPU was requested, but this PyTorch build cannot see an Intel GPU. " + "Install the XPU wheel and current Intel compute drivers." + ) + return choice + + if cuda_available: + return AcceleratorKind.CUDA + if xpu_available: + return AcceleratorKind.XPU + raise RuntimeError( + "No supported accelerator is available. FreeToken requires NVIDIA CUDA or Intel XPU." + ) + + +def detect_accelerator(requested: str = "auto") -> AcceleratorKind: + import torch + + xpu_available = hasattr(torch, "xpu") and torch.xpu.is_available() + return select_accelerator(requested, torch.cuda.is_available(), xpu_available) + + +def accelerator_runtime(kind: AcceleratorKind) -> ModuleType: + import torch + + match kind: + case AcceleratorKind.CUDA: + return torch.cuda + case AcceleratorKind.XPU: + return torch.xpu + + +def runtime_for_device(device_type: str) -> ModuleType: + try: + kind = AcceleratorKind(device_type) + except ValueError as exc: + raise RuntimeError(f"Unsupported accelerator device type {device_type!r}.") from exc + return accelerator_runtime(kind) + + +def apply_engine_accelerator_constraints( + config: EngineConfig, + kind: AcceleratorKind, + *, + model_config: ModelConfig, + is_moe: bool, +) -> None: + """Apply accelerator-specific EngineConfig backend constraints.""" + match kind: + case AcceleratorKind.CUDA: + return + case AcceleratorKind.XPU: + architectures = tuple(model_config.architectures) + quantized = any( + getattr(model_config, field, "none") != "none" + for field in ( + "dense_quant", + "attn_quant", + "lm_head_quant", + "checkpoint_quantization", + ) + ) + if is_moe or quantized or architectures != ("LlamaForCausalLM",): + raise RuntimeError( + "Intel XPU currently supports unquantized dense LlamaForCausalLM " + "checkpoints only. Quantized, MoE, and other architectures still " + "contain CUDA-specific kernels." + ) + if config.tp_info.size != 1: + raise RuntimeError( + "Intel XPU currently supports tensor parallel size 1 only." + ) + if config.attention_backend not in ("auto", "torch"): + raise RuntimeError( + "Intel XPU requires --attention-backend torch for the portable path." + ) + object.__setattr__(config, "attention_backend", "torch") + object.__setattr__(config, "cuda_graph_max_bs", 0) + object.__setattr__(config, "cuda_graph_bs", []) + object.__setattr__(config, "cache_type", "naive") + object.__setattr__(config, "use_pynccl", False) + object.__setattr__(config, "moe_prefill_hit_d2d", False) + + +def uses_expandable_segments(kind: AcceleratorKind) -> bool: + """Return whether the accelerator supports CUDA expandable segments.""" + return kind is AcceleratorKind.CUDA + + +def distributed_backend(kind: AcceleratorKind) -> Literal["nccl", "xccl"]: + """Return the torch distributed backend for an accelerator.""" + match kind: + case AcceleratorKind.CUDA: + return "nccl" + case AcceleratorKind.XPU: + return "xccl" + + +def create_accelerator_graph(device_type: str, runtime: ModuleType): + """Construct the graph implementation for an accelerator device type.""" + match AcceleratorKind(device_type): + case AcceleratorKind.CUDA: + return runtime.CUDAGraph() + case AcceleratorKind.XPU: + return runtime.XPUGraph() diff --git a/python/freetoken/attention/torch_native.py b/python/freetoken/attention/torch_native.py new file mode 100644 index 000000000..99a35efca --- /dev/null +++ b/python/freetoken/attention/torch_native.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import torch + +from freetoken.core import Batch +from freetoken.kernel.torch_ops import paged_attention + +from .base import AttentionSpec +from .triton import TritonAttentionBackend, TritonMetadata + + +class TorchAttentionBackend(TritonAttentionBackend): + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer_id: int, + batch: Batch, + attn_spec: AttentionSpec | None = None, + ) -> torch.Tensor: + metadata = batch.attn_metadata + assert isinstance(metadata, TritonMetadata) + self.kvcache.store_kv(k, v, batch.out_loc, layer_id) + k_raw = self.kvcache.k_cache(layer_id) + v_raw = self.kvcache.v_cache(layer_id) + kv_heads, head_dim = k_raw.shape[-2:] + spec = attn_spec or AttentionSpec() + if spec.sinks is not None: + raise RuntimeError("Torch XPU attention does not support attention sinks.") + return paged_attention( + q, + k_raw.view(-1, kv_heads, head_dim), + v_raw.view(-1, kv_heads, head_dim), + metadata.indices, + query_lens=tuple(req.extend_len for req in batch.padded_reqs), + kv_lens=tuple(req.device_len for req in batch.padded_reqs), + query_positions=metadata.q_positions, + scale=spec.sm_scale if spec.sm_scale is not None else head_dim**-0.5, + sliding_window=spec.sliding_window, + ) + + +__all__ = ["TorchAttentionBackend"] diff --git a/python/freetoken/benchmark/xpu.py b/python/freetoken/benchmark/xpu.py new file mode 100644 index 000000000..6524f2750 --- /dev/null +++ b/python/freetoken/benchmark/xpu.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import argparse +import json +import os +import time +from collections.abc import Callable +from dataclasses import asdict, dataclass + +import torch +from torch.nn import functional + + +@dataclass(frozen=True, slots=True) +class XpuBenchmarkResult: + device: str + torch_version: str + level_zero_v2: str + queue_mode: str + dtype: str + memory_mib: int + prefill_tokens: int + decode_context: int + heads: int + head_dim: int + warmup: int + iterations: int + memory_copy_gbps: float + prefill_tokens_per_second: float + decode_steps_per_second: float + + +def _positive_int(raw: str) -> int: + value = int(raw) + if value < 1: + raise argparse.ArgumentTypeError("expected a positive integer") + return value + + +def _nonnegative_int(raw: str) -> int: + value = int(raw) + if value < 0: + raise argparse.ArgumentTypeError("expected a non-negative integer") + return value + + +def _elapsed( + operation: Callable[[], None], + synchronize: Callable[[], None], + iterations: int, +) -> float: + synchronize() + started = time.perf_counter() + for _ in range(iterations): + operation() + synchronize() + return time.perf_counter() - started + + +@torch.inference_mode() +def benchmark_xpu_workloads( + *, + device: torch.device, + synchronize: Callable[[], None], + device_name: str, + dtype: torch.dtype, + memory_mib: int, + prefill_tokens: int, + decode_context: int, + heads: int, + head_dim: int, + warmup: int, + iterations: int, +) -> XpuBenchmarkResult: + element_size = torch.empty((), dtype=dtype).element_size() + elements = memory_mib * (1 << 20) // element_size + source = torch.randn(elements, dtype=dtype, device=device) + destination = torch.empty_like(source) + + prefill_q = torch.randn( + 1, heads, prefill_tokens, head_dim, dtype=dtype, device=device + ) + prefill_k = torch.randn_like(prefill_q) + prefill_v = torch.randn_like(prefill_q) + decode_q = torch.randn(1, heads, 1, head_dim, dtype=dtype, device=device) + decode_k = torch.randn( + 1, heads, decode_context, head_dim, dtype=dtype, device=device + ) + decode_v = torch.randn_like(decode_k) + + def copy_memory() -> None: + destination.copy_(source) + + def prefill_attention() -> None: + functional.scaled_dot_product_attention( + prefill_q, prefill_k, prefill_v, is_causal=True + ) + + def decode_attention() -> None: + functional.scaled_dot_product_attention(decode_q, decode_k, decode_v) + + for _ in range(warmup): + copy_memory() + prefill_attention() + decode_attention() + + memory_seconds = _elapsed(copy_memory, synchronize, iterations) + prefill_seconds = _elapsed(prefill_attention, synchronize, iterations) + decode_seconds = _elapsed(decode_attention, synchronize, iterations) + return XpuBenchmarkResult( + device=device_name, + torch_version=torch.__version__, + level_zero_v2=os.getenv("SYCL_UR_USE_LEVEL_ZERO_V2", "default"), + queue_mode=os.getenv( + "SYCL_PI_LEVEL_ZERO_USE_IMMEDIATE_COMMANDLISTS", "default" + ), + dtype=str(dtype).removeprefix("torch."), + memory_mib=memory_mib, + prefill_tokens=prefill_tokens, + decode_context=decode_context, + heads=heads, + head_dim=head_dim, + warmup=warmup, + iterations=iterations, + memory_copy_gbps=(memory_mib * (1 << 20) * iterations) / memory_seconds / 1e9, + prefill_tokens_per_second=prefill_tokens * iterations / prefill_seconds, + decode_steps_per_second=iterations / decode_seconds, + ) + + +def format_result(result: XpuBenchmarkResult, *, as_json: bool) -> str: + if as_json: + return json.dumps(asdict(result), sort_keys=True) + return "\n".join( + ( + f"Device: {result.device}", + f"PyTorch: {result.torch_version}", + f"Level Zero V2: {result.level_zero_v2}", + f"Level Zero immediate command lists: {result.queue_mode}", + f"Memory copy: {result.memory_copy_gbps:.2f} GB/s", + f"Prefill SDPA: {result.prefill_tokens_per_second:.1f} tokens/s", + f"Decode SDPA: {result.decode_steps_per_second:.1f} steps/s", + ) + ) + + +def main(argv: list[str] | None = None, prog: str = "ft bench xpu") -> int: + parser = argparse.ArgumentParser(prog=prog) + parser.add_argument("--memory-mib", type=_positive_int, default=256) + parser.add_argument("--prefill-tokens", type=_positive_int, default=512) + parser.add_argument("--decode-context", type=_positive_int, default=4096) + parser.add_argument("--heads", type=_positive_int, default=16) + parser.add_argument("--head-dim", type=_positive_int, default=128) + parser.add_argument("--warmup", type=_nonnegative_int, default=3) + parser.add_argument("--iterations", type=_positive_int, default=10) + parser.add_argument("--json", action="store_true") + args = parser.parse_args(argv) + if not hasattr(torch, "xpu") or not torch.xpu.is_available(): + parser.error("Intel XPU is not available in this PyTorch environment") + device = torch.device("xpu:0") + result = benchmark_xpu_workloads( + device=device, + synchronize=torch.xpu.synchronize, + device_name=torch.xpu.get_device_name(0), + dtype=torch.bfloat16, + memory_mib=args.memory_mib, + prefill_tokens=args.prefill_tokens, + decode_context=args.decode_context, + heads=args.heads, + head_dim=args.head_dim, + warmup=args.warmup, + iterations=args.iterations, + ) + print(format_result(result, as_json=args.json)) + return 0 + + +__all__ = ["XpuBenchmarkResult", "benchmark_xpu_workloads", "format_result", "main"] diff --git a/python/freetoken/cli.py b/python/freetoken/cli.py index 4e6deff23..d8d14875b 100644 --- a/python/freetoken/cli.py +++ b/python/freetoken/cli.py @@ -17,6 +17,7 @@ def _print_help(file: TextIO) -> None: launch Configure and launch an agent against a FreeToken server checkpoint Convert an HF safetensors checkpoint to FTW bench Run a micro-benchmark (e.g. "bench bw" = CPU vs PCIe bandwidth) + hardware Detect CUDA/Intel XPU readiness and print tuning guidance Use "ft --help" for command-specific options. Use "ft --version" to print the FreeToken version.""", @@ -61,12 +62,19 @@ def _run_daemon(argv: list[str]) -> int: return main(argv, prog="ft daemon") +def _run_hardware(argv: list[str]) -> int: + from freetoken.hardware import main + + return main(argv) + + def _print_bench_help(file: TextIO) -> None: print( """usage: ft bench [args] Subcommands: bw Benchmark CPU vs PCIe bandwidth and pick the MoE backend (hybrid/offload) + xpu Benchmark Intel XPU memory, prefill SDPA, and decode SDPA throughput Use "ft bench --help" for subcommand-specific options.""", file=file, @@ -85,6 +93,10 @@ def _run_bench(argv: list[str]) -> int: from freetoken.moe.benchbw import main return main(argv[1:], prog="ft bench bw") + if sub == "xpu": + from freetoken.benchmark.xpu import main + + return main(argv[1:], prog="ft bench xpu") print(f"unknown ft bench subcommand: {sub}", file=sys.stderr) _print_bench_help(sys.stderr) return 2 @@ -98,6 +110,7 @@ def _run_bench(argv: list[str]) -> int: "launch": "_run_launch", "checkpoint": "_run_checkpoint", "bench": "_run_bench", + "hardware": "_run_hardware", } diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index 4f2025025..a6082b19b 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Dict, List import torch +from freetoken.accelerator import create_accelerator_graph, runtime_for_device from freetoken.core import Batch, Req, get_global_ctx from freetoken.distributed import get_tp_info from freetoken.utils import init_logger, mem_GB @@ -88,7 +89,7 @@ def _determine_cuda_graph_bs( def get_free_memory(device: torch.device) -> int: - return torch.cuda.mem_get_info(device)[0] + return runtime_for_device(device.type).mem_get_info(device)[0] class GraphRunner: @@ -118,6 +119,7 @@ def __init__( self.moe_offload_cache = moe_offload_cache self.stream = stream self.device = device + self.accelerator = runtime_for_device(device.type) self._capture_graphs(max_seq_len, vocab_size, model) def _reset_moe_offload_cache(self) -> None: @@ -131,15 +133,15 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel # reads it as an indeterminate phase and animates the bar. Must precede the # graphs-disabled early return so that config gets the phase too. emit_progress("Capturing CUDA graphs / warming up", 0, 0) - self.graph_map: Dict[int, torch.cuda.CUDAGraph] = {} + self.graph_map = {} 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) - torch.cuda.synchronize(self.device) - torch.cuda.empty_cache() - torch.cuda.reset_peak_memory_stats(self.device) + self.accelerator.synchronize(self.device) + self.accelerator.empty_cache() + self.accelerator.reset_peak_memory_stats(self.device) logger.info_rank0(f"Start capturing CUDA graphs with sizes: {self.graph_bs_list}") free_memory = get_free_memory(self.device) @@ -159,7 +161,7 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel 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() + graph = create_accelerator_graph(self.device.type, self.accelerator) batch = Batch(reqs=[self.dummy_req] * bs, phase="decode") batch.padded_reqs = batch.reqs self.attn_backend.prepare_for_capture(batch) @@ -175,7 +177,7 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel self.buffer.logits[:bs] = model.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): + with self.accelerator.graph(graph, pool=pool, stream=self.stream): self.buffer.logits[:bs] = model.forward() self._reset_moe_offload_cache() if pool is None: diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index 01d14b1aa..0545b5911 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -18,6 +18,8 @@ class BatchSamplingArgs: def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + if device.type != "cuda": + return torch.tensor(data, dtype=dtype, device=device) return torch.tensor(data, dtype=dtype, pin_memory=True).to(device, non_blocking=True) @@ -27,6 +29,10 @@ def sample_impl( top_k: torch.Tensor | int | None, top_p: torch.Tensor | float | None, ) -> torch.Tensor: + if logits.device.type != "cuda": + from freetoken.kernel.torch_ops import sample + + return sample(logits, temperatures, top_k, top_p) from freetoken.kernel.backend import is_flashinfer_installed if is_flashinfer_installed(): @@ -74,7 +80,6 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: @nvtx_annotate("Sampler") def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: - with torch.cuda.nvtx.range("Sampler"): - if args.temperatures is None: # greedy sampling - return torch.argmax(logits, dim=-1) - return sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) + if args.temperatures is None: # greedy sampling + return torch.argmax(logits, dim=-1) + return sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) diff --git a/python/freetoken/hardware.py b/python/freetoken/hardware.py new file mode 100644 index 000000000..eaacd14ab --- /dev/null +++ b/python/freetoken/hardware.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol + +from freetoken.accelerator import AcceleratorKind, select_accelerator + + +class DeviceRuntime(Protocol): + def is_available(self) -> bool: ... + + def get_device_name(self, device: int = 0) -> str: ... + + +@dataclass(frozen=True) +class HardwareReport: + accelerator: AcceleratorKind + name: str + ready: bool + recommendations: tuple[str, ...] + + +def inspect_hardware( + requested: str, + cuda: DeviceRuntime, + xpu: DeviceRuntime, +) -> HardwareReport: + kind = select_accelerator(requested, cuda.is_available(), xpu.is_available()) + match kind: + case AcceleratorKind.CUDA: + return HardwareReport(kind, cuda.get_device_name(), True, ()) + case AcceleratorKind.XPU: + name = xpu.get_device_name() + is_b60 = "arc" in name.lower() and "b60" in name.lower() + recommendations = ( + "Enable Resizable BAR in firmware.", + "Use current Intel compute drivers and the PyTorch XPU wheel.", + "Keep the Arc B-series Level Zero V2 default; benchmark the legacy adapter only if needed.", + ) + return HardwareReport(kind, name, is_b60, recommendations) + + +def format_hardware_report(report: HardwareReport) -> str: + status = ( + "dense Llama eager path available" + if report.ready + else "driver detected; device model not validated" + ) + lines = [ + f"Accelerator: {report.accelerator.value}", + f"Device: {report.name}", + f"Intel Arc Pro B60 status: {status}" + if report.accelerator is AcceleratorKind.XPU + else f"Status: {status}", + ] + lines.extend(f"Tuning: {item}" for item in report.recommendations) + return "\n".join(lines) + + +def main(argv: list[str]) -> int: + import argparse + + import torch + + parser = argparse.ArgumentParser(prog="ft hardware") + parser.add_argument( + "--accelerator", choices=("auto", "cuda", "xpu"), default="auto" + ) + args = parser.parse_args(argv) + report = inspect_hardware(args.accelerator, torch.cuda, torch.xpu) + print(format_hardware_report(report)) + return 0 diff --git a/python/freetoken/layers/activation.py b/python/freetoken/layers/activation.py index 93602b6c5..f3c6cff4e 100644 --- a/python/freetoken/layers/activation.py +++ b/python/freetoken/layers/activation.py @@ -7,6 +7,10 @@ def silu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None): + if x.device.type != "cuda": + from freetoken.kernel.torch_ops import silu_and_mul + + return silu_and_mul(x, out=out) from freetoken.kernel.backend import is_flashinfer_installed if is_flashinfer_installed(): diff --git a/python/freetoken/layers/embedding.py b/python/freetoken/layers/embedding.py index 76cd759bb..1a6a751ce 100644 --- a/python/freetoken/layers/embedding.py +++ b/python/freetoken/layers/embedding.py @@ -38,13 +38,23 @@ def __init__( @nvtx_annotate("Embedding") def forward(self, x: torch.Tensor) -> torch.Tensor: - from freetoken.kernel import indexing + if x.device.type != "cuda": + if self.tp_size > 1: + start, length = self.vocab_range + local = x - start + valid = (local >= 0) & (local < length) + y = F.embedding(local.clamp(0, length - 1), self.weight) + y.masked_fill_(~valid.unsqueeze(-1), 0) + else: + y = F.embedding(x.to(torch.long), self.weight) + else: + from freetoken.kernel import indexing - y = indexing( - weights=self.weight, - indices=x, - vocab_range=self.vocab_range if self.tp_size > 1 else None, - ) + y = indexing( + weights=self.weight, + indices=x, + vocab_range=self.vocab_range if self.tp_size > 1 else None, + ) if self.tp_size > 1: y = self._comm.all_reduce(y) @@ -122,4 +132,4 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: output_tensor = output_tensor.view((self.tp_size,) + input_shape) output_tensor = output_tensor.permute(1, 0, 2).contiguous() output_tensor = output_tensor.reshape(input_shape[:1] + (self.tp_size * input_shape[1],)) - return output_tensor[:, : self.num_embeddings] \ No newline at end of file + return output_tensor[:, : self.num_embeddings] diff --git a/python/freetoken/layers/norm.py b/python/freetoken/layers/norm.py index df248136f..8bbef7b03 100644 --- a/python/freetoken/layers/norm.py +++ b/python/freetoken/layers/norm.py @@ -9,7 +9,9 @@ class RMSNorm(BaseOP): def __init__(self, size: int, eps: float) -> None: from freetoken.kernel.backend import is_flashinfer_installed - if is_flashinfer_installed(): + if getattr(torch.version, "xpu", None) is not None: + from freetoken.kernel.torch_ops import rmsnorm + elif is_flashinfer_installed(): from flashinfer import rmsnorm else: from freetoken.kernel.triton.norm import rmsnorm @@ -19,9 +21,18 @@ def __init__(self, size: int, eps: float) -> None: self.rmsnorm = rmsnorm def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.device.type != "cuda": + from freetoken.kernel.torch_ops import rmsnorm + + return rmsnorm(x, self.weight, self.eps) return self.rmsnorm(x, self.weight, self.eps) def forward_inplace(self, x: torch.Tensor) -> None: + if x.device.type != "cuda": + from freetoken.kernel.torch_ops import rmsnorm + + rmsnorm(x, self.weight, self.eps, out=x) + return self.rmsnorm(x, self.weight, self.eps, out=x) @@ -151,7 +162,9 @@ class RMSNormFused(BaseOP): def __init__(self, size: int, eps: float) -> None: from freetoken.kernel.backend import is_flashinfer_installed - if is_flashinfer_installed(): + if getattr(torch.version, "xpu", None) is not None: + from freetoken.kernel.torch_ops import fused_add_rmsnorm, rmsnorm + elif is_flashinfer_installed(): from flashinfer import fused_add_rmsnorm, rmsnorm else: from freetoken.kernel.triton.norm import fused_add_rmsnorm, rmsnorm @@ -164,6 +177,13 @@ def __init__(self, size: int, eps: float) -> None: def forward( self, x: torch.Tensor, residual: torch.Tensor | None = None ) -> Tuple[torch.Tensor, torch.Tensor]: + if x.device.type != "cuda": + from freetoken.kernel.torch_ops import fused_add_rmsnorm, rmsnorm + + if residual is None: + return rmsnorm(x, self.weight, self.eps), x + fused_add_rmsnorm(x, residual, self.weight, self.eps) + return x, residual if residual is None: return self.rmsnorm(x, self.weight, self.eps), x self.fused_add_rmsnorm(x, residual, self.weight, self.eps) diff --git a/python/freetoken/layers/rotary.py b/python/freetoken/layers/rotary.py index 3756a299b..0ea4b8444 100644 --- a/python/freetoken/layers/rotary.py +++ b/python/freetoken/layers/rotary.py @@ -61,7 +61,11 @@ def __init__( from freetoken.kernel.backend import is_flashinfer_installed - if is_flashinfer_installed(): + if getattr(torch.version, "xpu", None) is not None: + from freetoken.kernel.torch_ops import apply_rope_inplace + + apply_rope_with_cos_sin_cache_inplace = apply_rope_inplace + elif is_flashinfer_installed(): from flashinfer import apply_rope_with_cos_sin_cache_inplace else: from freetoken.kernel.triton.rope import apply_rope_with_cos_sin_cache_inplace @@ -74,6 +78,18 @@ def forward( query: torch.Tensor, key: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: + if query.device.type != "cuda": + from freetoken.kernel.torch_ops import apply_rope_inplace + + apply_rope_inplace( + positions, + query, + key, + self.head_size, + self._cos_sin_cache, + is_neox=self.is_neox, + ) + return query, key self.apply_rope_with_cos_sin_cache_inplace( positions=positions, query=query, diff --git a/python/freetoken/llm/llm.py b/python/freetoken/llm/llm.py index 1d4aaef0a..325d46000 100644 --- a/python/freetoken/llm/llm.py +++ b/python/freetoken/llm/llm.py @@ -129,7 +129,9 @@ def generate( self.mm_embeds_map[uid] = self.encode_images( mm["pixel_values"], mm["image_position_ids"] ) - torch.cuda.synchronize(self.device) + from freetoken.accelerator import runtime_for_device + + runtime_for_device(self.device.type).synchronize(self.device) try: self.run_forever() except RequestAllFinished: diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index f6105e1f8..315a8f2b8 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -243,6 +243,7 @@ class ModelConfig: # it bf16. Separate from dense_quant because only some NVFP4 checkpoints quantize lm_head # (modelopt MIXED_PRECISION does; pure NVFP4 leaves it bf16). lm_head_quant: str = "none" + checkpoint_quantization: str = "none" shared_expert_intermediate_size: int = 0 use_qk_norm: bool = False # ----- DeepSeek/GLM-style MoE extensions (default keeps other models intact) ----- diff --git a/python/freetoken/models/llama/config.py b/python/freetoken/models/llama/config.py index c8cc633a2..bd1768c66 100644 --- a/python/freetoken/models/llama/config.py +++ b/python/freetoken/models/llama/config.py @@ -5,6 +5,19 @@ from freetoken.models.config import ModelConfig, RotaryConfig +def _checkpoint_quantization(hf_config: Any) -> str: + quantization = getattr(hf_config, "quantization_config", None) + if quantization is None: + return "none" + if isinstance(quantization, dict): + method = quantization.get("quant_method") or quantization.get("quant_algo") + else: + method = getattr(quantization, "quant_method", None) or getattr( + quantization, "quant_algo", None + ) + return str(method or "quantized").lower() + + def parse_config(hf_config: Any) -> ModelConfig: num_kv_heads = getattr(hf_config, "num_key_value_heads", hf_config.num_attention_heads) head_dim = ( @@ -42,6 +55,7 @@ def parse_config(hf_config: Any) -> ModelConfig: norm_topk_prob=False, model_type=getattr(hf_config, "model_type", "llama"), architectures=getattr(hf_config, "architectures", ["LlamaForCausalLM"]), + checkpoint_quantization=_checkpoint_quantization(hf_config), ) diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index debea2823..5ea1fd3a2 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -6,7 +6,7 @@ from typing import Iterator import torch -from flashlib.kernels.slot_cache import N_STATS, Stat +from freetoken.moe.offload_stats import N_STATS, Stat # Fuse the per-bank expert copies into a single multi-bank launch (one per copy_missing # instead of one per bank). Set FREETOKEN_FUSED_COPY=0 to force the legacy per-bank path diff --git a/python/freetoken/moe/offload_stats.py b/python/freetoken/moe/offload_stats.py new file mode 100644 index 000000000..58fd80080 --- /dev/null +++ b/python/freetoken/moe/offload_stats.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +from enum import IntEnum + + +class Stat(IntEnum): + ACTIVE = 0 + MISS = 1 + CALLS = 2 + + +N_STATS = len(Stat) diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 355411617..a16f86734 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -65,9 +65,10 @@ def __init__(self, config: SchedulerConfig): # use another stream to overlap metadata processing with computation self.device = self.engine.device - self.stream = torch.cuda.Stream(device=self.device) - self.engine_stream_ctx = torch.cuda.stream(self.engine.stream) - torch.cuda.set_stream(self.stream) + self.accelerator = self.engine.accelerator + self.stream = self.accelerator.Stream(device=self.device) + self.engine_stream_ctx = self.accelerator.stream(self.engine.stream) + self.accelerator.set_stream(self.stream) # initialize other managers self.table_manager = TableManager(config.max_running_req, self.engine.page_table) @@ -160,7 +161,7 @@ def rebuild_cache( """ assert not self.prefill_manager.runnable, "rebuild requires no pending prefill" assert not self.decode_manager.runnable, "rebuild requires no running decode" - torch.cuda.synchronize(self.device) + self.accelerator.synchronize(self.device) if self.config.tp_info.size > 1: self.sync_all_ranks() self.engine.rebuild_runtime_cache( @@ -286,13 +287,13 @@ def run_forever(self) -> NoReturn: while True: self.normal_loop() else: - assert torch.cuda.current_stream() == self.stream + assert self.accelerator.current_stream() == self.stream data = None while True: data = self.overlap_loop(data) def shutdown(self) -> None: - torch.cuda.synchronize(self.device) + self.accelerator.synchronize(self.device) self.sync_all_ranks() self.engine.shutdown() @@ -467,9 +468,7 @@ def _swa_token_usage(self) -> Tuple[int, int] | None: def _gpu_mem_bytes(self) -> int: """Bytes this engine process holds on the GPU (torch's reserved caching-allocator pool: weights + KV + MoE cache + graphs). 0 on CPU. Cheap, no device sync.""" - if self.device.type != "cuda": - return 0 - return torch.cuda.memory_reserved(self.device) + return self.accelerator.memory_reserved(self.device) def _process_one_msg(self, msg: BaseBackendMsg) -> None: if isinstance(msg, BatchBackendMsg): diff --git a/python/freetoken/utils/torch_utils.py b/python/freetoken/utils/torch_utils.py index 9422b9e7d..938653ed2 100644 --- a/python/freetoken/utils/torch_utils.py +++ b/python/freetoken/utils/torch_utils.py @@ -21,6 +21,7 @@ def torch_dtype(dtype: torch.dtype): def nvtx_annotate(name: str, layer_id_field: str | None = None): + import torch import torch.cuda.nvtx as nvtx def decorator(fn): @@ -29,6 +30,11 @@ def wrapper(self, *args, **kwargs): display_name = name if layer_id_field and hasattr(self, layer_id_field): display_name = name.format(getattr(self, layer_id_field)) + device = getattr(self, "device", None) + if not torch.cuda.is_available() or ( + device is not None and getattr(device, "type", "cuda") != "cuda" + ): + return fn(self, *args, **kwargs) with nvtx.range(display_name): return fn(self, *args, **kwargs) diff --git a/requirements-xpu.txt b/requirements-xpu.txt new file mode 100644 index 000000000..16d01c00a --- /dev/null +++ b/requirements-xpu.txt @@ -0,0 +1,21 @@ +--extra-index-url https://download.pytorch.org/whl/xpu + +apache-tvm-ffi==0.1.13.post3 +einops>=0.8,<1 +fastapi>=0.115,<1 +gguf>=0.19,<1 +huggingface_hub>=1.5,<2 +msgpack>=1.1,<2 +modelscope>=1.37,<2 +numpy>=2.0,<2.5 +openai>=2.0,<3 +partial-json-parser>=0.2,<1 +prompt_toolkit>=3.0,<4 +pydantic>=2.9,<3 +pyzmq>=27,<28 +safetensors>=0.6,<1 +torch==2.11.0+xpu +tqdm>=4.66,<5 +transformers>=5.5,<6 +triton-xpu==3.7.0 +uvicorn>=0.30,<1 diff --git a/setup.py b/setup.py index cfe41b7d8..cfd0c5320 100644 --- a/setup.py +++ b/setup.py @@ -1,6 +1,7 @@ from __future__ import annotations import importlib.util +import os from pathlib import Path from setuptools import setup @@ -31,12 +32,14 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: return [str(cuda_home / "include")], library_dirs -cuda_include_dirs, cuda_library_dirs = _cuda_runtime_paths() -_check_toolchain() +build_xpu = os.environ.get("FREETOKEN_ACCELERATOR") == "xpu" +cuda_include_dirs, cuda_library_dirs = ([], []) if build_xpu else _cuda_runtime_paths() +if not build_xpu: + _check_toolchain() setup( - ext_modules=[ + ext_modules=[] if build_xpu else [ CppExtension( name="freetoken.kernel._pinned_tensor", sources=[ diff --git a/tests/models/test_llama_config.py b/tests/models/test_llama_config.py new file mode 100644 index 000000000..5472ccee1 --- /dev/null +++ b/tests/models/test_llama_config.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from types import SimpleNamespace + +import pytest + +from freetoken.accelerator import AcceleratorKind, apply_engine_accelerator_constraints +from freetoken.models.llama.config import parse_config + + +@dataclass(frozen=True) +class FakeTPInfo: + size: int = 1 + + +@dataclass(frozen=True) +class FakeEngineConfig: + attention_backend: str = "auto" + use_pynccl: bool = True + moe_prefill_hit_d2d: bool = True + cuda_graph_max_bs: int | None = None + tp_info: FakeTPInfo = field(default_factory=FakeTPInfo) + + +def _hf_config(quantization_config=None) -> SimpleNamespace: + return SimpleNamespace( + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + hidden_size=64, + vocab_size=32, + intermediate_size=128, + rms_norm_eps=1e-5, + max_position_embeddings=128, + hidden_act="silu", + architectures=["LlamaForCausalLM"], + model_type="llama", + quantization_config=quantization_config, + ) + + +@pytest.mark.parametrize("method", ["awq", "gptq"]) +def test_xpu_rejects_quantized_llama_metadata(method: str) -> None: + # Given: real Llama parser input carrying checkpoint quantization metadata. + model_config = parse_config(_hf_config({"quant_method": method})) + # When/Then: XPU admission rejects it before model allocation. + with pytest.raises(RuntimeError, match="unquantized"): + apply_engine_accelerator_constraints( + FakeEngineConfig(), + AcceleratorKind.XPU, + model_config=model_config, + is_moe=False, + ) + + +def test_xpu_accepts_unquantized_llama_metadata() -> None: + # Given: real Llama parser input without quantization metadata. + model_config = parse_config(_hf_config()) + # When: XPU admission evaluates the parsed model. + config = FakeEngineConfig() + apply_engine_accelerator_constraints( + config, + AcceleratorKind.XPU, + model_config=model_config, + is_moe=False, + ) + # Then: the portable attention path is selected. + assert config.attention_backend == "torch" diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 3caab91b9..7aa1ff92f 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -404,18 +404,23 @@ def test_lru_gpu_cache_assigns_unique_slots_for_large_miss_batch(): assert cache.src_indices[:256].tolist() == list(range(256)) -def test_adjust_config_converts_moe_cache_rate_to_cache_size(): +def test_adjust_config_converts_moe_cache_rate_to_cache_size(monkeypatch): from types import SimpleNamespace from freetoken.distributed import DistributedInfo from freetoken.engine.config import EngineConfig from freetoken.engine.engine import _adjust_config + monkeypatch.setattr( + "freetoken.engine.engine.accelerator_runtime", + lambda _kind: SimpleNamespace(get_device_name=lambda _device: "test gpu"), + ) + config = EngineConfig( model_path="/tmp/freetoken-test-model", tp_info=DistributedInfo(rank=0, size=1), dtype=torch.float16, - attention_backend="fi", + attention_backend="triton", moe_cache_rate=0.3, ) object.__setattr__( @@ -443,6 +448,8 @@ def test_adjust_config_converts_moe_cache_rate_to_cache_size(): def test_graph_capture_reuses_warm_offload_cache_before_capture(monkeypatch): + from types import SimpleNamespace + import freetoken.core as core from freetoken.core import Context, Req, get_global_ctx from freetoken.engine.graph import GraphRunner @@ -478,11 +485,17 @@ class FakeOffloadCache: def reset(self): events.append("reset") - monkeypatch.setattr("torch.cuda.CUDAGraph", FakeGraph) - monkeypatch.setattr("torch.cuda.graph", fake_cuda_graph) - monkeypatch.setattr("torch.cuda.synchronize", lambda device=None: None) - monkeypatch.setattr("torch.cuda.empty_cache", lambda: None) - monkeypatch.setattr("torch.cuda.reset_peak_memory_stats", lambda device=None: None) + runtime = SimpleNamespace( + synchronize=lambda device=None: None, + empty_cache=lambda: None, + reset_peak_memory_stats=lambda device=None: None, + graph=fake_cuda_graph, + ) + monkeypatch.setattr("freetoken.engine.graph.runtime_for_device", lambda device: runtime) + monkeypatch.setattr( + "freetoken.engine.graph.create_accelerator_graph", + lambda device, selected_runtime: FakeGraph(), + ) monkeypatch.setattr("freetoken.engine.graph.get_free_memory", lambda device: 1024) dummy_req = Req( diff --git a/tests/scheduler/test_cache_rebuild.py b/tests/scheduler/test_cache_rebuild.py index 0c0580708..bbd6c2a6a 100644 --- a/tests/scheduler/test_cache_rebuild.py +++ b/tests/scheduler/test_cache_rebuild.py @@ -143,12 +143,11 @@ def test_rebuild_cache_refreshes_prefill_budget(monkeypatch): from freetoken.scheduler.scheduler import Scheduler - monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) - sched = Scheduler.__new__(Scheduler) sched.prefill_manager = SimpleNamespace(runnable=False) sched.decode_manager = SimpleNamespace(runnable=False) sched.device = torch.device("cpu") + sched.accelerator = SimpleNamespace(synchronize=lambda *args: None) sched.config = SimpleNamespace(tp_info=SimpleNamespace(size=1), max_extend_tokens=100_000) sched.engine = SimpleNamespace( rebuild_runtime_cache=lambda **kw: None, num_pages=32, page_table=None diff --git a/tests/test_accelerator.py b/tests/test_accelerator.py new file mode 100644 index 000000000..44514d5b9 --- /dev/null +++ b/tests/test_accelerator.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from types import ModuleType + +import pytest +from freetoken.accelerator import ( + AcceleratorKind, + apply_engine_accelerator_constraints, + create_accelerator_graph, + distributed_backend, + runtime_for_device, + select_accelerator, + uses_expandable_segments, +) +from freetoken.hardware import format_hardware_report, inspect_hardware + + +class FakeRuntime: + def __init__(self, available: bool, name: str) -> None: + self.available = available + self.name = name + + def is_available(self) -> bool: + return self.available + + def get_device_name(self, device: int = 0) -> str: + return self.name + + +@dataclass(frozen=True) +class FakeTPInfo: + size: int = 1 + + +@dataclass(frozen=True) +class FakeEngineConfig: + attention_backend: str = "auto" + moe_backend: str = "auto" + use_pynccl: bool = True + moe_prefill_hit_d2d: bool = True + cuda_graph_max_bs: int | None = None + tp_info: FakeTPInfo = field(default_factory=lambda: FakeTPInfo()) + + +@dataclass(frozen=True) +class FakeModelConfig: + architectures: tuple[str, ...] = ("LlamaForCausalLM",) + + +@pytest.mark.parametrize( + ("requested", "cuda_available", "xpu_available", "expected"), + [ + ("auto", True, False, AcceleratorKind.CUDA), + ("auto", False, True, AcceleratorKind.XPU), + ("cuda", True, True, AcceleratorKind.CUDA), + ("xpu", True, True, AcceleratorKind.XPU), + ], +) +def test_select_accelerator_when_backend_is_available( + requested: str, + cuda_available: bool, + xpu_available: bool, + expected: AcceleratorKind, +) -> None: + # Given: CUDA/XPU availability and a user selection. + # When: the accelerator is resolved. + actual = select_accelerator(requested, cuda_available, xpu_available) + # Then: explicit selection wins and auto preserves CUDA precedence. + assert actual is expected + + +@pytest.mark.parametrize( + ("requested", "cuda_available", "xpu_available", "message"), + [ + ("auto", False, False, "No supported accelerator"), + ("cuda", False, True, "CUDA was requested"), + ("xpu", True, False, "Intel XPU was requested"), + ("bogus", True, True, "Unsupported accelerator"), + ], +) +def test_select_accelerator_when_backend_is_unavailable( + requested: str, + cuda_available: bool, + xpu_available: bool, + message: str, +) -> None: + # Given: a missing or invalid requested accelerator. + # When/Then: selection fails before model allocation with actionable context. + with pytest.raises(RuntimeError, match=message): + select_accelerator(requested, cuda_available, xpu_available) + + +def test_hardware_report_when_arc_pro_b60_is_available() -> None: + # Given: an XPU runtime exposing the Arc Pro B60 product name. + xpu = FakeRuntime(True, "Intel(R) Arc(TM) Pro B60 Graphics") + # When: readiness is inspected and rendered through the CLI boundary. + report = inspect_hardware("auto", FakeRuntime(False, ""), xpu) + rendered = format_hardware_report(report) + # Then: the report identifies the B60 and its constrained dense-Llama path. + assert report.ready + assert "Accelerator: xpu" in rendered + assert "Intel(R) Arc(TM) Pro B60 Graphics" in rendered + assert "dense Llama eager path available" in rendered + assert "Resizable BAR" in rendered + assert "Level Zero V2" in rendered + + +def test_apply_engine_accelerator_constraints_when_xpu_llama() -> None: + # Given: equivalent dense Llama configurations for CUDA and Intel XPU. + xpu_config = FakeEngineConfig() + cuda_config = FakeEngineConfig() + # When: accelerator constraints are applied. + apply_engine_accelerator_constraints( + xpu_config, + AcceleratorKind.XPU, + model_config=FakeModelConfig(), + is_moe=False, + ) + apply_engine_accelerator_constraints( + cuda_config, + AcceleratorKind.CUDA, + model_config=FakeModelConfig(), + is_moe=False, + ) + # Then: XPU selects the eager backend while CUDA remains unchanged. + assert xpu_config.attention_backend == "torch" + assert xpu_config.cuda_graph_max_bs == 0 + assert not xpu_config.use_pynccl + assert cuda_config == FakeEngineConfig() + + +@pytest.mark.parametrize( + "model_config", + [ + FakeModelConfig(("QwenForCausalLM",)), + FakeModelConfig(("LlamaForCausalLM",)), + ], +) +def test_apply_engine_accelerator_constraints_when_xpu_model_unsupported( + model_config: FakeModelConfig, +) -> None: + # Given: either an unported architecture or an otherwise supported MoE architecture. + is_moe = model_config.architectures == ("LlamaForCausalLM",) + # When/Then: configuration is rejected before model allocation. + with pytest.raises(RuntimeError, match="dense LlamaForCausalLM"): + apply_engine_accelerator_constraints( + FakeEngineConfig(), + AcceleratorKind.XPU, + model_config=model_config, + is_moe=is_moe, + ) + + +def test_apply_engine_accelerator_constraints_when_xpu_tensor_parallel() -> None: + # Given: a dense Llama configuration requesting two XPU ranks. + config = FakeEngineConfig(tp_info=FakeTPInfo(size=2)) + # When/Then: the unvalidated multi-device path is rejected before allocation. + with pytest.raises(RuntimeError, match="tensor parallel size 1"): + apply_engine_accelerator_constraints( + config, + AcceleratorKind.XPU, + model_config=FakeModelConfig(), + is_moe=False, + ) + + +def test_accelerator_dispatch_helpers() -> None: + # Given: a graph-capable runtime with distinct graph constructors. + runtime = ModuleType("fake_runtime") + calls: list[str] = [] + runtime.CUDAGraph = lambda: calls.append("cuda") or "cuda graph" + runtime.XPUGraph = lambda: calls.append("xpu") or "xpu graph" + # When: backend-owned dispatch decisions are resolved. + cuda_graph = create_accelerator_graph("cuda", runtime) + xpu_graph = create_accelerator_graph("xpu", runtime) + # Then: each backend selects only its applicable policy and graph constructor. + assert uses_expandable_segments(AcceleratorKind.CUDA) + assert not uses_expandable_segments(AcceleratorKind.XPU) + assert distributed_backend(AcceleratorKind.CUDA) == "nccl" + assert distributed_backend(AcceleratorKind.XPU) == "xccl" + assert (cuda_graph, xpu_graph, calls) == ( + "cuda graph", + "xpu graph", + ["cuda", "xpu"], + ) + + +def test_runtime_for_device_when_unsupported() -> None: + # Given: a non-accelerator device type. + # When: its runtime is requested. + with pytest.raises(RuntimeError) as exc_info: + runtime_for_device("cpu") + # Then: the established boundary error remains unchanged. + assert str(exc_info.value) == "Unsupported accelerator device type 'cpu'." diff --git a/tests/test_xpu_benchmark.py b/tests/test_xpu_benchmark.py new file mode 100644 index 000000000..3de45bd1e --- /dev/null +++ b/tests/test_xpu_benchmark.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +import json + +import pytest +import torch +from freetoken.benchmark.xpu import benchmark_xpu_workloads, format_result, main + + +def test_xpu_workload_benchmark_reports_all_hot_paths() -> None: + # Given: a tiny CPU reference workload using the same PyTorch operations as XPU. + # When: memory, prefill, and decode benchmarks execute. + result = benchmark_xpu_workloads( + device=torch.device("cpu"), + synchronize=lambda: None, + device_name="reference", + dtype=torch.float32, + memory_mib=1, + prefill_tokens=16, + decode_context=32, + heads=2, + head_dim=8, + warmup=0, + iterations=1, + ) + # Then: every performance signal is finite and positive. + assert result.memory_copy_gbps > 0 + assert result.prefill_tokens_per_second > 0 + assert result.decode_steps_per_second > 0 + + +def test_xpu_benchmark_json_is_machine_comparable() -> None: + # Given: a deterministic benchmark result. + result = benchmark_xpu_workloads( + device=torch.device("cpu"), + synchronize=lambda: None, + device_name="reference", + dtype=torch.float32, + memory_mib=1, + prefill_tokens=8, + decode_context=16, + heads=1, + head_dim=8, + warmup=0, + iterations=1, + ) + # When: the result is rendered as JSON. + payload = json.loads(format_result(result, as_json=True)) + # Then: environment and workload geometry accompany the measurements. + assert payload["device"] == "reference" + assert payload["queue_mode"] in {"default", "0", "1"} + assert payload["level_zero_v2"] in {"default", "0", "1"} + assert payload["prefill_tokens"] == 8 + assert payload["decode_context"] == 16 + assert payload["warmup"] == 0 + + +def test_xpu_benchmark_rejects_zero_iterations(capsys) -> None: + # Given: a benchmark invocation with no measurable iterations. + # When/Then: argument parsing rejects it before probing hardware. + with pytest.raises(SystemExit): + main(["--iterations", "0"]) + assert "positive integer" in capsys.readouterr().err + + +def test_xpu_benchmark_rejects_negative_iterations(capsys) -> None: + # Given: a benchmark invocation with a negative iteration count. + # When/Then: argument parsing rejects it before probing hardware. + with pytest.raises(SystemExit): + main(["--iterations", "-1"]) + assert "positive integer" in capsys.readouterr().err diff --git a/tests/test_xpu_llama_eager.py b/tests/test_xpu_llama_eager.py new file mode 100644 index 000000000..f24aadcd8 --- /dev/null +++ b/tests/test_xpu_llama_eager.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import sys + +import torch + +import freetoken.core as core +import freetoken.distributed.info as distributed_info +from freetoken.attention.torch_native import TorchAttentionBackend +from freetoken.core import Batch, Context, Req, SamplingParams +from freetoken.distributed import set_tp_info +from freetoken.kvcache.mha_pool import MHAKVCache +from freetoken.layers.rotary import set_rope_device +from freetoken.models.config import ModelConfig, RotaryConfig +from freetoken.models.llama.model import LlamaForCausalLM + + +def test_dense_llama_prefill_and_decode_run_through_portable_backend(monkeypatch) -> None: + # Given: a one-layer dense Llama model, paged KV cache, and two-token request. + monkeypatch.setattr(distributed_info, "_TP_INFO", None) + monkeypatch.setattr(core, "_GLOBAL_CTX", None) + set_tp_info(0, 1) + config = ModelConfig( + num_layers=1, + num_qo_heads=1, + num_kv_heads=1, + head_dim=64, + hidden_size=64, + vocab_size=32, + intermediate_size=128, + rms_norm_eps=1e-5, + rotary_config=RotaryConfig(64, 64, 32, 10_000.0, None), + hidden_act="silu", + tie_word_embeddings=False, + num_experts=0, + num_experts_per_tok=0, + moe_intermediate_size=0, + norm_topk_prob=False, + model_type="llama", + architectures=["LlamaForCausalLM"], + ) + ctx = Context(page_size=1) + ctx.kv_cache = MHAKVCache(1, 1, 64, 4, 1, torch.float32, torch.device("cpu")) + ctx.page_table = torch.zeros((2, 4), dtype=torch.int32) + ctx.page_table[0, :2] = torch.tensor([0, 1], dtype=torch.int32) + core.set_global_ctx(ctx) + ctx.attn_backend = TorchAttentionBackend(config) + req = Req( + input_ids=torch.tensor([1, 2], dtype=torch.int32), + table_idx=0, + cached_len=0, + output_len=1, + uid=1, + sampling_params=SamplingParams(), + cache_handle=None, + ) + batch = Batch([req], "prefill") + batch.padded_reqs = batch.reqs + batch.input_ids = req.input_ids + batch.positions = torch.tensor([0, 1], dtype=torch.int64) + batch.out_loc = torch.tensor([0, 1], dtype=torch.int32) + ctx.attn_backend.prepare_metadata(batch) + set_rope_device(torch.device("cpu")) + model = LlamaForCausalLM(config) + for tensor in model.state_dict().values(): + tensor.normal_(mean=0.0, std=0.02) + # When: the model executes prefill followed by one cached decode step. + with ctx.forward_batch(batch): + prefill_logits = model.forward() + req.complete_one() + req.append_host(torch.tensor([3], dtype=torch.int32)) + ctx.page_table[0, 2] = 2 + decode = Batch([req], "decode") + decode.padded_reqs = decode.reqs + decode.input_ids = req.input_ids[-1:] + decode.positions = torch.tensor([2], dtype=torch.int64) + decode.out_loc = torch.tensor([2], dtype=torch.int32) + ctx.attn_backend.prepare_metadata(decode) + with ctx.forward_batch(decode): + decode_logits = model.forward() + # Then: both phases produce finite next-token logits without a CUDA kernel. + assert prefill_logits.shape == decode_logits.shape == (1, 32) + assert torch.isfinite(prefill_logits).all() + assert torch.isfinite(decode_logits).all() + assert "freetoken.kernel.triton.norm" not in sys.modules + assert "freetoken.kernel.triton.rope" not in sys.modules diff --git a/tests/test_xpu_torch_ops.py b/tests/test_xpu_torch_ops.py new file mode 100644 index 000000000..620a2d7d1 --- /dev/null +++ b/tests/test_xpu_torch_ops.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import torch + +from freetoken.kernel.torch_ops import ( + apply_rope_inplace, + fused_add_rmsnorm, + paged_attention, + rmsnorm, + sample, + silu_and_mul, + store_cache, +) + + +def test_rmsnorm_matches_reference() -> None: + # Given: a small activation and scale vector. + x = torch.tensor([[1.0, 2.0, 3.0, 4.0]]) + weight = torch.tensor([0.5, 1.0, 1.5, 2.0]) + # When: the portable RMSNorm implementation runs. + actual = rmsnorm(x, weight, 1e-6) + # Then: it matches the direct PyTorch definition. + expected = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + 1e-6) * weight + torch.testing.assert_close(actual, expected) + + +def test_fused_add_rmsnorm_updates_both_buffers() -> None: + # Given: activation and residual buffers. + x = torch.tensor([[1.0, 2.0]]) + residual = torch.tensor([[3.0, 4.0]]) + weight = torch.ones(2) + # When: portable fused residual normalization runs in place. + fused_add_rmsnorm(x, residual, weight, 1e-6) + # Then: the residual owns the sum and x owns its normalized value. + expected_residual = torch.tensor([[4.0, 6.0]]) + torch.testing.assert_close(residual, expected_residual) + torch.testing.assert_close(x, rmsnorm(expected_residual, weight, 1e-6)) + + +def test_silu_and_mul_matches_reference() -> None: + # Given: un-interleaved gate and value halves. + x = torch.tensor([[1.0, -2.0, 3.0, 4.0]]) + # When: the portable gated activation runs. + actual = silu_and_mul(x) + # Then: it matches SiLU(gate) multiplied by the value half. + torch.testing.assert_close(actual, torch.nn.functional.silu(x[:, :2]) * x[:, 2:]) + + +def test_store_cache_uses_tensor_indices() -> None: + # Given: empty paged cache rows and two non-contiguous destinations. + k_cache = torch.zeros(4, 1, 2) + v_cache = torch.zeros(4, 1, 2) + indices = torch.tensor([3, 1]) + k = torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]]) + v = k + 10 + # When: portable cache storage runs. + store_cache(k_cache, v_cache, indices, k, v) + # Then: only the selected rows contain the new keys and values. + torch.testing.assert_close(k_cache[indices], k) + torch.testing.assert_close(v_cache[indices], v) + + +def test_rope_rotates_neox_pairs_in_place() -> None: + # Given: one query/key head and a quarter-turn cache entry. + query = torch.tensor([[1.0, 2.0, 3.0, 4.0]]) + key = query.clone() + cache = torch.tensor([[0.0, 0.0, 1.0, 1.0]]) + # When: portable NeoX RoPE is applied. + apply_rope_inplace(torch.tensor([0]), query, key, 4, cache, is_neox=True) + # Then: the first and second halves rotate as complex pairs. + expected = torch.tensor([[-3.0, -4.0, 1.0, 2.0]]) + torch.testing.assert_close(query, expected) + torch.testing.assert_close(key, expected) + + +def test_sampling_respects_per_row_top_k() -> None: + # Given: deterministic logits and top-k one for every row. + logits = torch.tensor([[1.0, 4.0, 2.0], [5.0, 0.0, 1.0]]) + temperatures = torch.ones(2) + # When: portable sampling runs. + actual = sample(logits, temperatures, torch.ones(2, dtype=torch.int32), None) + # Then: top-k one is equivalent to greedy selection. + torch.testing.assert_close(actual, torch.tensor([1, 0])) + + +def test_paged_attention_applies_causal_mask() -> None: + # Given: two zero queries whose cache values are ten and twenty. + q = torch.zeros(2, 1, 1) + k_cache = torch.zeros(2, 1, 1) + v_cache = torch.tensor([[[10.0]], [[20.0]]]) + # When: portable paged attention processes a two-token prefill. + actual = paged_attention( + q, + k_cache, + v_cache, + torch.tensor([0, 1]), + query_lens=(2,), + kv_lens=(2,), + query_positions=torch.tensor([0, 1]), + scale=1.0, + ) + # Then: token zero sees itself and token one sees both cache values. + torch.testing.assert_close(actual, torch.tensor([[[10.0]], [[15.0]]])) From 3f81bbac968f7173d1129916471db13c7e20bd13 Mon Sep 17 00:00:00 2001 From: Jackson57279 Date: Sun, 6 Sep 2026 15:23:44 -0500 Subject: [PATCH 2/2] feat(kvcache): decouple KV cache dtype from compute dtype; Modal deployment Adds `--kv-cache-dtype` so the paged KV cache can be stored narrower than the compute dtype, and a Modal deployment for M.O.G.-SEC-27B-1M-CTX-NVFP4 built around it. Why: single-stream decode is memory-bandwidth bound. Each token re-reads every active weight plus the entire KV cache, so tok/s ~= bandwidth / bytes-per-token. For the 27B hybrid at 1M context that is 91.6 GB/token with a bf16 cache -- 9.2 TB/s to hit 100 tok/s, which no single GPU has. An fp8 cache halves the KV half of that (65.5 -> 32.8 GB), bringing a single B200 from ~61 to ~95-108 tok/s and putting the target in reach at $6.25/hr. Engine changes: - EngineConfig.kv_cache_dtype (None = follow --dtype, so existing behaviour is bit-identical) with a resolved `kv_dtype` property used as the single source of truth by both the pool allocation and the byte budget, so capacity planning and the real allocation cannot disagree. - spec_kv_bytes_per_token budgets off the KV dtype rather than the compute dtype. - Context carries `compute_dtype`, letting the FlashInfer backend plan `q_data_type` at compute width while `kv_data_type` follows the pool -- queries are never quantized. - MHAKVCache.store_kv narrows K/V on write: store_cache is a templated raw byte-copy and cannot convert. The torch fallback aliases fp8 through uint8 because index_copy_ has no CPU kernel for float8. - BackendInfo.supports_quantized_kv gates the feature. A narrow cache on a backend that assumes q.dtype == cache.dtype is rejected at config time instead of failing inside a kernel launch, and `--attention-backend auto` skips such backends. Measured on L4/Qwen3-0.6B, same 18.44 GiB KV budget: bf16 stores 172,637 tokens at 114,688 B/token; fp8_e5m2 stores 345,271 at 57,344 -- exactly 2x the context, with a greedy continuation byte-identical to bf16 (271/271 chars). fp8_e4m3 is NOT recommended and warns: the pool stores K/V by a straight cast with no per-tensor scale and e4m3 saturates at +/-448, so activations clip and generation degenerates (observed collapsing into repeated punctuation after 14 chars). e5m2 spends the same byte on exponent range instead of mantissa and survives an unscaled cast; being the same width, e4m3 buys no extra bandwidth anyway. Also restores an availability guard in _adjust_config: the accelerator refactor replaced `torch.cuda.get_device_name(0) if torch.cuda.is_available() else None` with an unconditional call, which raised out of a driver probe on CPU-only hosts and broke config-time runs (and the test suite off-GPU). deploy/ contains the Modal app (CPU-built image, volume-cached weights and JIT artifacts, CPU-only download/validate gates before any GPU spend) and capacity.py, the bytes-per-token cost model whose numbers back the claims above. Co-Authored-By: Claude Opus 5 (1M context) --- deploy/__init__.py | 0 deploy/capacity.py | 186 +++++++++++ deploy/modal_app.py | 429 +++++++++++++++++++++++++ docs/modal-deployment.md | 137 ++++++++ python/freetoken/attention/__init__.py | 22 ++ python/freetoken/attention/fi.py | 12 +- python/freetoken/core.py | 4 + python/freetoken/engine/config.py | 18 ++ python/freetoken/engine/engine.py | 119 +++++-- python/freetoken/kernel/torch_ops.py | 162 ++++++++++ python/freetoken/kvcache/base.py | 6 +- python/freetoken/kvcache/mha_pool.py | 14 +- python/freetoken/server/args.py | 37 +++ tests/engine/test_cache_budget.py | 16 +- tests/kvcache/test_kv_cache_dtype.py | 146 +++++++++ 15 files changed, 1273 insertions(+), 35 deletions(-) create mode 100644 deploy/__init__.py create mode 100644 deploy/capacity.py create mode 100644 deploy/modal_app.py create mode 100644 docs/modal-deployment.md create mode 100644 python/freetoken/kernel/torch_ops.py create mode 100644 tests/kvcache/test_kv_cache_dtype.py diff --git a/deploy/__init__.py b/deploy/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/deploy/capacity.py b/deploy/capacity.py new file mode 100644 index 000000000..785a461a7 --- /dev/null +++ b/deploy/capacity.py @@ -0,0 +1,186 @@ +"""Decode-throughput / cost model for M.O.G.-SEC-27B-1M-CTX-NVFP4 on Modal. + +Autoregressive decode at batch size 1 is memory-bandwidth bound, not compute bound: to +emit one token the GPU must stream every active weight plus the whole KV cache through +the SMs exactly once. So + + tokens/sec ~= achievable_HBM_bandwidth / bytes_read_per_token + +and the engineering problem is entirely "shrink bytes_read_per_token". This module keeps +that arithmetic in one auditable place; `modal run deploy/modal_app.py::plan` prints it. + +Numbers for the model come from the checkpoint's own safetensors headers (measured, not +guessed) and its config.json. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +GB = 1e9 + +# --- checkpoint facts (measured from the safetensors headers) ---------------------- +BYTES_TOTAL = 29.390_348_384 * GB +BYTES_VISION = 0.921 * GB # SigLIP-style tower; dead weight for text-only serving +BYTES_LM_HEAD = 2.543 * GB # [248320, 5120] bf16 +BYTES_EMBED = 2.543 * GB # not read during decode (it is a gather, not a matmul) +BYTES_LINEAR_ATTN = 11.124 * GB # 48 GatedDeltaNet layers, left unquantized by the release +BYTES_MLP = 9.626 * GB # 64 layers, NVFP4 +BYTES_SELF_ATTN = 2.632 * GB # 16 full-attention layers; q/k/v bf16, o_proj NVFP4 + +# --- architecture (config.json) ---------------------------------------------------- +N_LAYERS = 64 +N_FULL_ATTN = 16 # layer_types: every 4th layer +N_LINEAR_ATTN = 48 +N_KV_HEADS = 4 +HEAD_DIM = 256 +LIN_V_HEADS, LIN_V_DIM, LIN_K_HEADS, LIN_K_DIM = 48, 128, 16, 128 +CONV_KERNEL = 4 + + +@dataclass(frozen=True) +class Gpu: + name: str + mem_gb: float + bw_tbs: float # peak HBM bandwidth, TB/s + usd_per_hour: float + + +# Modal published pricing (per-second x 3600). +GPUS = [ + Gpu("L40S", 48, 0.864, 0.000542 * 3600), + Gpu("A100-80GB", 80, 2.039, 0.000694 * 3600), + Gpu("RTX PRO 6000", 96, 1.792, 0.000842 * 3600), + Gpu("H100 SXM5", 80, 3.350, 0.001097 * 3600), + Gpu("H200 SXM", 141, 4.800, 0.001261 * 3600), + Gpu("B200", 180, 8.000, 0.001736 * 3600), + Gpu("B300", 288, 8.000, 0.001972 * 3600), +] + +# Fraction of peak HBM bandwidth a well-tuned decode loop actually sustains. 0.80 is +# optimistic-but-real for a fused CUDA-graph decode; 0.70 is the conservative planning +# number used for the headline claim. +BW_EFFICIENCY = 0.70 + + +def kv_bytes_per_token(kv_bits: int) -> float: + """KV bytes per token, counting only the 16 full-attention layers.""" + return 2 * N_KV_HEADS * HEAD_DIM * N_FULL_ATTN * (kv_bits / 8) + + +def linear_state_bytes(ssm_bits: int = 32) -> float: + """GatedDeltaNet recurrent state. Constant in context length -- the whole point.""" + per_layer_ssm = LIN_V_HEADS * LIN_V_DIM * LIN_K_DIM * (ssm_bits / 8) + conv_dim = 2 * LIN_K_HEADS * LIN_K_DIM + LIN_V_HEADS * LIN_V_DIM + per_layer_conv = conv_dim * CONV_KERNEL * 2 + return N_LINEAR_ATTN * (per_layer_ssm + per_layer_conv) + + +def weight_bytes(*, drop_vision: bool = True, lm_head_bits: int = 16, + linear_attn_bits: int = 16) -> float: + """Active weight bytes streamed per decode step. The model is dense, so 'active' + means all of them -- there is no MoE sparsity to exploit here.""" + total = BYTES_TOTAL + if drop_vision: + total -= BYTES_VISION + total -= BYTES_EMBED # embedding lookup is a gather, not a stream + total -= BYTES_LM_HEAD + total += BYTES_LM_HEAD * (lm_head_bits / 16) + total -= BYTES_LINEAR_ATTN + total += BYTES_LINEAR_ATTN * (linear_attn_bits / 16) + return total + + +def bytes_per_token(ctx: int, *, kv_bits: int = 8, lm_head_bits: int = 16, + linear_attn_bits: int = 16, drop_vision: bool = True) -> dict: + w = weight_bytes(drop_vision=drop_vision, lm_head_bits=lm_head_bits, + linear_attn_bits=linear_attn_bits) + kv = kv_bytes_per_token(kv_bits) * ctx + st = linear_state_bytes() + return {"weights": w, "kv": kv, "state": st, "total": w + kv + st} + + +def tps(gpu: Gpu, ctx: int, *, tp: int = 1, efficiency: float = BW_EFFICIENCY, + tp_scaling: float = 0.85, **kw) -> float: + """Predicted decode tok/s. With tensor parallelism each GPU reads its own shard, so + per-GPU bytes fall ~linearly while collectives eat `tp_scaling` of the win.""" + b = bytes_per_token(ctx, **kw)["total"] / tp + eff_bw = gpu.bw_tbs * 1e12 * efficiency * (tp_scaling if tp > 1 else 1.0) + return eff_bw / b + + +def footprint_gb(ctx: int, *, tp: int = 1, workspace_gb: float = 4.0, **kw) -> float: + """Resident VRAM per GPU: weights + KV + state + activation/workspace slack.""" + b = bytes_per_token(ctx, **kw) + resident = b["weights"] + b["kv"] + b["state"] + BYTES_EMBED + return resident / tp / GB + workspace_gb + + +def report(ctx: int = 1_000_000, target_tps: float = 100.0) -> str: + lines: list[str] = [] + A = lines.append + A(f"M.O.G.-SEC-27B-1M-CTX-NVFP4 decode model @ {ctx:,} ctx (target {target_tps:.0f} tok/s)") + A("=" * 94) + A("") + A("Hybrid architecture is what makes this tractable:") + A(f" {N_FULL_ATTN} full-attention layers carry the KV cache; {N_LINEAR_ATTN} GatedDeltaNet layers") + A(f" hold a fixed {linear_state_bytes()/GB:.3f} GB state regardless of context length.") + A(f" A same-size all-full-attention model would need {N_LAYERS/N_FULL_ATTN:.0f}x the KV bandwidth.") + A("") + + configs = [ + ("bf16 KV (today)", dict(kv_bits=16)), + ("fp8 KV", dict(kv_bits=8)), + ("fp8 KV+head", dict(kv_bits=8, lm_head_bits=8)), + ("fp8 all", dict(kv_bits=8, lm_head_bits=8, linear_attn_bits=8)), + ("fp4 KV/fp8 all", dict(kv_bits=4, lm_head_bits=8, linear_attn_bits=8)), + ] + A(f"{'configuration':<20}{'weights':>10}{'KV@1M':>10}{'total':>10} (bytes read per decode token)") + A("-" * 94) + for name, kw in configs: + b = bytes_per_token(ctx, **kw) + A(f"{name:<20}{b['weights']/GB:9.1f}G{b['kv']/GB:9.1f}G{b['total']/GB:9.1f}G") + A("") + + A(f"Predicted decode tok/s (at {BW_EFFICIENCY:.0%} of peak HBM bandwidth):") + A("") + hdr = f"{'GPU':<15}{'$/hr':>7}{'mem':>6}" + A(hdr + "".join(f"{n:>16}" for n, _ in configs)) + A("-" * 94) + for g in GPUS: + row = f"{g.name:<15}{g.usd_per_hour:7.2f}{g.mem_gb:5.0f}G" + for _name, kw in configs: + t = tps(g, ctx, **kw) + fits = footprint_gb(ctx, **kw) <= g.mem_gb + row += f"{(f'{t:.0f}' if fits else 'OOM'):>16}" + A(row) + A("") + + # Cheapest configuration that clears the target. + A(f"Cheapest options clearing {target_tps:.0f} tok/s:") + A("") + winners = [] + for g in GPUS: + for tp in (1, 2, 4): + for name, kw in configs: + t = tps(g, ctx, tp=tp, **kw) + if t >= target_tps and footprint_gb(ctx, tp=tp, **kw) <= g.mem_gb: + winners.append((g.usd_per_hour * tp, g.name, tp, name, t)) + winners.sort() + seen = set() + for cost, gname, tp, cname, t in winners: + if (gname, tp) in seen: + continue + seen.add((gname, tp)) + per_mtok = cost / (t * 3600) * 1e6 + A(f" ${cost:6.2f}/hr {gname:<14} TP={tp} {cname:<28} {t:6.0f} tok/s " + f"(${per_mtok:.2f}/M output tok)") + if len(seen) >= 6: + break + if not winners: + A(" none -- need a bigger lever (speculative decoding, or TP>4)") + return "\n".join(lines) + + +if __name__ == "__main__": + print(report()) diff --git a/deploy/modal_app.py b/deploy/modal_app.py new file mode 100644 index 000000000..9bf07c1bf --- /dev/null +++ b/deploy/modal_app.py @@ -0,0 +1,429 @@ +"""Serve M.O.G.-SEC-27B-1M-CTX-NVFP4 with FreeToken on Modal. + +Why this model fits a 1M-token context at interactive speed: it is a Qwen3.5-family +*hybrid*. Of its 64 decoder layers only 16 (every 4th) run full GQA attention; the other +48 are GatedDeltaNet linear-attention layers whose recurrent state is O(1) in context +length. So the KV cache that has to be re-read on every decode step grows 4x slower than +a conventional transformer's, and the linear layers contribute a flat ~150 MB regardless +of whether the context is 1k or 1M tokens. + +Decode is memory-bandwidth bound, so tokens/sec is set by bytes-read-per-token: + + bytes/token = weights(28.5 GB, dense -- every param is active) + + KV for the 16 full-attention layers (32 KiB/token at fp8) + + GatedDeltaNet state (~150 MB, constant) + +At 1M context with an fp8 KV cache that is ~61 GB/token, which needs ~6.1 TB/s to hit +100 tok/s. See `plan` below for the arithmetic against real Modal SKUs. + +Usage: + modal run deploy/modal_app.py::plan # cost/feasibility math, no cloud spend + modal run deploy/modal_app.py::download # ~29.4 GB onto the models volume + modal run deploy/modal_app.py::validate # CPU-only config gate + modal run deploy/modal_app.py::smoke # GPU: load + generate + modal run deploy/modal_app.py::bench + modal serve deploy/modal_app.py # OpenAI-compatible endpoint +""" + +from __future__ import annotations + +import os +import pathlib + +import modal + +# --- model ------------------------------------------------------------------------ +MODEL_ID = "Blackfrost-Research/M.O.G.-SEC-27B-1M-CTX-NVFP4" +MODEL_REVISION = "465f8de9584413f087625953d41bf4f6a2a898c0" +MODEL_DIR = "/models/mog-sec-27b-nvfp4" + +CUDA_TAG = "13.0.1-devel-ubuntu24.04" +REPO_ROOT = pathlib.Path(__file__).parent.parent + +# GPU for the serving/bench functions. B200 (8 TB/s HBM3e, 180 GB) is the cheapest SKU +# on Modal that can clear 100 tok/s at 1M context; override with FT_GPU when probing +# smaller/cheaper hardware (the CPU-only entrypoints ignore it entirely). +GPU = os.environ.get("FT_GPU", "B200") + +models_volume = modal.Volume.from_name("freetoken-models", create_if_missing=True) +kernels_volume = modal.Volume.from_name("freetoken-kernels", create_if_missing=True) +VOLUMES = {"/models": models_volume, "/kernel-cache": kernels_volume} + +_IGNORE = [ + "**/.git", "**/.git/**", "**/__pycache__", "**/*.pyc", "**/.pytest_cache", + "**/.ruff_cache", "**/freetoken-kernel-cache", "**/.venv", "**/uv.lock", + "**/assets", "**/.codegraph", "**/*.so", "**/build", +] + +image = ( + modal.Image.from_registry(f"nvidia/cuda:{CUDA_TAG}", add_python="3.12") + .apt_install("git", "build-essential", "ninja-build", "curl") + .env( + { + "CUDA_HOME": "/usr/local/cuda", + # The CUDA 13 base image points CXX at clang++, which it does not ship; + # torch.utils.cpp_extension honours CC/CXX, so pin them to the gcc that + # build-essential installs. + "CC": "gcc", + "CXX": "g++", + # B200/B300 are sm_100; keep 9.0 so the same image still runs on H100/H200. + "TORCH_CUDA_ARCH_LIST": "9.0;10.0", + "HF_HOME": "/models/hf-cache", + # Every JIT artifact lands on the shared volume, so a cold start neither + # re-downloads 29 GB nor recompiles Triton/flashinfer kernels. + "FREETOKEN_KERNEL_CACHE_DIR": "/kernel-cache/freetoken", + "TRITON_CACHE_DIR": "/kernel-cache/triton", + "FLASHINFER_WORKSPACE_BASE": "/kernel-cache/flashinfer", + "TORCHINDUCTOR_CACHE_DIR": "/kernel-cache/inductor", + } + ) + .pip_install("torch==2.11.0", "wheel", "setuptools>=77", "ninja") + .pip_install("hf_transfer", "huggingface_hub>=1.5,<2", "pytest>=6.0") + .add_local_dir(REPO_ROOT, "/opt/freetoken", copy=True, ignore=_IGNORE) + # No GPU on the build step: both ext_modules are plain CppExtension (.cpp, no .cu) + # and setup.py's _check_toolchain() only compares nvcc's CUDA major against torch's. + # Building on CPU keeps image builds off the B200 meter. + .run_commands("cd /opt/freetoken && pip install --no-build-isolation -e '.[accel]'") +) + +app = modal.App("freetoken-mog-sec-27b") + + +def _kv_dtype(name: str): + """CLI spelling -> torch dtype for the KV cache. 'auto' -> None -> follow --dtype. + Imported lazily: torch only exists inside the image, not on the launching machine.""" + import torch + + if name in ("auto", None): + return None + return { + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "fp8_e4m3": torch.float8_e4m3fn, + "fp8_e5m2": torch.float8_e5m2, + }[name] + + +# --- 0. planning (no cloud spend) -------------------------------------------------- +@app.local_entrypoint() +def plan(): + """Bytes-per-decode-token vs. Modal SKU bandwidth. Runs locally, costs nothing.""" + from deploy.capacity import report # noqa: PLC0415 + + print(report()) + + +# --- 1. weights ------------------------------------------------------------------- +@app.function(image=image, volumes=VOLUMES, timeout=60 * 60, cpu=8.0, memory=16384) +def download(force: bool = False) -> dict: + import os + import shutil + + from huggingface_hub import snapshot_download + + if force and os.path.isdir(MODEL_DIR): + shutil.rmtree(MODEL_DIR) + + path = snapshot_download( + MODEL_ID, revision=MODEL_REVISION, local_dir=MODEL_DIR, max_workers=16 + ) + models_volume.commit() + + files = [] + for root, _dirs, names in os.walk(path): + for n in names: + fp = os.path.join(root, n) + if not os.path.islink(fp): + files.append((os.path.relpath(fp, path), os.path.getsize(fp))) + files.sort(key=lambda x: -x[1]) + for n, sz in files[:12]: + print(f" {sz/1e9:8.3f} GB {n}") + total = sum(s for _, s in files) + print(f"TOTAL {total/1e9:.2f} GB in {len(files)} files -> {path}") + return {"path": path, "total_bytes": total, "num_files": len(files)} + + +# --- 2. CPU-only config gate ------------------------------------------------------ +@app.function(image=image, volumes=VOLUMES, timeout=60 * 20, cpu=4.0, memory=16384) +def validate() -> dict: + """Prove the checkpoint resolves through FreeToken's own config path before paying + for a GPU to find out. CPU container, cents.""" + import json + + from freetoken.models.register import get_model_spec + from freetoken.utils.hf import cached_load_hf_config + + hf_config = cached_load_hf_config(MODEL_DIR) + arch = hf_config.architectures[0] + spec = get_model_spec(arch) + print(f"architecture : {arch}") + print(f"registered : {spec.module}.{spec.model_cls}") + + text = getattr(hf_config, "text_config", hf_config) + layer_types = list(getattr(text, "layer_types", []) or []) + full = [i for i, t in enumerate(layer_types) if t == "full_attention"] + linear = [i for i, t in enumerate(layer_types) if t == "linear_attention"] + print(f"layers : {len(full)} full-attention, {len(linear)} linear-attention") + + out = { + "arch": arch, + "module": spec.module, + "num_hidden_layers": int(text.num_hidden_layers), + "hidden_size": int(text.hidden_size), + "num_attention_heads": int(text.num_attention_heads), + "num_key_value_heads": int(text.num_key_value_heads), + "head_dim": int(text.head_dim), + "vocab_size": int(text.vocab_size), + "max_position_embeddings": int(text.max_position_embeddings), + "num_full_attention_layers": len(full), + "num_linear_attention_layers": len(linear), + } + for name, nbytes in (("bf16", 2.0), ("fp8", 1.0), ("fp4", 0.5)): + per_tok = 2 * out["num_key_value_heads"] * out["head_dim"] * len(full) * nbytes + out[f"kv_bytes_per_token_{name}"] = per_tok + print(f" KV @{name:>4}: {per_tok/1024:7.1f} KiB/token -> {per_tok*1e6/1e9:6.2f} GB @ 1M ctx") + + # Now the real gate: FreeToken's own ModelConfig, which is what the engine consumes. + try: + import torch + + from freetoken.distributed import DistributedInfo + from freetoken.engine.config import EngineConfig + + cfg = EngineConfig( + model_path=MODEL_DIR, + tp_info=DistributedInfo(0, 1), + dtype=torch.bfloat16, + kv_cache_dtype=torch.float8_e4m3fn, + ) + mc = cfg.model_config + print(f"\nModelConfig OK: {mc.num_layers} layers, " + f"head_dim={mc.head_dim}, kv_heads={mc.num_kv_heads}") + print(f" compute dtype : {cfg.dtype}") + print(f" KV dtype : {cfg.kv_dtype}") + linear = mc.linear_attention_group() + print(f" linear group : {len(linear.layer_ids) if linear else 0} GDN layers") + specs = mc.kv_cache_group_specs() + print(f" KV specs : {[(s.num_layers, s.num_kv_heads, s.head_dim) for s in specs]}") + + from freetoken.kvcache.base import spec_kv_bytes_per_token + per_tok = sum(spec_kv_bytes_per_token(sp, cfg) for sp in specs if not sp.is_swa) + print(f" budgeted KV : {per_tok/1024:.1f} KiB/token " + f"({per_tok*1e6/1e9:.2f} GB @ 1M ctx, kv_dtype={cfg.kv_dtype})") + out["model_config_ok"] = True + out["budgeted_kv_bytes_per_token_fp8"] = per_tok + except Exception as exc: # noqa: BLE001 - reporting the gap is the point + import traceback + + traceback.print_exc() + out["model_config_ok"] = False + out["model_config_error"] = f"{type(exc).__name__}: {exc}" + + print(json.dumps(out, indent=2, default=str)) + return out + + +# --- 3. unit tests on the image ---------------------------------------------------- +@app.function(image=image, volumes=VOLUMES, timeout=60 * 20, cpu=4.0, memory=16384) +def tests(target: str = "tests/kvcache tests/engine") -> int: + """Run the KV-dtype/budget tests. CPU-only -- they never touch a device.""" + import subprocess + + rc = subprocess.run( + f"cd /opt/freetoken && python -m pytest -q {target}", + shell=True, check=False, + ).returncode + print(f"pytest exit={rc}") + return rc + + +@app.function(image=image, volumes=VOLUMES, gpu="L4", timeout=60 * 30) +def tests_gpu(target: str = "tests/kvcache tests/engine") -> int: + """Same suite with a real device attached. Some existing tests reach for + torch.cuda.get_device_name() through the MoE bench-profile lookup and can only pass + on a GPU container; running both tells a genuine regression apart from that.""" + import subprocess + + rc = subprocess.run( + f"cd /opt/freetoken && python -m pytest -q {target}", shell=True, check=False + ).returncode + print(f"pytest(gpu) exit={rc}") + return rc + + +# --- 4. GPU smoke test ------------------------------------------------------------- +@app.function(image=image, volumes=VOLUMES, gpu=GPU, timeout=60 * 45) +def smoke(kv_cache_dtype: str = "auto", max_seq_len: int = 32768) -> dict: + """Load the model on a GPU and generate. Small context so it is cheap; the point is + to prove the weight loader, the hybrid GDN path and the sampler all work.""" + import time + + import torch + + from freetoken.core import SamplingParams + from freetoken.llm import LLM + + kv = _kv_dtype(kv_cache_dtype) + print(f"loading {MODEL_DIR} kv_cache_dtype={kv_cache_dtype} max_seq_len={max_seq_len}") + t0 = time.time() + llm = LLM( + MODEL_DIR, + dtype=torch.bfloat16, + kv_cache_dtype=kv, + max_seq_len_override=max_seq_len, + max_running_req=1, + ) + load_s = time.time() - t0 + print(f"loaded in {load_s:.1f}s") + + out = llm.generate( + ["List three common categories of web application vulnerability."], + SamplingParams(temperature=0.0, max_tokens=64), + ) + text = out[0]["text"] if isinstance(out[0], dict) else str(out[0]) + print("--- completion ---") + print(text) + return {"load_s": load_s, "text": text, "kv_cache_dtype": kv_cache_dtype} + + +@app.function(image=image, volumes=VOLUMES, gpu=GPU, timeout=60 * 90) +def bench( + ctx: int = 131072, + gen: int = 64, + kv_cache_dtype: str = "auto", + max_seq_len: int | None = None, +) -> dict: + """Measure prefill and *decode* tok/s at a given context length. + + Decode rate is isolated from prefill by timing two runs -- one that stops after a + single token (prefill + 1 decode step) and one that generates `gen` tokens -- and + differencing them. That removes load, tokenization and prefill from the number, so + what comes out is the steady-state interactive rate a user would feel. + """ + import time + + import torch + + from freetoken.core import SamplingParams + from freetoken.llm import LLM + + kv = _kv_dtype(kv_cache_dtype) + max_seq_len = max_seq_len or (ctx + gen + 1024) + + t0 = time.time() + llm = LLM( + MODEL_DIR, + dtype=torch.bfloat16, + kv_cache_dtype=kv, + max_seq_len_override=max_seq_len, + max_running_req=1, + ) + load_s = time.time() - t0 + print(f"loaded in {load_s:.1f}s (ctx={ctx:,} kv={kv_cache_dtype})") + + # A synthetic prompt of exactly `ctx` token ids, passed as ids so tokenization + # cannot perturb the length we are measuring against. + prompt = [(i * 2654435761) % 200_000 + 1000 for i in range(ctx)] + + def run(n_tokens: int) -> float: + torch.cuda.synchronize() + t = time.time() + llm.generate([prompt], SamplingParams(temperature=0.0, max_tokens=n_tokens, + ignore_eos=True)) + torch.cuda.synchronize() + return time.time() - t + + t_prefill = run(1) # prefill + 1 decode step + t_full = run(gen) # prefill + gen decode steps + + decode_s = t_full - t_prefill + decode_tps = (gen - 1) / decode_s if decode_s > 0 else float("nan") + prefill_tps = ctx / t_prefill + + free, total = torch.cuda.mem_get_info() + result = { + "ctx": ctx, "gen": gen, "kv_cache_dtype": kv_cache_dtype, "gpu": GPU, + "load_s": round(load_s, 1), + "prefill_s": round(t_prefill, 2), "prefill_tps": round(prefill_tps, 1), + "decode_s": round(decode_s, 3), "decode_tps": round(decode_tps, 2), + "vram_used_gb": round((total - free) / 1e9, 1), + "vram_total_gb": round(total / 1e9, 1), + } + print("\n=== RESULT ===") + for k, v in result.items(): + print(f" {k:<16} {v}") + return result + + +# --- 5. fp8-KV correctness check on a small model ---------------------------------- +# The 27B target needs a B200-class card, but the --kv-cache-dtype plumbing it depends +# on is model-agnostic: pool allocation, the narrowing store, the flashinfer q/kv width +# split and the backend capability gate are all exercised by any GQA model. An L4 is +# sm_89 (native fp8) and inside Modal's no-payment-method tier, so this validates the +# engine change end to end for a few cents. +SMALL_MODEL = "Qwen/Qwen3-0.6B" + + +# max_inputs=1: the engine asserts the accelerator is not already initialized, so each +# probe needs a container that has never built an Engine -- a warm reused one would fail. +@app.function(image=image, volumes=VOLUMES, gpu="L4", timeout=60 * 40, max_inputs=1) +def probe_kv(kv_cache_dtype: str = "auto", model: str = SMALL_MODEL, + prompt: str | None = None) -> dict: + """Load `model` under one KV dtype and report the pool geometry + a greedy sample. + + One config per container on purpose: the engine asserts the accelerator is not yet + initialized, so two LLMs cannot share a process. + """ + import torch + + from freetoken.core import SamplingParams + from freetoken.llm import LLM + + prompt = prompt or "Explain why long-context LLM decoding is limited by memory bandwidth." + llm = LLM( + model, + dtype=torch.bfloat16, + kv_cache_dtype=_kv_dtype(kv_cache_dtype), + max_seq_len_override=4096, + max_running_req=1, + ) + pool = llm.engine.kv_cache + out = llm.generate([prompt], SamplingParams(temperature=0.0, max_tokens=48)) + text = out[0]["text"] if isinstance(out[0], dict) else str(out[0]) + result = { + "kv_cache_dtype": kv_cache_dtype, + "kv_pool_dtype": str(pool.dtype), + "kv_bytes_per_token": pool.unit_bytes()[0], + "attention_backend": llm.engine.config.attention_backend, + "text": text, + } + print(f"[{kv_cache_dtype}] pool={pool.dtype} {result['kv_bytes_per_token']} B/token " + f"backend={result['attention_backend']}") + print(f"[{kv_cache_dtype}] {text[:180]}") + return result + + +@app.local_entrypoint() +def verify_fp8_kv(model: str = SMALL_MODEL, kv: str = "fp8_e4m3"): + """A/B the same greedy continuation with a bf16 vs a narrowed KV cache.""" + bf16 = probe_kv.remote("auto", model) + fp8 = probe_kv.remote(kv, model) + + halved = bf16["kv_bytes_per_token"] == 2 * fp8["kv_bytes_per_token"] + common = 0 + for a, b in zip(bf16["text"], fp8["text"]): + if a != b: + break + common += 1 + + print("\n=== VERDICT ===") + print(f" backend : {bf16['attention_backend']} / {fp8['attention_backend']}") + print(f" pool dtype : {bf16['kv_pool_dtype']} -> {fp8['kv_pool_dtype']}") + print(f" KV bytes/token : {bf16['kv_bytes_per_token']} -> " + f"{fp8['kv_bytes_per_token']} halved={halved}") + # fp8 KV is lossy, so identical text is a bonus, not the contract: what must hold is + # that the cache is genuinely half as wide and the model still generates fluently. + print(f" identical text : {bf16['text'] == fp8['text']}") + print(f" common prefix : {common} chars") + print(f"\n bf16: {bf16['text'][:160]}") + print(f" fp8 : {fp8['text'][:160]}") diff --git a/docs/modal-deployment.md b/docs/modal-deployment.md new file mode 100644 index 000000000..35c50bfca --- /dev/null +++ b/docs/modal-deployment.md @@ -0,0 +1,137 @@ +# Serving M.O.G.-SEC-27B-1M-CTX-NVFP4 on Modal + +Deploying [`Blackfrost-Research/M.O.G.-SEC-27B-1M-CTX-NVFP4`](https://huggingface.co/Blackfrost-Research/M.O.G.-SEC-27B-1M-CTX-NVFP4) +with FreeToken on [Modal](https://modal.com), targeting **100 decode tok/s at a +1,000,000-token context**. + +Everything lives in [`deploy/modal_app.py`](../deploy/modal_app.py); the throughput and +cost arithmetic lives in [`deploy/capacity.py`](../deploy/capacity.py). + +## Why this model can do 1M context at all + +It is a Qwen3.5-family **hybrid**, and that is the whole story: + +| | count | per-token KV cost | +|---|---|---| +| full-attention (GQA) layers | 16 (every 4th) | 32 KiB/token at fp8 | +| GatedDeltaNet linear-attention layers | 48 | **0** — fixed 155 MB recurrent state | + +`config.json` sets `full_attention_interval: 4`, so only a quarter of the 64 layers carry +a KV cache. An all-full-attention model of the same size would need 4x the KV bandwidth, +and 1M context would be hopeless. + +FreeToken already registers this architecture — `Qwen3_5ForConditionalGeneration` maps to +`freetoken.models.qwen3_5_moe.Qwen3_5MoEForCausalLM`, which routes the dense MLP through +`Qwen3_5DenseMLP` and reads the compressed-tensors NVFP4 layout. No new model code was +needed. + +## The actual bottleneck: bytes read per decode token + +Single-stream decode is memory-bandwidth bound, not compute bound. To emit one token the +GPU streams every active weight plus the entire KV cache through the SMs exactly once: + +``` +tokens/sec ~= achievable_HBM_bandwidth / bytes_read_per_token +``` + +The model is **dense** (27B, no MoE sparsity), so "active weights" means all of them. +Measured from the checkpoint's own safetensors headers: + +| configuration | weights | KV @ 1M | total/token | +|---|---|---|---| +| bf16 KV (before this work) | 25.9 GB | 65.5 GB | **91.6 GB** | +| fp8 KV | 25.9 GB | 32.8 GB | **58.8 GB** | +| fp8 KV + fp8 lm_head | 24.7 GB | 32.8 GB | **57.6 GB** | +| fp8 KV + fp8 head + fp8 GDN | 19.1 GB | 32.8 GB | **52.0 GB** | + +At 100 tok/s those totals demand 9.2 / 5.9 / 5.8 / 5.2 TB/s respectively. + +## Which GPU clears 100 tok/s + +Predicted at 70% of peak HBM bandwidth (`python deploy/capacity.py` regenerates this): + +| GPU | $/hr | mem | bf16 KV | fp8 KV | fp8 all | +|---|---|---|---|---|---| +| A100-80GB | 2.50 | 80G | OOM | 24 | 27 | +| H100 SXM5 | 3.95 | 80G | OOM | 40 | 45 | +| H200 SXM | 4.54 | 141G | 37 | 57 | 65 | +| **B200** | **6.25** | **180G** | 61 | 95 | **108** | + +**A single B200 at $6.25/hr is the cheapest configuration that reaches the target**, and +only once the KV cache is fp8 — bf16 KV tops out around 61 tok/s no matter what else you +do. That is why this work added `--kv-cache-dtype`. + +Note B200 has plenty of *capacity* to spare (~63 GB resident of 180 GB); the binding +constraint is bandwidth, not memory. The leftover capacity is better spent on concurrent +requests than on a wider cache. + +## `--kv-cache-dtype` + +The KV cache storage width is now decoupled from the compute dtype: + +```bash +ft serve --model --dtype bfloat16 --kv-cache-dtype fp8_e5m2 +``` + +- `auto` (default) follows `--dtype`, preserving previous behaviour exactly. +- `fp8_e5m2` halves both the cache footprint and the bytes re-read each step. **This is + the one to use.** +- Queries and activations stay at the compute dtype; only the paged cache narrows. +- Requires an attention backend that plans query and KV widths separately. `fi` + (FlashInfer) does — its `plan()` takes `q_data_type` and `kv_data_type` independently. + Backends that assume `q.dtype == cache.dtype` are rejected at config time rather than + failing inside a kernel launch, and `--attention-backend auto` skips them. + +### Use e5m2, not e4m3 + +Measured on an L4 with Qwen3-0.6B (`modal run deploy/modal_app.py::verify_fp8_kv`), same +18.44 GiB KV budget both times: + +| | KV bytes/token | context capacity | output vs bf16 | +|---|---|---|---| +| bf16 | 114,688 | 172,637 tokens | — | +| `fp8_e5m2` | 57,344 | **345,271 tokens** | **byte-identical, 271/271 chars** | +| `fp8_e4m3` | 57,344 | 345,271 tokens | degenerates after 14 chars | + +Exactly 2x the context in the same memory, and with `e5m2` the greedy continuation was +*identical* to bf16 — not merely close. + +`e4m3` fails because the pool stores K/V by a straight cast with no per-tensor scale, and +e4m3 saturates at ±448; anything past that clips, and generation collapses into repeated +punctuation. `e5m2` spends the same byte on 5 exponent bits instead of 4 mantissa bits, so +it has the dynamic range to survive an unscaled cast. Since both are one byte wide, e4m3 +buys no extra bandwidth here anyway — it would only be preferable with calibrated +`k_scale`/`v_scale` plumbed through to the kernel, which is not implemented. Selecting it +logs a warning. + +## Running it + +```bash +modal run deploy/modal_app.py::plan # cost/feasibility math, no cloud spend +modal run deploy/modal_app.py::download # ~29.4 GB onto the models volume +modal run deploy/modal_app.py::validate # CPU-only config gate (cents) +modal run deploy/modal_app.py::tests # unit tests on the image +modal run deploy/modal_app.py::verify_fp8_kv # fp8-KV A/B on a small model (L4) +modal run deploy/modal_app.py::smoke # GPU: load + generate +modal run deploy/modal_app.py::bench --ctx 1000000 --kv-cache-dtype fp8_e5m2 +``` + +Weights and every JIT artifact (Triton, FlashInfer, Inductor) live on Modal volumes, so a +cold start neither re-downloads 29 GB nor recompiles kernels. `FT_GPU` overrides the GPU +for the device-backed functions. + +### Cost discipline + +- The image builds on CPU — both `ext_modules` are plain `CppExtension` (`.cpp`, no + `.cu`), and `setup.py` only compares nvcc's CUDA major to torch's. No GPU on the build. +- `download` and `validate` are CPU containers, so the checkpoint and the config gate cost + cents before any accelerator is touched. +- At 108 tok/s a single B200 works out to roughly **$16/M output tokens**. + +## Prerequisite: a Modal payment method + +B200 (and H200/H100/A100/L40S) are gated behind a payment method. Without one, Modal +accepts only **T4, L4 and A10**, none of which can hold a 29 GB checkpoint, let alone a +1M-token context. The `bench`/`smoke` entrypoints are ready to run the moment billing is +enabled; `verify_fp8_kv` is deliberately sized to run on an L4 inside the free tier so the +engine change can be validated without it. diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py index 746c04c4b..3f7fb9033 100644 --- a/python/freetoken/attention/__init__.py +++ b/python/freetoken/attention/__init__.py @@ -37,6 +37,12 @@ class BackendInfo: # linear layers bypass the backend entirely, but a backend whose metadata or # graph machinery assumes layer 0 is an attention layer can opt out here. hybrid_linear_ok: bool = True + # Whether forward() can read a paged KV cache stored NARROWER than the compute + # dtype (--kv-cache-dtype fp8_e4m3/fp8_e5m2). Backends that hand the cache + # straight to a kernel taking a separate kv_data_type can; the ones that assume + # q.dtype == cache.dtype cannot, and config-time validation rejects the pairing + # rather than letting it fail inside a kernel launch. + supports_quantized_kv: bool = False SUPPORTED_ATTENTION_BACKENDS = Registry[BackendCreator]("Attention Backend") @@ -62,6 +68,9 @@ def create_trtllm_backend(config: ModelConfig): BackendInfo( supported_types=frozenset({AttnType.FULL}), requires_flashinfer=True, + # plan() takes q_data_type and kv_data_type separately, so an fp8 paged cache + # is dequantized inside the kernel against bf16 queries. + supports_quantized_kv=True, ), ) def create_fi_backend(config: ModelConfig): @@ -96,6 +105,19 @@ def create_triton_backend(config: ModelConfig): return TritonAttentionBackend(config) +@SUPPORTED_ATTENTION_BACKENDS.register( + "torch", + BackendInfo( + supported_types=frozenset({AttnType.FULL}), + consumes_attn_spec=True, + ), +) +def create_torch_backend(config: ModelConfig): + from .torch_native import TorchAttentionBackend + + return TorchAttentionBackend(config) + + @SUPPORTED_ATTENTION_BACKENDS.register( "dsv4_sparse", BackendInfo(supported_types=frozenset({AttnType.DSV4})), diff --git a/python/freetoken/attention/fi.py b/python/freetoken/attention/fi.py index c9e58538f..404861f99 100644 --- a/python/freetoken/attention/fi.py +++ b/python/freetoken/attention/fi.py @@ -57,7 +57,8 @@ class FIMetadata(BaseAttnMetadata): page_size: Literal[1] # currently only support page_size=1 pos_encoding_mode: str seq_lens_cpu: torch.Tensor # on cpu - dtype: torch.dtype + dtype: torch.dtype # KV-cache STORAGE dtype (may be fp8) + q_dtype: torch.dtype # query/compute dtype (never quantized) wrapper: BatchPrefillWithPagedKVCacheWrapper | BatchDecodeWithPagedKVCacheWrapper initialized: bool = False # fmt: on @@ -86,6 +87,8 @@ def __init__(self, config: ModelConfig) -> None: self.config = config self.kvcache = get_global_ctx().kv_cache + # Queries stay at compute width even when the paged cache is fp8. + self.q_dtype = get_global_ctx().compute_dtype self.device = self.kvcache.device # fa2 split-KV prefill needs ``tmp_v <= qo_heads_local * padded_batch_size * # cta_tile_q * head_dim * 4`` bytes of scratch, where flashinfer's scheduler @@ -165,8 +168,8 @@ def _initialize_metadata_once(self, metadata: FIMetadata) -> None: page_size=metadata.page_size, pos_encoding_mode=metadata.pos_encoding_mode, seq_lens=metadata.seq_lens_cpu, - data_type=metadata.dtype, - q_data_type=metadata.dtype, + data_type=metadata.q_dtype, + q_data_type=metadata.q_dtype, kv_data_type=metadata.dtype, non_blocking=True, ) @@ -182,7 +185,7 @@ def _initialize_metadata_once(self, metadata: FIMetadata) -> None: page_size=metadata.page_size, pos_encoding_mode=metadata.pos_encoding_mode, seq_lens=metadata.seq_lens_cpu, - q_data_type=metadata.dtype, + q_data_type=metadata.q_dtype, kv_data_type=metadata.dtype, non_blocking=True, causal=True, @@ -256,6 +259,7 @@ def prepare_metadata(self, batch: Batch) -> None: pos_encoding_mode="NONE", seq_lens_cpu=seq_len_cpu, dtype=self.kvcache.dtype, + q_dtype=self.q_dtype, wrapper=self.decode_wrappers if batch.is_decode else self.prefill_wrapper, ) diff --git a/python/freetoken/core.py b/python/freetoken/core.py index ef0a539cb..77fd1e6ab 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -167,6 +167,10 @@ def padded_size(self) -> int: @dataclass class Context: page_size: int + # Model compute dtype. Distinct from `kv_cache.dtype`, which is the KV *storage* + # width and may be narrower (fp8); attention backends plan queries at this dtype + # while reading the cache at the pool's. + compute_dtype: torch.dtype = torch.bfloat16 # NOTE: this table always treat page_size = 1 page_table: torch.Tensor = field(init=False) attn_backend: BaseAttnBackend = field(init=False) diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 543012f39..ba0561845 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -18,6 +18,14 @@ class EngineConfig: model_path: str tp_info: DistributedInfo dtype: torch.dtype + # Storage dtype of the paged KV cache, decoupled from the compute dtype above. + # None -> follow `dtype` (the historical behaviour). Decode at long context is + # bandwidth-bound on re-reading the whole cache every step, so halving its width + # is close to a linear speedup: on a 1M-token hybrid context an fp8 cache reads + # 32.8 GB/token instead of 65.5 GB. Only backends whose BackendInfo sets + # `supports_quantized_kv` may pair with a narrower dtype than `dtype`. + kv_cache_dtype: torch.dtype | None = None + accelerator: str = "auto" max_running_req: int = 4 attention_backend: str = "auto" moe_backend: str = "auto" @@ -91,6 +99,16 @@ def model_config(self) -> ModelConfig: parse_config = _load_attr(spec.module, spec.parse_config) return parse_config(self.hf_config) + @property + def kv_dtype(self) -> torch.dtype: + """Resolved KV-cache storage dtype (falls back to the compute dtype). + + This is the single source of truth: the pool allocates with it, and + `spec_kv_bytes_per_token` budgets with it, so capacity planning and the real + allocation can never disagree. + """ + return self.kv_cache_dtype if self.kv_cache_dtype is not None else self.dtype + @property def max_seq_len(self) -> int: if self.max_seq_len_override is not None: diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index ff3c985f5..b7a27708d 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -7,6 +7,14 @@ from typing import Any, Dict, Iterable, NamedTuple, Tuple import torch +from freetoken.accelerator import ( + AcceleratorKind, + accelerator_runtime, + apply_engine_accelerator_constraints, + detect_accelerator, + distributed_backend, + uses_expandable_segments, +) from freetoken.attention import AttnType, attention_backend_info, create_attention_backend from freetoken.core import Batch, Context, Req, set_global_ctx from freetoken.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info @@ -114,7 +122,7 @@ def _backend_requirements_met(name: str) -> bool: def _resolve_auto_attention_backend( - required: frozenset[AttnType], hybrid_linear: bool + required: frozenset[AttnType], hybrid_linear: bool, quantized_kv: bool = False ) -> str: """First candidate (in per-type priority order) whose arch condition holds, whose packages are installed, and whose every comma part serves ALL required @@ -141,6 +149,13 @@ def _resolve_auto_attention_backend( continue if not _backend_parts_serve(name, required): continue + # A narrower-than-compute KV cache narrows the candidate set to backends that + # plan q and kv widths independently; skip the rest rather than auto-selecting + # one that would reject the pairing a few lines later in validation. + if quantized_kv and not all( + attention_backend_info(p).supports_quantized_kv for p in name.split(",") + ): + continue if hybrid_linear and not all( attention_backend_info(p).hybrid_linear_ok for p in name.split(",") ): @@ -154,6 +169,13 @@ def _resolve_auto_attention_backend( ) +def _kv_is_quantized(config) -> bool: + """True when the KV cache is stored narrower than the compute dtype.""" + kv = getattr(config, "kv_cache_dtype", None) + compute = getattr(config, "dtype", None) + return kv is not None and compute is not None and kv != compute + + def _validate_attention_backend_choice(config, override, required: frozenset[AttnType]) -> None: """Config-time type x backend capability check for the resolved (or explicit) backend string: every comma part must serve every required type and have its @@ -189,6 +211,31 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att f"backend {part!r} does not support hybrid-linear (GDN/mamba) models, " f"got {config.attention_backend!r}." ) + # getattr: duck-typed test configs omit the field, same as the _dtype probe above. + if getattr(config, "kv_cache_dtype", None) is torch.float8_e4m3fn: + # e4m3 saturates at +/-448, and this pool stores K/V by a straight cast with + # no per-tensor scale, so any activation past that range clips. Measured on + # Qwen3-0.6B: e4m3 degenerates into repeated punctuation after ~14 chars, + # while e5m2 (same 1 byte/elem, 5 exponent bits) reproduced the bf16 + # continuation exactly. Until calibrated k_scale/v_scale are plumbed through + # to the kernel, e5m2 is the correct unscaled choice -- and since both are + # one byte wide, e4m3 buys no extra bandwidth here anyway. + logger.warning_rank0( + "--kv-cache-dtype fp8_e4m3 stores K/V unscaled and clips above +/-448, " + "which can silently degrade generation quality. Prefer fp8_e5m2, which " + "is the same width and needs no calibration." + ) + if _kv_is_quantized(config) and not info.supports_quantized_kv: + supporting = [ + name + for name in ("fi", "fa", "trtllm", "triton") + if attention_backend_info(name).supports_quantized_kv + ] + raise ValueError( + f"--kv-cache-dtype {config.kv_cache_dtype} stores the paged KV cache " + f"narrower than --dtype {config.dtype}, which backend {part!r} cannot " + f"read; use {' or '.join(supporting) or 'a compute-width KV cache'}." + ) if AttnType.SWA in required and not info.consumes_attn_spec: # SWA models drive window/sinks/sm_scale through the per-call AttentionSpec; # a backend that drops it would attend with the wrong window silently. @@ -291,23 +338,28 @@ class ForwardOutput(NamedTuple): class Engine: def __init__(self, config: EngineConfig): - assert not torch.cuda.is_initialized() + self.accelerator_kind = detect_accelerator(config.accelerator) + self.accelerator = accelerator_runtime(self.accelerator_kind) + assert not self.accelerator.is_initialized() set_tp_info(rank=config.tp_info.rank, size=config.tp_info.size) - _ensure_expandable_segments() # before the first CUDA allocation below - _adjust_config(config) + if uses_expandable_segments(self.accelerator_kind): + _ensure_expandable_segments() + _adjust_config(config, self.accelerator_kind) - self.device = torch.device(f"cuda:{config.tp_info.rank}") - torch.cuda.set_device(self.device) + self.device = torch.device(f"{self.accelerator_kind.value}:{config.tp_info.rank}") + self.accelerator.set_device(self.device) torch.manual_seed(42) - self.stream = torch.cuda.Stream() - torch.cuda.set_stream(self.stream) + self.stream = self.accelerator.Stream() + self.accelerator.set_stream(self.stream) self.dtype = config.dtype + # KV storage width, decoupled from compute width (see EngineConfig.kv_dtype). + self.kv_dtype = config.kv_dtype self.config = config # retained for runtime cache rebuild (rebuild_runtime_cache) # KV pool family fixed at construction from the model config: its classmethods own the # page-token geometry and cost arithmetic the engine needs BEFORE the pool exists # (num_pages sizing, --moe-cache-auto); the instance owns rebuild/validation after. self._pool_cls = resolve_pool_class(config.model_config) - self.ctx = Context(config.page_size) + self.ctx = Context(config.page_size, compute_dtype=config.dtype) set_global_ctx(self.ctx) self.tp_cpu_group = self._init_communication(config) @@ -345,7 +397,7 @@ def __init__(self, config: EngineConfig): self.num_pages = self._pool_cls.solve_num_pages(config, available_memory) num_tokens = self.num_pages * config.page_size self.ctx.kv_cache = self.kv_cache = create_kv_pool( - config, self.num_pages, device=self.device, dtype=self.dtype + config, self.num_pages, device=self.device, dtype=self.kv_dtype ) # ======================= Linear (GatedDeltaNet) state initialization ======================== @@ -436,10 +488,11 @@ def _init_communication(self, config: EngineConfig) -> torch.distributed.Process max_bytes = ( config.max_forward_len * config.model_config.hidden_size * self.dtype.itemsize ) - enable_pynccl_distributed(config.tp_info, tp_cpu_group, max_bytes) + if config.use_pynccl: + enable_pynccl_distributed(config.tp_info, tp_cpu_group, max_bytes) else: torch.distributed.init_process_group( - backend="nccl", + backend=distributed_backend(self.accelerator_kind), rank=config.tp_info.rank, world_size=config.tp_info.size, timeout=timedelta(seconds=config.distributed_timeout), @@ -601,7 +654,7 @@ def _resolve_hybrid_fetch(self, config: EngineConfig, cache) -> None: return # explicit fixed cap from freetoken.moe.bench_profile import load_hybrid_fetch_fraction - gpu_name = torch.cuda.get_device_name(self.device) if torch.cuda.is_available() else None + gpu_name = self.accelerator.get_device_name(self.device) fraction = load_hybrid_fetch_fraction(cache.quant_format, gpu_name=gpu_name) if fraction is None: cache.hybrid_max_fetch = 1 @@ -654,9 +707,9 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: def _sync_get_memory(self) -> Tuple[int, int]: """Get the min and max free memory across TP ranks.""" - torch.cuda.synchronize(self.device) - torch.cuda.empty_cache() - torch.cuda.reset_peak_memory_stats(self.device) + self.accelerator.synchronize(self.device) + self.accelerator.empty_cache() + self.accelerator.reset_peak_memory_stats(self.device) free_memory = get_free_memory(self.device) free_mem_tensor = torch.tensor([free_memory, -free_memory], device="cpu", dtype=torch.int64) torch.distributed.all_reduce( @@ -802,7 +855,7 @@ def rebuild_runtime_cache( ), ) - torch.cuda.synchronize(self.device) + self.accelerator.synchronize(self.device) # Preserve the CUDA-graph batch-size set resolved at startup. The auto heuristic keys # off free memory, which is far smaller now that the caches are resident (post-cache # free << startup pre-load free), so re-deriving it here would silently drop large @@ -861,7 +914,7 @@ def rebuild_runtime_cache( ) def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: - assert torch.cuda.current_stream() == self.stream + assert self.accelerator.current_stream() == self.stream with self.ctx.forward_batch(batch): if self.graph_runner.can_use_cuda_graph(batch): logits = self.graph_runner.replay(batch) @@ -878,7 +931,7 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: batch_logits = logits[: batch.size] next_tokens_gpu = self.sampler.sample(batch_logits, args).to(torch.int32) next_tokens_cpu = next_tokens_gpu.to("cpu", non_blocking=True) - copy_done_event = torch.cuda.Event() + copy_done_event = self.accelerator.Event() copy_done_event.record(self.stream) return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event) @@ -903,8 +956,8 @@ def _warmup_prefill(self) -> None: dummy_row = self.page_table[self.dummy_req.table_idx] dummy_slot = int(dummy_row[0].item()) - started = torch.cuda.Event(enable_timing=True) - ended = torch.cuda.Event(enable_timing=True) + started = self.accelerator.Event(enable_timing=True) + ended = self.accelerator.Event(enable_timing=True) started.record(self.stream) try: for length in warmup_lens: @@ -933,7 +986,7 @@ def _warmup_prefill(self) -> None: if self.moe_offload_cache is not None: self.moe_offload_cache.reset() ended.record(self.stream) - torch.cuda.synchronize(self.device) + self.accelerator.synchronize(self.device) logger.info_rank0( f"Prefill warmup complete for lengths {warmup_lens} " f"in {started.elapsed_time(ended) / 1000.0:.3f} s" @@ -1086,7 +1139,10 @@ def _resolve_cpu_layers(config: EngineConfig, num_moe_layers: int) -> frozenset[ } -def _adjust_config(config: EngineConfig): +def _adjust_config( + config: EngineConfig, + accelerator_kind: AcceleratorKind = AcceleratorKind.CUDA, +) -> None: def override(attr: str, value: Any): # this is dangerous, use with caution object.__setattr__(config, attr, value) @@ -1097,6 +1153,12 @@ def override(attr: str, value: Any): # this is dangerous, use with caution has_linear_attention = getattr(model_config, "has_linear_attention", False) is_moe = getattr(model_config, "is_moe", False) expert_quant = getattr(model_config, "expert_quant", "none") + apply_engine_accelerator_constraints( + config, + accelerator_kind, + model_config=model_config, + is_moe=is_moe, + ) if not is_moe: # A dense model has no routed experts: the MoE knobs are inert, and the offload family @@ -1186,7 +1248,9 @@ def override(attr: str, value: Any): # this is dangerous, use with caution if config.attention_backend == "auto": override( "attention_backend", - _resolve_auto_attention_backend(required_attn_types, has_linear_attention), + _resolve_auto_attention_backend( + required_attn_types, has_linear_attention, _kv_is_quantized(config) + ), ) logger.info_rank0(f"Auto-selected attention backend: {config.attention_backend}") _validate_attention_backend_choice(config, override, required_attn_types) @@ -1245,7 +1309,12 @@ def override(attr: str, value: Any): # this is dangerous, use with caution bench_fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16") from freetoken.moe.bench_profile import load_backend_recommendation - gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else None + # Guard on availability: load_backend_recommendation treats None as "no profile", + # which is the right answer on a CPU-only host. get_device_name(0) would instead + # raise out of a driver probe (this ran under `torch.cuda.is_available()` before + # the accelerator refactor, and dropping the check broke CPU-only config runs). + _rt = accelerator_runtime(accelerator_kind) + gpu_name = _rt.get_device_name(0) if _rt.is_available() else None if load_backend_recommendation(bench_fmt, gpu_name=gpu_name) == "hybrid": from freetoken.moe.cpu_executor import compiled_extension_supports diff --git a/python/freetoken/kernel/torch_ops.py b/python/freetoken/kernel/torch_ops.py new file mode 100644 index 000000000..1e0bf11d3 --- /dev/null +++ b/python/freetoken/kernel/torch_ops.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +import torch +from torch.nn import functional + + +def rmsnorm( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, + out: torch.Tensor | None = None, +) -> torch.Tensor: + normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps) + result = (normalized * weight.float()).to(x.dtype) + if out is None: + return result + out.copy_(result) + return out + + +def fused_add_rmsnorm( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> None: + residual.add_(x) + rmsnorm(residual, weight, eps, out=x) + + +def silu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor: + gate, value = x.chunk(2, dim=-1) + result = functional.silu(gate) * value + if out is None: + return result + out.copy_(result) + return out + + +def store_cache( + k_cache: torch.Tensor, + v_cache: torch.Tensor, + indices: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, +) -> None: + locations = indices.to(torch.long) + row_shape = (locations.numel(), *k_cache.shape[1:]) + k, v = k.reshape(row_shape), v.reshape(row_shape) + # index_copy_ has no CPU kernel for the float8 dtypes, and this is a pure relocation + # of bits into the cache -- no arithmetic. Alias both sides as uint8 (same width, same + # layout) so a narrow KV cache stores through the fallback path exactly as it does + # through the CUDA byte-copy kernel. + if k_cache.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + k_cache, v_cache = k_cache.view(torch.uint8), v_cache.view(torch.uint8) + k, v = k.view(torch.uint8), v.view(torch.uint8) + k_cache.index_copy_(0, locations, k) + v_cache.index_copy_(0, locations, v) + + +def apply_rope_inplace( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + head_size: int, + cos_sin_cache: torch.Tensor, + *, + is_neox: bool, +) -> None: + rotary_dim = cos_sin_cache.shape[-1] + half = rotary_dim // 2 + cache = cos_sin_cache.index_select(0, positions.to(torch.long)) + cos = cache[:, :half].unsqueeze(1) + sin = cache[:, half:].unsqueeze(1) + for tensor in (query, key): + heads = tensor.view(tensor.shape[0], -1, head_size) + rotary = heads[..., :rotary_dim] + if is_neox: + first, second = rotary[..., :half], rotary[..., half:] + rotated = torch.cat((first * cos - second * sin, second * cos + first * sin), dim=-1) + else: + pairs = rotary.view(*rotary.shape[:-1], half, 2) + first, second = pairs.unbind(-1) + rotated = torch.stack( + (first * cos - second * sin, second * cos + first * sin), dim=-1 + ).flatten(-2) + rotary.copy_(rotated) + + +def sample( + logits: torch.Tensor, + temperatures: torch.Tensor, + top_k: torch.Tensor | int | None, + top_p: torch.Tensor | float | None, +) -> torch.Tensor: + scaled = logits / temperatures.unsqueeze(-1) + sorted_logits, sorted_indices = scaled.sort(dim=-1, descending=True) + ranks = torch.arange(scaled.shape[-1], device=scaled.device).unsqueeze(0) + if top_k is not None: + k = torch.as_tensor(top_k, device=scaled.device).reshape(-1, 1) + sorted_logits = sorted_logits.masked_fill(ranks >= k, -torch.inf) + sorted_probs = torch.softmax(sorted_logits, dim=-1) + if top_p is not None: + p = torch.as_tensor(top_p, device=scaled.device).reshape(-1, 1) + remove = sorted_probs.cumsum(-1) - sorted_probs > p + sorted_probs = sorted_probs.masked_fill(remove, 0) + sorted_probs = sorted_probs / sorted_probs.sum(-1, keepdim=True) + selected = torch.multinomial(sorted_probs, 1) + return sorted_indices.gather(-1, selected).squeeze(-1) + + +def paged_attention( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + indices: torch.Tensor, + *, + query_lens: tuple[int, ...], + kv_lens: tuple[int, ...], + query_positions: torch.Tensor, + scale: float, + sliding_window: int | None = None, +) -> torch.Tensor: + outputs: list[torch.Tensor] = [] + query_offset = 0 + kv_offset = 0 + for query_len, kv_len in zip(query_lens, kv_lens, strict=True): + query = q[query_offset : query_offset + query_len] + locations = indices[kv_offset : kv_offset + kv_len].to(torch.long) + keys = k_cache.index_select(0, locations) + values = v_cache.index_select(0, locations) + groups = query.shape[1] // keys.shape[1] + if groups > 1: + keys = keys.repeat_interleave(groups, dim=1) + values = values.repeat_interleave(groups, dim=1) + positions = query_positions[query_offset : query_offset + query_len] + key_positions = torch.arange(kv_len, device=q.device) + allowed = key_positions.unsqueeze(0) <= positions.unsqueeze(1) + if sliding_window is not None: + allowed &= key_positions.unsqueeze(0) > positions.unsqueeze(1) - sliding_window + output = functional.scaled_dot_product_attention( + query.transpose(0, 1).unsqueeze(0), + keys.transpose(0, 1).unsqueeze(0), + values.transpose(0, 1).unsqueeze(0), + attn_mask=allowed.unsqueeze(0).unsqueeze(0), + scale=scale, + ) + outputs.append(output.squeeze(0).transpose(0, 1)) + query_offset += query_len + kv_offset += kv_len + return torch.cat(outputs, dim=0) + + +__all__ = [ + "apply_rope_inplace", + "fused_add_rmsnorm", + "paged_attention", + "rmsnorm", + "sample", + "silu_and_mul", + "store_cache", +] diff --git a/python/freetoken/kvcache/base.py b/python/freetoken/kvcache/base.py index ae8cf9ecb..0f21603a5 100644 --- a/python/freetoken/kvcache/base.py +++ b/python/freetoken/kvcache/base.py @@ -18,7 +18,9 @@ class CacheRebuildRejected(Exception): def spec_kv_bytes_per_token(spec, config) -> int: """One paged-KV group's bytes per token: (1|2 slabs) x head_dim x local kv heads x dtype - x layers, plus the bf16 DSA index-key slab when the spec carries indexer dims. Pure + x layers, plus the bf16 DSA index-key slab when the spec carries indexer dims. The + per-element width is the *KV* dtype (``config.kv_dtype``, which falls back to the + compute dtype), so an fp8 cache budgets half the bytes it would at bf16. Pure per-spec arithmetic -- pool families compose it over THEIR OWN groups; no family branching here. (2 bytes/elem == the torch.bfloat16 dsa_pool.DSAKVCache._alloc hardcodes; keep the two in lockstep if the slab dtype ever changes.)""" @@ -26,7 +28,7 @@ def spec_kv_bytes_per_token(spec, config) -> int: (1 if spec.mla else 2) # MLA latent groups store one slab (V aliases K) * spec.head_dim * div_even(spec.num_kv_heads, config.tp_info.size, allow_replicate=True) - * config.dtype.itemsize + * getattr(config, "kv_dtype", config.dtype).itemsize * spec.num_layers ) return per_token + spec.index_head_dim * spec.num_index_layers * 2 diff --git a/python/freetoken/kvcache/mha_pool.py b/python/freetoken/kvcache/mha_pool.py index 8ed280b96..f10d2499b 100644 --- a/python/freetoken/kvcache/mha_pool.py +++ b/python/freetoken/kvcache/mha_pool.py @@ -124,9 +124,21 @@ def store_kv( out_loc: torch.Tensor, layer_id: int, ) -> None: - from freetoken.kernel import store_cache + if self._device.type != "cuda": + from freetoken.kernel.torch_ops import store_cache + else: + from freetoken.kernel import store_cache dense = self._dense(layer_id) + # store_cache is a templated raw byte-copy (element_size is in BYTES), so it + # cannot convert on the way in: narrow the incoming compute-dtype K/V here when + # the pool stores a different width (e.g. bf16 attention into an fp8 cache). + # The cast is O(tokens x kv_heads x head_dim) -- negligible next to the read + # traffic this quantization removes from every subsequent decode step. + cache_dtype = self._kv_buffer.dtype + if k.dtype != cache_dtype: + k = k.to(cache_dtype) + v = v.to(cache_dtype) store_cache( k_cache=self._k_buffer[dense].view(self._storage_shape), v_cache=self._v_buffer[dense].view(self._storage_shape), diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index b2857b752..6ccc91854 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -213,6 +213,27 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="Data type for model weights and activations. 'auto' will use FP16 for FP32/FP16 models and BF16 for BF16 models.", ) + parser.add_argument( + "--kv-cache-dtype", + type=str, + default="auto", + choices=["auto", "bfloat16", "float16", "fp8_e4m3", "fp8_e5m2"], + help=( + "Storage dtype of the paged KV cache, independent of --dtype. 'auto' follows " + "--dtype. fp8_e4m3/fp8_e5m2 halve the cache's footprint AND the bytes re-read " + "on every decode step, which is the dominant cost at long context -- on a " + "1M-token hybrid context that is 32.8 GB/token instead of 65.5 GB. Needs an " + "attention backend that plans query and KV widths separately (fi)." + ), + ) + + parser.add_argument( + "--accelerator", + choices=("auto", "cuda", "xpu"), + default=ServerArgs.accelerator, + help="Accelerator runtime. Auto prefers CUDA, then Intel XPU.", + ) + parser.add_argument( "--tensor-parallel-size", "--tp-size", @@ -677,6 +698,22 @@ def _infer_reasoning_parser(model_path: str) -> str | None: "float32": torch.float32, } kwargs["dtype"] = DTYPE_MAP[dtype_str] if isinstance(dtype_str, str) else dtype_str + + # KV cache width. "auto" -> None -> EngineConfig.kv_dtype falls back to the compute + # dtype, preserving the historical behaviour for every existing invocation. + KV_DTYPE_MAP = { + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "fp8_e4m3": torch.float8_e4m3fn, + "fp8_e5m2": torch.float8_e5m2, + } + kv_dtype_str = kwargs.get("kv_cache_dtype", "auto") + if kv_dtype_str in ("auto", None): + kwargs["kv_cache_dtype"] = None + elif isinstance(kv_dtype_str, str): + kwargs["kv_cache_dtype"] = KV_DTYPE_MAP[kv_dtype_str] + else: + kwargs["kv_cache_dtype"] = kv_dtype_str kwargs["tp_info"] = DistributedInfo(0, kwargs["tensor_parallel_size"]) del kwargs["tensor_parallel_size"] diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index a164f0b4d..b1464bee8 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -215,7 +215,7 @@ class Cfg: cuda_graph_bs = [1, 2] max_seq_len = 1024 page_size = 1 - attention_backend = "fi" + attention_backend = "triton" nvfp4_backend = "auto" num_page_override = None num_token_override = 5000 @@ -347,7 +347,7 @@ def _offload_engine_config(**overrides): model_path="/tmp/freetoken-test-model", tp_info=DistributedInfo(rank=0, size=1), dtype=torch.bfloat16, - attention_backend="fi", + attention_backend="triton", **overrides, ) object.__setattr__( @@ -382,7 +382,7 @@ def test_guard_raises_actionable_error_when_too_small(): assert "128" in msg and "moe-cache" in msg -def test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backend(): +def test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backend(monkeypatch): """Bare `ft serve `: no --moe-backend, no --moe-cache-* flags at all. args.py's parse-time default only fires when the backend is *already* @@ -396,6 +396,16 @@ def test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backend from freetoken.engine.engine import _adjust_config from freetoken.moe import is_offload_moe_backend + monkeypatch.setattr( + "freetoken.engine.engine.accelerator_runtime", + # Mirrors the torch.cuda/torch.xpu surface _adjust_config uses: it probes + # is_available() before asking for a device name, so a CPU-only host reports + # "no profile" instead of raising out of a driver probe. + lambda _kind: SimpleNamespace( + get_device_name=lambda _device: "test gpu", is_available=lambda: True + ), + ) + config = _offload_engine_config() _adjust_config(config) diff --git a/tests/kvcache/test_kv_cache_dtype.py b/tests/kvcache/test_kv_cache_dtype.py new file mode 100644 index 000000000..8a8b13ca8 --- /dev/null +++ b/tests/kvcache/test_kv_cache_dtype.py @@ -0,0 +1,146 @@ +"""KV-cache storage dtype decoupled from the compute dtype (--kv-cache-dtype). + +Decode at long context is bandwidth-bound on re-reading the whole KV cache every step, +so storing it at fp8 while computing at bf16 is close to a linear decode speedup. These +tests pin the three places that must agree: the resolved dtype, the byte budget, and the +backend capability gate. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.attention import attention_backend_info +from freetoken.kvcache.base import spec_kv_bytes_per_token + + +def _spec(**kw): + base = dict(head_dim=256, num_kv_heads=4, num_layers=16, mla=False, + index_head_dim=0, num_index_layers=0, is_swa=False) + base.update(kw) + return SimpleNamespace(**base) + + +def _cfg(dtype=torch.bfloat16, kv_cache_dtype=None, tp_size=1): + # Mirrors EngineConfig.kv_dtype's fallback without dragging in the full dataclass. + return SimpleNamespace( + dtype=dtype, + kv_cache_dtype=kv_cache_dtype, + kv_dtype=kv_cache_dtype if kv_cache_dtype is not None else dtype, + tp_info=SimpleNamespace(size=tp_size), + ) + + +# --- EngineConfig.kv_dtype resolution --------------------------------------------- +def test_kv_dtype_defaults_to_compute_dtype(): + from freetoken.distributed import DistributedInfo + from freetoken.engine.config import EngineConfig + + cfg = EngineConfig(model_path="x", tp_info=DistributedInfo(0, 1), dtype=torch.bfloat16) + assert cfg.kv_cache_dtype is None + assert cfg.kv_dtype is torch.bfloat16 + + +def test_kv_dtype_honours_explicit_override(): + from freetoken.distributed import DistributedInfo + from freetoken.engine.config import EngineConfig + + cfg = EngineConfig( + model_path="x", tp_info=DistributedInfo(0, 1), + dtype=torch.bfloat16, kv_cache_dtype=torch.float8_e4m3fn, + ) + assert cfg.kv_dtype is torch.float8_e4m3fn + # The compute dtype is untouched -- queries and activations stay bf16. + assert cfg.dtype is torch.bfloat16 + + +# --- budget arithmetic ------------------------------------------------------------- +def test_fp8_kv_halves_bytes_per_token(): + spec = _spec() + bf16 = spec_kv_bytes_per_token(spec, _cfg()) + fp8 = spec_kv_bytes_per_token(spec, _cfg(kv_cache_dtype=torch.float8_e4m3fn)) + assert bf16 == 2 * 4 * 256 * 2 * 16 # 2 slabs x kv_heads x head_dim x 2B x layers + assert fp8 * 2 == bf16 + + +def test_budget_uses_kv_dtype_not_compute_dtype(): + """A duck-typed config without kv_dtype must still budget off the compute dtype.""" + legacy = SimpleNamespace(dtype=torch.bfloat16, tp_info=SimpleNamespace(size=1)) + assert spec_kv_bytes_per_token(_spec(), legacy) == spec_kv_bytes_per_token(_spec(), _cfg()) + + +def test_hybrid_model_budgets_only_full_attention_layers(): + """The 27B hybrid carries KV on 16 of 64 layers; the other 48 hold an O(1) GDN state. + That 4x is what makes a 1M-token context affordable.""" + per_token = spec_kv_bytes_per_token(_spec(num_layers=16), _cfg(kv_cache_dtype=torch.float8_e4m3fn)) + assert per_token == 32 * 1024 # 32 KiB/token + assert per_token * 1_000_000 / 1e9 == pytest.approx(32.77, abs=0.1) # GB at 1M ctx + + +# --- backend capability gate ------------------------------------------------------- +def test_only_separate_width_backends_advertise_quantized_kv(): + # flashinfer's plan() takes q_data_type and kv_data_type independently. + assert attention_backend_info("fi").supports_quantized_kv is True + # triton's kernels assume q.dtype == cache.dtype. + assert attention_backend_info("triton").supports_quantized_kv is False + + +def test_quantized_kv_rejected_on_unsupporting_backend(): + from freetoken.engine.engine import _kv_is_quantized + + assert _kv_is_quantized(_cfg(kv_cache_dtype=torch.float8_e4m3fn)) is True + assert _kv_is_quantized(_cfg()) is False + assert _kv_is_quantized(_cfg(kv_cache_dtype=torch.bfloat16)) is False + + +# --- the store path narrows on write ---------------------------------------------- +@pytest.mark.parametrize("kv_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) +def test_store_kv_casts_into_a_narrower_pool(kv_dtype): + """store_cache is a raw byte-copy, so MHAKVCache.store_kv must narrow bf16 K/V itself.""" + from freetoken.kvcache.mha_pool import MHAKVCache + + pool = MHAKVCache( + num_kv_heads=2, num_layers=1, head_dim=8, num_pages=4, page_size=1, + dtype=kv_dtype, device=torch.device("cpu"), + ) + assert pool.dtype is kv_dtype + k = torch.ones(2, 2, 8, dtype=torch.bfloat16) + v = torch.full((2, 2, 8), 2.0, dtype=torch.bfloat16) + pool.store_kv(k, v, torch.tensor([0, 1], dtype=torch.int32), 0) + + assert pool.k_cache(0).dtype is kv_dtype + assert pool.k_cache(0)[:2].float().eq(1.0).all() + assert pool.v_cache(0)[:2].float().eq(2.0).all() + + +def test_fp8_pool_allocates_half_the_bytes(): + from freetoken.kvcache.mha_pool import MHAKVCache + + def unit(dtype): + return MHAKVCache( + num_kv_heads=4, num_layers=16, head_dim=256, num_pages=8, page_size=1, + dtype=dtype, device=torch.device("cpu"), + ).unit_bytes()[0] + + assert unit(torch.float8_e4m3fn) * 2 == unit(torch.bfloat16) + + +# --- e5m2 vs e4m3 ------------------------------------------------------------------ +def test_e5m2_is_the_unscaled_fp8_choice(): + """Both fp8 dtypes are one byte, so they buy identical bandwidth; only e5m2 has the + dynamic range to survive the pool's unscaled cast. + + Measured on L4/Qwen3-0.6B: e5m2 reproduced the bf16 continuation exactly (271/271 + chars) while e4m3 degenerated after 14. e4m3's finite max is ~448, so ordinary K/V + activations clip; e5m2 reaches ~57344. + """ + assert torch.finfo(torch.float8_e4m3fn).max == pytest.approx(448.0) + assert torch.finfo(torch.float8_e5m2).max == pytest.approx(57344.0) + assert ( + torch.finfo(torch.float8_e5m2).max > torch.finfo(torch.float8_e4m3fn).max + ), "e5m2 must be the wider-range option" + # Same width -> same bytes read per decode step, so preferring e5m2 costs no speed. + assert torch.finfo(torch.float8_e4m3fn).bits == torch.finfo(torch.float8_e5m2).bits == 8