diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 00000000..c6688ab3 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,9 @@ +.git +.github +.pytest_cache +.venv +**/__pycache__ +**/*.pyc +build +dist +tests diff --git a/benchmarks/bench_qwen4_qsa.py b/benchmarks/bench_qwen4_qsa.py new file mode 100644 index 00000000..502358d6 --- /dev/null +++ b/benchmarks/bench_qwen4_qsa.py @@ -0,0 +1,81 @@ +"""Microbenchmark Qwen4-Exp QSA selection and exact sparse attention.""" + +from __future__ import annotations + +import argparse +import statistics + +import torch + +from freetoken.attention.qsa import select_qsa_logical_rows +from freetoken.kernel.triton.qsa import qsa_sparse_gqa + + +def _median_ms(fn, warmup: int, repeats: int) -> float: + for _ in range(warmup): + fn() + samples = [] + for _ in range(repeats): + start = torch.cuda.Event(enable_timing=True) + stop = torch.cuda.Event(enable_timing=True) + start.record() + fn() + stop.record() + stop.synchronize() + samples.append(start.elapsed_time(stop)) + return statistics.median(samples) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--context", type=int, default=262_144) + parser.add_argument("--repeats", type=int, default=50) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if args.context % 4: + raise ValueError("context must divide by QSA's compression ratio of four") + + device = torch.device("cuda") + query = torch.randn(1, 4, 128, dtype=torch.bfloat16, device=device) + compressed = torch.randn( + args.context // 4, 1, 128, dtype=torch.bfloat16, device=device + ) + position = torch.tensor([args.context - 1], dtype=torch.int64, device=device) + selection_ms = _median_ms( + lambda: select_qsa_logical_rows( + query, + compressed, + position, + compress_ratio=4, + token_budget=2048, + ), + 5, + args.repeats, + ) + + attention_query = torch.randn(1, 24, 256, dtype=torch.bfloat16, device=device) + keys = torch.randn(2048, 2, 256, dtype=torch.bfloat16, device=device) + values = torch.randn_like(keys) + rows = torch.arange(2048, dtype=torch.int32, device=device).view(1, -1) + counts = torch.tensor([2048], dtype=torch.int32, device=device) + attention_ms = _median_ms( + lambda: qsa_sparse_gqa( + attention_query, keys, values, rows, counts, 256**-0.5 + ), + 5, + args.repeats, + ) + print( + { + "gpu": torch.cuda.get_device_name(device), + "context": args.context, + "selection_ms_per_layer": round(selection_ms, 3), + "attention_ms_per_layer": round(attention_ms, 3), + "qsa_ms_for_12_layers": round(12 * (selection_ms + attention_ms), 3), + } + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/qwen38-flash-next-rtx3090-128k.json b/benchmarks/qwen38-flash-next-rtx3090-128k.json new file mode 100644 index 00000000..ffe61659 --- /dev/null +++ b/benchmarks/qwen38-flash-next-rtx3090-128k.json @@ -0,0 +1,19 @@ +{ + "date": "2026-08-26", + "platform": "Windows, NVIDIA RTX 3090", + "model": "Qwen3.8-Flash-Next NVFP4 FTW", + "context_length": 131072, + "kv_cache_gib": 3.09, + "moe_cache_size": 2048, + "gpu_used_by_load_gib": 19.974, + "prompt_tokens": 37, + "output_tokens": 127, + "overall_output_tokens_per_second": 3.574, + "steady_decode_tokens_per_second": 5.791, + "moe_cache_miss_rate": 0.1651, + "notes": [ + "This result uses a short text prompt and a single request.", + "Steady decode excludes model load, prefill, and time to first token.", + "The 128K KV allocation permits a 2048-slot expert cache on a 24 GiB GPU." + ] +} diff --git a/benchmarks/qwen38-flash-next-rtx3090-bw.json b/benchmarks/qwen38-flash-next-rtx3090-bw.json new file mode 100644 index 00000000..d82e2f18 --- /dev/null +++ b/benchmarks/qwen38-flash-next-rtx3090-bw.json @@ -0,0 +1,50 @@ +{ + "version": 4, + "timestamp": "2026-08-26T18:36:33-04:00", + "epoch": 1787783793, + "gpu": { + "index": 0, + "name": "NVIDIA GeForce RTX 3090" + }, + "cpu": { + "physical_cores": 32, + "threads_used": 32 + }, + "threshold": 2.0, + "ceilings": { + "cpu_stream_read_gbs": 46.31, + "pcie_linear_h2d_gbs": 12.31, + "pcie_linear_d2h_gbs": 13.16 + }, + "dtypes": {}, + "dtype_kernels": {}, + "workloads": { + "qwen3.8-flash-next": { + "model": { + "name": "qwen3.8-flash-next", + "hidden": 2560, + "inter": 640, + "experts": 512, + "top_k": 10 + }, + "kernels": { + "nvfp4": { + "expert_bytes": 2772480, + "synth_experts": 512, + "cpu_moe_gbs": 0.19, + "cpu_moe_isa": "avx2+vnni(nvfp4-w4a8)", + "isa_sweep": null, + "pcie_gather_gbs": 9.73, + "cpu_moe_overlap_gbs": 0.13, + "pcie_gather_overlap_gbs": 9.68, + "ratio": 0.02, + "recommended": "offload", + "note": null + } + }, + "recommended_moe_backend": { + "nvfp4": "offload" + } + } + } +} diff --git a/benchmarks/qwen38-flash-next-server-validation.json b/benchmarks/qwen38-flash-next-server-validation.json new file mode 100644 index 00000000..48e5cdaa --- /dev/null +++ b/benchmarks/qwen38-flash-next-server-validation.json @@ -0,0 +1,32 @@ +{ + "date": "2026-08-26", + "platform": "Windows, NVIDIA RTX 3090", + "model": "Qwen3.8-Flash-Next NVFP4 FTW", + "served_model_name": "qwen3.8-flash-next-nvfp4", + "context_length": 262144, + "kv_cache_gib": 6.19, + "api": { + "health": "passed", + "models": "passed", + "chat_completions_text": { + "expected": "FREETOKEN_SERVER_OK", + "actual": "FREETOKEN_SERVER_OK", + "prompt_tokens": 21, + "completion_tokens": 6 + }, + "chat_completions_vision": { + "image": "assets/desktop-console.png", + "resolution": "2920x1944", + "expected": "FreeToken Desktop", + "actual": "FreeToken Desktop", + "prompt_tokens": 5589, + "completion_tokens": 6, + "reported_prefill_tokens_per_second": 143.53 + } + }, + "notes": [ + "OpenAI image_url data URI input passed through the multiprocess server.", + "Pixel transport uses BF16 because the vision patch projection converts pixels to BF16 before its first operation.", + "The validation server was stopped after the checks." + ] +} diff --git a/benchmarks/run_qwen4_smoke.py b/benchmarks/run_qwen4_smoke.py new file mode 100644 index 00000000..3936e6a8 --- /dev/null +++ b/benchmarks/run_qwen4_smoke.py @@ -0,0 +1,280 @@ +"""Run one repeatable Qwen3.8-Flash-Next text or image smoke test. + +This harness is intentionally offline. It loads the source checkout while +allowing the installed FreeToken wheel to supply unchanged native extensions. +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import time +from pathlib import Path + +import torch + + +def _nvidia_memory() -> tuple[int, int]: + output = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=memory.free,memory.total", + "--format=csv,noheader,nounits", + ], + text=True, + ) + free_mib, total_mib = (int(value.strip()) for value in output.splitlines()[0].split(",")) + return free_mib * 2**20, total_mib * 2**20 + + +def _bootstrap_native_extensions() -> None: + import freetoken.kernel + + override = os.getenv("FREETOKEN_INSTALLED_KERNEL_DIR") + candidate = ( + Path(override) + if override + else Path(sys.executable).resolve().parent.parent + / "Lib" + / "site-packages" + / "freetoken" + / "kernel" + ) + if candidate.is_dir() and str(candidate) not in freetoken.kernel.__path__: + freetoken.kernel.__path__.append(str(candidate)) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True) + parser.add_argument("--prompt", default="Reply with exactly: FREETOKEN_QWEN4_OK") + parser.add_argument("--synthetic-context-tokens", type=int) + parser.add_argument("--needle", default="FREETOKEN_CONTEXT_48291") + parser.add_argument("--image") + parser.add_argument("--disable-thinking", action="store_true") + parser.add_argument("--max-tokens", type=int, default=32) + parser.add_argument("--max-seq-len", type=int, default=8192) + parser.add_argument("--num-tokens", type=int, default=8192) + parser.add_argument("--max-prefill", type=int, default=2048) + parser.add_argument("--moe-cache-size", type=int, default=1024) + parser.add_argument("--moe-collect-stats", action="store_true") + parser.add_argument("--moe-backend", choices=("offload", "hybrid", "cpu"), default="offload") + parser.add_argument("--moe-cpu-layers") + parser.add_argument("--expert-load", choices=("auto", "serial", "parallel"), default="serial") + return parser.parse_args() + + +def _prepare_prompt( + model_path: str, + prompt: str, + image_path: str | None, + disable_thinking: bool = False, +): + from transformers import AutoProcessor, AutoTokenizer + + if image_path is None: + tokenizer = AutoTokenizer.from_pretrained(model_path) + encoded = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize=True, + add_generation_prompt=True, + enable_thinking=not disable_thinking, + ) + ids = encoded["input_ids"] if hasattr(encoded, "keys") else encoded + if isinstance(ids, torch.Tensor): + ids = ids.reshape(-1).tolist() + elif ids and isinstance(ids[0], list): + ids = ids[0] + ids = [int(token_id) for token_id in ids] + return ids, None + + from PIL import Image + + processor = AutoProcessor.from_pretrained(model_path) + assistant_prefix = "<|im_start|>assistant\n" + if disable_thinking: + assistant_prefix += "\n\n\n\n" + text = ( + "<|im_start|>user\n" + "<|vision_start|><|image_pad|><|vision_end|>" + f"{prompt}<|im_end|>\n{assistant_prefix}" + ) + with Image.open(image_path) as image: + encoded = processor(text=[text], images=[image.convert("RGB")], return_tensors="pt") + ids = encoded["input_ids"][0].tolist() + mm = { + "pixel_values": encoded["pixel_values"], + "image_grid_thw": encoded["image_grid_thw"], + "mm_token_type_ids": encoded["mm_token_type_ids"][0], + } + return ids, mm + + +def _prepare_synthetic_context( + model_path: str, + target_tokens: int, + needle: str, + disable_thinking: bool, +) -> tuple[list[int], None]: + """Build a repeatable needle-recall prompt close to ``target_tokens`` long.""" + from transformers import AutoTokenizer + + if target_tokens < 256: + raise ValueError("--synthetic-context-tokens must be at least 256") + tokenizer = AutoTokenizer.from_pretrained(model_path) + filler = "This sentence is filler and does not contain the verification code. " + prefix = "Read the full document and remember the verification code.\n\n" + needle_text = f"The verification code is {needle}.\n\n" + suffix = "\nWhat is the verification code? Reply with only the code." + + def encode(repetitions: int) -> list[int]: + before = repetitions // 2 + prompt = prefix + filler * before + needle_text + filler * (repetitions - before) + suffix + encoded = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize=True, + add_generation_prompt=True, + enable_thinking=not disable_thinking, + ) + ids = encoded["input_ids"] if hasattr(encoded, "keys") else encoded + if isinstance(ids, torch.Tensor): + ids = ids.reshape(-1).tolist() + elif ids and isinstance(ids[0], list): + ids = ids[0] + return [int(token_id) for token_id in ids] + + low, high = 0, target_tokens + best = encode(0) + while low <= high: + mid = (low + high) // 2 + candidate = encode(mid) + if len(candidate) <= target_tokens: + best = candidate + low = mid + 1 + else: + high = mid - 1 + return best, None + + +def main() -> None: + args = _parse_args() + _bootstrap_native_extensions() + + from freetoken.core import SamplingParams + from freetoken.llm import LLM + + if args.synthetic_context_tokens is not None: + if args.image is not None: + raise ValueError("--synthetic-context-tokens cannot be combined with --image") + prompt_ids, mm = _prepare_synthetic_context( + args.model, + args.synthetic_context_tokens, + args.needle, + disable_thinking=args.disable_thinking, + ) + else: + prompt_ids, mm = _prepare_prompt( + args.model, + args.prompt, + args.image, + disable_thinking=args.disable_thinking, + ) + # Engine requires CUDA to be uninitialized when it selects the process GPU. + # nvidia-smi gives us the baseline without creating a CUDA context. + free_before, total = _nvidia_memory() + load_start = time.perf_counter() + llm = LLM( + args.model, + dtype=torch.bfloat16, + max_running_req=1, + attention_backend="auto", + moe_backend=args.moe_backend, + nvfp4_backend="triton", + expert_load=args.expert_load, + moe_cache_size=args.moe_cache_size, + moe_collect_stats=args.moe_collect_stats, + moe_cpu_layers=args.moe_cpu_layers, + cache_type="naive", + max_seq_len_override=args.max_seq_len, + num_token_override=args.num_tokens, + max_extend_tokens=args.max_prefill, + ) + torch.cuda.synchronize() + load_seconds = time.perf_counter() - load_start + free_loaded, _ = torch.cuda.mem_get_info() + + from freetoken.message import DetokenizeMsg + + token_times: list[float] = [] + original_send_result = llm.send_result + + def timed_send_result(reply): + now = time.perf_counter() + for msg in reply: + if isinstance(msg, DetokenizeMsg) and not ( + msg.finished and msg.next_token in llm.eos_token_ids + ): + token_times.append(now) + original_send_result(reply) + + llm.send_result = timed_send_result + generation_start = time.perf_counter() + result = llm.generate( + [prompt_ids], + SamplingParams(max_tokens=args.max_tokens, temperature=0.0), + mm_inputs=[mm] if mm is not None else None, + )[0] + torch.cuda.synchronize() + generation_seconds = time.perf_counter() - generation_start + output_tokens = len(result["token_ids"]) + cache_stats = None + if args.moe_collect_stats and llm.engine.moe_offload_cache is not None: + cache = llm.engine.moe_offload_cache + cache_stats = cache.decode_miss_stats() + rates = [ + layer["miss_rate"] + for layer in cache.decode_miss_stats_per_layer()["per_layer"] + if layer["steps"] + ] + if rates: + cache_stats["layer_miss_rate_min"] = min(rates) + cache_stats["layer_miss_rate_max"] = max(rates) + ttft_seconds = token_times[0] - generation_start if token_times else None + decode_seconds = token_times[-1] - token_times[0] if len(token_times) > 1 else None + steady_decode_tps = ( + (len(token_times) - 1) / decode_seconds + if decode_seconds is not None and decode_seconds > 0 + else None + ) + print( + "QWEN4_SMOKE_RESULT " + + json.dumps( + { + "prompt_tokens": len(prompt_ids), + "output_tokens": output_tokens, + "load_seconds": round(load_seconds, 3), + "generation_seconds": round(generation_seconds, 3), + "overall_output_tps": round(output_tokens / generation_seconds, 3), + "ttft_seconds": round(ttft_seconds, 3) if ttft_seconds is not None else None, + "decode_seconds": round(decode_seconds, 3) if decode_seconds is not None else None, + "steady_decode_tps": ( + round(steady_decode_tps, 3) if steady_decode_tps is not None else None + ), + "gpu_total_gib": round(total / 2**30, 3), + "gpu_used_by_load_gib": round((free_before - free_loaded) / 2**30, 3), + "moe_backend": args.moe_backend, + "moe_cpu_layers": args.moe_cpu_layers, + "moe_cache_size": args.moe_cache_size, + "moe_cache_stats": cache_stats, + "text": result["text"], + }, + ensure_ascii=False, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/docker/Dockerfile.cuda13 b/docker/Dockerfile.cuda13 new file mode 100644 index 00000000..bcf9f521 --- /dev/null +++ b/docker/Dockerfile.cuda13 @@ -0,0 +1,47 @@ +# syntax=docker/dockerfile:1.7 +ARG UV_VERSION=0.9.3 +ARG CUDA_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 +FROM ghcr.io/astral-sh/uv:${UV_VERSION} AS uv + +FROM ${CUDA_IMAGE} + +ARG INSTALL_CUDA_COMPAT=0 +COPY --from=uv /uv /uvx /usr/local/bin/ + +ENV CUDA_HOME=/usr/local/cuda \ + DEBIAN_FRONTEND=noninteractive \ + LD_LIBRARY_PATH=/usr/local/cuda-13.0/compat:/usr/local/nvidia/lib:/usr/local/nvidia/lib64 \ + MAX_JOBS=8 \ + PATH=/opt/freetoken/.venv/bin:${PATH} \ + PYTHONUNBUFFERED=1 \ + UV_COMPILE_BYTECODE=1 \ + UV_LINK_MODE=copy \ + UV_PYTHON_DOWNLOADS=never + +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + ninja-build \ + python3 \ + python3-dev \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /opt/freetoken +COPY pyproject.toml uv.lock setup.py README.md LICENSE ./ +COPY python ./python + +# The frozen lock keeps torch/CUDA and native-kernel versions reproducible. +# Ampere uses FlashInfer for paged attention; the remaining kernels are Triton. +RUN --mount=type=cache,target=/root/.cache/uv \ + uv sync --frozen --no-dev --extra fi --no-editable + +RUN if [ "${INSTALL_CUDA_COMPAT}" = "1" ]; then \ + apt-get update \ + && apt-get install -y --no-install-recommends cuda-compat-13-0 \ + && rm -rf /var/lib/apt/lists/*; \ + fi + +EXPOSE 1919 +ENTRYPOINT ["ft"] +CMD ["--help"] diff --git a/docker/README.md b/docker/README.md new file mode 100644 index 00000000..5aed2e89 --- /dev/null +++ b/docker/README.md @@ -0,0 +1,27 @@ +# CUDA 13 image + +Build the locked FreeToken runtime from the repository root: + +```bash +docker build -f docker/Dockerfile.cuda13 -t freetoken:cuda13 . +``` + +The image uses CUDA 13.0.2, the exact dependency versions in `uv.lock`, and +FlashInfer's CUDA 13 wheels. Run it with an explicit model and a persistent +Hugging Face cache: + +```bash +docker run --gpus all -p 1919:1919 \ + -v "$HF_HOME:/models" -e HF_HOME=/models \ + freetoken:cuda13 serve --model Qwen/Qwen3.8-Flash-Next-FP8 \ + --host 0.0.0.0 --port 1919 +``` + +CUDA 13 normally requires driver 580 or newer. On hardware supported by +NVIDIA's forward-compatibility package, a driver 570 deployment can opt in: + +```bash +docker build -f docker/Dockerfile.cuda13 \ + --build-arg INSTALL_CUDA_COMPAT=1 \ + -t freetoken:cuda13-compat . +``` diff --git a/docs/models.md b/docs/models.md index e4850a12..cb9d831f 100644 --- a/docs/models.md +++ b/docs/models.md @@ -11,6 +11,7 @@ for them; other checkpoints of the same architectures work too. | GLM-4.7 | [nvidia/GLM-4.7-NVFP4](https://huggingface.co/nvidia/GLM-4.7-NVFP4) | | Qwen3.6 / Qwen3.5 MoE | [Qwen/Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B) ([-FP8](https://huggingface.co/Qwen/Qwen3.6-35B-A3B-FP8)), [nvidia/Qwen3.6-35B-A3B-NVFP4](https://huggingface.co/nvidia/Qwen3.6-35B-A3B-NVFP4), [Qwen/Qwen3.5-35B-A3B](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) ([-FP8](https://huggingface.co/Qwen/Qwen3.5-35B-A3B-FP8)) | | Qwen3.6 dense | [Qwen/Qwen3.6-27B](https://huggingface.co/Qwen/Qwen3.6-27B) ([-FP8](https://huggingface.co/Qwen/Qwen3.6-27B-FP8)), [nvidia/Qwen3.6-27B-NVFP4](https://huggingface.co/nvidia/Qwen3.6-27B-NVFP4) | +| Qwen3.8 Flash Next | [Qwen/Qwen3.8-Flash-Next-FP8](https://huggingface.co/Qwen/Qwen3.8-Flash-Next-FP8) (text-only; exact QSA prefix through 2,048 tokens; host-mapped PLE) | | Qwen3-MoE | [Qwen/Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | | gpt-oss | [openai/gpt-oss-120b](https://huggingface.co/openai/gpt-oss-120b), [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) | | Gemma-4 | [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it), [nvidia/Gemma-4-26B-A4B-NVFP4](https://huggingface.co/nvidia/Gemma-4-26B-A4B-NVFP4), [google/gemma-4-12B-it](https://huggingface.co/google/gemma-4-12B-it), [nvidia/Gemma-4-31B-IT-NVFP4](https://huggingface.co/nvidia/Gemma-4-31B-IT-NVFP4) .. | @@ -38,3 +39,8 @@ for them; other checkpoints of the same architectures work too. - DeepSeek-V4 checkpoints must keep the `inference/config.json` subdir — the authoritative model args are read from there. - Multimodal checkpoints are served text-only. +- Qwen3.8 Flash Next requires an offload-family MoE backend. FreeToken + automatically disables CUDA graphs and selects the naive cache because PLE + performs host-side gathers and owns per-request convolution state. The FP8 + checkpoint occupies about 173 GiB on disk; routed-expert banks use about + 113 GiB of pinned host RAM while the 48 GiB PLE table remains mmap-backed. diff --git a/pyproject.toml b/pyproject.toml index 8bd653f8..9391ebb7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,7 @@ dependencies = [ "numpy>=2.0,<2.5", "openai>=2.0,<3", "partial-json-parser>=0.2,<1", + "Pillow>=11,<13", "prompt_toolkit>=3.0,<4", "pydantic>=2.9,<3", "pyzmq>=27,<28", @@ -55,6 +56,7 @@ dependencies = [ # PyPI's torch 2.11.0 wheel is itself the cu130 build, so plain pip resolves # correctly from PyPI alone; uv additionally pins the index below. "torch>=2.11,<2.12", + "torchvision>=0.26,<0.27", "tqdm>=4.66,<5", "transformers>=5.5,<6", "triton==3.6.0; platform_system == 'Linux'", @@ -91,6 +93,7 @@ accel = ["freetoken[fi,sgl]"] # which would otherwise shadow PyPI under uv's first-index strategy). [tool.uv.sources] torch = { index = "pytorch-cu130" } +torchvision = { index = "pytorch-cu130" } sglang-kernel = { index = "sglang-cu130" } [[tool.uv.index]] diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py index 746c04c4..657868bc 100644 --- a/python/freetoken/attention/__init__.py +++ b/python/freetoken/attention/__init__.py @@ -132,6 +132,22 @@ def create_m3_sparse_backend(config: ModelConfig): return M3SparseAttnBackend(config) +@SUPPORTED_ATTENTION_BACKENDS.register( + "qsa", + BackendInfo( + supported_types=frozenset({AttnType.QSA}), + # A 64-token full page becomes a 16-row compressed page at ratio 4. + # Keeping both page sizes aligned makes full_row // ratio exact. + page_sizes=(64,), + hybrid_linear_ok=True, + ), +) +def create_qsa_backend(config: ModelConfig): + from .qsa import QSAAttnBackend + + return QSAAttnBackend(config) + + def attention_backend_info(name: str) -> BackendInfo: return SUPPORTED_ATTENTION_BACKENDS.info(name) diff --git a/python/freetoken/attention/base.py b/python/freetoken/attention/base.py index eb39d472..a35c685e 100644 --- a/python/freetoken/attention/base.py +++ b/python/freetoken/attention/base.py @@ -24,6 +24,9 @@ class AttnType(str, Enum): # GQA block-sparse (MiniMax-M3): paged GQA K/V + a per-sparse-layer index-key # slab; the indexer picks top-k 128-token blocks per query -> BSAKVCache BSA = "bsa" + # Qwen sparse attention: paged GQA K/V plus one compressed index key per + # ``indexer_compress_ratio`` tokens for every sparse-attention layer. + QSA = "qsa" @property def backend_driven(self) -> bool: diff --git a/python/freetoken/attention/qsa.py b/python/freetoken/attention/qsa.py new file mode 100644 index 00000000..17cbf203 --- /dev/null +++ b/python/freetoken/attention/qsa.py @@ -0,0 +1,392 @@ +"""Qwen compressed sparse-attention backend. + +QSA keeps exact full-resolution K/V. Its small four-head indexer scores one +compressed key for each four-token group, selects 512 groups, expands them to +2048 token rows, and appends the current incomplete group. The final attention +therefore uses the model's original K/V values without approximation. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, List + +import torch +from freetoken.core import Batch, get_global_ctx + +from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata + +if TYPE_CHECKING: + from freetoken.models import ModelConfig + +_SCORE_WORKSPACE_BYTES = 128 << 20 + + +@dataclass +class QSAMetadata(BaseAttnMetadata): + cu_seqlens_q: torch.Tensor + cu_seqlens_q_host: tuple[int, ...] + logical_positions: torch.Tensor + last_indices: torch.Tensor + compressed_rows: tuple[torch.Tensor, ...] + + def get_last_indices(self, bs: int) -> torch.Tensor: + return self.last_indices[:bs] + + +def _compact_expanded_selection( + block_indices: torch.Tensor, + query_positions: torch.Tensor, + *, + compress_ratio: int, + token_budget: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Expand compressed blocks and append the visible incomplete-group tail.""" + if block_indices.is_cuda: + from freetoken.kernel.triton.qsa import compact_qsa_blocks + + return compact_qsa_blocks( + block_indices, + query_positions, + compress_ratio=compress_ratio, + token_budget=token_budget, + ) + rows = block_indices.shape[0] + device = block_indices.device + offsets = torch.arange(compress_ratio, device=device, dtype=torch.long) + expanded = block_indices.long().unsqueeze(-1) * compress_ratio + offsets + expanded = torch.where( + block_indices.long().unsqueeze(-1) >= 0, + expanded, + torch.full_like(expanded, -1), + ).reshape(rows, -1)[:, :token_budget] + positions = query_positions.to(device=device, dtype=torch.long) + expanded = torch.where( + (expanded >= 0) & (expanded <= positions.unsqueeze(1)), + expanded, + torch.full_like(expanded, -1), + ) + + tail_offsets = torch.arange(compress_ratio - 1, device=device, dtype=torch.long) + visible = positions + 1 + tail_start = torch.div(visible, compress_ratio, rounding_mode="floor") * compress_ratio + tail_count = visible - tail_start + tail = tail_start.unsqueeze(1) + tail_offsets.unsqueeze(0) + tail = torch.where( + tail_offsets.unsqueeze(0) < tail_count.unsqueeze(1), + tail, + torch.full_like(tail, -1), + ) + result = torch.full( + (rows, token_budget + compress_ratio - 1), + -1, + dtype=torch.long, + device=device, + ) + result[:, :token_budget].copy_(expanded) + # topk is sorted, so every finite block precedes its -inf/-1 padding. + # Insert the incomplete tail directly after those expanded blocks instead + # of sorting all 2,051 columns on every layer and decode step. + block_counts = (block_indices >= 0).sum(dim=1) + tail_columns = block_counts.unsqueeze(1) * compress_ratio + tail_offsets.unsqueeze(0) + tail_live = tail_offsets.unsqueeze(0) < tail_count.unsqueeze(1) + row_ids = torch.arange(rows, device=device).unsqueeze(1).expand_as(tail_columns) + result[row_ids[tail_live], tail_columns[tail_live]] = tail[tail_live] + counts = (block_counts * compress_ratio + tail_count).to(torch.int32) + return result.to(torch.int32), counts + + +def select_qsa_logical_rows( + index_q: torch.Tensor, + compressed_keys: torch.Tensor, + query_positions: torch.Tensor, + *, + compress_ratio: int, + token_budget: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Score compressed keys and return exact logical token selections. + + Index scores are ``sum_h(relu(q_h dot k)) / sqrt(dim)``. Row tiling + bounds the FP32 score workspace without changing the per-row top-k result. + """ + rows, heads, dim = index_q.shape + blocks = compressed_keys.shape[0] + block_budget = token_budget // compress_ratio + output_blocks = torch.full( + (rows, block_budget), -1, dtype=torch.int32, device=index_q.device + ) + if rows == 0: + return _compact_expanded_selection( + output_blocks, + query_positions, + compress_ratio=compress_ratio, + token_budget=token_budget, + ) + if blocks: + # One four-head dot tensor and one reduced logits matrix are live at once. + bytes_per_row = max(blocks * torch.float32.itemsize * (heads + 1), 1) + row_chunk = max(1, min(rows, _SCORE_WORKSPACE_BYTES // bytes_per_row)) + keys = compressed_keys[:, 0].transpose(0, 1) + columns = torch.arange(blocks, device=index_q.device).unsqueeze(0) + for start in range(0, rows, row_chunk): + stop = min(start + row_chunk, rows) + queries = index_q[start:stop].reshape((stop - start) * heads, dim) + if queries.is_cuda and queries.dtype in (torch.bfloat16, torch.float16): + # Tensor-core BF16/FP16 inputs with FP32 accumulation/output. + # This avoids materializing a full FP32 copy of the long index + # cache on every sparse layer and decode step. + dots = torch.mm(queries, keys, out_dtype=torch.float32) + else: + dots = queries.float() @ keys.float() + dots = dots.view(stop - start, heads, blocks) + logits = torch.relu_(dots).sum(dim=1) + logits.mul_(dim**-0.5) + visible_blocks = torch.div( + query_positions[start:stop].to(torch.long) + 1, + compress_ratio, + rounding_mode="floor", + ) + logits.masked_fill_(columns >= visible_blocks.unsqueeze(1), -float("inf")) + width = min(block_budget, blocks) + if width: + scores, picks = torch.topk(logits, width, dim=1) + picks = torch.where( + torch.isfinite(scores), picks, torch.full_like(picks, -1) + ) + output_blocks[start:stop, :width] = picks.to(torch.int32) + return _compact_expanded_selection( + output_blocks, + query_positions, + compress_ratio=compress_ratio, + token_budget=token_budget, + ) + + +class QSAAttnBackend(BaseAttnBackend): + def __init__(self, config: ModelConfig) -> None: + from freetoken.kvcache.qsa_pool import QSAKVCache + + self.config = config + self.args = config.qwen4_args + self.kvcache = get_global_ctx().kv_cache + if not isinstance(self.kvcache, QSAKVCache): + raise TypeError(f"qsa backend needs QSAKVCache, got {type(self.kvcache).__name__}") + self.device = self.kvcache.device + self.compress_ratio = int(self.args.indexer_compress_ratio) + self.token_budget = int(self.args.indexer_budget) + if self.token_budget % self.compress_ratio: + raise ValueError("QSA token budget must divide by its compression ratio") + + def prepare_metadata(self, batch: Batch) -> None: + reqs = batch.padded_reqs if hasattr(batch, "padded_reqs") else batch.reqs + lengths = [int(req.extend_len) for req in reqs] + cu_host = [0] + for length in lengths: + cu_host.append(cu_host[-1] + length) + cu = torch.tensor(cu_host, dtype=torch.int32, device=self.device) + logical = torch.cat( + [ + torch.arange(req.cached_len, req.device_len, device=self.device) + for req in reqs + if req.extend_len + ], + dim=0, + ) if sum(lengths) else torch.empty(0, dtype=torch.int64, device=self.device) + page_table = get_global_ctx().page_table + compressed_rows: list[torch.Tensor] = [] + for req in reqs: + complete_blocks = int(req.device_len) // self.compress_ratio + if complete_blocks: + starts = torch.arange( + complete_blocks, device=self.device, dtype=torch.long + ).mul_(self.compress_ratio) + full_rows = page_table[int(req.table_idx)].index_select(0, starts) + rows = torch.div( + full_rows.to(torch.int64), + self.compress_ratio, + rounding_mode="floor", + ) + else: + rows = torch.empty(0, dtype=torch.int64, device=self.device) + compressed_rows.append(rows) + batch.attn_metadata = QSAMetadata( + cu_seqlens_q=cu, + cu_seqlens_q_host=tuple(cu_host), + logical_positions=logical, + last_indices=cu[1:] - 1, + compressed_rows=tuple(compressed_rows), + ) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer_id: int, + batch: Batch, + attn_spec: AttentionSpec | None = None, + ) -> torch.Tensor: + raise NotImplementedError("Qwen4-Exp QSA layers call qsa_forward()") + + def _compress_current_keys(self, indexer, index_k, layer_id: int, batch: Batch) -> None: + md = batch.attn_metadata + assert isinstance(md, QSAMetadata) + reqs = batch.padded_reqs if hasattr(batch, "padded_reqs") else batch.reqs + if reqs: + self.kvcache.ensure_pending_capacity(max(int(req.table_idx) for req in reqs) + 1) + page_table = get_global_ctx().page_table + pooled: list[torch.Tensor] = [] + rope_positions: list[torch.Tensor] = [] + compressed_rows: list[torch.Tensor] = [] + cu = md.cu_seqlens_q_host + ratio = self.compress_ratio + + for req_id, req in enumerate(reqs): + begin, stop = cu[req_id], cu[req_id + 1] + if begin == stop: + continue + start, end = int(req.cached_len), int(req.device_len) + request_row = int(req.table_idx) + current = index_k[begin:stop] + batch_rope = getattr(batch, "rope_positions", None) + if batch_rope is None: + current_rope = md.logical_positions[begin:stop].view(-1, 1).expand(-1, 3) + else: + current_rope = batch_rope[:, begin:stop].transpose(0, 1) + if start == 0: + self.kvcache.clear_pending(layer_id, request_row) + + first_end = ((start + ratio) // ratio) * ratio - 1 + for group_end in range(first_end, end, ratio): + group_start = group_end - ratio + 1 + if group_start < start: + prior_pos = torch.arange(group_start, start, device=self.device) + prior = self.kvcache.pending_group(layer_id, request_row, prior_pos) + first_rope = self.kvcache.pending_rope_group( + layer_id, request_row, prior_pos[:1] + )[0] + current_part = current[: group_end - start + 1] + members = torch.cat((prior, current_part), dim=0) + else: + lo = group_start - start + members = current[lo : lo + ratio] + first_rope = current_rope[lo] + if members.shape[0] != ratio: + raise RuntimeError("QSA compression received an incomplete key group") + pooled.append(members.float().mean(dim=0).to(index_k.dtype)) + rope_positions.append(first_rope) + full_row = page_table[request_row, group_start].to(torch.int64) + compressed_rows.append(torch.div(full_row, ratio, rounding_mode="floor")) + + # Only the newest occurrence of each modulo slot is needed. This + # avoids duplicate-index writes for a long prefill chunk. + keep_start = max(start, end - ratio) + keep_positions = torch.arange(keep_start, end, device=self.device) + self.kvcache.store_pending( + layer_id, + request_row, + keep_positions, + current[keep_start - start :], + current_rope[keep_start - start :], + ) + + if pooled: + pooled_tensor = torch.stack(pooled) + positions = torch.stack(rope_positions).transpose(0, 1).contiguous() + normalized = indexer.normalize_compressed_keys(pooled_tensor, positions) + rows = torch.stack(compressed_rows).to(torch.int64) + self.kvcache.store_compressed_k(normalized, rows, layer_id) + + def _select_physical_rows( + self, index_q: torch.Tensor, layer_id: int, batch: Batch + ) -> tuple[torch.Tensor, torch.Tensor]: + md = batch.attn_metadata + assert isinstance(md, QSAMetadata) + reqs = batch.padded_reqs if hasattr(batch, "padded_reqs") else batch.reqs + page_table = get_global_ctx().page_table + cu = md.cu_seqlens_q_host + selections: list[torch.Tensor] = [] + counts: list[torch.Tensor] = [] + ratio = self.compress_ratio + compressed_pool = self.kvcache.compressed_k_cache(layer_id) + + for req_id, req in enumerate(reqs): + begin, stop = cu[req_id], cu[req_id + 1] + if begin == stop: + continue + positions = md.logical_positions[begin:stop] + compressed_rows = md.compressed_rows[req_id] + if compressed_rows.numel(): + compressed_keys = compressed_pool.index_select(0, compressed_rows) + else: + compressed_keys = compressed_pool[:0] + logical, live = select_qsa_logical_rows( + index_q[begin:stop], + compressed_keys, + positions, + compress_ratio=ratio, + token_budget=self.token_budget, + ) + safe = logical.clamp_min(0).long() + physical = page_table[int(req.table_idx)].index_select( + 0, safe.reshape(-1) + ).reshape_as(logical) + physical = torch.where(logical >= 0, physical, torch.full_like(physical, -1)) + selections.append(physical.to(torch.int32)) + counts.append(live) + if not selections: + width = self.token_budget + ratio - 1 + return ( + torch.empty((0, width), dtype=torch.int32, device=self.device), + torch.empty(0, dtype=torch.int32, device=self.device), + ) + return torch.cat(selections, dim=0), torch.cat(counts, dim=0) + + def qsa_forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + index_q: torch.Tensor, + index_k: torch.Tensor, + indexer, + layer_id: int, + batch: Batch, + ) -> torch.Tensor: + from freetoken.kernel.triton.qsa import qsa_sparse_gqa + + self.kvcache.store_kv(k, v, batch.out_loc, layer_id) + self._compress_current_keys(indexer, index_k, layer_id, batch) + selected, counts = self._select_physical_rows(index_q, layer_id, batch) + # k is flattened [N, kv_heads * head_dim], so take the pool geometry + # directly instead of inferring a head count from the flattened input. + k_raw = self.kvcache.k_cache(layer_id) + v_raw = self.kvcache.v_cache(layer_id) + k_rows = k_raw.view(-1, k_raw.shape[-2], k_raw.shape[-1]) + v_rows = v_raw.view(-1, v_raw.shape[-2], v_raw.shape[-1]) + return qsa_sparse_gqa( + q, + k_rows, + v_rows, + selected, + counts, + q.shape[-1] ** -0.5, + ) + + def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: + # Qwen4-Exp disables CUDA graphs because PLE owns per-request recurrent state. + return None + + def prepare_for_capture(self, batch: Batch) -> None: + self.prepare_metadata(batch) + + def prepare_for_replay(self, batch: Batch) -> None: + self.prepare_metadata(batch) + + +__all__ = [ + "QSAAttnBackend", + "QSAMetadata", + "select_qsa_logical_rows", +] diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py index 420faf2c..2f643bca 100644 --- a/python/freetoken/checkpoint/convert.py +++ b/python/freetoken/checkpoint/convert.py @@ -17,6 +17,7 @@ import glob import hashlib +import json import os import shutil import threading @@ -64,6 +65,53 @@ def _source_fingerprint(model_path: str, model_config, *, device) -> str: _SKIP_NAMES = ("model.safetensors.index.json",) # indexes shards the FTW replaces # .freetoken_expert_cache: the legacy per-bank cache (can be tens of GB of stale .bin) _SKIP_DIRS = (".git", ".cache", ".freetoken_expert_cache") +_QWEN4_PLE_FRAGMENT = ".ple.ple_embedding.ngram_embedding." + + +def _copy_host_mapped_weights(model_path: str, out_dir: str) -> list[str]: + """Preserve large CPU-mapped weights that are intentionally absent from FTW. + + Qwen4-Exp PLE tables are random-access host embeddings. They must remain + memory-mapped safetensors instead of being copied into resident RAM or GPU + weights. Copy only the shards named by those keys and write a small index + that contains only the preserved host weights. + """ + index_path = os.path.join(model_path, "model.safetensors.index.json") + if not os.path.isfile(index_path): + return [] + with open(index_path, encoding="utf-8") as index_file: + index = json.load(index_file) + source_map = index.get("weight_map", {}) + weight_map = { + name: filename + for name, filename in source_map.items() + if _QWEN4_PLE_FRAGMENT in name + } + if not weight_map: + return [] + + copied: list[str] = [] + total_size = 0 + for rel in sorted(set(weight_map.values())): + src = os.path.join(model_path, rel) + if not os.path.isfile(src): + raise FileNotFoundError(f"host-mapped checkpoint shard is missing: {src}") + dst = os.path.join(out_dir, rel) + os.makedirs(os.path.dirname(dst), exist_ok=True) + shutil.copy2(src, dst) + copied.append(rel) + total_size += os.path.getsize(src) + + slim_index = { + "metadata": {"total_size": total_size, "freetoken_host_mapped_only": True}, + "weight_map": weight_map, + } + dst_index = os.path.join(out_dir, "model.safetensors.index.json") + with open(dst_index, "w", encoding="utf-8") as index_file: + json.dump(slim_index, index_file, indent=2, sort_keys=True) + index_file.write("\n") + copied.append("model.safetensors.index.json") + return copied def _copy_metadata(model_path: str, out_dir: str) -> list[str]: @@ -100,6 +148,7 @@ def _copy_metadata(model_path: str, out_dir: str) -> list[str]: os.makedirs(os.path.dirname(dst), exist_ok=True) shutil.copy2(src, dst) copied.append(rel) + copied.extend(_copy_host_mapped_weights(model_path, out_dir)) return copied diff --git a/python/freetoken/checkpoint/ftw.py b/python/freetoken/checkpoint/ftw.py index a365e3d5..dcd6fe5b 100644 --- a/python/freetoken/checkpoint/ftw.py +++ b/python/freetoken/checkpoint/ftw.py @@ -37,6 +37,7 @@ import mmap import os import re +import sys import threading from concurrent.futures import ThreadPoolExecutor @@ -53,7 +54,6 @@ DEFAULT_SHARD_LIMIT = 8 << 30 # 8 GiB; must be a multiple of ALIGN _SHARD_FMT = "freetoken-{:05d}.ftw" _DEFAULT_CHUNK = 8 << 20 -_BANK_CONCURRENCY = 4 _ALPHA_NAMES = ("gate_up_alpha", "down_alpha") # Per-layer expert-bank entry name (converter streaming path, see checkpoint/convert.py): # each layer of a bank is its own FTW tensor instead of one flat [num_layers*E, ...] region. @@ -261,7 +261,10 @@ def _map(self, file: str) -> memoryview: if entry is None: fd = os.open(os.path.join(self.dir, file), os.O_RDONLY) try: - m = mmap.mmap(fd, 0, prot=mmap.PROT_READ) + if sys.platform == "win32": + m = mmap.mmap(fd, 0, access=mmap.ACCESS_READ) + else: + m = mmap.mmap(fd, 0, prot=mmap.PROT_READ) finally: os.close(fd) # the mapping keeps its own reference to the file try: @@ -276,10 +279,21 @@ def close(self) -> None: for fd in self._fds.values(): os.close(fd) self._fds.clear() - for m, mv in self._maps.values(): + self.drop_maps() + + def drop_maps(self) -> None: + """Close cached source mappings after all active reads have finished. + + Windows counts touched file mappings in the process working set. Keeping + every FTW shard mapped while also filling the 60+ GiB pinned destination + bank can exhaust RAM and force paging. Streaming callers use this method + between entries so only the current source shard stays mapped. + """ + maps = self._maps + self._maps = {} + for m, mv in maps.values(): mv.release() m.close() - self._maps.clear() def _pieces(self, global_off: int, nbytes: int): """Yield (file, file_off, dest_off, length) covering [global_off, +nbytes), @@ -382,6 +396,8 @@ def _producer(): for e in entries: buf = _transient_buffer(e["nbytes"]) reader.read_into(memoryview(buf), e, workers=workers, chunk=chunk) + if sys.platform == "win32": + reader.drop_maps() dt = _dtype_of(e["dtype"]) t = torch.frombuffer(buf, dtype=dt, count=e["nbytes"] // _elsize(dt)) if not _put((e["name"], t.view(*e["shape"]) if e["shape"] else t, buf, e["nbytes"])): @@ -568,12 +584,27 @@ def _read_layer(job): pins.submit(bank, residency[layer_id]) bar.update(entry["nbytes"]) - with ThreadPoolExecutor(min(max(_BANK_CONCURRENCY, 16), max(n_jobs, 1))) as ex: - futures = [ex.submit(_read_alpha, e) for e in alpha_entries] - futures += [ex.submit(_read_row, job) for job in row_jobs] - futures += [ex.submit(_read_layer, job) for job in layer_jobs] - for f in futures: - f.result() + if sys.platform == "win32": + # Windows has no O_DIRECT path. Stream one physically ordered + # entry at a time and unmap its source shard immediately. The + # previous four-entry queue retained touched source mappings while + # the 63+ GiB pinned destination stayed resident, which drove a + # 128 GiB system into paging. read_into still uses its inner copy + # pool, and PinPipeline still overlaps registration with reads. + jobs = [(e["global_off"], _read_alpha, e) for e in alpha_entries] + jobs += [(job[2], _read_row, job) for job in row_jobs] + jobs += [(job[2]["global_off"], _read_layer, job) for job in layer_jobs] + for _off, read, job in sorted(jobs, key=lambda item: item[0]): + read(job) + reader.drop_maps() + else: + ordered_layer_jobs = sorted(layer_jobs, key=lambda job: job[2]["global_off"]) + with ThreadPoolExecutor(min(16, max(n_jobs, 1))) as ex: + futures = [ex.submit(_read_alpha, e) for e in alpha_entries] + futures += [ex.submit(_read_row, job) for job in row_jobs] + futures += [ex.submit(_read_layer, job) for job in ordered_layer_jobs] + for f in futures: + f.result() finally: bar.close() reader.close() diff --git a/python/freetoken/core.py b/python/freetoken/core.py index ef0a539c..51717c94 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -43,6 +43,11 @@ class Req: # Optional precomputed multimodal soft-token embeddings (GPU, [num_image_tokens, # hidden]) scattered at image-token positions during this request's prefill. mm_embeds: torch.Tensor | None = None + # Optional Qwen multimodal RoPE coordinates for the original prompt, + # shaped [3, prompt_tokens] on CPU. Generated tokens use + # ``logical_position + mrope_position_delta`` on all three axes. + mrope_position_ids: torch.Tensor | None = None + mrope_position_delta: int = 0 # --- hybrid-radix (GDN linear-state) per-request slots; None for non-hybrid models or # until allocated from LinearStatePool. Set by the scheduler (P2). --- @@ -116,6 +121,9 @@ class Batch: # these fields should be set by scheduler input_ids: torch.Tensor = field(init=False) positions: torch.Tensor = field(init=False) + # Model-space rotary coordinates. None means ordinary 1-D ``positions``. + # Qwen VL uses [3, tokens] temporal/height/width coordinates. + rope_positions: torch.Tensor | None = field(default=None, init=False) out_loc: torch.Tensor | None = field(init=False) # Per-(padded-)request table_idx as a GPU int64 tensor, used by GatedDeltaNet # decode to gather/scatter recurrent+conv state without host-side loops (so the diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index cd6505d2..dd9a8499 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -128,6 +128,8 @@ def _resolve_auto_attention_backend( candidates.append(("dsa", True)) if AttnType.BSA in required: candidates.append(("m3_sparse", True)) + if AttnType.QSA in required: + candidates.append(("qsa", True)) if AttnType.SWA in required: candidates.append(("triton", True)) if AttnType.FULL in required: @@ -176,7 +178,10 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att if missing: valid = [ name - for name in ("fa", "fi", "trtllm", "triton", "dsa", "dsv4_sparse", "m3_sparse") + for name in ( + "fa", "fi", "trtllm", "triton", "dsa", "dsv4_sparse", + "m3_sparse", "qsa", + ) if required <= attention_backend_info(name).supported_types ] missing_names = "/".join(sorted(t.value for t in missing)) @@ -323,6 +328,11 @@ def __init__(self, config: EngineConfig): with torch.device("meta"), torch_dtype(config.dtype): self.model = create_model(config.model_config) self.model.load_state_dict(self._load_weight_state_dict(config)) + if hasattr(self.model, "load_host_weights"): + self.model.load_host_weights( + config.model_path, + dummy=config.use_dummy_weight, + ) post_weights_free = self._sync_get_memory()[0] self._weights_bytes = self._baseline_free - post_weights_free # Pool-budget baseline for the desktop cache sliders: free VRAM after the weights are @@ -1161,14 +1171,57 @@ def _cpu_moe_executor_viable(model_config) -> bool: def _pin_budget_bytes() -> int | None: """Bytes this process can safely cudaHostRegister, or None when the platform does not cap pinning (plain Linux). - WSL's WDDM-backed CUDA caps pinning near half of RAM, shared across processes -- budget 40%. FREETOKEN_PIN_BUDGET_GB overrides anywhere.""" + Native Windows and WSL both use WDDM-backed CUDA, which caps registered host + memory near half of physical RAM and shares that pool across processes. Keep + 20% headroom by budgeting 40%. ``FREETOKEN_PIN_BUDGET_GB`` overrides this on + every platform.""" if env := os.environ.get("FREETOKEN_PIN_BUDGET_GB"): return int(float(env) * 2**30) - if not hasattr(os, "uname") or "microsoft" not in os.uname().release.lower(): # WSL kernel tag + + if os.name == "nt": + total = _windows_total_physical_memory() + return int(total * 0.4) if total is not None else None + + if not hasattr(os, "uname") or "microsoft" not in os.uname().release.lower(): return None return int(os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") * 0.4) +def _windows_total_physical_memory() -> int | None: + """Return native-Windows physical RAM through ``GlobalMemoryStatusEx``. + + This stays stdlib-only because the server must make the pin-budget decision + before optional monitoring packages are available. ``None`` is a defensive + fallback for an unexpected Win32 API failure; callers then retain the old + explicit-override behavior through ``FREETOKEN_PIN_BUDGET_GB``. + """ + if os.name != "nt": + return None + + import ctypes + + class _MemoryStatusEx(ctypes.Structure): + _fields_ = [ + ("dwLength", ctypes.c_ulong), + ("dwMemoryLoad", ctypes.c_ulong), + ("ullTotalPhys", ctypes.c_ulonglong), + ("ullAvailPhys", ctypes.c_ulonglong), + ("ullTotalPageFile", ctypes.c_ulonglong), + ("ullAvailPageFile", ctypes.c_ulonglong), + ("ullTotalVirtual", ctypes.c_ulonglong), + ("ullAvailVirtual", ctypes.c_ulonglong), + ("ullAvailExtendedVirtual", ctypes.c_ulonglong), + ] + + status = _MemoryStatusEx() + status.dwLength = ctypes.sizeof(_MemoryStatusEx) + try: + ok = ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)) + except (AttributeError, OSError): + return None + return int(status.ullTotalPhys) if ok else None + + def _auto_cpu_layers(config: EngineConfig, num_moe_layers: int) -> frozenset[int]: """Pick CPU (locked) MoE layers automatically when the banks exceed the pin budget. @@ -1220,6 +1273,8 @@ def override(attr: str, value: Any): # this is dangerous, use with caution model_config = config.model_config single_stream_only = getattr(model_config, "single_stream_only", False) + requires_naive_cache = getattr(model_config, "requires_naive_cache", False) + supports_cuda_graph = getattr(model_config, "supports_cuda_graph", True) is_dsv4 = getattr(model_config, "dsv4_args", None) is not None has_swa_attention = getattr(model_config, "has_swa_attention", False) has_linear_attention = getattr(model_config, "has_linear_attention", False) @@ -1259,6 +1314,14 @@ def override(attr: str, value: Any): # this is dangerous, use with caution override("cuda_graph_bs", [1]) override("cuda_graph_max_bs", 1) + if not supports_cuda_graph: + override("cuda_graph_bs", []) + override("cuda_graph_max_bs", 0) + logger.info_rank0( + f"CUDA graphs disabled for {getattr(model_config, 'model_type', 'model')}: " + "the model requires host-side work during forward" + ) + if config.cuda_graph_max_bs is None: override("cuda_graph_max_bs", config.max_running_req) @@ -1281,6 +1344,14 @@ def override(attr: str, value: Any): # this is dangerous, use with caution ) override("cache_type", "swa_radix") + if requires_naive_cache and getattr(config, "cache_type", "radix") != "naive": + override("cache_type", "naive") + logger.warning_rank0( + f"Cache type overridden to 'naive' for " + f"{getattr(model_config, 'model_type', 'model')}: model-owned runtime state " + "cannot be restored from radix prefixes" + ) + if has_linear_attention: override( "cache_type", @@ -1292,11 +1363,11 @@ def override(attr: str, value: Any): # this is dangerous, use with caution # comma part must serve every required type, with packages/arch available. required_attn_types = _required_attn_types(model_config) _dtype = getattr(config, "dtype", None) # duck-typed test configs omit it - if AttnType.BSA in required_attn_types and _dtype is not None and _dtype.itemsize != 2: + if required_attn_types & {AttnType.BSA, AttnType.QSA} and _dtype is not None and _dtype.itemsize != 2: # Reject at config time: the BSA pool's own assert only fires after the # model is resident (and not at all under `python -O`). raise ValueError( - f"--dtype {config.dtype}: block-sparse attention serves 16-bit " + f"--dtype {config.dtype}: sparse attention serves 16-bit " "compute only (the index slab budgets 2 bytes/token); use bfloat16 " "or float16." ) diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index 4f202502..1215fd9c 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -24,6 +24,7 @@ class GraphCaptureBuffer: input_ids: torch.Tensor out_loc: torch.Tensor positions: torch.Tensor + rope_positions: torch.Tensor logits: torch.Tensor table_idx: torch.Tensor # per-request slot id for GatedDeltaNet state gather/scatter # Decode GDN query indptr = arange(bs+1); a constant per captured bs, filled once. @@ -35,6 +36,7 @@ def init(cls, bs: int, vocab_size: int, device: torch.device) -> GraphCaptureBuf input_ids=torch.zeros(bs, dtype=torch.int32, device=device), out_loc=torch.zeros(bs, dtype=torch.int32, device=device), positions=torch.zeros(bs, dtype=torch.int32, device=device), + rope_positions=torch.zeros((3, bs), dtype=torch.int64, device=device), logits=torch.empty(bs, vocab_size, dtype=torch.float32, device=device), table_idx=torch.zeros(bs, dtype=torch.int32, device=device), fla_cu_seqlens=torch.arange(bs + 1, dtype=torch.int32, device=device), @@ -48,6 +50,7 @@ def set_batch(self, batch: Batch) -> None: batch.input_ids = self.input_ids[_slice] batch.out_loc = self.out_loc[_slice] batch.positions = self.positions[_slice] + batch.rope_positions = self.rope_positions[:, _slice] batch.linear_table_idx = self.table_idx[_slice] # Decode GDN metadata reads the persistent cu_seqlens (constant arange) and the # persistent table_idx slot map, so the captured kernels see stable addresses. @@ -61,6 +64,10 @@ def copy_from(self, batch: Batch) -> None: if batch.out_loc is not None: self.out_loc[_slice] = batch.out_loc self.positions[_slice] = batch.positions + rope_positions = batch.rope_positions + if rope_positions is None: + rope_positions = batch.positions.to(torch.int64).expand(3, -1) + self.rope_positions[:, _slice] = rope_positions if batch.linear_table_idx is not None: self.table_idx[_slice] = batch.linear_table_idx diff --git a/python/freetoken/kernel/__init__.py b/python/freetoken/kernel/__init__.py index 5dc51571..7da52812 100644 --- a/python/freetoken/kernel/__init__.py +++ b/python/freetoken/kernel/__init__.py @@ -1,3 +1,17 @@ +import os as _os + + +_installed_kernel_dir = _os.environ.get("FREETOKEN_INSTALLED_KERNEL_DIR") +if ( + _installed_kernel_dir + and _os.path.isdir(_installed_kernel_dir) + and _installed_kernel_dir not in __path__ +): + # A source checkout can reuse the matching native extensions from an + # installed FreeToken runtime. Spawned server workers import this package + # again, so the path must be applied here instead of only in a launcher. + __path__.append(_installed_kernel_dir) + from .index import indexing from .fast_index_copy import fast_index_copy_jit, update_copy_flag_jit from .moe_impl import ( diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py index a00154d0..c0c786b2 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/python/freetoken/kernel/aot_models.py @@ -128,6 +128,16 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int # https://huggingface.co//raw/main/config.json. SUPPORTED_MODELS: tuple[AotModel, ...] = ( # ---- MoE checkpoints (offload expert banks) ---- + AotModel( + name="Qwen/Qwen3.8-Flash-Next", + architecture="Qwen4ExpForConditionalGeneration", + hidden_size=2560, + kv_groups=((2, 256),), + top_k=10, + moe_intermediate_size=640, + expert_formats=("nvfp4",), + aliases=("RadixArk/Qwen3.8-Flash-Next-NVFP4",), + ), AotModel( name="Qwen/Qwen3-30B-A3B", architecture="Qwen3MoeForCausalLM", diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index 8e917832..be89047b 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -1,5 +1,9 @@ #pragma once +#if defined(_WIN32) && !defined(__always_inline) +#define __always_inline __forceinline__ +#endif + #include #include diff --git a/python/freetoken/kernel/csrc/include/freetoken/warp.cuh b/python/freetoken/kernel/csrc/include/freetoken/warp.cuh index d03063c3..00fbed30 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/warp.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/warp.cuh @@ -1,7 +1,9 @@ #pragma once #include +#ifndef _WIN32 #include +#endif #include diff --git a/python/freetoken/kernel/fla/layernorm_gated.py b/python/freetoken/kernel/fla/layernorm_gated.py index 55c2f00b..23d31a8c 100644 --- a/python/freetoken/kernel/fla/layernorm_gated.py +++ b/python/freetoken/kernel/fla/layernorm_gated.py @@ -47,6 +47,7 @@ def _layer_norm_fwd_1pass_kernel( NORM_BEFORE_GATE: tl.constexpr, IS_RMS_NORM: tl.constexpr, ACTIVATION: tl.constexpr, + WEIGHT_PLUS_ONE: tl.constexpr, ): # Map the program id to the starting row of X and Y it should compute. row_start = tl.program_id(0) * ROWS_PER_BLOCK @@ -106,6 +107,8 @@ def _layer_norm_fwd_1pass_kernel( w_offsets = cols + group * N w_mask = cols < N w = tl.load(W + w_offsets, mask=w_mask, other=0.0).to(tl.float32) + if WEIGHT_PLUS_ONE: + w += 1.0 if HAS_BIAS: b = tl.load(B + w_offsets, mask=w_mask, other=0.0).to(tl.float32) @@ -155,6 +158,7 @@ def _layer_norm_fwd( norm_before_gate=True, is_rms_norm=False, activation: str = "swish", + weight_plus_one: bool = False, ): M, N = x.shape if group_size is None: @@ -216,6 +220,7 @@ def _layer_norm_fwd( IS_RMS_NORM=is_rms_norm, num_warps=num_warps, ACTIVATION=activation, + WEIGHT_PLUS_ONE=weight_plus_one, ) return out, mean, rstd @@ -231,8 +236,13 @@ def rms_norm_gated( norm_before_gate=True, is_rms_norm=False, activation: str = "swish", + weight_plus_one: bool = False, ): - """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))""" + """Apply grouped (RMS)Norm and an optional gate. + + ``weight_plus_one`` keeps centered RMSNorm checkpoint weights unmodified and + performs the ``1 + weight`` operation in fp32 inside the kernel. + """ x_shape_og = x.shape # reshape input data into 2D tensor @@ -257,6 +267,7 @@ def rms_norm_gated( norm_before_gate=norm_before_gate, is_rms_norm=is_rms_norm, activation=activation, + weight_plus_one=weight_plus_one, ) return y.reshape(x_shape_og) diff --git a/python/freetoken/kernel/index.py b/python/freetoken/kernel/index.py index 95e61ff4..a5ba1875 100644 --- a/python/freetoken/kernel/index.py +++ b/python/freetoken/kernel/index.py @@ -1,15 +1,19 @@ from __future__ import annotations import functools +import sys +import warnings from typing import TYPE_CHECKING, Tuple +import torch + from .utils import KernelConfig, load_jit, make_cpp_args if TYPE_CHECKING: - import torch from tvm_ffi import Module DEFAULT_INDEX_KERNEL_CONFIG = KernelConfig(num_threads=128, max_occupancy=1, use_pdl=False) +_TORCH_FALLBACK_KEYS: set[tuple[int, int]] = set() @functools.cache @@ -49,6 +53,33 @@ def indexing( output = weights.new_empty(indices.shape[0], weights.shape[1]) element_size = weights.shape[1] * weights.element_size() - module = _jit_index_module(element_size, num_splits=num_splits_for(element_size)) - module.launch(weights, indices, output, vocab_range) + num_splits = num_splits_for(element_size) + key = (element_size, num_splits) + module = None + if key not in _TORCH_FALLBACK_KEYS: + try: + module = _jit_index_module(element_size, num_splits=num_splits) + except RuntimeError as exc: + if sys.platform != "win32": + raise + _TORCH_FALLBACK_KEYS.add(key) + warnings.warn( + f"Falling back to torch.index_select for {element_size}-byte embedding rows " + f"because the Windows CUDA index kernel is unavailable: {exc}", + RuntimeWarning, + stacklevel=2, + ) + if module is not None: + module.launch(weights, indices, output, vocab_range) + return output + + if vocab_range is None: + torch.index_select(weights, 0, indices.to(torch.int64), out=output) + return output + + start, length = vocab_range + valid = (indices >= start) & (indices < start + length) + local_indices = (indices - start).clamp(0, max(0, length - 1)).to(torch.int64) + torch.index_select(weights, 0, local_indices, out=output) + output.masked_fill_(~valid[:, None], 0) return output diff --git a/python/freetoken/kernel/triton/qsa.py b/python/freetoken/kernel/triton/qsa.py new file mode 100644 index 00000000..0a859ccf --- /dev/null +++ b/python/freetoken/kernel/triton/qsa.py @@ -0,0 +1,264 @@ +"""Exact gathered-row GQA attention for Qwen compressed sparse attention.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +BLOCK_H = 16 +BLOCK_T = 32 + + +@triton.jit +def _compact_qsa_blocks_kernel( + blocks_ptr, + positions_ptr, + output_ptr, + counts_ptr, + stride_bn, + stride_bt, + stride_on, + stride_ot, + BLOCK_COUNT: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + TOKEN_BUDGET: tl.constexpr, + OUTPUT_WIDTH: tl.constexpr, + BLOCK_BLOCKS: tl.constexpr, + BLOCK_OUTPUT: tl.constexpr, +): + row = tl.program_id(0) + block_offsets = tl.arange(0, BLOCK_BLOCKS) + chosen = tl.load( + blocks_ptr + row * stride_bn + block_offsets * stride_bt, + mask=block_offsets < BLOCK_COUNT, + other=-1, + ) + live_blocks = tl.sum((chosen >= 0).to(tl.int32), axis=0) + position = tl.load(positions_ptr + row).to(tl.int64) + visible = position + 1 + tail_start = (visible // COMPRESS_RATIO) * COMPRESS_RATIO + tail_count = visible - tail_start + live_tokens = live_blocks * COMPRESS_RATIO + + columns = tl.arange(0, BLOCK_OUTPUT) + block_slot = columns // COMPRESS_RATIO + block_value = tl.load( + blocks_ptr + row * stride_bn + block_slot * stride_bt, + mask=(columns < live_tokens) & (block_slot < BLOCK_COUNT), + other=-1, + ).to(tl.int64) + expanded = block_value * COMPRESS_RATIO + columns % COMPRESS_RATIO + tail_offset = columns - live_tokens + value = tl.where( + columns < live_tokens, + expanded, + tl.where(tail_offset < tail_count, tail_start + tail_offset, -1), + ) + tl.store( + output_ptr + row * stride_on + columns * stride_ot, + value.to(tl.int32), + mask=columns < OUTPUT_WIDTH, + ) + tl.store(counts_ptr + row, (live_tokens + tail_count).to(tl.int32)) + + +def compact_qsa_blocks( + block_indices: torch.Tensor, + query_positions: torch.Tensor, + *, + compress_ratio: int, + token_budget: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused CUDA expansion, tail insertion, and live-count calculation.""" + if not block_indices.is_cuda or not query_positions.is_cuda: + raise ValueError("compact_qsa_blocks is a CUDA kernel") + rows, blocks = block_indices.shape + width = token_budget + compress_ratio - 1 + output = torch.empty((rows, width), dtype=torch.int32, device=block_indices.device) + counts = torch.empty(rows, dtype=torch.int32, device=block_indices.device) + if not rows: + return output, counts + blocks_i32 = block_indices.to(torch.int32).contiguous() + positions = query_positions.to(torch.int64).contiguous() + _compact_qsa_blocks_kernel[(rows,)]( + blocks_i32, + positions, + output, + counts, + blocks_i32.stride(0), + blocks_i32.stride(1), + output.stride(0), + output.stride(1), + BLOCK_COUNT=blocks, + COMPRESS_RATIO=compress_ratio, + TOKEN_BUDGET=token_budget, + OUTPUT_WIDTH=width, + BLOCK_BLOCKS=triton.next_power_of_2(blocks), + BLOCK_OUTPUT=triton.next_power_of_2(width), + num_warps=8, + num_stages=1, + ) + return output, counts + + +@triton.jit +def _qsa_sparse_gqa_kernel( + q_ptr, k_ptr, v_ptr, rows_ptr, counts_ptr, out_ptr, + scale, + H, KVH, D, TOPK, GQA, + stride_qn, stride_qh, stride_qd, + stride_kr, stride_kh, stride_kd, + stride_vr, stride_vh, stride_vd, + stride_rn, stride_rt, + stride_on, stride_oh, stride_od, + BLOCK_D: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_T: tl.constexpr, +): + row = tl.program_id(0) + kv_head = tl.program_id(1) + + head_offsets = tl.arange(0, BLOCK_H) + heads = kv_head * GQA + head_offsets + head_mask = (head_offsets < GQA) & (heads < H) + dims = tl.arange(0, BLOCK_D) + dim_mask = dims < D + + q = tl.load( + q_ptr + row * stride_qn + heads[:, None] * stride_qh + dims[None, :] * stride_qd, + mask=head_mask[:, None] & dim_mask[None, :], + other=0.0, + ) + running_max = tl.full((BLOCK_H,), -float("inf"), tl.float32) + running_sum = tl.zeros((BLOCK_H,), tl.float32) + acc = tl.zeros((BLOCK_H, BLOCK_D), tl.float32) + active = tl.load(counts_ptr + row) + row_base = rows_ptr + row * stride_rn + + for tile in range(0, tl.cdiv(active, BLOCK_T)): + token_offsets = tile * BLOCK_T + tl.arange(0, BLOCK_T) + token_mask = token_offsets < active + physical = tl.load( + row_base + token_offsets * stride_rt, mask=token_mask, other=-1 + ) + valid = token_mask & (physical >= 0) + physical = tl.maximum(physical, 0) + k = tl.load( + k_ptr + + physical[None, :] * stride_kr + + kv_head * stride_kh + + dims[:, None] * stride_kd, + mask=dim_mask[:, None] & valid[None, :], + other=0.0, + ) + scores = tl.dot(q, k) * scale + scores = tl.where(valid[None, :], scores, -float("inf")) + new_max = tl.maximum(running_max, tl.max(scores, axis=1)) + alpha = tl.where(new_max == -float("inf"), 1.0, tl.exp(running_max - new_max)) + probs = tl.where(valid[None, :], tl.exp(scores - new_max[:, None]), 0.0) + running_sum = running_sum * alpha + tl.sum(probs, axis=1) + acc *= alpha[:, None] + v = tl.load( + v_ptr + + physical[:, None] * stride_vr + + kv_head * stride_vh + + dims[None, :] * stride_vd, + mask=valid[:, None] & dim_mask[None, :], + other=0.0, + ) + acc += tl.dot(probs.to(v.dtype), v) + running_max = new_max + + output = tl.where(running_sum[:, None] > 0, acc / running_sum[:, None], 0.0) + tl.store( + out_ptr + row * stride_on + heads[:, None] * stride_oh + dims[None, :] * stride_od, + output.to(out_ptr.dtype.element_ty), + mask=head_mask[:, None] & dim_mask[None, :], + ) + + +def _torch_qsa_sparse_gqa( + q: torch.Tensor, + k_rows: torch.Tensor, + v_rows: torch.Tensor, + selected_rows: torch.Tensor, + counts: torch.Tensor, + sm_scale: float, +) -> torch.Tensor: + """Small CPU/reference implementation used by unit tests.""" + output = torch.zeros_like(q) + gqa = q.shape[1] // k_rows.shape[1] + for row in range(q.shape[0]): + count = int(counts[row]) + if count == 0: + continue + indices = selected_rows[row, :count].long() + keys = k_rows.index_select(0, indices) + values = v_rows.index_select(0, indices) + for kv_head in range(k_rows.shape[1]): + heads = slice(kv_head * gqa, (kv_head + 1) * gqa) + scores = torch.einsum( + "hd,td->ht", q[row, heads].float(), keys[:, kv_head].float() + ) * sm_scale + probs = torch.softmax(scores, dim=-1).to(values.dtype) + output[row, heads] = torch.einsum( + "ht,td->hd", probs, values[:, kv_head] + ).to(output.dtype) + return output + + +@torch.no_grad() +def qsa_sparse_gqa( + q: torch.Tensor, + k_rows: torch.Tensor, + v_rows: torch.Tensor, + selected_rows: torch.Tensor, + counts: torch.Tensor, + sm_scale: float, +) -> torch.Tensor: + """Attend to an arbitrary physical-row list for each query. + + Shapes are ``q[N,H,D]``, ``k/v[R,KVH,D]``, and + ``selected_rows[N,K]``. Valid entries are compact and ``counts[N]`` gives + the live width of each row. + """ + if q.ndim != 3 or k_rows.ndim != 3 or v_rows.shape != k_rows.shape: + raise ValueError("QSA expects q [N,H,D] and matching k/v [R,KVH,D]") + if q.shape[-1] != k_rows.shape[-1] or q.shape[1] % k_rows.shape[1]: + raise ValueError("QSA requires matching head dimensions and integral GQA groups") + if selected_rows.shape[0] != q.shape[0] or counts.numel() != q.shape[0]: + raise ValueError("QSA selection rows must match query rows") + if not q.is_cuda: + return _torch_qsa_sparse_gqa(q, k_rows, v_rows, selected_rows, counts, sm_scale) + + q = q.contiguous() + k_rows = k_rows.contiguous() + v_rows = v_rows.contiguous() + selected_rows = selected_rows.to(torch.int32).contiguous() + counts = counts.to(torch.int32).contiguous() + output = torch.empty_like(q) + n, heads, dim = q.shape + kv_heads = k_rows.shape[1] + gqa = heads // kv_heads + if gqa > BLOCK_H: + raise ValueError(f"QSA Triton kernel supports a GQA group up to {BLOCK_H}, got {gqa}") + _qsa_sparse_gqa_kernel[(n, kv_heads)]( + q, k_rows, v_rows, selected_rows, counts, output, + float(sm_scale), + heads, kv_heads, dim, selected_rows.shape[1], gqa, + q.stride(0), q.stride(1), q.stride(2), + k_rows.stride(0), k_rows.stride(1), k_rows.stride(2), + v_rows.stride(0), v_rows.stride(1), v_rows.stride(2), + selected_rows.stride(0), selected_rows.stride(1), + output.stride(0), output.stride(1), output.stride(2), + BLOCK_D=triton.next_power_of_2(dim), + BLOCK_H=BLOCK_H, + BLOCK_T=BLOCK_T, + num_warps=8, + num_stages=2, + ) + return output + + +__all__ = ["compact_qsa_blocks", "qsa_sparse_gqa"] diff --git a/python/freetoken/kernel/triton/rope.py b/python/freetoken/kernel/triton/rope.py index 0284c82d..ba131beb 100644 --- a/python/freetoken/kernel/triton/rope.py +++ b/python/freetoken/kernel/triton/rope.py @@ -91,6 +91,77 @@ def _rope_tiled( tl.store(K + base_k + d1 * stride_kd, ok1.to(K.dtype.element_ty), mask=kmask) +@triton.jit(do_not_specialize=["nnz"]) +def _mrope_tiled( + Q, K, POS, CACHE, + stride_qbs, stride_qh, stride_qd, + stride_kbs, stride_kh, stride_kd, + stride_pa, stride_ps, + nnz, + HEAD_Q, HEAD_K, rotary_dim, half, + H_SPAN: tl.constexpr, + W_SPAN: tl.constexpr, + HAS_K: tl.constexpr, + INTERLEAVE: tl.constexpr, + BLOCK_SEQ: tl.constexpr, + BLOCK_HEAD: tl.constexpr, + BLOCK_DHALF: tl.constexpr, +): + """Qwen interleaved temporal/height/width RoPE in one fused launch.""" + seq_pid = tl.program_id(0) + head_pid = tl.program_id(1) + seq_range = seq_pid * BLOCK_SEQ + tl.arange(0, BLOCK_SEQ) + head_range = head_pid * BLOCK_HEAD + tl.arange(0, BLOCK_HEAD) + d = tl.arange(0, BLOCK_DHALF) + seq_mask = seq_range < nnz + dmask = d < half + + axis = tl.zeros((BLOCK_DHALF,), dtype=tl.int32) + axis = tl.where((d % 3 == 1) & (d < H_SPAN), 1, axis) + axis = tl.where((d % 3 == 2) & (d < W_SPAN), 2, axis) + pos = tl.load( + POS + axis[None, :] * stride_pa + seq_range[:, None] * stride_ps, + mask=seq_mask[:, None] & dmask[None, :], + other=0, + ).to(tl.int64) + cache_row = CACHE + pos * rotary_dim + cs_mask = seq_mask[:, None] & dmask[None, :] + cos = tl.load(cache_row + d[None, :], mask=cs_mask, other=0.0)[:, None, :] + sin = tl.load(cache_row + half + d[None, :], mask=cs_mask, other=0.0)[:, None, :] + + if INTERLEAVE: + d0 = 2 * d + d1 = 2 * d + 1 + else: + d0 = d + d1 = half + d + d0 = d0[None, None, :] + d1 = d1[None, None, :] + + qmask = ( + seq_mask[:, None, None] + & (head_range[None, :, None] < HEAD_Q) + & dmask[None, None, :] + ) + base_q = seq_range[:, None, None] * stride_qbs + head_range[None, :, None] * stride_qh + q0 = tl.load(Q + base_q + d0 * stride_qd, mask=qmask, other=0.0).to(tl.float32) + q1 = tl.load(Q + base_q + d1 * stride_qd, mask=qmask, other=0.0).to(tl.float32) + tl.store(Q + base_q + d0 * stride_qd, q0 * cos - q1 * sin, mask=qmask) + tl.store(Q + base_q + d1 * stride_qd, q1 * cos + q0 * sin, mask=qmask) + + if HAS_K: + kmask = ( + seq_mask[:, None, None] + & (head_range[None, :, None] < HEAD_K) + & dmask[None, None, :] + ) + base_k = seq_range[:, None, None] * stride_kbs + head_range[None, :, None] * stride_kh + k0 = tl.load(K + base_k + d0 * stride_kd, mask=kmask, other=0.0).to(tl.float32) + k1 = tl.load(K + base_k + d1 * stride_kd, mask=kmask, other=0.0).to(tl.float32) + tl.store(K + base_k + d0 * stride_kd, k0 * cos - k1 * sin, mask=kmask) + tl.store(K + base_k + d1 * stride_kd, k1 * cos + k0 * sin, mask=kmask) + + def apply_rope_with_cos_sin_cache_inplace( positions: torch.Tensor, query: torch.Tensor, @@ -140,4 +211,62 @@ def apply_rope_with_cos_sin_cache_inplace( ) -__all__ = ["apply_rope_with_cos_sin_cache_inplace"] +def apply_mrope_with_cos_sin_cache_inplace( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + head_size: int, + cos_sin_cache: torch.Tensor, + mrope_section: tuple[int, int, int], + is_neox: bool = True, +) -> None: + """Apply Qwen interleaved 3-axis RoPE to query and key in place.""" + if str(cos_sin_cache.dtype) != "torch.float32": + raise ValueError("cos_sin_cache should be float32") + assert query.is_cuda and key.is_cuda and positions.is_cuda + assert cos_sin_cache.is_contiguous() + if positions.ndim != 2 or positions.shape[0] != 3: + raise ValueError(f"MRoPE positions must have shape [3, tokens], got {positions.shape}") + nnz = query.shape[0] + if positions.shape[1] != nnz: + raise ValueError("MRoPE position count must match query rows") + if nnz == 0: + return + rotary_dim = cos_sin_cache.shape[1] + half = rotary_dim // 2 + if sum(mrope_section) != half: + raise ValueError( + f"MRoPE sections {mrope_section} must sum to rotary_dim / 2 ({half})" + ) + block_dhalf = triton.next_power_of_2(half) + head_q = query.shape[1] // head_size + head_k = key.shape[1] // head_size + qv = query.view(nnz, head_q, head_size) + kv = key.view(nnz, head_k, head_size) + max_head = max(head_q, head_k) + grid = lambda meta: ( + triton.cdiv(nnz, meta["BLOCK_SEQ"]), + triton.cdiv(max_head, meta["BLOCK_HEAD"]), + ) + _mrope_tiled[grid]( + qv, kv, positions, cos_sin_cache, + qv.stride(0), qv.stride(1), qv.stride(2), + kv.stride(0), kv.stride(1), kv.stride(2), + positions.stride(0), positions.stride(1), + nnz, head_q, head_k, rotary_dim, half, + H_SPAN=int(mrope_section[1]) * 3, + W_SPAN=int(mrope_section[2]) * 3, + HAS_K=True, + INTERLEAVE=not is_neox, + BLOCK_DHALF=block_dhalf, + BLOCK_SEQ=16, + BLOCK_HEAD=1, + num_warps=4, + num_stages=1, + ) + + +__all__ = [ + "apply_mrope_with_cos_sin_cache_inplace", + "apply_rope_with_cos_sin_cache_inplace", +] diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index 7a0164b5..f649f401 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -4,6 +4,7 @@ import os import pathlib import re +import sys from typing import TYPE_CHECKING, List, NamedTuple, Tuple, TypeAlias, Union if TYPE_CHECKING: @@ -17,7 +18,9 @@ DISABLE_JIT_ENV = "FREETOKEN_DISABLE_JIT" _TRUE_VALUES = {"1", "true", "yes", "on"} DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")] -DEFAULT_CFLAGS = ["-std=c++20", "-O3"] +DEFAULT_CFLAGS = ( + ["/std:c++20", "/O2"] if sys.platform == "win32" else ["-std=c++20", "-O3"] +) DEFAULT_CUDA_CFLAGS = ["-std=c++20", "-O3", "--expt-relaxed-constexpr"] DEFAULT_LDFLAGS = [] @@ -145,6 +148,18 @@ def _kernel_cache_dir() -> pathlib.Path | None: return pathlib.Path(get_jit_cache_dir()).expanduser() +def _prebuilt_library_path( + cache_dir: pathlib.Path, + name: str, + *, + platform: str | None = None, +) -> pathlib.Path: + """Return the platform-native cached kernel library path.""" + platform = sys.platform if platform is None else platform + suffix = ".dll" if platform == "win32" else ".so" + return cache_dir / name / f"{name}{suffix}" + + def _load_prebuilt(name: str) -> Module | None: cache_dir = _kernel_cache_dir() if cache_dir is None: @@ -155,16 +170,16 @@ def _load_prebuilt(name: str) -> Module | None: ) return None - so_path = cache_dir / name / f"{name}.so" - if so_path.exists(): + library_path = _prebuilt_library_path(cache_dir, name) + if library_path.exists(): import tvm_ffi - return tvm_ffi.load_module(str(so_path)) + return tvm_ffi.load_module(str(library_path)) if _env_enabled(DISABLE_JIT_ENV): raise RuntimeError( "JIT compilation is disabled by FREETOKEN_DISABLE_JIT, " - f"but prebuilt kernel {name!r} was not found at {so_path}" + f"but prebuilt kernel {name!r} was not found at {library_path}" ) return None diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py index 1bb352b4..d182b3b1 100644 --- a/python/freetoken/kvcache/__init__.py +++ b/python/freetoken/kvcache/__init__.py @@ -61,6 +61,10 @@ def resolve_pool_class(model_config: ModelConfig) -> type[BaseKVCachePool]: from .bsa_pool import BSAKVCache return BSAKVCache + if AttnType.QSA in types: + from .qsa_pool import QSAKVCache + + return QSAKVCache from .mha_pool import MHAKVCache return MHAKVCache @@ -174,6 +178,24 @@ def create_kvcache_pool( num_index_layers=spec.num_index_layers, ) + if len(kv_specs) == 1 and kv_specs[0].attn_type == _AttnType.QSA: + from .qsa_pool import QSAKVCache + + spec = kv_specs[0] + return QSAKVCache( + num_kv_heads=spec.num_kv_heads, + num_layers=model_config.num_layers, + head_dim=spec.head_dim, + num_pages=num_pages, + page_size=page_size, + dtype=dtype, + device=device, + index_num_kv_heads=spec.index_num_kv_heads, + index_head_dim=spec.index_head_dim, + compress_ratio=spec.index_compress_ratio, + layer_ids=spec.layer_ids, + ) + if len(kv_specs) == 1 and kv_specs[0].mla: from .dsa_pool import DSAKVCache, MLAKVCache diff --git a/python/freetoken/kvcache/base.py b/python/freetoken/kvcache/base.py index ae8cf9ec..aedab15d 100644 --- a/python/freetoken/kvcache/base.py +++ b/python/freetoken/kvcache/base.py @@ -29,7 +29,13 @@ def spec_kv_bytes_per_token(spec, config) -> int: * config.dtype.itemsize * spec.num_layers ) - return per_token + spec.index_head_dim * spec.num_index_layers * 2 + index_bytes = spec.index_head_dim * spec.num_index_layers * 2 + if spec.attn_type.value == "qsa": + index_bytes *= spec.index_num_kv_heads + if index_bytes % spec.index_compress_ratio: + raise ValueError("QSA index bytes must divide evenly by its compression ratio") + index_bytes //= spec.index_compress_ratio + return per_token + index_bytes class BaseKVCachePool(ABC): diff --git a/python/freetoken/kvcache/qsa_pool.py b/python/freetoken/kvcache/qsa_pool.py new file mode 100644 index 00000000..d0379347 --- /dev/null +++ b/python/freetoken/kvcache/qsa_pool.py @@ -0,0 +1,187 @@ +"""Paged KV storage for Qwen compressed sparse attention (QSA). + +The ordinary K/V cache keeps one row per token. The QSA index cache keeps +one row per ``compress_ratio`` tokens. A full KV page therefore maps to one +smaller, page-aligned QSA page. This keeps the translation exact and cheap: +the compressed row for a complete token group is ``full_row // ratio``. +""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from .mha_pool import MHAKVCache + + +class QSAKVCache(MHAKVCache): + """MHA/GQA K/V plus compressed QSA index keys and a small pending ring.""" + + def __init__( + self, + num_kv_heads: int, + num_layers: int, + head_dim: int, + num_pages: int, + page_size: int, + dtype: torch.dtype, + device: torch.device, + index_num_kv_heads: int, + index_head_dim: int, + compress_ratio: int, + layer_ids: Sequence[int], + ) -> None: + if compress_ratio < 2 or page_size % compress_ratio: + raise ValueError( + "QSA needs a compression ratio >= 2 that divides the KV page size" + ) + if dtype.itemsize != 2: + raise ValueError(f"QSA index keys require a 2-byte compute dtype, got {dtype}") + self._page_size = int(page_size) + self._compress_ratio = int(compress_ratio) + self._compressed_page_size = self._page_size // self._compress_ratio + self._index_num_kv_heads = int(index_num_kv_heads) + self._index_head_dim = int(index_head_dim) + self._num_index_layers = len(layer_ids) + self._pending_k: torch.Tensor | None = None + self._pending_pos: torch.Tensor | None = None + self._pending_rope: torch.Tensor | None = None + super().__init__( + num_kv_heads=num_kv_heads, + num_layers=num_layers, + head_dim=head_dim, + num_pages=num_pages, + page_size=page_size, + dtype=dtype, + device=device, + layer_ids=layer_ids, + ) + self._alloc_compressed(num_pages) + + def _alloc_compressed(self, num_pages: int) -> None: + self._compressed_k = torch.empty( + self._num_index_layers, + num_pages, + self._compressed_page_size, + self._index_num_kv_heads, + self._index_head_dim, + dtype=self.dtype, + device=self.device, + ) + + def rebuild(self, num_pages: int) -> None: + self._compressed_k = None + super().rebuild(num_pages) + self._alloc_compressed(num_pages) + + def unit_bytes(self) -> tuple[int, int]: + kv, swa = super().unit_bytes() + full_tokens = int(self._kv_buffer.shape[2]) * self._page_size + index_bytes = int(self._compressed_k.numel() * self._compressed_k.element_size()) + return kv + index_bytes // full_tokens, swa + + @property + def compress_ratio(self) -> int: + return self._compress_ratio + + @property + def compressed_page_size(self) -> int: + return self._compressed_page_size + + def compressed_k_cache(self, layer_id: int) -> torch.Tensor: + """Return row-flat compressed keys ``[rows, kv_heads, index_dim]``.""" + return self._compressed_k[self._dense(layer_id)].view( + -1, self._index_num_kv_heads, self._index_head_dim + ) + + def store_compressed_k( + self, keys: torch.Tensor, compressed_rows: torch.Tensor, layer_id: int + ) -> None: + self.compressed_k_cache(layer_id)[compressed_rows.long()] = keys + + def ensure_pending_capacity(self, request_rows: int) -> None: + """Allocate or grow the per-request incomplete-group ring. + + The ring is tiny compared with the paged cache. It stores at most + ``ratio - 1`` useful raw keys per active request and layer. + """ + current = 0 if self._pending_k is None else int(self._pending_k.shape[1]) + if current >= request_rows: + return + new_rows = max(request_rows, max(16, current * 2)) + shape = ( + self._num_index_layers, + new_rows, + self._compress_ratio, + self._index_num_kv_heads, + self._index_head_dim, + ) + pending = torch.empty(shape, dtype=self.dtype, device=self.device) + positions = torch.full( + shape[:3], -1, dtype=torch.int64, device=self.device + ) + rope = torch.full( + (*shape[:3], 3), -1, dtype=torch.int64, device=self.device + ) + if self._pending_k is not None: + pending[:, :current].copy_(self._pending_k) + positions[:, :current].copy_(self._pending_pos) + rope[:, :current].copy_(self._pending_rope) + self._pending_k = pending + self._pending_pos = positions + self._pending_rope = rope + + def clear_pending(self, layer_id: int, request_row: int) -> None: + self._pending_pos[self._dense(layer_id), request_row].fill_(-1) + + def pending_group( + self, layer_id: int, request_row: int, positions: torch.Tensor + ) -> torch.Tensor: + dense = self._dense(layer_id) + slots = torch.remainder(positions, self._compress_ratio).long() + actual = self._pending_pos[dense, request_row].index_select(0, slots) + expected = positions.to(device=actual.device, dtype=actual.dtype) + if not torch.equal(actual, expected): + raise RuntimeError( + "QSA pending-key state is missing; use the naive cache and do not " + "resume a prefix without its QSA state" + ) + return self._pending_k[dense, request_row].index_select(0, slots) + + def pending_rope_group( + self, layer_id: int, request_row: int, positions: torch.Tensor + ) -> torch.Tensor: + """Return stored [tokens, 3] rotary coordinates after state validation.""" + self.pending_group(layer_id, request_row, positions) + dense = self._dense(layer_id) + slots = torch.remainder(positions, self._compress_ratio).long() + return self._pending_rope[dense, request_row].index_select(0, slots) + + def store_pending( + self, + layer_id: int, + request_row: int, + positions: torch.Tensor, + keys: torch.Tensor, + rope_positions: torch.Tensor | None = None, + ) -> None: + dense = self._dense(layer_id) + slots = torch.remainder(positions, self._compress_ratio).long() + self._pending_k[dense, request_row].index_copy_(0, slots, keys) + self._pending_pos[dense, request_row].index_copy_( + 0, slots, positions.to(device=self.device, dtype=torch.int64) + ) + if rope_positions is None: + rope_positions = positions.to(device=self.device, dtype=torch.int64).view(-1, 1).expand(-1, 3) + if rope_positions.shape != (positions.numel(), 3): + raise ValueError( + "QSA pending RoPE positions must have shape [tokens, 3], got " + f"{tuple(rope_positions.shape)}" + ) + self._pending_rope[dense, request_row].index_copy_( + 0, slots, rope_positions.to(device=self.device, dtype=torch.int64) + ) + + +__all__ = ["QSAKVCache"] diff --git a/python/freetoken/llm/llm.py b/python/freetoken/llm/llm.py index 1d4aaef0..f6b60780 100644 --- a/python/freetoken/llm/llm.py +++ b/python/freetoken/llm/llm.py @@ -40,11 +40,12 @@ def __init__(self, model_path: str, dtype: torch.dtype = torch.bfloat16, **kwarg self.pending_requests: List[Tuple[List[int] | str, SamplingParams]] = [] self.status_map: Dict[int, RequestStatus] = {} self.mm_embeds_map: Dict[int, torch.Tensor] = {} + self.mrope_map: Dict[int, Tuple[torch.Tensor, int]] = {} self.counter = 0 @torch.inference_mode() def encode_images( - self, pixel_values: torch.Tensor, image_position_ids: torch.Tensor + self, pixel_values: torch.Tensor, image_metadata: torch.Tensor ) -> torch.Tensor: """Run the vision tower + projector on processor outputs, returning the ``[num_image_tokens, hidden]`` soft-token embeddings (on device).""" @@ -52,7 +53,7 @@ def encode_images( if not hasattr(model, "encode_images"): raise RuntimeError(f"{type(model).__name__} does not support image inputs") return model.encode_images( - pixel_values.to(self.device), image_position_ids.to(self.device) + pixel_values.to(self.device), image_metadata.to(self.device) ) def _tokenize_one(self, prompt: List[int] | str) -> torch.Tensor: @@ -78,6 +79,12 @@ def offline_receive_msg(self, blocking: bool = False) -> List[BaseBackendMsg]: input_ids=input_ids, sampling_params=sampling_params, mm_embeds=self.mm_embeds_map.get(uid), + mrope_position_ids=( + self.mrope_map[uid][0] if uid in self.mrope_map else None + ), + mrope_position_delta=( + self.mrope_map[uid][1] if uid in self.mrope_map else 0 + ), ) ) self.status_map[uid] = RequestStatus( @@ -112,12 +119,14 @@ def generate( ``mm_inputs`` (optional) is aligned with ``prompts``; each entry is either ``None`` (text-only) or a dict with ``pixel_values`` ``[N, P, 3*patch**2]`` and - ``image_position_ids`` ``[N, P, 2]`` from the HF processor. For multimodal + either ``image_position_ids`` (Gemma) or ``image_grid_thw`` (Qwen) from + the HF processor. For multimodal prompts pass token-id ``prompts`` containing ``image_token_id`` placeholders. """ self.pending_requests = [] self.status_map = {} self.mm_embeds_map = {} + self.mrope_map = {} self.counter = 0 if isinstance(sampling_params, SamplingParams): sampling_params = [sampling_params] * len(prompts) @@ -126,9 +135,30 @@ def generate( if mm_inputs is not None: for uid, mm in enumerate(mm_inputs): if mm is not None: + metadata = mm.get("image_grid_thw", mm.get("image_position_ids")) + if metadata is None: + raise ValueError( + "multimodal input needs image_grid_thw or image_position_ids" + ) self.mm_embeds_map[uid] = self.encode_images( - mm["pixel_values"], mm["image_position_ids"] + mm["pixel_values"], metadata ) + if "image_grid_thw" in mm: + mm_types = mm.get("mm_token_type_ids") + if mm_types is None: + raise ValueError( + "Qwen multimodal input needs mm_token_type_ids from the processor" + ) + from freetoken.models.qwen4_exp.mrope import build_mrope_positions + + prompt_ids = self._tokenize_one(prompts[uid]) + merge_size = int(self.config.model_config.vision_config.spatial_merge_size) + self.mrope_map[uid] = build_mrope_positions( + prompt_ids, + mm_types, + mm["image_grid_thw"], + merge_size, + ) torch.cuda.synchronize(self.device) try: self.run_forever() diff --git a/python/freetoken/message/backend.py b/python/freetoken/message/backend.py index c42ecc5a..e8f5d0d5 100644 --- a/python/freetoken/message/backend.py +++ b/python/freetoken/message/backend.py @@ -37,6 +37,17 @@ class UserMsg(BaseBackendMsg): # Optional precomputed multimodal soft-token embeddings (GPU tensor). Only used by # the in-process offline path; remains None for the (serialized) online path. mm_embeds: torch.Tensor | None = None + # Online multimodal requests carry processor outputs from the CPU tokenizer + # worker to the GPU scheduler. The scheduler turns these into mm_embeds before + # admission. Pixel rows are BF16 because the vision patch projection casts to + # its BF16 weights immediately; this halves local transport without changing + # the model input. + mm_pixel_values: torch.Tensor | None = None + mm_image_grid_thw: torch.Tensor | None = None + mm_token_type_ids: torch.Tensor | None = None + # Optional CPU [3, prompt_tokens] Qwen multimodal RoPE coordinates. + mrope_position_ids: torch.Tensor | None = None + mrope_position_delta: int = 0 @dataclass diff --git a/python/freetoken/message/utils.py b/python/freetoken/message/utils.py index ee92adf5..c807c9e1 100644 --- a/python/freetoken/message/utils.py +++ b/python/freetoken/message/utils.py @@ -32,10 +32,13 @@ def serialize_type(self) -> Dict: serialized = {} if isinstance(self, torch.Tensor): - assert self.dim() == 1, "we can only serialize 1D tensor for now" + tensor = self.detach().to(device="cpu").contiguous() serialized["__type__"] = "Tensor" - serialized["buffer"] = self.numpy().tobytes() - serialized["dtype"] = str(self.dtype) + # A byte view supports every torch dtype, including bfloat16, which NumPy + # cannot represent directly on all supported versions. + serialized["buffer"] = tensor.view(torch.uint8).numpy().tobytes() + serialized["dtype"] = str(tensor.dtype) + serialized["shape"] = list(tensor.shape) return serialized # normal type @@ -68,10 +71,12 @@ def deserialize_type(cls_map: Dict[str, Type], data: Dict) -> Any: if type_name == "Tensor": buffer = data["buffer"] dtype_str = data["dtype"].replace("torch.", "") - np_dtype = getattr(np, dtype_str) assert isinstance(buffer, bytes) - np_tensor = np.frombuffer(buffer, dtype=np_dtype) - return torch.from_numpy(np_tensor.copy()) + shape = tuple(data.get("shape", ())) + torch_dtype = getattr(torch, dtype_str) + raw = torch.from_numpy(np.frombuffer(buffer, dtype=np.uint8).copy()) + tensor = raw.view(torch_dtype) + return tensor.reshape(shape) if "shape" in data else tensor cls = cls_map.get(type_name) if cls is None: diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index f6105e1f..6b4350f8 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -97,6 +97,9 @@ class KVCacheGroupSpec: mla: bool = False index_head_dim: int = 0 num_index_layers: int = 0 + index_num_kv_heads: int = 1 + index_compress_ratio: int = 1 + index_token_budget: int = 0 # Attention-type taxonomy value for this group; drives the backend capability # matrix and (with the pool factory) selects the KV pool family. attn_type: AttnType = AttnType.FULL @@ -147,6 +150,28 @@ class SWAAttentionGroupConfig(BaseAttentionGroupConfig): sliding_window: int +@dataclass(frozen=True) +class QSAAttentionGroupConfig(BaseAttentionGroupConfig): + """Qwen compressed sparse-attention group. + + Full K/V remains paged at token resolution. The indexer keeps one BF16 + key per ``index_compress_ratio`` tokens, then selects + ``index_token_budget`` original tokens for exact sparse GQA. + """ + + kind: ClassVar[Literal["qsa"]] = "qsa" + cache_kind: ClassVar[Literal["qsa_paged"]] = "qsa_paged" + + num_kv_heads: int + head_dim: int + rotary_config: RotaryConfig + index_num_heads: int + index_num_kv_heads: int + index_head_dim: int + index_token_budget: int + index_compress_ratio: int + + @dataclass(frozen=True) class LinearGatedDeltaGroupConfig(BaseAttentionGroupConfig): kind: ClassVar[Literal["linear_gated_delta"]] = "linear_gated_delta" @@ -177,6 +202,7 @@ class on purpose: subclassing SWAAttentionGroupConfig would flip has_swa_attenti AttentionGroupConfig: TypeAlias = ( FullAttentionGroupConfig + | QSAAttentionGroupConfig | SWAAttentionGroupConfig | LinearGatedDeltaGroupConfig | DSV4AttentionGroupConfig @@ -288,9 +314,14 @@ class ModelConfig: # swigluoai/dense-MLP scalars the model module needs. Opaque to model-agnostic engine # code; None for every other model. m3_args: Any | None = None + # Qwen4-Exp payload: hyper-connections, PLE host embedding geometry, and the + # QSA exact-context ceiling. Opaque outside the qwen4_exp model package. + qwen4_args: Any | None = None # Generic execution-path capability flags (set by a model's parse_config) so the engine and # factories stay model-agnostic instead of branching on dsv4_args: single_stream_only: bool = False # model runs one sequence at a time -> force bs=1 + requires_naive_cache: bool = False # model owns host/runtime state radix cannot snapshot + supports_cuda_graph: bool = True # False when forward performs host-side dynamic work @property def is_moe(self) -> bool: @@ -382,6 +413,8 @@ def attn_type_for_layer(self, layer_id: int) -> AttnType: return AttnType.SWA if isinstance(group, DSV4AttentionGroupConfig): return AttnType.DSV4 + if isinstance(group, QSAAttentionGroupConfig): + return AttnType.QSA return _full_group_attn_type(group) def kv_cache_group_specs(self) -> Tuple[KVCacheGroupSpec, ...]: @@ -412,6 +445,22 @@ def kv_cache_group_specs(self) -> Tuple[KVCacheGroupSpec, ...]: attn_type=_full_group_attn_type(group), ) ) + elif isinstance(group, QSAAttentionGroupConfig): + specs.append( + KVCacheGroupSpec( + name=group.name, + layer_ids=group.layer_ids, + num_kv_heads=group.num_kv_heads, + head_dim=group.head_dim, + sliding_window=None, + index_head_dim=group.index_head_dim, + num_index_layers=len(group.layer_ids), + index_num_kv_heads=group.index_num_kv_heads, + index_compress_ratio=group.index_compress_ratio, + index_token_budget=group.index_token_budget, + attn_type=AttnType.QSA, + ) + ) elif isinstance(group, SWAAttentionGroupConfig): specs.append( KVCacheGroupSpec( diff --git a/python/freetoken/models/loader.py b/python/freetoken/models/loader.py index 49419364..22737592 100644 --- a/python/freetoken/models/loader.py +++ b/python/freetoken/models/loader.py @@ -55,10 +55,16 @@ def iter_weight_files(model_path: str) -> list[str]: def drop_page_cache(path: str) -> None: """drop a file's page cache: banks + full checkpoint cache don't both fit in host RAM (OOM).""" + posix_fadvise = getattr(os, "posix_fadvise", None) + dontneed = getattr(os, "POSIX_FADV_DONTNEED", None) + if posix_fadvise is None or dontneed is None: + # Windows has no POSIX page-cache advisory API. The hint is optional; + # skipping it is correct and lets the serial expert loader run there. + return try: fd = os.open(path, os.O_RDONLY) try: - os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED) + posix_fadvise(fd, 0, 0, dontneed) finally: os.close(fd) except OSError: diff --git a/python/freetoken/models/qwen3_5_moe/attention.py b/python/freetoken/models/qwen3_5_moe/attention.py index 2421264e..c911ea44 100644 --- a/python/freetoken/models/qwen3_5_moe/attention.py +++ b/python/freetoken/models/qwen3_5_moe/attention.py @@ -60,10 +60,11 @@ def __init__(self, config: ModelConfig, layer_id: int): ) self.o_proj = make_replicated(config, self.qo_attn_dim, config.hidden_size, has_bias=False) - def _project(self, x: torch.Tensor): + def _project(self, x: torch.Tensor, positions: torch.Tensor | None = None): """Returns (q, k, v, gate): q [N, num_q, head_dim] post qk-norm+rope, k [N, num_kv*head_dim] post norm+rope, v [N, num_kv*head_dim], gate [N, num_q*head_dim].""" - positions = get_global_ctx().batch.positions + if positions is None: + positions = get_global_ctx().batch.positions qkv = self.qkv_proj.forward(x) qg, k, v = torch.split(qkv, self._qkv_split, dim=-1) qg = qg.view(-1, self.num_q, self.head_dim * 2) diff --git a/python/freetoken/models/qwen4_exp/__init__.py b/python/freetoken/models/qwen4_exp/__init__.py new file mode 100644 index 00000000..c4d5c731 --- /dev/null +++ b/python/freetoken/models/qwen4_exp/__init__.py @@ -0,0 +1,19 @@ +from .config import parse_config +from .model import Qwen4ExpForCausalLM +from .weight import ( + iter_weights, + iter_weights_parallel, + load_nvfp4_expert_sources, + load_nvfp4_expert_sources_parallel, + setup_offload_expert_banks, +) + +__all__ = [ + "Qwen4ExpForCausalLM", + "iter_weights", + "iter_weights_parallel", + "load_nvfp4_expert_sources", + "load_nvfp4_expert_sources_parallel", + "parse_config", + "setup_offload_expert_banks", +] diff --git a/python/freetoken/models/qwen4_exp/args.py b/python/freetoken/models/qwen4_exp/args.py new file mode 100644 index 00000000..535cc883 --- /dev/null +++ b/python/freetoken/models/qwen4_exp/args.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class Qwen4VisionConfig: + depth: int + hidden_size: int + intermediate_size: int + num_heads: int + num_position_embeddings: int + out_hidden_size: int + patch_size: int + spatial_merge_size: int + temporal_patch_size: int + in_channels: int + hidden_act: str + deepstack_visual_indexes: tuple[int, ...] + + +@dataclass(frozen=True) +class Qwen4ExpArgs: + hc_count: int + hc_lowrank: int + ple_layer_ids: tuple[int, ...] + ple_embed_dim: int + ple_conv_kernel_size: int + ngram_size: int + heads_per_ngram: int + ngram_vocab_size_base: int + split_ngram_parts: int + eos_token_id: int + indexer_n_heads: int + indexer_kv_heads: int + indexer_head_dim: int + indexer_budget: int + indexer_compress_ratio: int + output_gate_type: str + mrope_section: tuple[int, int, int] + mrope_interleaved: bool + + +__all__ = ["Qwen4ExpArgs", "Qwen4VisionConfig"] diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py new file mode 100644 index 00000000..b53811a9 --- /dev/null +++ b/python/freetoken/models/qwen4_exp/config.py @@ -0,0 +1,189 @@ +from __future__ import annotations + +from typing import Any + +from freetoken.models.config import ( + LinearGatedDeltaGroupConfig, + ModelConfig, + QSAAttentionGroupConfig, + RotaryConfig, + detect_expert_quant, +) + +from .args import Qwen4ExpArgs, Qwen4VisionConfig + + +def parse_config(hf_config: Any) -> ModelConfig: + text = hf_config.text_config + layer_types = list(text.layer_types) + sparse_attention_types = {"full_attention", "qwen_sparse_attention"} + unsupported = sorted(set(layer_types) - {"linear_attention", *sparse_attention_types}) + if unsupported: + raise ValueError(f"Unsupported Qwen4-Exp layer types: {unsupported}") + + head_dim = int(text.head_dim) + rope = text.rope_parameters + rotary_dim = round(head_dim * float(rope.get("partial_rotary_factor", 1.0))) + indexer_budget = int(text.indexer_budget) + rotary = RotaryConfig( + head_dim=head_dim, + rotary_dim=rotary_dim, + max_position=int(text.max_position_embeddings), + base=float(rope["rope_theta"]), + scaling=None, + ) + + full_ids = tuple( + i for i, layer_type in enumerate(layer_types) if layer_type in sparse_attention_types + ) + linear_ids = tuple( + i for i, layer_type in enumerate(layer_types) if layer_type == "linear_attention" + ) + groups = ( + LinearGatedDeltaGroupConfig( + name="linear", + layer_ids=linear_ids, + num_key_heads=int(text.linear_num_key_heads), + num_value_heads=int(text.linear_num_value_heads), + key_head_dim=int(text.linear_key_head_dim), + value_head_dim=int(text.linear_value_head_dim), + conv_kernel_dim=int(text.linear_conv_kernel_dim), + output_gate=True, + ), + QSAAttentionGroupConfig( + name="qsa", + layer_ids=full_ids, + num_kv_heads=int(text.num_key_value_heads), + head_dim=head_dim, + rotary_config=rotary, + index_num_heads=int(text.indexer_n_heads), + index_num_kv_heads=int(text.indexer_kv_heads), + index_head_dim=int(text.indexer_head_dim), + index_token_budget=indexer_budget, + index_compress_ratio=int(text.indexer_compress_ratio), + ), + ) + + eos_token_id = text.eos_token_id + if isinstance(eos_token_id, list): + eos_token_id = eos_token_id[0] + qwen4_args = Qwen4ExpArgs( + hc_count=int(text.hc_count), + hc_lowrank=int(text.hc_lowrank), + ple_layer_ids=tuple(int(layer_id) - 1 for layer_id in text.ple_layer_ids), + ple_embed_dim=int(text.ple_embed_dim), + ple_conv_kernel_size=int(text.ple_conv_kernel_size), + ngram_size=int(text.ngram_size), + heads_per_ngram=int(text.heads_per_ngram), + ngram_vocab_size_base=int(text.ngram_vocab_size_base), + split_ngram_parts=int(text.split_ngram_parts), + eos_token_id=int(eos_token_id), + indexer_n_heads=int(text.indexer_n_heads), + indexer_kv_heads=int(text.indexer_kv_heads), + indexer_head_dim=int(text.indexer_head_dim), + indexer_budget=indexer_budget, + indexer_compress_ratio=int(text.indexer_compress_ratio), + output_gate_type=str(text.output_gate_type or text.hidden_act), + mrope_section=tuple(int(value) for value in rope["mrope_section"]), + mrope_interleaved=bool(rope.get("mrope_interleaved", False)), + ) + if not qwen4_args.mrope_interleaved: + raise ValueError("Qwen4-Exp requires interleaved MRoPE") + if sum(qwen4_args.mrope_section) * 2 != rotary_dim: + raise ValueError( + "Qwen4-Exp mrope_section must cover the rotary dimension: " + f"{qwen4_args.mrope_section} vs {rotary_dim}" + ) + raw_vision = getattr(hf_config, "vision_config", None) + vision_config = None + if raw_vision is not None: + vision_config = Qwen4VisionConfig( + depth=int(raw_vision.depth), + hidden_size=int(raw_vision.hidden_size), + intermediate_size=int(raw_vision.intermediate_size), + num_heads=int(raw_vision.num_heads), + num_position_embeddings=int(raw_vision.num_position_embeddings), + out_hidden_size=int(raw_vision.out_hidden_size), + patch_size=int(raw_vision.patch_size), + spatial_merge_size=int(raw_vision.spatial_merge_size), + temporal_patch_size=int(raw_vision.temporal_patch_size), + in_channels=int(raw_vision.in_channels), + hidden_act=str(raw_vision.hidden_act), + deepstack_visual_indexes=tuple( + int(index) for index in raw_vision.deepstack_visual_indexes + ), + ) + + quant = getattr(hf_config, "quantization_config", None) + get_quant = ( + quant.get + if isinstance(quant, dict) + else (lambda key, default=None: getattr(quant, key, default)) + ) + detected_quant = detect_expert_quant(hf_config) + if detected_quant == "fp8": + raw_block_size = get_quant("weight_block_size") + block_size = ( + tuple(int(value) for value in raw_block_size) if raw_block_size is not None else None + ) + if block_size != (128, 128): + raise ValueError( + "Qwen4-Exp block-FP8 checkpoints require a 128x128 weight block size" + ) + expert_quant = "fp8_block" + elif detected_quant == "nvfp4": + # RadixArk's ModelOpt checkpoint quantizes only the routed experts. The + # attention, GDN, mHC, shared experts, router, embeddings, lm_head, and + # vision tensors remain BF16; PLE remains its source FP8 format. + block_size = None + expert_quant = "nvfp4" + else: + raise ValueError( + "Qwen4-Exp requires routed experts in 128x128 block-FP8 or ModelOpt NVFP4; " + f"detected {detected_quant!r}" + ) + + return ModelConfig( + num_layers=int(text.num_hidden_layers), + num_qo_heads=int(text.num_attention_heads), + num_kv_heads=int(text.num_key_value_heads), + head_dim=head_dim, + hidden_size=int(text.hidden_size), + vocab_size=int(text.vocab_size), + intermediate_size=int(getattr(text, "intermediate_size", 0) or 0), + hidden_act=str(text.hidden_act), + rms_norm_eps=float(text.rms_norm_eps), + tie_word_embeddings=bool(getattr(text, "tie_word_embeddings", False)), + rotary_config=rotary, + num_experts=int(text.num_experts), + num_experts_per_tok=int(text.num_experts_per_tok), + moe_intermediate_size=int(text.moe_intermediate_size), + shared_expert_intermediate_size=int(text.shared_expert_intermediate_size), + # The released Qwen3.8-Flash-Next configs omit this older Qwen MoE + # field. Omission means that the router weights are not renormalized. + norm_topk_prob=bool(getattr(text, "norm_topk_prob", False)), + model_type=str(hf_config.model_type), + architectures=list(hf_config.architectures), + moe_enabled=True, + expert_quant=expert_quant, + weight_block_size=block_size, + # Only routed experts and PLE are FP8 in the official checkpoint. All + # attention, hyper-connection, and shared-expert projections stay BF16. + attn_quant="none", + dense_quant="none", + lm_head_quant="none", + use_qk_norm=True, + # Qwen3.8-Flash-Next is a VL checkpoint. Vision is part of this model, + # not an optional text-only add-on. + vision_config=vision_config, + image_token_id=getattr(hf_config, "image_token_id", None), + attention_groups=groups, + qwen4_args=qwen4_args, + # PLE keeps per-request dilated-convolution state outside the generic + # radix cache and performs mmap-backed CPU gathers during every forward. + requires_naive_cache=True, + supports_cuda_graph=False, + ) + + +__all__ = ["parse_config"] diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py new file mode 100644 index 00000000..d0ba393a --- /dev/null +++ b/python/freetoken/models/qwen4_exp/model.py @@ -0,0 +1,715 @@ +from __future__ import annotations + +import json +import math +import os +from dataclasses import replace +from typing import TYPE_CHECKING + +import safetensors +import torch +import torch.nn.functional as F +from freetoken.core import get_global_ctx +from freetoken.layers import ( + BaseOP, + GemmaPlusOneRMSNorm, + LinearColParallelMerged, + LinearReplicated, + LinearRowParallel, + OPList, + ParallelLMHead, + VocabParallelEmbedding, + make_moe_layer, + silu_and_mul, + StateLessOP, + get_rope, +) +from freetoken.models.blocks import BaseLLMModel +from freetoken.models.qwen3_5_moe.attention import Qwen3_5Attention +from freetoken.models.qwen3_5_moe.gdn import Qwen3_5GatedDeltaNet +from freetoken.utils import download_hf_weight, nvtx_annotate + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + + from .args import Qwen4ExpArgs + + +class _Qwen4MRoPE(StateLessOP): + """Partial, interleaved temporal/height/width RoPE for Qwen4-Exp.""" + + def __init__(self, config: ModelConfig): + rotary = config.rotary_config + self._base = get_rope( + head_dim=rotary.head_dim, + rotary_dim=rotary.rotary_dim, + max_position=rotary.max_position, + base=rotary.base, + ) + self.mrope_section = tuple(config.qwen4_args.mrope_section) + + @property + def head_size(self) -> int: + return self._base.head_size + + @property + def rotary_dim(self) -> int: + return self._base.rotary_dim + + @property + def is_neox(self) -> bool: + return self._base.is_neox + + @property + def _cos_sin_cache(self) -> torch.Tensor: + return self._base._cos_sin_cache + + @_cos_sin_cache.setter + def _cos_sin_cache(self, value: torch.Tensor) -> None: + self._base._cos_sin_cache = value + + def apply_inplace( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + head_size: int | None = None, + ) -> None: + head_size = self.head_size if head_size is None else int(head_size) + if positions.ndim == 1: + self._base.apply_rope_with_cos_sin_cache_inplace( + positions=positions, + query=query, + key=key, + head_size=head_size, + cos_sin_cache=self._cos_sin_cache, + is_neox=self.is_neox, + ) + return + if query.is_cuda: + from freetoken.kernel.triton.rope import ( + apply_mrope_with_cos_sin_cache_inplace, + ) + + apply_mrope_with_cos_sin_cache_inplace( + positions=positions, + query=query, + key=key, + head_size=head_size, + cos_sin_cache=self._cos_sin_cache, + mrope_section=self.mrope_section, + is_neox=self.is_neox, + ) + return + + # CPU reference path for exact unit tests and configuration checks. + if positions.ndim != 2 or positions.shape != (3, query.shape[0]): + raise ValueError( + f"MRoPE positions must have shape (3, {query.shape[0]}), got " + f"{tuple(positions.shape)}" + ) + half = self.rotary_dim // 2 + pair = torch.arange(half, device=positions.device) + axis = torch.zeros(half, dtype=torch.long, device=positions.device) + axis[(pair % 3 == 1) & (pair < self.mrope_section[1] * 3)] = 1 + axis[(pair % 3 == 2) & (pair < self.mrope_section[2] * 3)] = 2 + selected = positions.long().transpose(0, 1)[:, axis] + dim = pair.view(1, -1).expand_as(selected) + cos = self._cos_sin_cache[:, :half][selected, dim] + sin = self._cos_sin_cache[:, half:][selected, dim] + for tensor in (query, key): + heads = tensor.shape[1] // head_size + view = tensor.view(tensor.shape[0], heads, head_size) + first = view[..., :half].float().clone() + second = view[..., half : self.rotary_dim].float().clone() + view[..., :half].copy_((first * cos[:, None] - second * sin[:, None]).to(view.dtype)) + view[..., half : self.rotary_dim].copy_( + (second * cos[:, None] + first * sin[:, None]).to(view.dtype) + ) + + def forward( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + self.apply_inplace(positions, query, key) + return query, key + + +class _GroupedRMSNorm(BaseOP): + def __init__(self, size: int, group_size: int, eps: float): + if size % group_size: + raise ValueError(f"RMSNorm size {size} is not divisible by group size {group_size}") + self.weight = torch.empty(size) + self.group_size = group_size + self.eps = eps + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + from freetoken.kernel.fla import rms_norm_gated + + return rms_norm_gated( + x=hidden, + weight=self.weight, + bias=None, + eps=self.eps, + group_size=self.group_size, + is_rms_norm=True, + weight_plus_one=True, + ) + + +class _GatedRMSNorm(BaseOP): + def __init__(self, size: int, eps: float, activation: str): + self.weight = torch.empty(size) + self.eps = eps + self.activation = activation + + def forward(self, hidden: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + from freetoken.kernel.fla import rms_norm_gated + + return rms_norm_gated( + x=hidden, + weight=self.weight, + bias=None, + z=gate, + eps=self.eps, + is_rms_norm=True, + norm_before_gate=True, + activation=self.activation, + ) + + +class _GatedResidual(BaseOP): + def __init__(self, config: ModelConfig, combine: bool = True): + args: Qwen4ExpArgs = config.qwen4_args + self.hc_count = args.hc_count + self.hidden_size = config.hidden_size + hc_size = self.hc_count * self.hidden_size + self.hc_norm = _GroupedRMSNorm(hc_size, self.hidden_size, config.rms_norm_eps) + self.input_mix_weight_down = LinearReplicated(hc_size, args.hc_lowrank, has_bias=False) + self.input_mix_weight_up = LinearReplicated(args.hc_lowrank, hc_size, has_bias=False) + self.block_inject_weight = ( + LinearReplicated(hc_size, self.hc_count, has_bias=False) if combine else None + ) + + def forward(self, hyper_input: torch.Tensor): + normalized = self.hc_norm.forward(hyper_input) + mix = F.silu(self.input_mix_weight_down.forward(normalized) / self.hc_count) + mix = torch.sigmoid(self.input_mix_weight_up.forward(mix)) + mix = mix.view(-1, self.hc_count, self.hidden_size) + mixed = (mix * normalized.view(-1, self.hc_count, self.hidden_size)).mean(dim=1) + if self.block_inject_weight is None: + return mixed + inject = 2 * torch.sigmoid(self.block_inject_weight.forward(normalized) / self.hc_count) + return mixed, hyper_input, inject + + +class _SharedExpert(BaseOP): + def __init__(self, config: ModelConfig): + width = config.shared_expert_intermediate_size + self.gate_up_proj = LinearColParallelMerged( + config.hidden_size, [width, width], has_bias=False + ) + self.down_proj = LinearRowParallel(width, config.hidden_size, has_bias=False) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + return self.down_proj.forward(silu_and_mul(self.gate_up_proj.forward(hidden))) + + +class _SparseMoE(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + weight_format = "fp8_block" if config.expert_quant == "fp8_block" else "bf16" + self.experts = make_moe_layer( + config, + layer_id=layer_id, + renormalize=bool(config.norm_topk_prob), + weight_format=weight_format, + ) + self.gate = LinearReplicated(config.hidden_size, config.num_experts, has_bias=False) + self.shared_expert = _SharedExpert(config) + self.shared_expert_gate = LinearReplicated(config.hidden_size, 1, has_bias=False) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + router_logits = self.gate.forward(hidden) + shared = self.shared_expert.forward(hidden) + shared *= torch.sigmoid(self.shared_expert_gate.forward(hidden)) + return self.experts.forward(hidden_states=hidden, router_logits=router_logits) + shared + + +def _shift_right_ignore_eos(tokens: torch.Tensor, shift: int, eos_token_id: int) -> torch.Tensor: + if shift == 0: + return tokens + positions = torch.arange(tokens.numel(), dtype=torch.long) + eos_positions = torch.where(tokens == eos_token_id, positions, -1) + previous_eos_inclusive = torch.cummax(eos_positions, dim=0).values + previous_eos = torch.cat([eos_positions.new_full((1,), -1), previous_eos_inclusive[:-1]]) + segment_start = previous_eos + 1 + source_positions = positions - shift + shifted = tokens[source_positions.clamp_min(0)] + valid = (positions - segment_start >= shift) & (source_positions >= 0) + return torch.where(valid, shifted, tokens.new_full((), eos_token_id)) + + +def build_ngram_ids( + tokens: torch.Tensor, + *, + ngram_size: int, + heads_per_ngram: int, + eos_token_id: int, + multipliers: torch.Tensor, + vocab_sizes: torch.Tensor, + offsets: torch.Tensor, +) -> torch.Tensor: + tokens = tokens.to(dtype=torch.long, device="cpu") + shifted = [ + _shift_right_ignore_eos(tokens, shift, eos_token_id) for shift in range(ngram_size) + ] + blocks = [] + for ngram in range(2, ngram_size + 1): + start = (ngram - 2) * heads_per_ngram + stop = start + heads_per_ngram + mixed = shifted[0] * multipliers[0] + for position in range(1, ngram): + mixed = torch.bitwise_xor(mixed, shifted[position] * multipliers[position]) + sizes = vocab_sizes[start:stop] + heads = torch.remainder(mixed.unsqueeze(-1), sizes) + blocks.append(heads + offsets[start:stop]) + return torch.cat(blocks, dim=-1) + + +def _ple_request_tokens(req, forwarded_ids: torch.Tensor | None = None) -> torch.Tensor: + """Return the complete host token history visible to this forward. + + The overlap scheduler advances ``device_len`` before it drains the prior + sampled token to ``req.input_ids``. During decode, that one current token is + already present in ``batch.input_ids``. Join it to the committed host prefix + so PLE hashes the same history as a non-overlapped forward. + """ + host_len = req.input_ids.numel() + if host_len >= req.device_len: + return req.input_ids[: req.device_len] + if host_len != req.cached_len: + raise RuntimeError( + "Qwen4-Exp PLE host history has an unexpected gap: " + f"host={host_len}, cached={req.cached_len}, device={req.device_len}" + ) + if forwarded_ids is None or forwarded_ids.numel() != req.extend_len: + actual = 0 if forwarded_ids is None else forwarded_ids.numel() + raise RuntimeError( + "Qwen4-Exp PLE needs the current forwarded tokens: " + f"got {actual}, expected {req.extend_len}" + ) + return torch.cat((req.input_ids[: req.cached_len], forwarded_ids.to(device="cpu"))) + + +class _HostNGramEmbedding(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + args: Qwen4ExpArgs = config.qwen4_args + self.layer_id = layer_id + self.ngram_size = args.ngram_size + self.heads_per_ngram = args.heads_per_ngram + self.eos_token_id = args.eos_token_id + self.embedding_dim = args.ple_embed_dim + self.split_ngram_parts = args.split_ngram_parts + self.ngram_heads = (args.ngram_size - 1) * args.heads_per_ngram + self.head_dim = self.embedding_dim // self.ngram_heads + self.layer_multipliers = torch.empty(args.ngram_size, dtype=torch.long) + self.ngram_heads_vocab_sizes = torch.empty(self.ngram_heads, dtype=torch.long) + self.ngram_heads_offsets = torch.empty(self.ngram_heads, dtype=torch.long) + self._handles = [] + self._shards: list[torch.Tensor] = [] + self._shard_ends = torch.empty(0, dtype=torch.long) + self._scale = torch.tensor(1.0, dtype=torch.bfloat16) + self._host_constants: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None + self._dummy = False + + def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: + if dummy: + self._dummy = True + return + folder = download_hf_weight(model_path) + index_path = os.path.join(folder, "model.safetensors.index.json") + with open(index_path) as index_file: + weight_map = json.load(index_file)["weight_map"] + prefix = ( + f"model.language_model.layers.{self.layer_id}.ple.ple_embedding." + "ngram_embedding" + ) + shard_count = len([key for key in weight_map if key.startswith(prefix + ".shard_")]) + if shard_count != self.split_ngram_parts: + raise RuntimeError( + f"Qwen4-Exp PLE has {shard_count} shards, expected {self.split_ngram_parts}" + ) + shard_keys = [f"{prefix}.shard_{shard_id}.weight" for shard_id in range(shard_count)] + if not shard_keys or any(key not in weight_map for key in shard_keys): + raise RuntimeError(f"Incomplete Qwen4-Exp PLE shards under {prefix}") + + handles = {} + shards = [] + for key in shard_keys: + filename = weight_map[key] + handle = handles.get(filename) + if handle is None: + handle = safetensors.safe_open( + os.path.join(folder, filename), framework="pt", device="cpu" + ).__enter__() + handles[filename] = handle + shard = handle.get_tensor(key) + if shard.dtype != torch.float8_e4m3fn or shard.shape[1] != self.head_dim: + raise RuntimeError(f"Unexpected PLE shard {key}: {shard.dtype} {tuple(shard.shape)}") + shards.append(shard.view(torch.uint8)) + scale_key = prefix + ".weight_scale" + scale_handle = handles.get(weight_map[scale_key]) + if scale_handle is None: + scale_handle = safetensors.safe_open( + os.path.join(folder, weight_map[scale_key]), framework="pt", device="cpu" + ).__enter__() + handles[weight_map[scale_key]] = scale_handle + + self._handles = list(handles.values()) + self._shards = shards + self._shard_ends = torch.tensor([shard.shape[0] for shard in shards]).cumsum(0) + self._scale = scale_handle.get_tensor(scale_key).reshape(()) + self._host_constants = ( + self.layer_multipliers.cpu(), + self.ngram_heads_vocab_sizes.cpu(), + self.ngram_heads_offsets.cpu(), + ) + expected_rows = int(self._host_constants[1][-1] + self._host_constants[2][-1]) + if int(self._shard_ends[-1]) < expected_rows: + raise RuntimeError( + f"PLE table has {int(self._shard_ends[-1])} rows, needs {expected_rows}" + ) + + def _current_ngram_ids(self) -> torch.Tensor: + if self._host_constants is None: + raise RuntimeError("Qwen4-Exp PLE host weights are not loaded") + batch = get_global_ctx().batch + reqs = batch.padded_reqs if batch.is_decode else batch.reqs + multipliers, vocab_sizes, offsets = self._host_constants + pieces = [] + forwarded_host = None + forwarded_offset = 0 + for req in reqs: + extend_len = req.extend_len + forwarded = None + if req.input_ids.numel() < req.device_len: + if forwarded_host is None: + forwarded_host = batch.input_ids.detach().to(device="cpu") + forwarded = forwarded_host[ + forwarded_offset : forwarded_offset + extend_len + ] + tokens = _ple_request_tokens(req, forwarded) + all_ids = build_ngram_ids( + tokens, + ngram_size=self.ngram_size, + heads_per_ngram=self.heads_per_ngram, + eos_token_id=self.eos_token_id, + multipliers=multipliers, + vocab_sizes=vocab_sizes, + offsets=offsets, + ) + pieces.append(all_ids[req.cached_len : req.device_len]) + forwarded_offset += extend_len + result = torch.cat(pieces, dim=0) + if result.shape[0] != batch.input_ids.numel(): + raise RuntimeError( + f"PLE token count {result.shape[0]} does not match batch {batch.input_ids.numel()}" + ) + return result + + def forward(self, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + if self._dummy: + token_count = get_global_ctx().batch.input_ids.numel() + return torch.zeros(token_count, self.embedding_dim, device=device, dtype=dtype) + ngram_ids = self._current_ngram_ids().reshape(-1) + shard_ids = torch.bucketize(ngram_ids, self._shard_ends, right=True) + output = torch.empty( + ngram_ids.numel(), + self.head_dim, + dtype=torch.uint8, + pin_memory=torch.cuda.is_available(), + ) + starts = torch.cat([self._shard_ends.new_zeros(1), self._shard_ends[:-1]]) + for shard_id in shard_ids.unique().tolist(): + positions = torch.nonzero(shard_ids == shard_id, as_tuple=False).flatten() + local_ids = ngram_ids.index_select(0, positions) - starts[shard_id] + rows = self._shards[shard_id].index_select(0, local_ids) + output.index_copy_(0, positions, rows) + fp8 = output.to(device=device, non_blocking=True).view(torch.float8_e4m3fn) + embedded = fp8.to(dtype) * self._scale.to(device=device, dtype=dtype) + return embedded.view(-1, self.embedding_dim) + + +class _DepthwiseConv(BaseOP): + def __init__(self, channels: int, kernel_size: int): + self.weight = torch.empty(channels, 1, kernel_size) + + +class _PLELayer(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + args: Qwen4ExpArgs = config.qwen4_args + self.layer_id = layer_id + self.hidden_size = config.hidden_size + self.hc_count = args.hc_count + hc_size = self.hidden_size * self.hc_count + self.ple_embedding = _HostNGramEmbedding(config, layer_id) + self.key_proj = LinearReplicated(args.ple_embed_dim, hc_size, has_bias=False) + self.value_proj = LinearReplicated(args.ple_embed_dim, self.hidden_size, has_bias=False) + self.norm_key = _GroupedRMSNorm(hc_size, self.hidden_size, config.rms_norm_eps) + self.norm_query = _GroupedRMSNorm(hc_size, self.hidden_size, config.rms_norm_eps) + self.norm_conv = _GroupedRMSNorm(hc_size, self.hidden_size, config.rms_norm_eps) + self.conv1d = _DepthwiseConv(hc_size, args.ple_conv_kernel_size) + self.dilation = args.ngram_size + self.state_len = (args.ple_conv_kernel_size - 1) * self.dilation + self._conv_states: dict[int, torch.Tensor] = {} + + def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: + self.ple_embedding.load_host_weights(model_path, dummy=dummy) + + def _short_conv(self, hidden: torch.Tensor) -> torch.Tensor: + batch = get_global_ctx().batch + reqs = batch.padded_reqs if batch.is_decode else batch.reqs + outputs = [] + offset = 0 + weight = self.conv1d.weight + for req in reqs: + length = req.extend_len + current = hidden[offset : offset + length].transpose(0, 1).unsqueeze(0) + state = self._conv_states.get(req.table_idx) + if req.cached_len == 0: + state = current.new_zeros(1, current.shape[1], self.state_len) + elif state is None: + raise RuntimeError( + "Qwen4-Exp PLE state cannot resume a radix prefix; serve with --cache-type naive" + ) + combined = torch.cat([state, current], dim=-1) + convolved = F.conv1d( + combined, + weight, + groups=weight.shape[0], + dilation=self.dilation, + ) + outputs.append(F.silu(convolved).squeeze(0).transpose(0, 1)) + self._conv_states[req.table_idx] = combined[..., -self.state_len :].detach() + offset += length + return torch.cat(outputs, dim=0) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + embeddings = self.ple_embedding.forward(hidden.device, hidden.dtype) + key = self.norm_key.forward(self.key_proj.forward(embeddings)) + key = key.view(-1, self.hc_count, self.hidden_size) + value = self.value_proj.forward(embeddings) + query = self.norm_query.forward(hidden).view(-1, self.hc_count, self.hidden_size) + gate = (key * query).sum(dim=-1, keepdim=True) / math.sqrt(self.hidden_size) + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + gated = (torch.sigmoid(gate) * value.unsqueeze(1)).flatten(1) + normalized = self.norm_conv.forward(gated) + return gated + self._short_conv(normalized) + + +class _QSAIndexer(BaseOP): + """Qwen4-Exp's weight-free four-head compressed-key indexer.""" + + def __init__(self, config: ModelConfig, rotary): + args: Qwen4ExpArgs = config.qwen4_args + self.num_q_heads = args.indexer_n_heads + self.num_kv_heads = args.indexer_kv_heads + self.head_dim = args.indexer_head_dim + self.q_dim = self.num_q_heads * self.head_dim + self.k_dim = self.num_kv_heads * self.head_dim + self.index_qk_proj = LinearReplicated( + config.hidden_size, self.q_dim + self.k_dim, has_bias=False + ) + self.q_layernorm = GemmaPlusOneRMSNorm(self.head_dim, config.rms_norm_eps) + self.k_layernorm = GemmaPlusOneRMSNorm(self.head_dim, config.rms_norm_eps) + self.rotary = rotary + if self.rotary.rotary_dim > self.head_dim: + raise ValueError( + f"QSA index head {self.head_dim} is smaller than rotary dim " + f"{self.rotary.rotary_dim}" + ) + + def _apply_rope(self, tensor: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + if tensor.numel() == 0: + return tensor + shape = tensor.shape + flat = tensor.reshape(shape[0], -1).contiguous() + # The shared RoPE object was built for 256-wide main heads. Call its + # kernel with the 128-wide QSA head size instead of using forward(), + # which would interpret the fused index-query row with the wrong stride. + dummy_key = torch.zeros( + shape[0], self.head_dim, dtype=tensor.dtype, device=tensor.device + ) + self.rotary.apply_inplace( + positions=positions, + query=flat, + key=dummy_key, + head_size=self.head_dim, + ) + return flat.view(shape) + + def project(self, hidden: torch.Tensor, positions: torch.Tensor): + qk = self.index_qk_proj.forward(hidden) + q_raw, k_raw = torch.split(qk, (self.q_dim, self.k_dim), dim=-1) + q = q_raw.view(-1, self.num_q_heads, self.head_dim).contiguous() + k = k_raw.view(-1, self.num_kv_heads, self.head_dim).contiguous() + q = self.q_layernorm.forward(q) + q = self._apply_rope(q, positions) + return q, k + + def normalize_compressed_keys( + self, keys: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + keys = self.k_layernorm.forward(keys.contiguous()) + return self._apply_rope(keys, positions) + + +class Qwen4ExpAttention(Qwen3_5Attention): + def __init__(self, config: ModelConfig, layer_id: int): + super().__init__(config, layer_id) + self.rotary = _Qwen4MRoPE(config) + # Qwen4 stores centered q/k norm weights (effective scale is 1 + w). + self.q_norm = GemmaPlusOneRMSNorm(config.head_dim, config.rms_norm_eps) + self.k_norm = GemmaPlusOneRMSNorm(config.head_dim, config.rms_norm_eps) + self.indexer = _QSAIndexer(config, self.rotary) + + @nvtx_annotate("QSA") + def forward(self, x: torch.Tensor) -> torch.Tensor: + ctx = get_global_ctx() + rope_positions = ctx.batch.rope_positions + if rope_positions is None: + rope_positions = ctx.batch.positions + q, k, v, gate = self._project(x, rope_positions) + index_q, index_k = self.indexer.project(x, rope_positions) + output = ctx.attn_backend.qsa_forward( + q, + k, + v, + index_q, + index_k, + self.indexer, + self.layer_id, + ctx.batch, + ) + return self._combine(output, gate) + + +class Qwen4ExpDecoderLayer(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + self._layer_id = layer_id + self._is_linear = config.is_linear_layer(layer_id) + dense_config = replace(config, expert_quant="none", attn_quant="none") + if self._is_linear: + group = config.linear_attention_group() + assert group is not None + self.linear_attn = Qwen3_5GatedDeltaNet( + hidden_size=config.hidden_size, + num_k_heads=group.num_key_heads, + num_v_heads=group.num_value_heads, + head_k_dim=group.key_head_dim, + head_v_dim=group.value_head_dim, + conv_kernel_size=group.conv_kernel_dim, + rms_norm_eps=config.rms_norm_eps, + layer_id=layer_id, + expert_quant="none", + attn_quant="none", + ) + self.linear_attn.norm = _GatedRMSNorm( + group.value_head_dim, + config.rms_norm_eps, + config.qwen4_args.output_gate_type, + ) + else: + self.self_attn = Qwen4ExpAttention(dense_config, layer_id) + self.mlp = _SparseMoE(config, layer_id) + self.ple = ( + _PLELayer(config, layer_id) + if layer_id in config.qwen4_args.ple_layer_ids + else None + ) + self.attn_hyper_connection = _GatedResidual(config) + self.mlp_hyper_connection = _GatedResidual(config) + + @nvtx_annotate("Layer_{}", layer_id_field="_layer_id") + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + if self.ple is not None: + hidden = hidden + self.ple.forward(hidden) + mixed, residual, weights = self.attn_hyper_connection.forward(hidden) + mixed = ( + self.linear_attn.forward(mixed) + if self._is_linear + else self.self_attn.forward(mixed) + ) + hidden = residual + (mixed.unsqueeze(1) * weights.unsqueeze(-1)).flatten(1) + mixed, residual, weights = self.mlp_hyper_connection.forward(hidden) + mixed = self.mlp.forward(mixed) + return residual + (mixed.unsqueeze(1) * weights.unsqueeze(-1)).flatten(1) + + +class Qwen4ExpModel(BaseOP): + def __init__(self, config: ModelConfig): + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.layers = OPList( + [Qwen4ExpDecoderLayer(config, layer_id) for layer_id in range(config.num_layers)] + ) + self.hyper_connection_mixer = _GatedResidual(config, combine=False) + self.hc_count = config.qwen4_args.hc_count + self._image_token_id = config.image_token_id + + def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: + for layer in self.layers.op_list: + if layer.ple is not None: + layer.ple.load_host_weights(model_path, dummy=dummy) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + hidden = self.embed_tokens.forward(input_ids) + mm_embeds = getattr(get_global_ctx().batch, "mm_embeds", None) + if mm_embeds is not None and self._image_token_id is not None: + mask = input_ids == self._image_token_id + slots = int(mask.sum().item()) + if slots != mm_embeds.shape[0]: + raise ValueError( + f"image-token slots ({slots}) do not match vision features " + f"({mm_embeds.shape[0]})" + ) + hidden = hidden.masked_scatter(mask.unsqueeze(-1), mm_embeds.to(hidden.dtype)) + hidden = hidden.repeat(1, self.hc_count) + for layer in self.layers.op_list: + hidden = layer.forward(hidden) + return self.hyper_connection_mixer.forward(hidden) + + +class Qwen4ExpForCausalLM(BaseLLMModel): + def __init__(self, config: ModelConfig): + self.model = Qwen4ExpModel(config) + self.lm_head = ParallelLMHead( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + tie_word_embeddings=config.tie_word_embeddings, + tied_embedding=self.model.embed_tokens if config.tie_word_embeddings else None, + ) + if config.is_multimodal: + from .vision import Qwen4VisionModel + + self.visual = Qwen4VisionModel(config.vision_config) + super().__init__() + + @torch.inference_mode() + def encode_images( + self, pixel_values: torch.Tensor, image_grid_thw: torch.Tensor + ) -> torch.Tensor: + if not hasattr(self, "visual"): + raise RuntimeError("Qwen4-Exp vision weights are not loaded") + return self.visual.forward(pixel_values, image_grid_thw) + + def load_host_weights(self, model_path: str, *, dummy: bool = False) -> None: + self.model.load_host_weights(model_path, dummy=dummy) + + def forward(self) -> torch.Tensor: + hidden = self.model.forward(get_global_ctx().batch.input_ids) + return self.lm_head.forward(hidden) + + +__all__ = ["Qwen4ExpForCausalLM", "build_ngram_ids"] diff --git a/python/freetoken/models/qwen4_exp/mrope.py b/python/freetoken/models/qwen4_exp/mrope.py new file mode 100644 index 00000000..c7ba4c4a --- /dev/null +++ b/python/freetoken/models/qwen4_exp/mrope.py @@ -0,0 +1,80 @@ +"""Qwen VL multimodal rotary-position helpers.""" + +from __future__ import annotations + +import itertools + +import torch + + +def build_mrope_positions( + input_ids: torch.Tensor, + mm_token_type_ids: torch.Tensor, + image_grid_thw: torch.Tensor, + spatial_merge_size: int, +) -> tuple[torch.Tensor, int]: + """Build exact Qwen3-VL image/text MRoPE positions for one request. + + ``mm_token_type_ids`` uses 0 for text and 1 for image tokens. Video is + rejected until the server carries its timestamps and grid metadata. + """ + tokens = input_ids.detach().to(device="cpu", dtype=torch.int64).reshape(-1) + types = mm_token_type_ids.detach().to(device="cpu", dtype=torch.int64).reshape(-1) + grids = image_grid_thw.detach().to(device="cpu", dtype=torch.int64).reshape(-1, 3) + if types.numel() != tokens.numel(): + raise ValueError( + "mm_token_type_ids length must match input_ids: " + f"{types.numel()} != {tokens.numel()}" + ) + if spatial_merge_size < 1: + raise ValueError("spatial_merge_size must be positive") + + grid_index = 0 + current_position = 0 + pieces: list[torch.Tensor] = [] + for modality, group in itertools.groupby(enumerate(types.tolist()), lambda item: item[1]): + members = list(group) + group_len = len(members) + if modality == 0: + positions = torch.arange(current_position, current_position + group_len) + pieces.append(positions.view(1, -1).expand(3, -1)) + current_position += group_len + continue + if modality != 1: + raise NotImplementedError( + "Qwen4-Exp video MRoPE is not enabled; image input is supported" + ) + if grid_index >= grids.shape[0]: + raise ValueError("mm_token_type_ids contains more image groups than image_grid_thw") + grid_t, grid_h, grid_w = (int(value) for value in grids[grid_index].tolist()) + grid_index += 1 + if grid_h % spatial_merge_size or grid_w % spatial_merge_size: + raise ValueError("image grid height and width must divide by spatial_merge_size") + llm_t = grid_t + llm_h = grid_h // spatial_merge_size + llm_w = grid_w // spatial_merge_size + expected = llm_t * llm_h * llm_w + if group_len != expected: + raise ValueError( + "image-token group length does not match image_grid_thw: " + f"{group_len} != {expected}" + ) + temporal = torch.arange(llm_t) + height = torch.arange(llm_h) + current_position + width = torch.arange(llm_w) + current_position + t_grid, h_grid, w_grid = torch.meshgrid( + temporal, height, width, indexing="ij" + ) + vision = torch.stack((t_grid, h_grid, w_grid), dim=0).reshape(3, -1) + vision[0].add_(current_position) + pieces.append(vision) + current_position += max(llm_h, llm_w) + + if grid_index != grids.shape[0]: + raise ValueError("image_grid_thw contains more images than mm_token_type_ids") + positions = torch.cat(pieces, dim=1) if pieces else torch.empty((3, 0), dtype=torch.int64) + delta = int(positions.max().item() + 1 - tokens.numel()) if tokens.numel() else 0 + return positions.contiguous(), delta + + +__all__ = ["build_mrope_positions"] diff --git a/python/freetoken/models/qwen4_exp/vision.py b/python/freetoken/models/qwen4_exp/vision.py new file mode 100644 index 00000000..aec8d2f6 --- /dev/null +++ b/python/freetoken/models/qwen4_exp/vision.py @@ -0,0 +1,241 @@ +"""Qwen3-VL-compatible vision tower used by Qwen3.8-Flash-Next.""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F +from freetoken.layers import BaseOP, LinearReplicated, OPList + +from .args import Qwen4VisionConfig + + +class _LayerNorm(BaseOP): + def __init__(self, size: int, eps: float = 1e-6): + self.weight = torch.empty(size) + self.bias = torch.empty(size) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.layer_norm(x, (x.shape[-1],), self.weight, self.bias, self.eps) + + +class _Embedding(BaseOP): + def __init__(self, count: int, width: int): + self.weight = torch.empty(count, width) + + def forward(self, indices: torch.Tensor) -> torch.Tensor: + return F.embedding(indices, self.weight) + + +class _Conv3dPatchProjection(BaseOP): + def __init__(self, config: Qwen4VisionConfig): + self.weight = torch.empty( + config.hidden_size, + config.in_channels, + config.temporal_patch_size, + config.patch_size, + config.patch_size, + ) + self.bias = torch.empty(config.hidden_size) + + def forward(self, pixels: torch.Tensor) -> torch.Tensor: + # The processor already patchifies each Conv3d receptive field into one + # row. Flattening the kernel makes this exactly the stride==kernel Conv3d. + return F.linear(pixels.to(self.weight.dtype), self.weight.flatten(1), self.bias) + + +class Qwen4VisionPatchEmbed(BaseOP): + def __init__(self, config: Qwen4VisionConfig): + self.proj = _Conv3dPatchProjection(config) + + def forward(self, pixels: torch.Tensor) -> torch.Tensor: + return self.proj.forward(pixels) + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + half = x.shape[-1] // 2 + return torch.cat((-x[..., half:], x[..., :half]), dim=-1) + + +def _apply_vision_rope( + q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + q_dtype, k_dtype = q.dtype, k.dtype + cos = cos.unsqueeze(-2).float() + sin = sin.unsqueeze(-2).float() + qf, kf = q.float(), k.float() + q = qf * cos + _rotate_half(qf) * sin + k = kf * cos + _rotate_half(kf) * sin + return q.to(q_dtype), k.to(k_dtype) + + +class Qwen4VisionAttention(BaseOP): + def __init__(self, config: Qwen4VisionConfig): + self.num_heads = config.num_heads + self.head_dim = config.hidden_size // config.num_heads + self.qkv = LinearReplicated( + config.hidden_size, config.hidden_size * 3, has_bias=True + ) + self.proj = LinearReplicated( + config.hidden_size, config.hidden_size, has_bias=True + ) + + def forward( + self, + hidden: torch.Tensor, + segment_lengths: tuple[int, ...], + cos: torch.Tensor, + sin: torch.Tensor, + ) -> torch.Tensor: + tokens = hidden.shape[0] + q, k, v = ( + self.qkv.forward(hidden) + .view(tokens, 3, self.num_heads, self.head_dim) + .permute(1, 0, 2, 3) + .unbind(0) + ) + q, k = _apply_vision_rope(q, k, cos, sin) + outputs = [] + start = 0 + for length in segment_lengths: + stop = start + length + qs = q[start:stop].transpose(0, 1).unsqueeze(0) + ks = k[start:stop].transpose(0, 1).unsqueeze(0) + vs = v[start:stop].transpose(0, 1).unsqueeze(0) + out = F.scaled_dot_product_attention( + qs, ks, vs, is_causal=False, scale=self.head_dim**-0.5 + ) + outputs.append(out.squeeze(0).transpose(0, 1)) + start = stop + if start != tokens: + raise ValueError( + f"vision grid describes {start} patches but processor supplied {tokens}" + ) + return self.proj.forward(torch.cat(outputs, dim=0).reshape(tokens, -1)) + + +class Qwen4VisionMLP(BaseOP): + def __init__(self, config: Qwen4VisionConfig): + self.linear_fc1 = LinearReplicated( + config.hidden_size, config.intermediate_size, has_bias=True + ) + self.linear_fc2 = LinearReplicated( + config.intermediate_size, config.hidden_size, has_bias=True + ) + self.hidden_act = config.hidden_act + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + projected = self.linear_fc1.forward(hidden) + if self.hidden_act == "gelu_pytorch_tanh": + projected = F.gelu(projected, approximate="tanh") + elif self.hidden_act == "gelu": + projected = F.gelu(projected) + else: + raise ValueError(f"Unsupported Qwen vision activation {self.hidden_act!r}") + return self.linear_fc2.forward(projected) + + +class Qwen4VisionBlock(BaseOP): + def __init__(self, config: Qwen4VisionConfig): + self.norm1 = _LayerNorm(config.hidden_size) + self.norm2 = _LayerNorm(config.hidden_size) + self.attn = Qwen4VisionAttention(config) + self.mlp = Qwen4VisionMLP(config) + + def forward(self, hidden, segment_lengths, cos, sin): + hidden = hidden + self.attn.forward( + self.norm1.forward(hidden), segment_lengths, cos, sin + ) + return hidden + self.mlp.forward(self.norm2.forward(hidden)) + + +class Qwen4VisionPatchMerger(BaseOP): + def __init__(self, config: Qwen4VisionConfig, use_postshuffle_norm: bool = False): + merged = config.hidden_size * config.spatial_merge_size**2 + self.norm = _LayerNorm(merged if use_postshuffle_norm else config.hidden_size) + self.linear_fc1 = LinearReplicated(merged, merged, has_bias=True) + self.linear_fc2 = LinearReplicated(merged, config.out_hidden_size, has_bias=True) + self.use_postshuffle_norm = use_postshuffle_norm + self.merged = merged + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + if self.use_postshuffle_norm: + hidden = self.norm.forward(hidden.view(-1, self.merged)) + else: + hidden = self.norm.forward(hidden).view(-1, self.merged) + return self.linear_fc2.forward(F.gelu(self.linear_fc1.forward(hidden))) + + +class Qwen4VisionModel(BaseOP): + def __init__(self, config: Qwen4VisionConfig): + self.patch_embed = Qwen4VisionPatchEmbed(config) + self.pos_embed = _Embedding(config.num_position_embeddings, config.hidden_size) + self.blocks = OPList([Qwen4VisionBlock(config) for _ in range(config.depth)]) + self.merger = Qwen4VisionPatchMerger(config) + self.deepstack_merger_list = OPList( + [Qwen4VisionPatchMerger(config, True) for _ in config.deepstack_visual_indexes] + ) + self.spatial_merge_size = config.spatial_merge_size + self.num_grid_per_side = int(math.sqrt(config.num_position_embeddings)) + self.head_dim = config.hidden_size // config.num_heads + self.deepstack_visual_indexes = config.deepstack_visual_indexes + # The engine constructs the model under ``torch.device("meta")``. This + # value is derived, not loaded from the checkpoint, so keep only its + # shape here and materialize the tiny frequency vector on the real device. + self._inv_dim = self.head_dim // 2 + + def _position_data(self, grid_thw: torch.Tensor, dtype: torch.dtype): + from transformers.vision_utils import ( + get_vision_interpolation_indices_and_weights, + get_vision_position_ids, + ) + + indices, weights = get_vision_interpolation_indices_and_weights( + grid_thw, + num_grid_per_side=self.num_grid_per_side, + mode="bilinear", + align_corners=True, + spatial_merge_size=self.spatial_merge_size, + ) + pos = (self.pos_embed.forward(indices) * weights[:, :, None].to(dtype)).sum(1) + position_ids = get_vision_position_ids(grid_thw, self.spatial_merge_size) + inv = 1.0 / ( + 10000.0 + ** ( + torch.arange( + 0, + self._inv_dim, + 2, + dtype=torch.float32, + device=position_ids.device, + ) + / self._inv_dim + ) + ) + rotary = (position_ids.unsqueeze(-1).float() * inv).flatten(1) + rotary = torch.cat((rotary, rotary), dim=-1) + return pos, rotary.cos().to(dtype), rotary.sin().to(dtype) + + def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: + grid_thw = grid_thw.to(device=pixel_values.device, dtype=torch.long) + hidden = self.patch_embed.forward(pixel_values) + pos, cos, sin = self._position_data(grid_thw, hidden.dtype) + hidden = hidden + pos.to(hidden.dtype) + segment_lengths = tuple( + int(h) * int(w) + for t, h, w in grid_thw.detach().cpu().tolist() + for _ in range(int(t)) + ) + deepstack = [] + for layer_id, block in enumerate(self.blocks.op_list): + hidden = block.forward(hidden, segment_lengths, cos, sin) + if layer_id in self.deepstack_visual_indexes: + slot = self.deepstack_visual_indexes.index(layer_id) + deepstack.append(self.deepstack_merger_list.op_list[slot].forward(hidden)) + merged = self.merger.forward(hidden) + return torch.cat((merged, *deepstack), dim=-1) if deepstack else merged + + +__all__ = ["Qwen4VisionModel"] diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py new file mode 100644 index 00000000..65793983 --- /dev/null +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +from typing import Iterator + +import safetensors +import torch +from freetoken.distributed import get_tp_info +from freetoken.models.loader import iter_weight_files +from tqdm import tqdm + +from freetoken.models.qwen3_5_moe.weight import ( + iter_weights_parallel, + load_nvfp4_expert_sources, + load_nvfp4_expert_sources_parallel, + setup_offload_expert_banks, +) + + +_FUSIONS = { + ".self_attn.qkv_proj.weight": ( + ".self_attn.q_proj.weight", + ".self_attn.k_proj.weight", + ".self_attn.v_proj.weight", + ), + ".linear_attn.in_proj.weight": ( + ".linear_attn.in_proj_qkv.weight", + ".linear_attn.in_proj_z.weight", + ".linear_attn.in_proj_b.weight", + ".linear_attn.in_proj_a.weight", + ), + ".mlp.shared_expert.gate_up_proj.weight": ( + ".mlp.shared_expert.gate_proj.weight", + ".mlp.shared_expert.up_proj.weight", + ), +} + + +def _rename(raw_name: str) -> str | None: + if raw_name.startswith("mtp."): + return None + if raw_name.startswith("model.visual."): + return "visual." + raw_name[len("model.visual.") :] + if raw_name.startswith("visual."): + return raw_name + if ".ple.ple_embedding.ngram_embedding." in raw_name: + return None + if raw_name.startswith("model.language_model."): + return "model." + raw_name[len("model.language_model.") :] + if raw_name.startswith("language_model."): + return "model." + raw_name[len("language_model.") :] + return raw_name + + +def _try_fuse(name: str, tensor: torch.Tensor, buffers: dict): + for fused_suffix, parts in _FUSIONS.items(): + for index, part in enumerate(parts): + if name.endswith(part): + fused_name = name[: -len(part)] + fused_suffix + slots = buffers.setdefault(fused_name, {}) + slots[index] = tensor + if len(slots) == len(parts): + del buffers[fused_name] + return fused_name, torch.cat([slots[i] for i in range(len(parts))], dim=0) + return () + return None + + +def iter_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + if get_tp_info().size > 1: + raise NotImplementedError("Qwen4-Exp currently supports TP=1 only") + if include_moe_experts: + raise ValueError("Qwen4-Exp requires --moe-backend offload, cpu, or hybrid") + if not include_non_moe: + return + + buffers = {} + for filename in tqdm( + iter_weight_files(model_path), + desc="Loading Qwen4-Exp resident weights", + disable=not get_tp_info().is_primary(), + ): + with safetensors.safe_open(filename, framework="pt", device=str(device)) as handle: + for raw_name in handle.keys(): + name = _rename(raw_name) + if ( + name is None + or ".mlp.experts." in name + or raw_name.endswith(".weight_scale_inv") + ): + continue + tensor = handle.get_tensor(raw_name) + fused = _try_fuse(name, tensor, buffers) + if fused is not None: + if fused: + yield fused + continue + yield name, tensor + if buffers: + raise RuntimeError(f"Incomplete Qwen4-Exp projection fusions: {sorted(buffers)}") + + +__all__ = [ + "iter_weights", + "iter_weights_parallel", + "load_nvfp4_expert_sources", + "load_nvfp4_expert_sources_parallel", + "setup_offload_expert_banks", +] diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 0c033ca0..9626261a 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -65,6 +65,12 @@ class ModelSpec: "freetoken.models.qwen3_5_moe", "Qwen3_5MoEForCausalLM", ), + # Qwen3.8-Flash-Next / Qwen4-Exp text tower. Routed FP8 experts use the + # standard offload banks; its 51 GB PLE table is mmap-backed on the host. + "Qwen4ExpForConditionalGeneration": ModelSpec( + "freetoken.models.qwen4_exp", + "Qwen4ExpForCausalLM", + ), # Muse-Glimmer-30B (model_type muse_glimmer): multimodal wrapper config (text tower in # text_config, weights under model.language_model.); served text-only. Dense gated GQA # with a [SWA x3, full] pattern -- full layers are NoPE -- weightless qk norms, centered diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py index f3e5359a..e713ddd7 100644 --- a/python/freetoken/moe/benchbw.py +++ b/python/freetoken/moe/benchbw.py @@ -113,6 +113,9 @@ class Workload: # Preset workloads. Dims from the model configs / benchmarks/bench_offload_cache_copy.py. # E/top_k/H/I are what drive the per-expert byte size and thus the bandwidths. WORKLOADS: dict[str, Workload] = { + "qwen3.8-flash-next": Workload( + "qwen3.8-flash-next", 2560, 640, 512, 10, ("nvfp4",) + ), "qwen3.6-moe": Workload("qwen3.6-moe", 2048, 512, 256, 8, ("bf16", "nvfp4", "fp8_block")), "qwen3-30b": Workload("qwen3-30b", 2048, 768, 128, 8, ("bf16",)), "gemma4-26b": Workload("gemma4-26b", 2816, 704, 128, 8, ("bf16",), activation="gelu_tanh"), diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py index fe7e417d..cefd8b01 100644 --- a/python/freetoken/moe/fused.py +++ b/python/freetoken/moe/fused.py @@ -46,20 +46,30 @@ def fused_topk( from freetoken.kernel.backend import is_triton_kernels_installed - # triton_kernels ships no Windows wheel, and unlike flashinfer/sgl_kernel it is not one - # of the six ops the in-repo triton kernels cover -- so this router needs its own fallback. - if not is_triton_kernels_installed(): + # The external Triton kernel pads several internal dimensions to powers of two, but + # does not pad N_EXPTS_ACT. Its tl.arange therefore cannot compile for values such as + # Qwen3.8-Flash-Next's topk=10. Use the exact Torch path for those models as well as + # on systems (notably Windows) where triton_kernels has no wheel. + triton_topk_supported = topk > 0 and (topk & (topk - 1)) == 0 + triton_topk_installed = is_triton_kernels_installed() + if not triton_topk_installed or not triton_topk_supported: global _warned_torch_topk if not _warned_torch_topk: _warned_torch_topk = True - # Once, not per call: this runs every MoE forward. On Linux a missing - # triton_kernels used to fail fast with ImportError; keep the misconfiguration - # visible without giving up the fallback that Windows needs. - logger.warning_rank0( - "fused_topk: triton_kernels is not installed -> pure-torch router fallback " - "(numerically equivalent, slower). Expected on Windows (no wheel); on Linux " - "install triton_kernels to restore the fused router." - ) + if not triton_topk_supported: + logger.warning_rank0( + f"fused_topk: topk={topk} is not supported by triton_kernels -> " + "pure-torch router fallback (numerically equivalent, slower)." + ) + else: + # Once, not per call: this runs every MoE forward. On Linux a missing + # triton_kernels used to fail fast with ImportError; keep the + # misconfiguration visible without giving up the fallback Windows needs. + logger.warning_rank0( + "fused_topk: triton_kernels is not installed -> pure-torch router " + "fallback (numerically equivalent, slower). Expected on Windows " + "(no wheel); on Linux install triton_kernels to restore the fused router." + ) return _torch_fused_topk(gating_output, topk, renormalize, num_token_non_padded) from triton_kernels.topk import topk as triton_kernels_topk diff --git a/python/freetoken/moe/host_banks.py b/python/freetoken/moe/host_banks.py index d7af348a..700ac87c 100644 --- a/python/freetoken/moe/host_banks.py +++ b/python/freetoken/moe/host_banks.py @@ -147,7 +147,29 @@ def release(self) -> None: For buffers that are done being read (the converter). No-op for born-pinned banks: registered pages cannot be dropped.""" if self._pinned: return - self._buf.madvise(mmap.MADV_DONTNEED) + madvise = getattr(self._buf, "madvise", None) + dontneed = getattr(mmap, "MADV_DONTNEED", None) + if madvise is not None and dontneed is not None: + madvise(dontneed) + return + if os.name == "nt": + # Windows does not expose mmap.madvise. The streaming converter is + # finished with this bank, including every tensor alias kept in its + # source lists. Reset that shared Tensor object, then close the pagefile + # mapping so each completed layer releases physical memory immediately. + self.tensor.set_(torch.empty(0, dtype=self.tensor.dtype)) + buf = self._buf + try: + buf.close() + except BufferError: + # A backend-specific repacked view can still export the mapping. + # Correctness does not depend on discarding it; Windows can reclaim + # the pages under pressure and process exit releases the mapping. + return + try: + _LIVE_BUFFERS.remove(buf) + except ValueError: + pass def lock(self) -> None: """mlock the (now-filled) buffer: resident without CUDA pin quota, but no device address -- only the CPU executor can serve a locked layer. @@ -172,6 +194,10 @@ def lock(self) -> None: def _os_lock(addr: int, nbytes: int) -> None: + if os.name == "nt": + _windows_lock(addr, nbytes) + return + global _os_locked_total import resource @@ -197,6 +223,64 @@ def _os_lock(addr: int, nbytes: int) -> None: _os_locked_total += nbytes +def _windows_lock(addr: int, nbytes: int) -> None: + """Keep a host-bank range resident with ``VirtualLock`` on native Windows. + + Windows limits ``VirtualLock`` to the process working-set floor. Grow that + floor before locking each completed bank. This does not consume CUDA's WDDM + registered-memory quota and does not create a GPU device address, so these + banks remain CPU-executor-only exactly like Linux ``mlock`` banks. + """ + global _os_locked_total + import ctypes + from ctypes import wintypes + + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + process = kernel32.GetCurrentProcess() + + kernel32.GetProcessWorkingSetSize.argtypes = [ + wintypes.HANDLE, + ctypes.POINTER(ctypes.c_size_t), + ctypes.POINTER(ctypes.c_size_t), + ] + kernel32.GetProcessWorkingSetSize.restype = wintypes.BOOL + kernel32.SetProcessWorkingSetSize.argtypes = [ + wintypes.HANDLE, + ctypes.c_size_t, + ctypes.c_size_t, + ] + kernel32.SetProcessWorkingSetSize.restype = wintypes.BOOL + kernel32.VirtualLock.argtypes = [ctypes.c_void_p, ctypes.c_size_t] + kernel32.VirtualLock.restype = wintypes.BOOL + + current_min = ctypes.c_size_t() + current_max = ctypes.c_size_t() + if not kernel32.GetProcessWorkingSetSize( + process, ctypes.byref(current_min), ctypes.byref(current_max) + ): + err = ctypes.get_last_error() + raise OSError(err, f"GetProcessWorkingSetSize: {ctypes.FormatError(err)}") + + reserve = 256 << 20 + want = _os_locked_total + nbytes + reserve + new_min = max(current_min.value, want) + new_max = max(current_max.value, new_min + reserve) + if not kernel32.SetProcessWorkingSetSize(process, new_min, new_max): + err = ctypes.get_last_error() + raise OSError( + err, + f"SetProcessWorkingSetSize({new_min / 2**30:.1f} GiB minimum): " + f"{ctypes.FormatError(err)}", + ) + if not kernel32.VirtualLock(ctypes.c_void_p(addr), ctypes.c_size_t(nbytes)): + err = ctypes.get_last_error() + raise OSError( + err, + f"VirtualLock({nbytes / 2**30:.1f} GiB): {ctypes.FormatError(err)}", + ) + _os_locked_total += nbytes + + def alloc_banks(specs: dict[str, tuple[tuple[int, ...], torch.dtype]]) -> dict[str, HostBank]: """Allocate (lazy, unpinned) host banks from ``{name: (shape, dtype)}``.""" return {name: HostBank(shape, dtype) for name, (shape, dtype) in specs.items()} diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py index be84874c..4733705f 100644 --- a/python/freetoken/scheduler/prefill.py +++ b/python/freetoken/scheduler/prefill.py @@ -178,6 +178,8 @@ def _add_one_req( cache_handle=cache_handle, sampling_params=pending_req.sampling_params, mm_embeds=pending_req.mm_embeds, + mrope_position_ids=pending_req.mrope_position_ids, + mrope_position_delta=pending_req.mrope_position_delta, ) # Hybrid GDN per-request state slots (None for non-hybrid). On a fresh admit these are # freshly allocated; on a chunked continuation they are inherited from the prior chunk. @@ -238,7 +240,14 @@ class PrefillManager: def add_one_req(self, req: UserMsg) -> None: self.pending_list.append( - PendingReq(req.uid, req.input_ids, req.sampling_params, mm_embeds=req.mm_embeds) + PendingReq( + req.uid, + req.input_ids, + req.sampling_params, + mm_embeds=req.mm_embeds, + mrope_position_ids=req.mrope_position_ids, + mrope_position_delta=req.mrope_position_delta, + ) ) def schedule_next_batch(self, prefill_budget: int) -> Batch | None: diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 48923e3b..cf7cd75b 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -474,6 +474,36 @@ def _gpu_mem_bytes(self) -> int: return 0 return torch.cuda.memory_reserved(self.device) + @torch.inference_mode() + def _prepare_multimodal_request(self, msg: UserMsg) -> None: + """Convert online Qwen processor outputs into vision embeddings and mRoPE.""" + grid = msg.mm_image_grid_thw + token_types = msg.mm_token_type_ids + if grid is None or token_types is None: + raise ValueError("Qwen image input needs image_grid_thw and mm_token_type_ids") + model = self.engine.model + if not hasattr(model, "encode_images"): + raise ValueError(f"{type(model).__name__} does not support image inputs") + + msg.mm_embeds = model.encode_images( + msg.mm_pixel_values.to(self.device), grid.to(self.device) + ) + from freetoken.models.qwen4_exp.mrope import build_mrope_positions + + vision_config = self.config.model_config.vision_config + if vision_config is None: + raise ValueError("model vision configuration is not loaded") + msg.mrope_position_ids, msg.mrope_position_delta = build_mrope_positions( + msg.input_ids, + token_types, + grid, + int(vision_config.spatial_merge_size), + ) + # Release the large CPU transport tensor before prefill. + msg.mm_pixel_values = None + msg.mm_image_grid_thw = None + msg.mm_token_type_ids = None + def _process_one_msg(self, msg: BaseBackendMsg) -> None: if isinstance(msg, BatchBackendMsg): for msg in msg.data: @@ -489,6 +519,15 @@ def _process_one_msg(self, msg: BaseBackendMsg) -> None: "Dropping request %d because its abort arrived before admission", msg.uid ) return + if msg.mm_pixel_values is not None: + try: + self._prepare_multimodal_request(msg) + except Exception as exc: # noqa: BLE001 - fail this request, not the server + logger.warning_rank0("Image processing failed for request %d: %s", msg.uid, exc) + self.send_result( + [ErrorReplyMsg(uid=msg.uid, error=f"could not encode image: {exc}")] + ) + return input_len, max_seq_len = len(msg.input_ids), self.engine.max_seq_len max_output_len = max_seq_len - input_len if max_output_len <= 0: @@ -781,6 +820,7 @@ def _prepare_batch(self, batch: Batch) -> ForwardInput: if batch.is_prefill: self._gather_multimodal(batch) batch.positions = _make_positions(batch, self.device) + batch.rope_positions = _make_rope_positions(batch, self.device) input_mapping = _make_input_tuple(batch, self.device) write_mapping = _make_write_tuple(batch, self.device) batch.out_loc = self.engine.page_table[input_mapping] @@ -890,6 +930,38 @@ def _make_positions(batch: Batch, device: torch.device) -> torch.Tensor: return indices_host.to(device, non_blocking=True) +def _make_rope_positions(batch: Batch, device: torch.device) -> torch.Tensor | None: + """Build packed three-axis positions when a batch contains Qwen VL input.""" + if not any(req.mrope_position_ids is not None for req in batch.padded_reqs): + return None + needed_size = sum(req.extend_len for req in batch.padded_reqs) + host = torch.empty((3, needed_size), dtype=torch.int64, pin_memory=True) + offset = 0 + for req in batch.padded_reqs: + start, end = int(req.cached_len), int(req.device_len) + length = end - start + if not length: + continue + prompt_positions = req.mrope_position_ids + prompt_len = 0 if prompt_positions is None else int(prompt_positions.shape[1]) + prompt_stop = min(end, prompt_len) + copied = max(prompt_stop - start, 0) + if copied: + host[:, offset : offset + copied].copy_( + prompt_positions[:, start:prompt_stop] + ) + generated_start = start + copied + if generated_start < end: + generated = torch.arange( + generated_start, + end, + dtype=torch.int64, + ).add_(int(req.mrope_position_delta)) + host[:, offset + copied : offset + length].copy_(generated.expand(3, -1)) + offset += length + return host.to(device, non_blocking=True) + + def _make_input_tuple(batch: Batch, device: torch.device) -> Indice2D: mapping_host = torch.empty(len(batch.positions), dtype=torch.int64, pin_memory=True) offset = 0 diff --git a/python/freetoken/scheduler/utils.py b/python/freetoken/scheduler/utils.py index 0d6dd512..6f5a6b49 100644 --- a/python/freetoken/scheduler/utils.py +++ b/python/freetoken/scheduler/utils.py @@ -18,6 +18,8 @@ class PendingReq: sampling_params: SamplingParams chunked_req: ChunkedReq | None = None mm_embeds: torch.Tensor | None = None + mrope_position_ids: torch.Tensor | None = None + mrope_position_delta: int = 0 @property def input_len(self) -> int: diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 3e2acc85..c3fb5c0a 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -59,6 +59,25 @@ BACKEND_DEATH_EXIT_GRACE_S = 10.0 +def windows_selector_loop_factory(use_subprocess: bool = False) -> asyncio.AbstractEventLoop: + """Return the Windows loop required by pyzmq.asyncio. + + Uvicorn 0.36 and later bypasses the global event-loop policy and creates a + Proactor loop directly on Windows. pyzmq needs ``add_reader``, which only + the Selector loop provides without an optional Tornado dependency. + """ + del use_subprocess + return asyncio.SelectorEventLoop() + + +def _uvicorn_loop() -> str: + return ( + "freetoken.server.api_server:windows_selector_loop_factory" + if os.name == "nt" + else "auto" + ) + + def get_global_state() -> FrontendManager: global _GLOBAL_STATE assert _GLOBAL_STATE is not None, "Global state is not initialized" @@ -897,7 +916,9 @@ def _serve_and_run_shell(host: str, port: int) -> None: netloc = f"[{host}]:{port}" if ":" in host else f"{host}:{port}" origin = resolve_server_url(f"http://{netloc}").origin - server = uvicorn.Server(uvicorn.Config(app, host=host, port=port, access_log=False)) + server = uvicorn.Server( + uvicorn.Config(app, host=host, port=port, access_log=False, loop=_uvicorn_loop()) + ) thread = threading.Thread(target=server.run, name="freetoken-uvicorn", daemon=True) thread.start() _install_shell_stop_handlers() @@ -930,6 +951,12 @@ def run_api_server(config: ServerArgs, start_backend: Callable[[], "Any"], run_s global _GLOBAL_STATE, _MODEL_SAMPLING + if os.name == "nt": + # pyzmq.asyncio needs add_reader/add_writer. Python's default Windows + # Proactor loop does not provide them unless Tornado is installed. + # Select the native compatible loop before Uvicorn creates its loop. + asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) + if config.sampling_defaults == "model" and not config.use_dummy_weight: _MODEL_SAMPLING = load_generation_sampling(config.model_path) # Always surface the effective default sampling (model-recommended where available, @@ -1037,4 +1064,4 @@ def _on_meta(meta: dict) -> None: _serve_and_run_shell(host, port) return # uvicorn stays on the main thread (signal handling unchanged); ^C reaches the worker group. - uvicorn.run(app, host=host, port=port) + uvicorn.run(app, host=host, port=port, loop=_uvicorn_loop()) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index a71b6819..adef3e39 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -48,15 +48,39 @@ class ServerArgs(SchedulerConfig): def share_tokenizer(self) -> bool: return self.num_tokenizer == 0 + def _zmq_addr(self, channel: int) -> str: + if os.name == "nt": + # libzmq does not implement ipc:// on Windows. Keep every internal + # socket on loopback, after the HTTP port and the distributed port. + port = self.server_port + 2 + channel + if port > 65535: + raise ValueError( + f"server port {self.server_port} leaves no room for internal ports" + ) + return f"tcp://127.0.0.1:{port}" + return f"ipc:///tmp/freetoken_{channel}{self._unique_suffix}" + + @property + def zmq_backend_addr(self) -> str: + return self._zmq_addr(0) + + @property + def zmq_detokenizer_addr(self) -> str: + return self._zmq_addr(1) + + @property + def zmq_scheduler_broadcast_addr(self) -> str: + return self._zmq_addr(2) + @property def zmq_frontend_addr(self) -> str: - return "ipc:///tmp/freetoken_3" + self._unique_suffix + return self._zmq_addr(3) @property def zmq_tokenizer_addr(self) -> str: if self.share_tokenizer: return self.zmq_detokenizer_addr - result = "ipc:///tmp/freetoken_4" + self._unique_suffix + result = self._zmq_addr(4) assert result != self.zmq_detokenizer_addr return result @@ -188,7 +212,10 @@ def _infer_reasoning_parser(model_path: str) -> str | None: tag in marker for tag in ("v4", "deepseek_v4", "v3.2", "v32") ): return "deepseekv32" - if "qwen3" in marker or "qwen3.5" in marker or "qwen3_5" in marker: + if any( + tag in marker + for tag in ("qwen3", "qwen3.5", "qwen3_5", "qwen4_exp", "qwen4exp") + ): return "qwen3" if "glm" in marker: return "glm" diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index be05d908..aa7f4c23 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -188,9 +188,12 @@ def pick(value, key, framework): def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: - """Normalize OpenAI-shaped message dicts for the chat template: flatten text - content parts to a string and decode tool-call arguments from JSON. Raises - ValueError on a non-text content part (text-only server). Shared by all adapters.""" + """Normalize OpenAI-shaped messages for the chat template. + + Text-only part lists are flattened for broad template compatibility. Image + parts stay structured so multimodal templates can insert their image tokens + and the tokenizer worker can load the corresponding pixels. + """ return [_render_message(m) for m in messages] @@ -198,7 +201,10 @@ def _render_message(message: dict[str, Any]) -> dict[str, Any]: m = dict(message) content = m.get("content") if isinstance(content, list): - m["content"] = _flatten_text_parts(content) + if any(_is_image_part(part) for part in content): + m["content"] = [_normalize_content_part(part) for part in content] + else: + m["content"] = _flatten_text_parts(content) # Templates read different reasoning keys (reasoning_content: most; reasoning: # gemma4; thinking: gpt-oss) — accept any, emit both. reasoning = m.get("reasoning_content") or m.get("reasoning") or m.get("thinking") @@ -230,6 +236,23 @@ def _render_message(message: dict[str, Any]) -> dict[str, Any]: return m +def _is_image_part(part: Any) -> bool: + if not isinstance(part, dict): + return False + return part.get("type") in {"image", "image_url"} or "image" in part or "image_url" in part + + +def _normalize_content_part(part: Any) -> dict[str, Any]: + if not isinstance(part, dict): + raise ValueError("Message content parts must be objects") + ptype = part.get("type") + if ptype == "text": + return {"type": "text", "text": part.get("text") or ""} + if _is_image_part(part): + return dict(part) + raise ValueError(f"Unsupported content part type: {ptype}") + + def _flatten_text_parts(parts: list[Any]) -> str: texts: list[str] = [] for part in parts: @@ -237,7 +260,7 @@ def _flatten_text_parts(parts: list[Any]) -> str: if ptype == "text": texts.append((part.get("text") if isinstance(part, dict) else None) or "") else: - raise ValueError(f"Unsupported content part type for text-only server: {ptype}") + raise ValueError(f"Unsupported content part type: {ptype}") return "".join(texts) diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py index 530e862d..4d065068 100644 --- a/python/freetoken/tokenizer/server.py +++ b/python/freetoken/tokenizer/server.py @@ -1,7 +1,11 @@ from __future__ import annotations import multiprocessing as mp +import base64 +import io from typing import Any, List +from urllib.parse import unquote_to_bytes, urlparse +from urllib.request import Request, urlopen import torch from freetoken.message import ( @@ -82,9 +86,10 @@ def _send_generation_replies( def _tokenize_requests( tokenize_manager: Any, + multimodal_processor: Any, messages: List[TokenizeMsg], logger: Any, -) -> tuple[List[TokenizeMsg], List[torch.Tensor], List[UserReply]]: +) -> tuple[List[TokenizeMsg], List[torch.Tensor], List[dict[str, torch.Tensor] | None], List[UserReply]]: """Tokenize independently, returning backend work plus terminal frontend errors. Successful tokenization deliberately emits no prompt-token reply: accounting starts @@ -92,10 +97,11 @@ def _tokenize_requests( """ ok_msgs: List[TokenizeMsg] = [] ok_tensors: List[torch.Tensor] = [] + ok_multimodal: List[dict[str, torch.Tensor] | None] = [] errors: List[UserReply] = [] for msg in messages: try: - tokens = tokenize_manager.tokenize([msg])[0] + tokens, mm = multimodal_processor.encode(msg, tokenize_manager) except Exception as exc: # noqa: BLE001 — isolate, never crash the worker logger.warning(f"tokenization failed for request {msg.uid}: {exc!r}") errors.append( @@ -121,7 +127,95 @@ def _tokenize_requests( continue ok_msgs.append(msg) ok_tensors.append(tokens) - return ok_msgs, ok_tensors, errors + ok_multimodal.append(mm) + return ok_msgs, ok_tensors, ok_multimodal, errors + + +_MAX_IMAGE_BYTES = 64 * 1024 * 1024 + + +def _image_source(part: dict[str, Any]) -> str: + value = part.get("image_url", part.get("image")) + if isinstance(value, dict): + value = value.get("url") + if not isinstance(value, str) or not value: + raise ValueError("image content part needs a non-empty image URL") + return value + + +def _download_image(source: str) -> bytes: + if source.startswith("data:"): + header, separator, payload = source.partition(",") + if not separator: + raise ValueError("invalid image data URL") + try: + data = base64.b64decode(payload, validate=True) if ";base64" in header else unquote_to_bytes(payload) + except Exception as exc: + raise ValueError("invalid image data URL") from exc + elif urlparse(source).scheme in {"http", "https"}: + request = Request(source, headers={"User-Agent": "FreeToken/vision"}) + with urlopen(request, timeout=30) as response: # noqa: S310 - explicit http(s) check above + data = response.read(_MAX_IMAGE_BYTES + 1) + else: + raise ValueError("image URL must use data, http, or https") + if len(data) > _MAX_IMAGE_BYTES: + raise ValueError("image exceeds the 64 MiB input limit") + return data + + +def _message_image_sources(text: str | List[dict[str, Any]]) -> list[str]: + if not isinstance(text, list): + return [] + sources: list[str] = [] + for message in text: + content = message.get("content") if isinstance(message, dict) else None + if not isinstance(content, list): + continue + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") in {"image", "image_url"} or "image" in part or "image_url" in part: + sources.append(_image_source(part)) + return sources + + +class _MultimodalProcessor: + def __init__(self, model_path: str): + self.model_path = model_path + self.processor = None + + def encode( + self, msg: TokenizeMsg, tokenize_manager: Any + ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: + sources = _message_image_sources(msg.text) + if not sources: + return tokenize_manager.tokenize([msg])[0], None + + from PIL import Image + from transformers import AutoProcessor + + if self.processor is None: + self.processor = AutoProcessor.from_pretrained(self.model_path) + prompt = tokenize_manager.render_prompt(msg) + images = [] + try: + for source in sources: + with Image.open(io.BytesIO(_download_image(source))) as image: + images.append(image.convert("RGB")) + encoded = self.processor(text=[prompt], images=images, return_tensors="pt") + finally: + for image in images: + image.close() + + required = {"input_ids", "pixel_values", "image_grid_thw", "mm_token_type_ids"} + missing = sorted(required.difference(encoded)) + if missing: + raise ValueError(f"model processor did not return: {', '.join(missing)}") + return encoded["input_ids"][0].to(dtype=torch.int32), { + "pixel_values": encoded["pixel_values"].to(dtype=torch.bfloat16), + "image_grid_thw": encoded["image_grid_thw"].to(dtype=torch.int64), + "mm_token_type_ids": encoded["mm_token_type_ids"][0].to(dtype=torch.int32), + } @torch.inference_mode() @@ -148,6 +242,7 @@ def tokenize_worker( from .tokenize import TokenizeManager tokenize_manager = TokenizeManager(tokenizer) + multimodal_processor = _MultimodalProcessor(tokenizer_path) detokenize_manager = DetokenizeManager( tokenizer, load_eos_token_ids(tokenizer_path, tokenizer) ) @@ -245,18 +340,26 @@ def tokenize_worker( # Tokenize per-message so a single un-renderable request (e.g. a chat template # that rejects the message layout) becomes a terminal error reply for THAT uid # instead of an uncaught exception that kills the worker and bricks the server. - ok_msgs, ok_tensors, errors = _tokenize_requests( - tokenize_manager, tokenize_msg, logger + ok_msgs, ok_tensors, ok_multimodal, errors = _tokenize_requests( + tokenize_manager, multimodal_processor, tokenize_msg, logger ) if errors: send_frontend.put( errors[0] if len(errors) == 1 else BatchFrontendMsg(data=errors) ) if ok_msgs: - backend = [ - UserMsg(uid=msg.uid, input_ids=t, sampling_params=msg.sampling_params) - for msg, t in zip(ok_msgs, ok_tensors, strict=True) - ] + backend = [] + for msg, tokens, mm in zip(ok_msgs, ok_tensors, ok_multimodal, strict=True): + backend.append( + UserMsg( + uid=msg.uid, + input_ids=tokens, + sampling_params=msg.sampling_params, + mm_pixel_values=(mm or {}).get("pixel_values"), + mm_image_grid_thw=(mm or {}).get("image_grid_thw"), + mm_token_type_ids=(mm or {}).get("mm_token_type_ids"), + ) + ) send_backend.put(backend[0] if len(backend) == 1 else BatchBackendMsg(data=backend)) if len(abort_msg) > 0: batch_output = BatchBackendMsg( diff --git a/scripts/start-qwen38-flash-next.ps1 b/scripts/start-qwen38-flash-next.ps1 new file mode 100644 index 00000000..9d2671fa --- /dev/null +++ b/scripts/start-qwen38-flash-next.ps1 @@ -0,0 +1,51 @@ +[CmdletBinding()] +param( + [Parameter(Mandatory = $true)] + [string]$ModelPath, + [int]$Port = 1927, + [int]$MaxContext = 262144, + [int]$MoeCacheSize = 1024 +) + +$ErrorActionPreference = 'Stop' +$repoRoot = (Resolve-Path (Join-Path $PSScriptRoot '..')).Path +$python = Join-Path $env:LOCALAPPDATA 'FreeToken\venv\Scripts\python.exe' +$kernelDir = Join-Path $env:LOCALAPPDATA 'FreeToken\venv\Lib\site-packages\freetoken\kernel' + +foreach ($requiredPath in @($ModelPath, $python, $kernelDir)) { + if (-not (Test-Path -LiteralPath $requiredPath)) { + throw "Required FreeToken path does not exist: $requiredPath" + } +} + +$env:PYTHONPATH = Join-Path $repoRoot 'python' +$env:FREETOKEN_INSTALLED_KERNEL_DIR = $kernelDir +$env:FREETOKEN_DISABLE_KERNEL_CACHE_VERSION_CHECK = '1' +$env:FREETOKEN_LOAD_VISION = '1' +$env:HF_HUB_DISABLE_PROGRESS_BARS = '1' +$env:FREETOKEN_PIN_BUDGET_GB = '64' + +Write-Host "Starting Qwen3.8-Flash-Next NVFP4 on http://127.0.0.1:$Port" +Write-Host "OpenAI model: qwen3.8-flash-next-nvfp4; context: $MaxContext; vision: enabled" + +& $python -m freetoken.cli serve ` + --model $ModelPath ` + --host 127.0.0.1 ` + --port $Port ` + --served-model-name qwen3.8-flash-next-nvfp4 ` + --dtype bfloat16 ` + --max-running-requests 1 ` + --max-seq-len-override $MaxContext ` + --num-tokens $MaxContext ` + --max-prefill-length 8192 ` + --attention-backend auto ` + --cache-type naive ` + --moe-backend offload ` + --nvfp4-backend triton ` + --expert-load serial ` + --moe-cache-size $MoeCacheSize ` + --moe-cpu-layers 0 ` + --tool-call-parser qwen ` + --reasoning-parser qwen3 + +exit $LASTEXITCODE diff --git a/tests/checkpoint/__init__.py b/tests/checkpoint/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/tests/checkpoint/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/checkpoint/test_convert_metadata.py b/tests/checkpoint/test_convert_metadata.py new file mode 100644 index 00000000..a58b00b7 --- /dev/null +++ b/tests/checkpoint/test_convert_metadata.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import torch + +from freetoken.checkpoint.convert import _copy_metadata +from freetoken.checkpoint.ftw import FTWReader, FTWWriter, iter_ftw_weights + + +def test_copy_metadata_keeps_only_qwen4_host_mapped_shards(tmp_path: Path) -> None: + source = tmp_path / "source" + output = tmp_path / "output" + source.mkdir() + (source / "config.json").write_text("{}", encoding="utf-8") + (source / "model-plefp8-00000.safetensors").write_bytes(b"ple") + (source / "model-00001.safetensors").write_bytes(b"dense") + ple_name = ( + "model.language_model.layers.0.ple.ple_embedding." + "ngram_embedding.shard_0.weight" + ) + index = { + "metadata": {"total_size": 8}, + "weight_map": { + ple_name: "model-plefp8-00000.safetensors", + "model.embed_tokens.weight": "model-00001.safetensors", + }, + } + (source / "model.safetensors.index.json").write_text( + json.dumps(index), encoding="utf-8" + ) + + copied = _copy_metadata(str(source), str(output)) + + assert (output / "config.json").is_file() + assert (output / "model-plefp8-00000.safetensors").read_bytes() == b"ple" + assert not (output / "model-00001.safetensors").exists() + slim = json.loads((output / "model.safetensors.index.json").read_text()) + assert slim["weight_map"] == {ple_name: "model-plefp8-00000.safetensors"} + assert slim["metadata"]["freetoken_host_mapped_only"] is True + assert sorted(copied) == [ + "config.json", + "model-plefp8-00000.safetensors", + "model.safetensors.index.json", + ] + + +def test_ftw_buffered_reader_works_on_the_current_platform(tmp_path: Path) -> None: + output = tmp_path / "ftw" + writer = FTWWriter(str(output), shard_limit=4096) + expected = torch.arange(64, dtype=torch.int32).reshape(8, 8) + writer.add_tensor("weight", expected, kind="weight") + writer.finalize({}) + + loaded = list(iter_ftw_weights(str(output), workers=2)) + + assert len(loaded) == 1 + assert loaded[0][0] == "weight" + assert torch.equal(loaded[0][1], expected) + + +def test_ftw_reader_can_drop_and_reopen_source_maps(tmp_path: Path) -> None: + output = tmp_path / "ftw" + writer = FTWWriter(str(output), shard_limit=4096) + expected = torch.arange(64, dtype=torch.int32).reshape(8, 8) + writer.add_tensor("weight", expected, kind="weight") + writer.finalize({}) + + reader = FTWReader(str(output)) + reader._direct = 0 + reader._probed = True + entry = reader.entries("weight")[0] + destination = bytearray(4096) + reader.read_into(memoryview(destination), entry, workers=2) + assert reader._maps + + reader.drop_maps() + assert not reader._maps + destination[:] = b"\0" * len(destination) + reader.read_into(memoryview(destination), entry, workers=2) + actual = torch.frombuffer(destination, dtype=torch.int32, count=64).reshape(8, 8) + assert torch.equal(actual, expected) + reader.close() diff --git a/tests/engine/test_attention_backend_matrix.py b/tests/engine/test_attention_backend_matrix.py index f27640b9..eb6e01f8 100644 --- a/tests/engine/test_attention_backend_matrix.py +++ b/tests/engine/test_attention_backend_matrix.py @@ -67,6 +67,9 @@ def _model_config(kind): elif kind == "bsa": # MiniMax-M3 shape: one FULL-family group, mla=False + index dims -> BSA. specs = (_spec("full", AttnType.BSA, index_head_dim=128),) + elif kind == "qsa": + mc.has_linear_attention = True + specs = (_spec("qsa", AttnType.QSA, index_head_dim=128),) elif kind == "linear_hybrid": mc.has_linear_attention = True specs = (_spec("full", AttnType.FULL),) @@ -109,6 +112,7 @@ def _patch_env(monkeypatch, *, major=9, flashinfer=True, sgl=True): ("dsa", "dsa"), # MLA + DSA indexer (GLM-5.2 shape) ("dsv4", "dsv4_sparse"), ("bsa", "m3_sparse"), # MiniMax-M3 block-sparse GQA + ("qsa", "qsa"), # Qwen4 compressed sparse GQA ], ) def test_auto_resolves_per_type(monkeypatch, kind, expected): @@ -131,6 +135,15 @@ def test_auto_bsa_sets_block_page_size(monkeypatch): assert config.page_size == 128 +def test_auto_qsa_sets_aligned_page_size(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_env(monkeypatch) + config = _config("qsa", attention_backend="auto") + _adjust_config(config) + assert config.page_size == 64 + + def test_bsa_rejects_float32_dtype(monkeypatch): # --dtype float32 used to pass config validation and die on the pool's # itemsize==2 assert only after the model was resident. @@ -159,10 +172,13 @@ def test_auto_dsv4_sets_window_page_size(monkeypatch): ("full", "dsa"), ("full", "dsv4_sparse"), ("full", "m3_sparse"), + ("full", "qsa"), ("swa", "dsa"), # forward gates: generic backends on the BSA-locked model ("bsa", "fi"), ("bsa", "triton"), + ("qsa", "fi"), + ("qsa", "triton"), # forward gates: generic backends on type-locked models ("mla", "fi"), ("mla", "triton"), @@ -199,6 +215,7 @@ def test_illegal_combinations_rejected_at_config_time(monkeypatch, kind, backend ("swa", "triton"), ("full", "triton"), ("full", "fa,fi"), + ("qsa", "qsa"), ], ) def test_legal_explicit_combinations_pass(monkeypatch, kind, backend): @@ -245,6 +262,22 @@ def _info(name): _adjust_config(config) +def test_model_runtime_capabilities_force_safe_cache_and_graph(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_env(monkeypatch) + config = _config("linear_hybrid", attention_backend="auto") + object.__setattr__(config, "cache_type", "radix") + config.model_config.requires_naive_cache = True + config.model_config.supports_cuda_graph = False + + _adjust_config(config) + + assert config.cache_type == "naive" + assert config.cuda_graph_bs == [] + assert config.cuda_graph_max_bs == 0 + + def test_trtllm_page_size_coercion_is_part_aware(monkeypatch): from freetoken.engine.engine import _adjust_config diff --git a/tests/engine/test_moe_cpu_layers.py b/tests/engine/test_moe_cpu_layers.py index dee36a11..a3b0fc40 100644 --- a/tests/engine/test_moe_cpu_layers.py +++ b/tests/engine/test_moe_cpu_layers.py @@ -10,7 +10,9 @@ import pytest from freetoken.engine.engine import _parse_cpu_layers_spec as parse +from freetoken.engine.engine import _pin_budget_bytes as pin_budget from freetoken.engine.engine import _resolve_cpu_layers as resolve +from freetoken.engine import engine as engine_module L = 40 @@ -61,6 +63,20 @@ def test_resolve_backend_dispatch(): assert resolve(_cfg("fused", "8"), L) == frozenset() +def test_pin_budget_native_windows(monkeypatch): + monkeypatch.delenv("FREETOKEN_PIN_BUDGET_GB", raising=False) + monkeypatch.setattr(engine_module.os, "name", "nt") + monkeypatch.setattr(engine_module, "_windows_total_physical_memory", lambda: 128 << 30) + assert pin_budget() == int((128 << 30) * 0.4) + + +def test_pin_budget_env_overrides_platform(monkeypatch): + monkeypatch.setenv("FREETOKEN_PIN_BUDGET_GB", "47.5") + monkeypatch.setattr(engine_module.os, "name", "nt") + monkeypatch.setattr(engine_module, "_windows_total_physical_memory", lambda: 1) + assert pin_budget() == int(47.5 * 2**30) + + if __name__ == "__main__": import sys diff --git a/tests/kernels/test_index.py b/tests/kernels/test_index.py new file mode 100644 index 00000000..d4d1c51f --- /dev/null +++ b/tests/kernels/test_index.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import pytest +import torch + + +def test_windows_index_fallback_matches_index_select(monkeypatch): + from freetoken.kernel import index + + index._TORCH_FALLBACK_KEYS.clear() + monkeypatch.setattr(index.sys, "platform", "win32") + monkeypatch.setattr( + index, + "_jit_index_module", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("no matching kernel")), + ) + weights = torch.arange(30, dtype=torch.float32).view(10, 3) + indices = torch.tensor([7, 2], dtype=torch.int32) + + with pytest.warns(RuntimeWarning, match="torch.index_select"): + actual = index.indexing(weights, indices) + + torch.testing.assert_close(actual, weights[[7, 2]]) + assert (weights.shape[1] * weights.element_size(), 1) in index._TORCH_FALLBACK_KEYS + + +def test_windows_index_fallback_masks_remote_vocab_rows(monkeypatch): + from freetoken.kernel import index + + index._TORCH_FALLBACK_KEYS.clear() + monkeypatch.setattr(index.sys, "platform", "win32") + monkeypatch.setattr( + index, + "_jit_index_module", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("no matching kernel")), + ) + weights = torch.arange(12, dtype=torch.float32).view(4, 3) + indices = torch.tensor([4, 6, 9], dtype=torch.int64) + + with pytest.warns(RuntimeWarning): + actual = index.indexing(weights, indices, vocab_range=(4, 4)) + + torch.testing.assert_close(actual[0], weights[0]) + torch.testing.assert_close(actual[1], weights[2]) + torch.testing.assert_close(actual[2], torch.zeros(3)) diff --git a/tests/kernels/test_kernel_cache_version.py b/tests/kernels/test_kernel_cache_version.py index 3ff5b021..91f3a78c 100644 --- a/tests/kernels/test_kernel_cache_version.py +++ b/tests/kernels/test_kernel_cache_version.py @@ -4,9 +4,22 @@ version scheme (runtime `0.1.1+g`, cache `0.1.1+cu130.g`) introduced by scripts/build-release-wheels.sh.""" +from pathlib import Path + import pytest -from freetoken.kernel.utils import _kernel_cache_version_ok +from freetoken.kernel.utils import _kernel_cache_version_ok, _prebuilt_library_path + + +@pytest.mark.parametrize( + ("platform", "suffix"), + [("win32", ".dll"), ("linux", ".so")], +) +def test_prebuilt_library_path_uses_platform_suffix(platform: str, suffix: str) -> None: + name = "freetoken__store_1024_128_1_false" + assert _prebuilt_library_path(Path("cache"), name, platform=platform) == ( + Path("cache") / name / f"{name}{suffix}" + ) @pytest.mark.parametrize( diff --git a/tests/kernels/test_qsa.py b/tests/kernels/test_qsa.py new file mode 100644 index 00000000..fc3eb947 --- /dev/null +++ b/tests/kernels/test_qsa.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import pytest +import torch + +from freetoken.attention.qsa import _compact_expanded_selection, select_qsa_logical_rows +from freetoken.kernel.triton.qsa import qsa_sparse_gqa + + +def test_qsa_selection_is_dense_before_budget_and_keeps_tail(): + torch.manual_seed(1) + # Four complete groups plus a two-token tail at query position 17. + q = torch.randn(1, 4, 8) + keys = torch.randn(4, 1, 8) + selected, counts = select_qsa_logical_rows( + q, + keys, + torch.tensor([17]), + compress_ratio=4, + token_budget=2048, + ) + assert counts.tolist() == [18] + assert set(selected[0, :18].tolist()) == set(range(18)) + assert torch.all(selected[0, 18:] == -1) + + +def test_qsa_selection_obeys_query_causality_for_prefill_rows(): + torch.manual_seed(2) + q = torch.randn(4, 4, 8) + keys = torch.randn(2, 1, 8) + positions = torch.tensor([0, 3, 4, 7]) + selected, counts = select_qsa_logical_rows( + q, keys, positions, compress_ratio=4, token_budget=8 + ) + assert counts.tolist() == [1, 4, 5, 8] + for row, position in enumerate(positions.tolist()): + assert set(selected[row, : counts[row]].tolist()) == set(range(position + 1)) + + +def test_qsa_sparse_gqa_matches_explicit_attention_cpu(): + torch.manual_seed(3) + q = torch.randn(3, 4, 8, dtype=torch.bfloat16) + k = torch.randn(12, 2, 8, dtype=torch.bfloat16) + v = torch.randn(12, 2, 8, dtype=torch.bfloat16) + rows = torch.tensor( + [[0, 2, 4, -1], [1, 3, 5, 7], [8, 9, -1, -1]], dtype=torch.int32 + ) + counts = torch.tensor([3, 4, 2], dtype=torch.int32) + actual = qsa_sparse_gqa(q, k, v, rows, counts, 8**-0.5) + + expected = torch.zeros_like(q) + for row in range(3): + chosen = rows[row, : counts[row]].long() + for kv_head in range(2): + heads = slice(kv_head * 2, (kv_head + 1) * 2) + score = torch.einsum( + "hd,td->ht", q[row, heads].float(), k[chosen, kv_head].float() + ) * 8**-0.5 + expected[row, heads] = torch.einsum( + "ht,td->hd", + torch.softmax(score, dim=-1).to(v.dtype), + v[chosen, kv_head], + ).to(expected.dtype) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qsa_sparse_gqa_matches_reference_on_cuda(): + torch.manual_seed(17) + q = torch.randn(3, 24, 256, device="cuda", dtype=torch.bfloat16) + k = torch.randn(4096, 2, 256, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) + rows = torch.randint(0, 4096, (3, 131), device="cuda", dtype=torch.int32) + counts = torch.tensor([131, 97, 41], device="cuda", dtype=torch.int32) + actual = qsa_sparse_gqa(q, k, v, rows, counts, 256**-0.5) + reference = qsa_sparse_gqa( + q.cpu(), k.cpu(), v.cpu(), rows.cpu(), counts.cpu(), 256**-0.5 + ) + torch.testing.assert_close(actual.cpu(), reference, rtol=0.03, atol=0.01) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qsa_cuda_compaction_matches_cpu_with_short_tail(): + blocks = torch.tensor([[-1, -1], [1, 0], [0, -1]], dtype=torch.int32) + positions = torch.tensor([0, 11, 6], dtype=torch.int64) + expected_rows, expected_counts = _compact_expanded_selection( + blocks, + positions, + compress_ratio=4, + token_budget=8, + ) + actual_rows, actual_counts = _compact_expanded_selection( + blocks.cuda(), + positions.cuda(), + compress_ratio=4, + token_budget=8, + ) + assert torch.equal(actual_rows.cpu(), expected_rows) + assert torch.equal(actual_counts.cpu(), expected_counts) diff --git a/tests/kvcache/test_qsa_pool.py b/tests/kvcache/test_qsa_pool.py new file mode 100644 index 00000000..de41afea --- /dev/null +++ b/tests/kvcache/test_qsa_pool.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import torch + +from freetoken.distributed.info import DistributedInfo +from freetoken.kvcache.qsa_pool import QSAKVCache + + +def _pool(monkeypatch, pages=3): + monkeypatch.setattr( + "freetoken.kvcache.mha_pool.get_tp_info", + lambda: DistributedInfo(rank=0, size=1), + ) + return QSAKVCache( + num_kv_heads=2, + num_layers=6, + head_dim=16, + num_pages=pages, + page_size=64, + dtype=torch.bfloat16, + device=torch.device("cpu"), + index_num_kv_heads=1, + index_head_dim=8, + compress_ratio=4, + layer_ids=(1, 5), + ) + + +def test_qsa_pool_geometry_and_cost(monkeypatch): + pool = _pool(monkeypatch) + assert pool.k_cache(1).shape == (3, 64, 2, 16) + assert pool.k_cache(5).shape == (3, 64, 2, 16) + assert pool.compressed_k_cache(1).shape == (3 * 16, 1, 8) + kv, swa = pool.unit_bytes() + assert kv == 2 * 2 * 2 * 16 * 2 + 2 * 1 * 8 * 2 // 4 + assert swa == 0 + + +def test_qsa_pool_rebuild_and_compressed_rows(monkeypatch): + pool = _pool(monkeypatch) + keys = torch.arange(16, dtype=torch.bfloat16).view(2, 1, 8) + pool.store_compressed_k(keys, torch.tensor([0, 17]), layer_id=5) + assert torch.equal(pool.compressed_k_cache(5)[17], keys[1]) + pool.rebuild(5) + assert pool.k_cache(1).shape[0] == 5 + assert pool.compressed_k_cache(5).shape == (5 * 16, 1, 8) + + +def test_qsa_pending_ring_validates_logical_positions(monkeypatch): + pool = _pool(monkeypatch) + pool.ensure_pending_capacity(4) + positions = torch.tensor([5, 6, 7]) + keys = torch.randn(3, 1, 8, dtype=torch.bfloat16) + rope = torch.tensor([[5, 5, 5], [6, 7, 8], [7, 9, 11]]) + pool.store_pending(5, 2, positions, keys, rope) + assert torch.equal(pool.pending_group(5, 2, positions), keys) + assert torch.equal(pool.pending_rope_group(5, 2, positions), rope) diff --git a/tests/models/test_models_loader.py b/tests/models/test_models_loader.py index 7254739c..959b4206 100644 --- a/tests/models/test_models_loader.py +++ b/tests/models/test_models_loader.py @@ -6,6 +6,17 @@ import torch +def test_drop_page_cache_is_a_noop_without_posix_fadvise(tmp_path, monkeypatch): + from freetoken.models import loader + + path = tmp_path / "weights.safetensors" + path.write_bytes(b"weights") + monkeypatch.delattr(loader.os, "posix_fadvise", raising=False) + monkeypatch.delattr(loader.os, "POSIX_FADV_DONTNEED", raising=False) + + loader.drop_page_cache(str(path)) + + def test_shard_tensor_splits_vocab_with_ceil_partition(): from freetoken.models.loader import shard_tensor diff --git a/tests/models/test_qwen4_exp.py b/tests/models/test_qwen4_exp.py new file mode 100644 index 00000000..7d9912e4 --- /dev/null +++ b/tests/models/test_qwen4_exp.py @@ -0,0 +1,232 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +import freetoken.models.qwen4_exp as qwen4_exp +from freetoken.models.qwen4_exp.config import parse_config +from freetoken.models.qwen4_exp.model import _ple_request_tokens, build_ngram_ids +from freetoken.models.qwen4_exp.weight import _rename, _try_fuse +from freetoken.models.register import get_model_spec + + +def _config(quantization_config=None): + text = SimpleNamespace( + layer_types=[ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + ], + head_dim=256, + rope_parameters={ + "partial_rotary_factor": 0.25, + "rope_theta": 10_000_000, + "mrope_interleaved": True, + "mrope_section": [11, 11, 10], + }, + indexer_budget=2048, + indexer_n_heads=4, + indexer_kv_heads=1, + indexer_head_dim=128, + max_position_embeddings=262_144, + num_key_value_heads=2, + linear_num_key_heads=16, + linear_num_value_heads=48, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_conv_kernel_dim=4, + eos_token_id=248044, + hc_count=4, + hc_lowrank=320, + ple_layer_ids=[2], + ple_embed_dim=2560, + ple_conv_kernel_size=4, + ngram_size=3, + heads_per_ngram=8, + ngram_vocab_size_base=20_000_000, + split_ngram_parts=128, + indexer_compress_ratio=4, + output_gate_type="sigmoid", + hidden_act="silu", + num_hidden_layers=4, + num_attention_heads=24, + hidden_size=2560, + vocab_size=248320, + rms_norm_eps=1e-6, + num_experts=512, + num_experts_per_tok=10, + moe_intermediate_size=640, + shared_expert_intermediate_size=640, + norm_topk_prob=None, + tie_word_embeddings=False, + ) + return SimpleNamespace( + text_config=text, + quantization_config=( + quantization_config + if quantization_config is not None + else {"quant_method": "fp8", "weight_block_size": [128, 128]} + ), + model_type="qwen4_exp", + architectures=["Qwen4ExpForConditionalGeneration"], + image_token_id=248056, + ) + + +def test_qwen4_config_uses_exact_qsa_prefix(): + config = parse_config(_config()) + assert config.rotary_config.max_position == 262_144 + assert config.expert_quant == "fp8_block" + assert config.attn_quant == "none" + assert config.qwen4_args.ple_layer_ids == (1,) + assert config.qwen4_args.output_gate_type == "sigmoid" + assert config.requires_naive_cache + assert not config.supports_cuda_graph + assert config.is_linear_layer(0) + assert not config.is_linear_layer(3) + + +def test_qwen4_config_accepts_transformers_sparse_attention_alias(): + hf_config = _config() + hf_config.text_config.layer_types[-1] = "qwen_sparse_attention" + config = parse_config(hf_config) + assert not config.is_linear_layer(3) + assert config.attn_type_for_layer(3).value == "qsa" + spec = config.kv_cache_group_specs()[0] + assert spec.layer_ids == (3,) + assert spec.index_head_dim == 128 + assert spec.index_compress_ratio == 4 + assert spec.index_token_budget == 2048 + + +def test_qwen4_config_accepts_missing_norm_topk_prob(): + hf_config = _config() + del hf_config.text_config.norm_topk_prob + config = parse_config(hf_config) + assert not config.norm_topk_prob + + +def test_qwen4_config_accepts_routed_expert_nvfp4(): + config = parse_config( + _config( + { + "quant_method": "modelopt", + "quant_algo": "NVFP4", + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": {"num_bits": 4, "group_size": 16, "type": "float"}, + } + }, + "ignore": [ + "*.self_attn.*", + "*.linear_attn.*", + "*.mlp.shared_expert.*", + "*.ple.*", + "model.visual.*", + "lm_head", + ], + } + ) + ) + assert config.expert_quant == "nvfp4" + assert config.weight_block_size is None + assert config.attn_quant == "none" + assert config.dense_quant == "none" + assert config.lm_head_quant == "none" + + +def test_qwen4_config_rejects_unsupported_expert_quant(): + with pytest.raises(ValueError, match="requires routed experts"): + parse_config(_config({"quant_method": "gptq"})) + + +def test_qwen4_exports_nvfp4_loader_hooks(): + assert callable(qwen4_exp.load_nvfp4_expert_sources) + assert callable(qwen4_exp.load_nvfp4_expert_sources_parallel) + + +def test_qwen4_registry_entry(): + spec = get_model_spec("Qwen4ExpForConditionalGeneration") + assert spec.module == "freetoken.models.qwen4_exp" + assert spec.model_cls == "Qwen4ExpForCausalLM" + + +def test_qwen4_weight_names(): + assert _rename("model.language_model.layers.1.ple.key_proj.weight") == ( + "model.layers.1.ple.key_proj.weight" + ) + assert _rename("model.visual.blocks.0.attn.qkv.weight") == ( + "visual.blocks.0.attn.qkv.weight" + ) + assert _rename("model.language_model.layers.3.self_attn.indexer.q_layernorm.weight") == ( + "model.layers.3.self_attn.indexer.q_layernorm.weight" + ) + + +def test_qwen4_projection_fusion_order(): + buffers = {} + base = "model.layers.3.self_attn." + parts = [ + ("q_proj.weight", torch.full((2, 3), 1.0)), + ("k_proj.weight", torch.full((1, 3), 2.0)), + ("v_proj.weight", torch.full((1, 3), 3.0)), + ] + assert _try_fuse(base + parts[0][0], parts[0][1], buffers) == () + assert _try_fuse(base + parts[1][0], parts[1][1], buffers) == () + name, fused = _try_fuse(base + parts[2][0], parts[2][1], buffers) + assert name == base + "qkv_proj.weight" + assert fused[:, 0].tolist() == [1.0, 1.0, 2.0, 3.0] + + +def test_ngram_hash_resets_at_eos(): + tokens = torch.tensor([4, 5, 99, 6, 7]) + multipliers = torch.tensor([3, 5, 7]) + sizes = torch.tensor([101, 103]) + offsets = torch.tensor([0, 101]) + ids = build_ngram_ids( + tokens, + ngram_size=3, + heads_per_ngram=1, + eos_token_id=99, + multipliers=multipliers, + vocab_sizes=sizes, + offsets=offsets, + ) + assert ids.shape == (5, 2) + expected_bigram_after_eos = (6 * 3) ^ (99 * 5) + assert ids[3, 0].item() == expected_bigram_after_eos % 101 + + +def test_ple_request_tokens_uses_complete_prefill_history(): + req = SimpleNamespace( + input_ids=torch.tensor([11, 12, 13]), + cached_len=0, + device_len=3, + extend_len=3, + ) + assert _ple_request_tokens(req).tolist() == [11, 12, 13] + + +def test_ple_request_tokens_joins_overlap_decode_token(): + req = SimpleNamespace( + input_ids=torch.tensor([11, 12]), + cached_len=2, + device_len=3, + extend_len=1, + ) + assert _ple_request_tokens(req, torch.tensor([13])).tolist() == [11, 12, 13] + + +def test_ple_request_tokens_rejects_noncontiguous_host_history(): + req = SimpleNamespace( + input_ids=torch.tensor([11]), + cached_len=2, + device_len=3, + extend_len=1, + ) + with pytest.raises(RuntimeError, match="unexpected gap"): + _ple_request_tokens(req, torch.tensor([13])) diff --git a/tests/models/test_qwen4_mrope.py b/tests/models/test_qwen4_mrope.py new file mode 100644 index 00000000..305b69e0 --- /dev/null +++ b/tests/models/test_qwen4_mrope.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.models.config import RotaryConfig +from freetoken.models.qwen4_exp.mrope import build_mrope_positions +from freetoken.models.qwen4_exp.model import _Qwen4MRoPE + + +def test_mrope_position_builder_matches_transformers_reference(): + from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLModel + + class Reference: + config = SimpleNamespace(vision_config=SimpleNamespace(spatial_merge_size=2)) + get_vision_position_ids = Qwen3VLModel.get_vision_position_ids + get_rope_index = Qwen3VLModel.get_rope_index + + input_ids = torch.arange(9) + token_types = torch.tensor([0, 0, 1, 1, 1, 1, 0, 0, 0]) + grid = torch.tensor([[1, 4, 4]]) + actual, delta = build_mrope_positions(input_ids, token_types, grid, 2) + expected, expected_delta = Reference().get_rope_index( + input_ids.view(1, -1), + token_types.view(1, -1), + image_grid_thw=grid, + ) + assert torch.equal(actual, expected[:, 0]) + assert delta == int(expected_delta[0, 0]) + + +def _rotary() -> _Qwen4MRoPE: + config = SimpleNamespace( + rotary_config=RotaryConfig( + head_dim=256, + rotary_dim=64, + max_position=128, + base=10_000_000, + scaling=None, + ), + qwen4_args=SimpleNamespace(mrope_section=(11, 11, 10)), + ) + return _Qwen4MRoPE(config) + + +def _reference_rotate( + tensor: torch.Tensor, positions: torch.Tensor, cache: torch.Tensor, head_size: int +) -> torch.Tensor: + output = tensor.clone() + view = output.view(output.shape[0], -1, head_size) + half = cache.shape[1] // 2 + pair = torch.arange(half, device=positions.device) + axis = torch.zeros(half, dtype=torch.long, device=positions.device) + axis[(pair % 3 == 1) & (pair < 33)] = 1 + axis[(pair % 3 == 2) & (pair < 30)] = 2 + selected = positions.transpose(0, 1)[:, axis] + dim = pair.view(1, -1).expand_as(selected) + cos = cache[:, :half][selected, dim] + sin = cache[:, half:][selected, dim] + first = view[..., :half].float().clone() + second = view[..., half : 2 * half].float().clone() + view[..., :half] = (first * cos[:, None] - second * sin[:, None]).to(view.dtype) + view[..., half : 2 * half] = (second * cos[:, None] + first * sin[:, None]).to(view.dtype) + return output + + +def test_mrope_cpu_rotation_matches_axis_reference(): + torch.manual_seed(17) + rotary = _rotary() + positions = torch.tensor( + [[2, 3, 4, 9], [2, 3, 5, 9], [2, 4, 6, 9]], dtype=torch.int64 + ) + query = torch.randn(4, 2 * 256) + key = torch.randn(4, 256) + expected_q = _reference_rotate(query, positions, rotary._cos_sin_cache, 256) + expected_k = _reference_rotate(key, positions, rotary._cos_sin_cache, 256) + rotary.forward(positions, query, key) + torch.testing.assert_close(query, expected_q) + torch.testing.assert_close(key, expected_k) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_mrope_rtx_kernel_matches_cpu_reference(): + torch.manual_seed(23) + rotary = _rotary() + positions = torch.randint(0, 128, (3, 37), dtype=torch.int64) + query = torch.randn(37, 24 * 256, dtype=torch.bfloat16) + key = torch.randn(37, 2 * 256, dtype=torch.bfloat16) + expected_q = _reference_rotate(query, positions, rotary._cos_sin_cache, 256) + expected_k = _reference_rotate(key, positions, rotary._cos_sin_cache, 256) + rotary._cos_sin_cache = rotary._cos_sin_cache.cuda() + query_gpu = query.cuda() + key_gpu = key.cuda() + rotary.forward(positions.cuda(), query_gpu, key_gpu) + torch.testing.assert_close(query_gpu.cpu(), expected_q, rtol=0, atol=2e-3) + torch.testing.assert_close(key_gpu.cpu(), expected_k, rtol=0, atol=2e-3) diff --git a/tests/models/test_qwen4_vision.py b/tests/models/test_qwen4_vision.py new file mode 100644 index 00000000..0caf9e75 --- /dev/null +++ b/tests/models/test_qwen4_vision.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +import torch + +from freetoken.models.qwen4_exp.args import Qwen4VisionConfig +from freetoken.models.qwen4_exp.vision import Qwen4VisionModel + + +def _config() -> Qwen4VisionConfig: + return Qwen4VisionConfig( + depth=2, + hidden_size=32, + intermediate_size=64, + num_heads=4, + num_position_embeddings=16, + out_hidden_size=24, + patch_size=2, + spatial_merge_size=2, + temporal_patch_size=2, + in_channels=3, + hidden_act="gelu_pytorch_tanh", + deepstack_visual_indexes=(), + ) + + +def test_qwen4_vision_matches_transformers_qwen3_vl_reference(): + from transformers.models.qwen3_vl.configuration_qwen3_vl import ( + Qwen3VLVisionConfig, + ) + from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionModel + + torch.manual_seed(9) + ours_config = _config() + reference_config = Qwen3VLVisionConfig( + depth=ours_config.depth, + hidden_size=ours_config.hidden_size, + intermediate_size=ours_config.intermediate_size, + num_heads=ours_config.num_heads, + num_position_embeddings=ours_config.num_position_embeddings, + out_hidden_size=ours_config.out_hidden_size, + patch_size=ours_config.patch_size, + spatial_merge_size=ours_config.spatial_merge_size, + temporal_patch_size=ours_config.temporal_patch_size, + in_channels=ours_config.in_channels, + hidden_act=ours_config.hidden_act, + deepstack_visual_indexes=[], + _attn_implementation="sdpa", + ) + reference = Qwen3VLVisionModel(reference_config).eval().cpu() + with torch.device("cpu"): + ours = Qwen4VisionModel(ours_config) + ours.load_state_dict(dict(reference.state_dict())) + + grid = torch.tensor([[1, 4, 4]], dtype=torch.long, device="cpu") + pixels = torch.randn( + 16, + ours_config.in_channels + * ours_config.temporal_patch_size + * ours_config.patch_size + * ours_config.patch_size, + device="cpu", + ) + with torch.inference_mode(): + expected = reference(pixels, grid).pooler_output + actual = ours.forward(pixels, grid) + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + +def test_qwen4_vision_derived_rope_survives_meta_construction(): + torch.manual_seed(17) + config = _config() + with torch.device("cpu"): + reference = Qwen4VisionModel(config) + with torch.no_grad(): + for tensor in reference.state_dict().values(): + if tensor.is_floating_point(): + tensor.uniform_(-0.02, 0.02) + else: + tensor.zero_() + with torch.device("meta"): + model = Qwen4VisionModel(config) + with torch.device("cpu"): + model.load_state_dict(dict(reference.state_dict())) + + assert not hasattr(model, "_inv_freq") + assert model._inv_dim == config.hidden_size // config.num_heads // 2 + + grid = torch.tensor([[1, 4, 4]], dtype=torch.long, device="cpu") + pixels = torch.randn( + 16, + config.in_channels + * config.temporal_patch_size + * config.patch_size + * config.patch_size, + device="cpu", + ) + with torch.inference_mode(): + expected = reference.forward(pixels, grid) + actual = model.forward(pixels, grid) + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) diff --git a/tests/moe/test_fused_moe.py b/tests/moe/test_fused_moe.py index 1fd0f2e5..9185f1ca 100644 --- a/tests/moe/test_fused_moe.py +++ b/tests/moe/test_fused_moe.py @@ -58,6 +58,34 @@ def test_fused_topk_accepts_triton_kernel_tuple_output(): torch.testing.assert_close(weights, ref_weights, rtol=2e-4, atol=2e-4) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("renormalize", [False, True]) +def test_fused_topk_uses_exact_torch_fallback_for_non_power_of_two_topk( + monkeypatch, renormalize +): + from freetoken.kernel import backend + from freetoken.moe.fused import fused_topk + + monkeypatch.setattr(backend, "is_triton_kernels_installed", lambda: True) + torch.manual_seed(38) + logits = torch.randn((3, 512), device="cuda", dtype=torch.bfloat16) + hidden_states = torch.zeros((3, 64), device="cuda", dtype=torch.bfloat16) + + weights, ids = fused_topk( + hidden_states, + logits, + topk=10, + renormalize=renormalize, + ) + + probabilities = torch.softmax(logits.float(), dim=-1) + expected_weights, expected_ids = torch.topk(probabilities, 10, dim=-1) + if renormalize: + expected_weights = expected_weights / expected_weights.sum(dim=-1, keepdim=True) + torch.testing.assert_close(ids, expected_ids.to(torch.int32), rtol=0, atol=0) + torch.testing.assert_close(weights, expected_weights, rtol=0, atol=0) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16, 24]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 422ca867..8b62a907 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -868,3 +868,25 @@ def boom(addr, nbytes): with hb.PinPipeline() as pins: pins(1, {"gate_up": hb.HostBank((4,), torch.uint8)}) assert plan2.actual == {1: hb.HostResidency.PAGEABLE.value} + + +def test_os_lock_dispatches_to_virtual_lock_on_windows(monkeypatch): + import freetoken.moe.host_banks as hb + + calls = [] + monkeypatch.setattr(hb.os, "name", "nt") + monkeypatch.setattr(hb, "_windows_lock", lambda addr, nbytes: calls.append((addr, nbytes))) + hb._os_lock(0x1000, 8192) + assert calls == [(0x1000, 8192)] + + +def test_host_bank_release_is_portable() -> None: + import freetoken.moe.host_banks as hb + + bank = hb.HostBank((4096,), torch.uint8) + alias = bank.tensor + bank.tensor.fill_(1) + bank.release() + if hb.os.name == "nt": + assert alias.numel() == 0 + assert bank._buf.closed diff --git a/tests/scheduler/test_mrope_positions.py b/tests/scheduler/test_mrope_positions.py new file mode 100644 index 00000000..743e6d78 --- /dev/null +++ b/tests/scheduler/test_mrope_positions.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from freetoken.scheduler.scheduler import _make_rope_positions + + +def _req(start, end, positions=None, delta=0): + return SimpleNamespace( + cached_len=start, + device_len=end, + extend_len=end - start, + mrope_position_ids=positions, + mrope_position_delta=delta, + ) + + +def test_scheduler_packs_prompt_mrope_and_plain_text_positions(): + prompt = torch.tensor( + [[0, 1, 2, 2, 2, 3], [0, 1, 2, 2, 3, 3], [0, 1, 2, 3, 2, 3]] + ) + batch = SimpleNamespace( + padded_reqs=[_req(2, 6, prompt, -2), _req(0, 2)] + ) + actual = _make_rope_positions(batch, torch.device("cpu")) + expected = torch.cat((prompt[:, 2:6], torch.tensor([[0, 1], [0, 1], [0, 1]])), dim=1) + assert torch.equal(actual, expected) + + +def test_scheduler_uses_mrope_delta_for_generated_tokens(): + prompt = torch.arange(18, dtype=torch.int64).view(3, 6) + batch = SimpleNamespace(padded_reqs=[_req(6, 8, prompt, -2)]) + actual = _make_rope_positions(batch, torch.device("cpu")) + assert torch.equal(actual, torch.tensor([[4, 5], [4, 5], [4, 5]])) + + +def test_scheduler_skips_rope_tensor_for_text_only_batch(): + batch = SimpleNamespace(padded_reqs=[_req(3, 5)]) + assert _make_rope_positions(batch, torch.device("cpu")) is None diff --git a/tests/server/test_message_wire.py b/tests/server/test_message_wire.py index 3bd7cc6b..ce299cda 100644 --- a/tests/server/test_message_wire.py +++ b/tests/server/test_message_wire.py @@ -7,6 +7,8 @@ from __future__ import annotations +import torch + from freetoken.message import ( BaseBackendMsg, DetokenizeMsg, @@ -18,6 +20,7 @@ CacheRebuildResultMsg, PromptAdmittedMsg, TokenizeMsg, + UserMsg, UserReply, ) from freetoken.core import SamplingParams @@ -122,3 +125,24 @@ def test_client_dicts_with_the_wire_tag_key_survive_intact(): assert isinstance(out, TokenizeMsg) assert out.chat_template_kwargs == payload assert out.tools[0]["function"]["parameters"] == payload + + +def test_multidimensional_bfloat16_tensor_survives_backend_wire() -> None: + pixels = torch.arange(24, dtype=torch.float32).reshape(3, 8).to(torch.bfloat16) + msg = UserMsg( + uid=9, + input_ids=torch.tensor([1, 2, 3], dtype=torch.int32), + sampling_params=SamplingParams(max_tokens=4), + mm_pixel_values=pixels, + mm_image_grid_thw=torch.tensor([[1, 2, 4]], dtype=torch.int64), + mm_token_type_ids=torch.tensor([0, 1, 1], dtype=torch.int32), + ) + + out = BaseBackendMsg.decoder(msg.encoder()) + + assert isinstance(out, UserMsg) + assert out.mm_pixel_values.dtype == torch.bfloat16 + assert out.mm_pixel_values.shape == (3, 8) + assert torch.equal(out.mm_pixel_values, pixels) + assert torch.equal(out.mm_image_grid_thw, msg.mm_image_grid_thw) + assert torch.equal(out.mm_token_type_ids, msg.mm_token_type_ids) diff --git a/tests/server/test_multimodal_tokenizer.py b/tests/server/test_multimodal_tokenizer.py new file mode 100644 index 00000000..2d3aeee2 --- /dev/null +++ b/tests/server/test_multimodal_tokenizer.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +import base64 + +import pytest + +from freetoken.tokenizer.server import _download_image, _message_image_sources + + +def test_openai_image_data_url_is_extracted_and_decoded() -> None: + raw = b"small-image-payload" + source = "data:image/png;base64," + base64.b64encode(raw).decode("ascii") + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": source}}, + {"type": "text", "text": "describe"}, + ], + } + ] + + assert _message_image_sources(messages) == [source] + assert _download_image(source) == raw + + +def test_local_file_image_url_is_rejected() -> None: + with pytest.raises(ValueError, match="data, http, or https"): + _download_image("file:///C:/secret.png") diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index facd469b..a36f981e 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -145,6 +145,29 @@ def test_chat_request_accepts_tool_messages_and_assistant_tool_calls(): assert req.messages[0].tool_calls[0].function.arguments == '{"city":"Paris"}' +def test_chat_request_preserves_image_parts_for_multimodal_worker(): + req = ChatCompletionRequest( + model="client-model", + messages=[ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}}, + {"type": "text", "text": "Read the title"}, + ], + } + ], + max_tokens=8, + ) + + content = chat_request_to_genspec(req, {}).messages[0]["content"] + + assert isinstance(content, list) + assert content[0]["type"] == "image_url" + assert content[0]["image_url"]["url"].startswith("data:image/png;base64,") + assert content[1] == {"type": "text", "text": "Read the title"} + + def test_chat_request_reasoning_replay_field_aliases(): # Any replay field name in -> both template-read field names out. for field in ("reasoning_content", "reasoning", "thinking"): diff --git a/tests/server/test_windows_zmq_addresses.py b/tests/server/test_windows_zmq_addresses.py new file mode 100644 index 00000000..93fd0af7 --- /dev/null +++ b/tests/server/test_windows_zmq_addresses.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import os + +import torch + +from freetoken.distributed import DistributedInfo +from freetoken.server import api_server +from freetoken.server.args import ServerArgs + + +def test_server_internal_zmq_addresses_are_supported_on_current_platform() -> None: + args = ServerArgs( + model_path="test-model", + tp_info=DistributedInfo(rank=0, size=1), + dtype=torch.bfloat16, + server_port=1927, + num_tokenizer=1, + ) + addresses = [ + args.zmq_backend_addr, + args.zmq_detokenizer_addr, + args.zmq_scheduler_broadcast_addr, + args.zmq_frontend_addr, + args.zmq_tokenizer_addr, + ] + + assert len(set(addresses)) == len(addresses) + if os.name == "nt": + assert addresses == [f"tcp://127.0.0.1:{port}" for port in range(1929, 1934)] + else: + assert all(address.startswith("ipc:///tmp/freetoken_") for address in addresses) + + +def test_uvicorn_uses_a_zmq_compatible_loop_on_windows() -> None: + if os.name == "nt": + assert api_server._uvicorn_loop().endswith(":windows_selector_loop_factory") + loop = api_server.windows_selector_loop_factory() + try: + assert isinstance(loop, __import__("asyncio").SelectorEventLoop) + assert hasattr(loop, "add_reader") + finally: + loop.close() + else: + assert api_server._uvicorn_loop() == "auto"