Initial ZONOS2 ZeroGPU Space

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+zonos2/vendor/nemo_text_processing/text_normalization/en/data/number/cardinal_number_name.far filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.py[cod]
+.gradio/
+.venv/
+venv/
+*.wav
+.DS_Store

README.md

@@ -1,13 +1,18 @@
 ---
-title: Zonos2
-emoji: 📚
-colorFrom: purple
-colorTo: gray
+title: ZONOS2
+emoji: 🔊
+colorFrom: blue
+colorTo: indigo
 sdk: gradio
-sdk_version: 6.18.0
-python_version: '3.13'
+sdk_version: 6.10.0
 app_file: app.py
-pinned: false
+short_description: ZeroGPU text-to-speech with ZONOS2
+python_version: "3.10"
+startup_duration_timeout: 1h
+preload_from_hub:
+  - Zyphra/ZONOS2
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# ZONOS2 on ZeroGPU
+
+Gradio Space for [Zyphra/ZONOS2](https://huggingface.co/Zyphra/ZONOS2), using the upstream [ZONOS2](https://github.com/Zyphra/ZONOS2) inference code with small ZeroGPU compatibility patches.

app.py

@@ -0,0 +1,317 @@
+import os
+import tempfile
+import threading
+import time
+import wave
+from pathlib import Path
+
+os.environ.setdefault("HF_HOME", "/data/.cache/huggingface")
+os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
+os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
+os.environ.setdefault("ZONOS2_TTS_NORM_CACHE_DIR", "/tmp/zonos2-tts-norm")
+os.environ.setdefault("GRADIO_SSR_MODE", "false")
+os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+for cache_dir in (
+    os.environ["HF_HOME"],
+    os.environ["HF_MODULES_CACHE"],
+    os.environ["MPLCONFIGDIR"],
+    os.environ["ZONOS2_TTS_NORM_CACHE_DIR"],
+):
+    Path(cache_dir).mkdir(parents=True, exist_ok=True)
+
+import spaces
+import gradio as gr
+import numpy as np
+import torch
+
+from zonos2.message import TTSSamplingParams
+from zonos2.tokenizer.textnorm import SERVER_TO_NEMO_LANG, TTSTextNormalizer
+from zonos2.tts import TTSLLM
+
+MODEL_ID = "Zyphra/ZONOS2"
+SAMPLE_RATE = 44100
+LANGUAGES = [
+    ("English (US)", "en_us"),
+    ("English (UK)", "en_gb"),
+    ("French", "fr_fr"),
+    ("German", "de"),
+    ("Spanish", "es"),
+    ("Italian", "it"),
+    ("Portuguese (Brazil)", "pt_br"),
+    ("Japanese", "ja"),
+    ("Mandarin Chinese", "cmn"),
+    ("Korean", "ko"),
+]
+SPEAKING_RATE_BUCKETS = [
+    ("Default", "default"),
+    ("Very slow", "0"),
+    ("Slow", "1"),
+    ("Relaxed", "2"),
+    ("Natural", "3"),
+    ("Bright", "4"),
+    ("Fast", "5"),
+    ("Very fast", "6"),
+    ("Extreme", "7"),
+]
+
+torch.backends.cuda.matmul.allow_tf32 = True
+
+
+def _load_model() -> TTSLLM:
+    print(f"Loading {MODEL_ID} for ZeroGPU inference...", flush=True)
+    started = time.perf_counter()
+    model = TTSLLM(
+        model_path=MODEL_ID,
+        decode_audio=True,
+        cuda_graph_max_bs=0,
+        max_running_req=4,
+        max_extend_tokens=4096,
+        memory_ratio=0.75,
+        use_pynccl=False,
+    )
+    elapsed = time.perf_counter() - started
+    print(f"Loaded {MODEL_ID} in {elapsed:.1f}s", flush=True)
+    return model
+
+
+TTS = _load_model()
+TEXT_NORMALIZER = TTSTextNormalizer()
+TTS_LOCK = threading.Lock()
+
+
+def _estimate_duration(*args, **kwargs) -> int:
+    max_tokens = kwargs.get("max_tokens")
+    if max_tokens is None and len(args) > 4:
+        max_tokens = args[4]
+    try:
+        max_tokens = int(max_tokens)
+    except (TypeError, ValueError):
+        max_tokens = 768
+    return min(180, max(60, 45 + max_tokens // 12))
+
+
+def _pcm_float32_to_wav(audio_bytes: bytes, sample_rate: int = SAMPLE_RATE) -> str:
+    audio = np.frombuffer(audio_bytes, dtype=np.float32)
+    if audio.size == 0:
+        raise gr.Error("The model returned no audio. Try increasing max tokens.")
+    audio = np.nan_to_num(audio, nan=0.0, posinf=0.0, neginf=0.0)
+    audio = np.clip(audio, -1.0, 1.0)
+    audio_i16 = (audio * 32767.0).astype(np.int16)
+
+    handle = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+    handle.close()
+    with wave.open(handle.name, "wb") as wav:
+        wav.setnchannels(1)
+        wav.setsampwidth(2)
+        wav.setframerate(sample_rate)
+        wav.writeframes(audio_i16.tobytes())
+    return handle.name
+
+
+def _normalize_text(text: str, language: str, enabled: bool) -> str:
+    if not enabled:
+        return text
+    if language not in SERVER_TO_NEMO_LANG:
+        return text
+    return TEXT_NORMALIZER.normalize(text, language)
+
+
+def _speaking_rate_bucket(value: str) -> int | None:
+    if value == "default":
+        return None
+    return int(value)
+
+
+@spaces.GPU(duration=_estimate_duration)
+def synthesize(
+    text: str,
+    language: str,
+    text_normalization: bool,
+    speaking_rate: str,
+    max_tokens: int,
+    temperature: float,
+    topk: int,
+    top_p: float,
+    min_p: float,
+    repetition_penalty: float,
+    seed: int,
+):
+    text = (text or "").strip()
+    if not text:
+        raise gr.Error("Enter text to synthesize.")
+    if len(text) > 1200:
+        raise gr.Error("Keep the prompt under 1200 characters for this Space.")
+
+    normalized = _normalize_text(text, language, text_normalization)
+    params = TTSSamplingParams(
+        temperature=float(temperature),
+        topk=int(topk),
+        top_p=float(top_p),
+        min_p=float(min_p),
+        max_tokens=int(max_tokens),
+        repetition_window=50,
+        repetition_penalty=float(repetition_penalty),
+        repetition_codebooks=8,
+        seed=None if seed is None or int(seed) < 0 else int(seed),
+    )
+
+    started = time.perf_counter()
+    with TTS_LOCK:
+        torch.cuda.set_stream(TTS.stream)
+        result = TTS.generate_one(
+            normalized,
+            params,
+            decode_audio=True,
+            speaking_rate_bucket=_speaking_rate_bucket(speaking_rate),
+            quality_buckets=None,
+        )
+    elapsed = time.perf_counter() - started
+
+    wav_path = _pcm_float32_to_wav(result["audio"], result.get("sample_rate", SAMPLE_RATE))
+    frames = len(result.get("audio_tokens") or [])
+    eos_frame = result.get("eos_frame")
+    status = f"Generated {frames} frames in {elapsed:.1f}s"
+    if eos_frame is not None:
+        status += f" (EOS frame {eos_frame})"
+    if normalized != text:
+        status += f"\n\nNormalized text: {normalized}"
+    return wav_path, status
+
+
+CSS = """
+main, .gradio-container, .gradio-container > .fillable {
+    max-width: 1180px !important;
+    margin-inline: auto !important;
+}
+.compact-status textarea {
+    font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
+}
+"""
+
+
+with gr.Blocks(title="ZONOS2", css=CSS) as demo:
+    gr.Markdown("# ZONOS2")
+    with gr.Row():
+        with gr.Column(scale=5):
+            text = gr.Textbox(
+                label="Text",
+                value="In the quiet hum of the studio, ZONOS2 turns written words into natural speech.",
+                lines=6,
+                max_length=1200,
+            )
+            with gr.Row():
+                language = gr.Dropdown(
+                    choices=LANGUAGES,
+                    value="en_us",
+                    label="Language",
+                )
+                speaking_rate = gr.Dropdown(
+                    choices=SPEAKING_RATE_BUCKETS,
+                    value="default",
+                    label="Speaking rate",
+                )
+            text_normalization = gr.Checkbox(value=True, label="Text normalization")
+            generate = gr.Button("Generate", variant="primary")
+        with gr.Column(scale=4):
+            audio = gr.Audio(label="Audio", type="filepath", format="wav")
+            status = gr.Textbox(
+                label="Status",
+                lines=5,
+                interactive=False,
+                elem_classes=["compact-status"],
+            )
+
+    with gr.Accordion("Sampling", open=False):
+        with gr.Row():
+            max_tokens = gr.Slider(
+                minimum=128,
+                maximum=2048,
+                step=64,
+                value=768,
+                label="Max audio tokens",
+            )
+            seed = gr.Number(value=-1, precision=0, label="Seed (-1 random)")
+        with gr.Row():
+            temperature = gr.Slider(0.1, 2.0, value=1.15, step=0.05, label="Temperature")
+            topk = gr.Slider(1, 512, value=106, step=1, label="Top-k")
+        with gr.Row():
+            top_p = gr.Slider(0.0, 1.0, value=0.0, step=0.01, label="Top-p")
+            min_p = gr.Slider(0.0, 0.5, value=0.18, step=0.01, label="Min-p")
+            repetition_penalty = gr.Slider(
+                1.0,
+                2.0,
+                value=1.2,
+                step=0.05,
+                label="Repetition penalty",
+            )
+
+    gr.Examples(
+        examples=[
+            [
+                "The first explorers landed just after sunrise, carrying maps, coffee, and impossible optimism.",
+                "en_us",
+                True,
+                "default",
+                512,
+                1.15,
+                106,
+                0.0,
+                0.18,
+                1.2,
+                -1,
+            ],
+            [
+                "Le modèle parle avec une voix claire, expressive et naturellement rythmée.",
+                "fr_fr",
+                True,
+                "default",
+                512,
+                1.15,
+                106,
+                0.0,
+                0.18,
+                1.2,
+                -1,
+            ],
+        ],
+        inputs=[
+            text,
+            language,
+            text_normalization,
+            speaking_rate,
+            max_tokens,
+            temperature,
+            topk,
+            top_p,
+            min_p,
+            repetition_penalty,
+            seed,
+        ],
+    )
+
+    generate.click(
+        fn=synthesize,
+        inputs=[
+            text,
+            language,
+            text_normalization,
+            speaking_rate,
+            max_tokens,
+            temperature,
+            topk,
+            top_p,
+            min_p,
+            repetition_penalty,
+            seed,
+        ],
+        outputs=[audio, status],
+        api_name="generate",
+        concurrency_limit=1,
+    )
+
+
+if __name__ == "__main__":
+    demo.queue(max_size=8, default_concurrency_limit=1).launch()

requirements.txt

@@ -0,0 +1,19 @@
+gradio>=6.10.0
+spaces>=0.41.0
+numpy
+tqdm
+msgpack
+sgl_kernel>=0.3.17.post1
+descript-audio-codec==1.0.0
+transformers>=4.56.0,<=4.57.3
+huggingface_hub
+pyzmq
+apache-tvm-ffi>=0.1.4
+nvidia-cutlass-dsl==4.3.1
+flashinfer-cubin==0.5.3
+flashinfer-python>=0.5.3
+flashinfer-jit-cache @ https://github.com/flashinfer-ai/flashinfer/releases/download/v0.5.3/flashinfer_jit_cache-0.5.3+cu128-cp39-abi3-manylinux_2_28_x86_64.whl
+ninja>=1.13.0
+kernels>=0.12.1
+pynini==2.1.6
+sacremoses>=0.1.1

zonos2/__main__.py

@@ -0,0 +1,5 @@
+from .server import launch_server
+
+assert __name__ == "__main__"
+
+launch_server()

zonos2/attention/__init__.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol
+
+import torch
+from zonos2.utils import Registry, init_logger, is_sm90_supported, is_sm100_supported
+
+from .base import BaseAttnBackend, BaseAttnMetadata, HybridBackend
+
+if TYPE_CHECKING:
+    from zonos2.kvcache import BaseKVCache
+    from zonos2.models import ModelConfig
+
+logger = init_logger(__name__)
+
+
+class BackendCreator(Protocol):
+    def __call__(
+        self, config: ModelConfig, kvcache: BaseKVCache, page_table: torch.Tensor
+    ) -> BaseAttnBackend: ...
+
+
+SUPPORTED_ATTENTION_BACKENDS = Registry[BackendCreator]("Attention Backend")
+
+
+def resolve_auto_backend(config: ModelConfig) -> str:
+    """Determine the best attention backend based on the GPU architecture and model."""
+    if is_sm100_supported():  # blackwell
+        return "fi"
+    elif is_sm90_supported():  # hopper
+        return "fa,fi"
+    else:  # pre-hopper
+        return "fi"
+
+
+@SUPPORTED_ATTENTION_BACKENDS.register("fi")
+def create_fi_backend(config: ModelConfig, kvcache: BaseKVCache, page_table: torch.Tensor):
+    from .fi import FlashInferBackend
+
+    return FlashInferBackend(config, kvcache, page_table)
+
+
+@SUPPORTED_ATTENTION_BACKENDS.register("fa")
+def create_fa_backend(config: ModelConfig, kvcache: BaseKVCache, page_table: torch.Tensor):
+    from .fa import FlashAttentionBackend
+
+    return FlashAttentionBackend(config, kvcache, page_table)
+
+
+def validate_backend(backend: str):
+    if backend != "auto":
+        required_backends = backend.split(",") if "," in backend else [backend]
+        supported = SUPPORTED_ATTENTION_BACKENDS.supported_names()
+        for b in required_backends:
+            if b not in supported:
+                from argparse import ArgumentTypeError
+
+                raise ArgumentTypeError(
+                    f"Unsupported attention backend: {b}. Supported backends: {supported}"
+                )
+    return backend
+
+
+def create_attention_backend(
+    backend: str,
+    config: ModelConfig,
+    kvcache: BaseKVCache,
+    page_table: torch.Tensor,
+) -> BaseAttnBackend:
+    if backend == "auto":
+        backend = resolve_auto_backend(config)
+        logger.info(f"Auto-selected attention backend: {backend}")
+
+    if "," in backend:
+        assert backend.count(",") == 1, "Only one comma is allowed in hybrid backend"
+        p_backend, d_backend = backend.split(",", 1)
+        if p_backend != d_backend:
+            logger.info(f"Using hybrid attention backend: prefill={p_backend}, decode={d_backend}")
+            p_backend = create_attention_backend(p_backend, config, kvcache, page_table)
+            d_backend = create_attention_backend(d_backend, config, kvcache, page_table)
+            return HybridBackend(p_backend, d_backend)
+        backend = p_backend  # both are the same, fall through to single backend
+        logger.warning(f"P/D attention backends are the same: {backend}, using single backend.")
+
+    return SUPPORTED_ATTENTION_BACKENDS[backend](config, kvcache, page_table)
+
+
+__all__ = [
+    "BaseAttnMetadata",
+    "BaseAttnBackend",
+    "create_attention_backend",
+    "SUPPORTED_ATTENTION_BACKENDS",
+    "validate_backend",
+]

zonos2/attention/base.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List
+
+if TYPE_CHECKING:
+    import torch
+    from zonos2.core import TTSBatch
+
+
+@dataclass
+class BaseAttnMetadata(ABC):
+    positions: torch.Tensor
+
+    @abstractmethod
+    def get_last_indices(self, bs: int) -> torch.Tensor: ...
+
+
+class BaseAttnBackend(ABC):
+    @abstractmethod
+    def forward(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer_id: int, batch: TTSBatch
+    ) -> torch.Tensor: ...
+
+    @abstractmethod
+    def prepare_metadata(self, batch: TTSBatch) -> None: ...
+
+    @abstractmethod
+    def init_capture_graph(
+        self, max_seq_len: int, bs_list: List[int], frame_width: int = 1
+    ) -> None: ...
+
+    @abstractmethod
+    def prepare_for_capture(self, batch: TTSBatch) -> None: ...
+
+    @abstractmethod
+    def prepare_for_replay(self, batch: TTSBatch) -> None: ...
+
+
+class HybridBackend(BaseAttnBackend):
+    def __init__(
+        self,
+        prefill_backend: BaseAttnBackend,
+        decode_backend: BaseAttnBackend,
+    ) -> None:
+        self.prefill_backend = prefill_backend
+        self.decode_backend = decode_backend
+
+    def forward(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer_id: int, batch: TTSBatch
+    ) -> torch.Tensor:
+        backend = self.prefill_backend if batch.is_prefill else self.decode_backend
+        return backend.forward(q, k, v, layer_id, batch)
+
+    def prepare_metadata(self, batch: TTSBatch) -> None:
+        backend = self.prefill_backend if batch.is_prefill else self.decode_backend
+        return backend.prepare_metadata(batch)
+
+    def init_capture_graph(
+        self, max_seq_len: int, bs_list: List[int], frame_width: int = 1
+    ) -> None:
+        self.decode_backend.init_capture_graph(max_seq_len, bs_list, frame_width)
+
+    def prepare_for_capture(self, batch: TTSBatch) -> None:
+        self.decode_backend.prepare_for_capture(batch)
+
+    def prepare_for_replay(self, batch: TTSBatch) -> None:
+        self.decode_backend.prepare_for_replay(batch)

zonos2/attention/fa.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List, Tuple
+
+import torch
+
+from .base import BaseAttnBackend, BaseAttnMetadata
+from .utils import BaseCaptureData, make_positions
+
+if TYPE_CHECKING:
+    from zonos2.core import TTSBatch
+    from zonos2.kvcache import BaseKVCache
+    from zonos2.models import ModelConfig
+
+
+@dataclass
+class FACaptureData(BaseCaptureData):
+    pass
+
+
+@dataclass
+class FAMetadata(BaseAttnMetadata):
+    cu_seqlens_k: torch.Tensor
+    cu_seqlens_q: torch.Tensor
+    cache_seqlens: torch.Tensor
+    max_seqlen_k: int
+    max_seqlen_q: int
+
+    page_table: torch.Tensor
+
+    def get_positions(self) -> torch.Tensor:
+        return self.positions
+
+    def get_last_indices(self, bs: int) -> torch.Tensor:
+        return self.cu_seqlens_q[1 : 1 + bs] - 1
+
+
+class FlashAttentionBackend(BaseAttnBackend):
+    def __init__(self, config: ModelConfig, kvcache: BaseKVCache, page_table: torch.Tensor):
+        self.config = config
+        self.kvcache = kvcache
+        self.capture: FACaptureData | None = None
+        self.max_graph_bs = 0
+        self.capture_bs: List[int] = []
+        self.scale = config.head_dim**-0.5
+        self.page_table = page_table
+
+    def forward(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer_id: int, batch: TTSBatch
+    ) -> torch.Tensor:
+        metadata = batch.attn_metadata
+        assert isinstance(metadata, FAMetadata)
+        self.kvcache.store_kv(k, v, batch.out_loc, layer_id)
+        return _fa_sgl_impl(
+            q=q,
+            k_cache=self.kvcache.k_cache(layer_id),
+            v_cache=self.kvcache.v_cache(layer_id),
+            page_table=metadata.page_table,
+            cache_seqlens=metadata.cache_seqlens,
+            cu_seqlens_q=metadata.cu_seqlens_q,
+            cu_seqlens_k_new=metadata.cu_seqlens_k,
+            max_seqlen_q=metadata.max_seqlen_q,
+            softmax_scale=self.scale,
+        )
+
+    def prepare_metadata(self, batch: TTSBatch) -> None:
+        reqs = batch.padded_reqs
+
+        padded_size = len(reqs)
+        seqlens_q = [req.extend_len for req in reqs]
+        seqlens_k = [req.device_len for req in reqs]
+        cached_lens = [req.cached_len for req in reqs]
+        max_seqlen_k = max(seqlens_k)
+        max_seqlen_q = max(seqlens_q)
+        cpu_kwargs = {"device": "cpu", "dtype": torch.int32, "pin_memory": True}
+
+        device = self.kvcache.device
+        cache_seqlens = torch.tensor(seqlens_k, **cpu_kwargs)
+        cache_seqlens = cache_seqlens.to(device, non_blocking=True)
+        cu_seqlens_k = torch.tensor([0] + seqlens_k, **cpu_kwargs).cumsum_(dim=0)
+        cu_seqlens_k = cu_seqlens_k.to(device, non_blocking=True)
+
+        if max_seqlen_q == 1:
+            cu_seqlens_q = torch.arange(0, padded_size + 1, device=device, dtype=torch.int32)
+        elif all(l == 0 for l in cached_lens):  # prefill with no cache hit
+            cu_seqlens_q = cu_seqlens_k
+        else:  # normal extend prefill, with partial cache hit
+            cu_seqlens_q = torch.tensor([0] + seqlens_q, **cpu_kwargs).cumsum_(dim=0)
+            cu_seqlens_q = cu_seqlens_q.to(self.kvcache.device, non_blocking=True)
+
+        positions = make_positions(device, reqs)
+        page_table = self.page_table
+        new_page_table = torch.stack([page_table[req.table_idx, :max_seqlen_k] for req in reqs])
+
+        # copy from CPU to GPU
+        batch.attn_metadata = FAMetadata(
+            cu_seqlens_k=cu_seqlens_k,
+            cu_seqlens_q=cu_seqlens_q,
+            positions=positions,
+            cache_seqlens=cache_seqlens,
+            max_seqlen_k=max_seqlen_k,
+            max_seqlen_q=max_seqlen_q,
+            page_table=new_page_table,
+        )
+
+    def init_capture_graph(
+        self, max_seq_len: int, bs_list: List[int], frame_width: int = 1
+    ) -> None:
+        assert self.capture is None, "Capture already initialized."
+        max_bs = max(bs_list)
+        capture = FACaptureData.create(max_bs, max_seq_len, self.kvcache.device, frame_width)
+        self.max_graph_bs = max_bs
+        self.capture = capture
+        self.capture_bs = sorted(bs_list)
+
+    def prepare_for_capture(self, batch: TTSBatch) -> None:
+        assert (bs := batch.size) in self.capture_bs and self.capture
+        capture = self.capture
+        metadata = FAMetadata(
+            cu_seqlens_k=capture.cu_seqlens_k[: bs + 1],
+            cu_seqlens_q=capture.cu_seqlens_q[: bs + 1],
+            positions=capture.positions[:bs],
+            cache_seqlens=capture.seq_lens[:bs],
+            max_seqlen_k=capture.page_table.size(1),
+            max_seqlen_q=1,  # decode only
+            page_table=capture.page_table[:bs, :],
+        )
+        batch.attn_metadata = metadata
+        batch.input_ids = capture.input_ids[:bs]
+        batch.out_loc = capture.out_loc[:bs]
+
+    def prepare_for_replay(self, batch: TTSBatch) -> None:
+        metadata, bs = batch.attn_metadata, batch.padded_size
+        assert isinstance(metadata, FAMetadata)
+        assert self.capture is not None and bs in self.capture_bs
+        capture = self.capture
+
+        # Copy all dynamic tensors to capture buffers (fixed memory addresses for CUDA graph)
+        capture.input_ids[:bs].copy_(batch.input_ids)
+        capture.out_loc[:bs].copy_(batch.out_loc)
+        capture.positions[:bs].copy_(metadata.positions)
+        capture.cu_seqlens_k[: bs + 1].copy_(metadata.cu_seqlens_k)
+        capture.seq_lens[:bs].copy_(metadata.cache_seqlens)
+        capture.page_table[:bs, : metadata.max_seqlen_k].copy_(metadata.page_table)
+
+        # Update metadata and batch to point to capture buffers
+        # This ensures CUDA graph uses fixed memory addresses for all inputs
+        metadata.positions = capture.positions[:bs]
+        metadata.cu_seqlens_k = capture.cu_seqlens_k[: bs + 1]
+        metadata.cache_seqlens = capture.seq_lens[:bs]
+        metadata.page_table = capture.page_table[:bs, :]
+        batch.input_ids = capture.input_ids[:bs]
+        batch.out_loc = capture.out_loc[:bs]
+
+
+def _fa_sgl_impl(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    page_table: torch.Tensor,
+    cache_seqlens: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k_new: torch.Tensor,
+    max_seqlen_q: int,
+    softmax_scale: float,
+    sm_margin: int = 0,
+    window_size: Tuple[int, int] = (-1, -1),  # -1 means infinite context window
+    softcap: float = 0.0,  # 0.0 means deactivated
+    num_splits: int = 0,  # Can be tuned for speed
+    pack_gqa: bool | None = None,  # Can be tuned for speed
+    causal: bool = True,
+) -> torch.Tensor:
+    try:
+        from sgl_kernel.flash_attn import flash_attn_with_kvcache
+    except ImportError as e:
+        raise ImportError(
+            "sgl_kernel.flash_attn is not found. Please install it with `pip install sgl-kernel`.\n"
+            "If you're sure it's correctly installed, try `apt update && apt install libnuma1`."
+        ) from e
+
+    return flash_attn_with_kvcache(  # type: ignore
+        q=q,
+        k_cache=k_cache,
+        v_cache=v_cache,
+        page_table=page_table,
+        cache_seqlens=cache_seqlens,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k_new=cu_seqlens_k_new,
+        max_seqlen_q=max_seqlen_q,
+        softmax_scale=softmax_scale,
+        sm_margin=sm_margin,
+        window_size=window_size,
+        softcap=softcap,
+        num_splits=num_splits,
+        pack_gqa=pack_gqa,
+        causal=causal,
+        ver=3,  # TODO: support FA4 on blackwell
+    )

zonos2/attention/fi.py

@@ -0,0 +1,360 @@
+from __future__ import annotations
+
+import logging
+import math
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING, Dict, List, Literal
+
+import torch
+from zonos2.distributed import get_tp_info
+from zonos2.env import ENV
+from zonos2.utils import divide_even
+from zonos2.utils.logger import init_logger
+
+from .base import BaseAttnBackend, BaseAttnMetadata
+from .utils import BaseCaptureData, make_positions
+
+if TYPE_CHECKING:
+    from flashinfer import (
+        BatchDecodeWithPagedKVCacheWrapper,
+        BatchPrefillWithPagedKVCacheWrapper,
+        CUDAGraphBatchDecodeWithPagedKVCacheWrapper,
+    )
+    from zonos2.core import TTSBatch
+    from zonos2.kvcache import BaseKVCache
+    from zonos2.models import ModelConfig
+
+
+def _next_power_of_2(n: int) -> int:
+    if n <= 1:
+        return 1
+    return 1 << math.ceil(math.log2(n))
+
+
+logger = init_logger(__name__)
+
+_fi_debug_counter = 0
+
+
+@dataclass
+class FICaptureData(BaseCaptureData):
+    @property
+    def one_tensor(self) -> torch.Tensor:
+        return self.seq_lens
+
+    @property
+    def indices(self) -> torch.Tensor:
+        return self.page_table
+
+
+@dataclass
+class FIMetadata(BaseAttnMetadata):
+    # fmt: off
+    cu_seqlens_q_cpu:   torch.Tensor  # on cpu
+    cu_seqlens_k_cpu:   torch.Tensor  # on cpu
+    cu_seqlens_q_gpu:   torch.Tensor  # on gpu
+    indices:            torch.Tensor  # on gpu
+    last_page_len_cpu:  torch.Tensor  # on cpu
+    num_qo_heads:       int
+    num_kv_heads:       int
+    head_dim:           int
+    page_size:          Literal[1] # currently only support page_size=1
+    pos_encoding_mode:  str
+    seq_lens_cpu:       torch.Tensor  # on cpu
+    dtype:              torch.dtype
+    wrapper:            BatchPrefillWithPagedKVCacheWrapper | BatchDecodeWithPagedKVCacheWrapper
+    initialized:        bool = False
+    sm_scale:           float = 0.0  # softmax scale (1/sqrt(head_dim)), set during prepare_metadata
+    # fmt: on
+
+    def __post_init__(self) -> None:
+        assert self.page_size == 1, "Currently only page_size=1 is supported."
+        assert (
+            self.positions.is_cuda
+            and self.cu_seqlens_k_cpu.is_cpu
+            and self.cu_seqlens_q_cpu.is_cpu
+            and self.cu_seqlens_q_gpu.is_cuda
+            and self.indices.is_cuda
+            and self.last_page_len_cpu.is_cpu
+            and self.seq_lens_cpu.is_cpu
+        )
+
+    def get_positions(self) -> torch.Tensor:
+        return self.positions
+
+    def get_last_indices(self, bs: int) -> torch.Tensor:
+        return self.cu_seqlens_q_gpu[1 : 1 + bs] - 1
+
+
+class FlashInferBackend(BaseAttnBackend):
+    def __init__(
+        self,
+        config: ModelConfig,
+        kvcache: BaseKVCache,
+        page_table: torch.Tensor,
+    ) -> None:
+        from flashinfer import (
+            BatchDecodeWithPagedKVCacheWrapper,
+            BatchPrefillWithPagedKVCacheWrapper,
+        )
+
+        self.config = config
+        self.kvcache = kvcache
+        self.device = kvcache.device
+        self.float_workspace_buffer = torch.empty(
+            128 * 1024 * 1024, dtype=torch.uint8, device=self.device
+        )
+        self.prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
+            self.float_workspace_buffer,
+            kv_layout="NHD",
+            backend="fa2",  # flashinfer fa3 is buggy, use fa2 instead
+        )
+        self.decode_wrappers = BatchDecodeWithPagedKVCacheWrapper(
+            self.float_workspace_buffer,
+            kv_layout="NHD",
+        )
+
+        # NOTE: some hack to reuse the int_workspace_buffer
+        self.int_workspace_buffer = self.prefill_wrapper._int_workspace_buffer
+        self.decode_wrappers._int_workspace_buffer = self.int_workspace_buffer
+
+        # initialize some data members
+        tp_size = get_tp_info().size
+        self.qo_head_local = divide_even(self.config.num_qo_heads, tp_size)
+        self.kv_head_local = divide_even(self.config.num_kv_heads, tp_size)
+
+        self.cached_ones_cpu: torch.Tensor = torch.tensor([], dtype=torch.int32, pin_memory=True)
+        # for cuda graph
+        self.capture_bs: List[int] = []
+        self.max_graph_bs = 0
+        self.graph_wrappers: Dict[int, CUDAGraphBatchDecodeWithPagedKVCacheWrapper] = {}
+        self.capture: FICaptureData | None = None
+        self.page_table = page_table
+
+    @staticmethod
+    def _initialize_metadata_once(metadata: FIMetadata) -> None:
+        if metadata.initialized:
+            return
+
+        from flashinfer import (
+            BatchDecodeWithPagedKVCacheWrapper,
+            CUDAGraphBatchDecodeWithPagedKVCacheWrapper,
+        )
+
+        metadata.initialized = True
+
+        is_cuda_graph = isinstance(metadata.wrapper, CUDAGraphBatchDecodeWithPagedKVCacheWrapper)
+        is_decode = isinstance(metadata.wrapper, BatchDecodeWithPagedKVCacheWrapper) or is_cuda_graph
+
+        if is_cuda_graph and logger.isEnabledFor(logging.DEBUG):
+            global _fi_debug_counter
+            if _fi_debug_counter <= 5:
+                logger.debug("plan() for CUDAGraphWrapper: indices shape=%s", metadata.indices.shape)
+                logger.debug("  cu_seqlens_k_cpu=%s", metadata.cu_seqlens_k_cpu.tolist())
+
+        if is_decode:
+            # Decode wrapper (both regular and CUDA graph variants)
+            # sm_scale must be passed to plan() for decode, not run()
+            metadata.wrapper.plan(
+                indptr=metadata.cu_seqlens_k_cpu,
+                indices=metadata.indices,
+                last_page_len=metadata.last_page_len_cpu,
+                num_qo_heads=metadata.num_qo_heads,
+                num_kv_heads=metadata.num_kv_heads,
+                head_dim=metadata.head_dim,
+                page_size=metadata.page_size,
+                pos_encoding_mode=metadata.pos_encoding_mode,
+                seq_lens=metadata.seq_lens_cpu,
+                data_type=metadata.dtype,
+                q_data_type=metadata.dtype,
+                kv_data_type=metadata.dtype,
+                sm_scale=metadata.sm_scale,
+                non_blocking=True,
+            )
+        else:
+            # Prefill wrapper - sm_scale passed to plan()
+            metadata.wrapper.plan(
+                qo_indptr=metadata.cu_seqlens_q_cpu,
+                paged_kv_indptr=metadata.cu_seqlens_k_cpu,
+                paged_kv_indices=metadata.indices,
+                paged_kv_last_page_len=metadata.last_page_len_cpu,
+                num_qo_heads=metadata.num_qo_heads,
+                num_kv_heads=metadata.num_kv_heads,
+                head_dim_qk=metadata.head_dim,
+                page_size=metadata.page_size,
+                pos_encoding_mode=metadata.pos_encoding_mode,
+                seq_lens=metadata.seq_lens_cpu,
+                q_data_type=metadata.dtype,
+                kv_data_type=metadata.dtype,
+                sm_scale=metadata.sm_scale,
+                non_blocking=True,
+                causal=True,
+            )
+
+    def _get_ones_cpu(self, bs: int) -> torch.Tensor:
+        if bs <= len(self.cached_ones_cpu):
+            return self.cached_ones_cpu[:bs]
+        # padding to next pow of 2
+        next_len = _next_power_of_2(bs)
+        self.cached_ones_cpu = torch.ones(next_len, dtype=torch.int32, pin_memory=True)
+        return self.cached_ones_cpu[:bs]
+
+    def forward(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer_id: int, batch: TTSBatch
+    ) -> torch.Tensor:
+        metadata = batch.attn_metadata
+        assert isinstance(metadata, FIMetadata)
+        self._initialize_metadata_once(metadata)
+        self.kvcache.store_kv(k, v, batch.out_loc, layer_id)
+        kv_cache = (self.kvcache.k_cache(layer_id), self.kvcache.v_cache(layer_id))
+
+        # Debug: check attention inputs/outputs for first few steps at layer 0
+        global _fi_debug_counter
+        is_capturing = torch.cuda.is_current_stream_capturing()
+        if (logger.isEnabledFor(logging.DEBUG)
+                and layer_id == 0 and batch.is_decode and not is_capturing
+                and _fi_debug_counter < 10):
+            _fi_debug_counter += 1
+            step = _fi_debug_counter
+            logger.debug("ATTN Step %d Layer 0", step)
+            logger.debug("  Q shape: %s, norm: %.4f, mean: %.6f", q.shape, q.norm().item(), q.mean().item())
+            logger.debug("  K shape: %s, norm: %.4f, mean: %.6f", k.shape, k.norm().item(), k.mean().item())
+            logger.debug("  V shape: %s, norm: %.4f, mean: %.6f", v.shape, v.norm().item(), v.mean().item())
+            logger.debug("  out_loc: %s", batch.out_loc.tolist())
+            logger.debug("  positions: %s", metadata.positions.tolist())
+            logger.debug("  indices shape: %s, first 10: %s", metadata.indices.shape, metadata.indices[:10].tolist())
+            logger.debug("  cu_seqlens_k: %s", metadata.cu_seqlens_k_cpu.tolist())
+            k_cache, v_cache = kv_cache
+            if batch.out_loc.numel() > 0:
+                slot = batch.out_loc[0].item()
+                logger.debug("  KV cache slot %d: k_norm=%.4f", slot, k_cache.view(-1, k_cache.shape[-1])[slot].norm().item())
+
+        output = metadata.wrapper.run(q=q, paged_kv_cache=kv_cache)
+
+        if (logger.isEnabledFor(logging.DEBUG)
+                and layer_id == 0 and batch.is_decode and not is_capturing
+                and _fi_debug_counter <= 10):
+            logger.debug("  Output shape: %s, norm: %.4f, mean: %.6f", output.shape, output.norm().item(), output.mean().item())
+
+        return output
+
+    def prepare_metadata(self, batch: TTSBatch) -> None:
+        reqs = batch.padded_reqs
+
+        padded_size = len(reqs)
+        seqlens_q = [req.extend_len for req in reqs]
+        seqlens_k = [req.device_len for req in reqs]
+        cached_lens = [req.cached_len for req in reqs]
+        max_seqlen_q = max(seqlens_q)
+        cpu_kwargs = {"device": "cpu", "dtype": torch.int32, "pin_memory": True}
+
+        device = self.device
+        seq_len_cpu = torch.tensor(seqlens_k, **cpu_kwargs)
+        cu_seqlens_k_cpu = torch.tensor([0] + seqlens_k, **cpu_kwargs).cumsum_(dim=0)
+        if max_seqlen_q == 1:  # decode with all extend_len = 1
+            cu_seqlens_q_cpu = torch.arange(0, padded_size + 1, **cpu_kwargs)
+        elif all(l == 0 for l in cached_lens):  # prefill with no cache hit
+            cu_seqlens_q_cpu = cu_seqlens_k_cpu
+        else:  # normal extend prefill, with partial cache hit
+            cu_seqlens_q_cpu = torch.tensor([0] + seqlens_q, **cpu_kwargs).cumsum_(dim=0)
+        batch.attn_metadata = FIMetadata(
+            positions=make_positions(device, reqs),
+            cu_seqlens_q_cpu=cu_seqlens_q_cpu,
+            cu_seqlens_k_cpu=cu_seqlens_k_cpu,
+            cu_seqlens_q_gpu=cu_seqlens_q_cpu.to(device, non_blocking=True),
+            indices=torch.cat([self.page_table[req.table_idx, : req.device_len] for req in reqs]),
+            last_page_len_cpu=self._get_ones_cpu(padded_size),
+            num_qo_heads=self.qo_head_local,
+            num_kv_heads=self.kv_head_local,
+            head_dim=self.config.head_dim,
+            page_size=1,
+            pos_encoding_mode="NONE",
+            seq_lens_cpu=seq_len_cpu,
+            dtype=self.kvcache.dtype,
+            wrapper=self.decode_wrappers if batch.is_decode else self.prefill_wrapper,
+            sm_scale=1.0 / math.sqrt(self.config.head_dim),
+        )
+
+    def init_capture_graph(
+        self, max_seq_len: int, bs_list: List[int], frame_width: int = 1
+    ) -> None:
+        assert self.capture is None, "Capture already initialized."
+        max_bs = max(bs_list)
+        capture = FICaptureData.create(max_bs, max_seq_len, self.kvcache.device, frame_width)
+        capture.page_table = capture.page_table.view(-1)  # use 1D as ragged indices
+        self.max_graph_bs = max_bs
+        self.capture = capture
+        self.capture_bs = sorted(bs_list)
+
+    @cached_property
+    def use_tensor_cores(self) -> bool:
+        if (overriden_value := ENV.FLASHINFER_USE_TENSOR_CORES.value) is not None:
+            logger.warning(f"Overriding FlashInfer tensor core usage to {overriden_value}")
+            return overriden_value
+        GQA = self.config.num_qo_heads // self.config.num_kv_heads
+        return GQA >= 4
+
+    def prepare_for_capture(self, batch: TTSBatch) -> None:
+        from flashinfer import CUDAGraphBatchDecodeWithPagedKVCacheWrapper
+
+        bs = batch.size
+        assert bs in self.capture_bs and bs not in self.graph_wrappers and self.capture
+        batch.padded_reqs = batch.reqs
+        capture = self.capture
+        self.graph_wrappers[bs] = CUDAGraphBatchDecodeWithPagedKVCacheWrapper(
+            self.float_workspace_buffer,
+            kv_layout="NHD",
+            use_tensor_cores=self.use_tensor_cores,
+            indptr_buffer=capture.cu_seqlens_k[: bs + 1],
+            indices_buffer=capture.indices,
+            last_page_len_buffer=capture.one_tensor[:bs],
+        )
+        self.graph_wrappers[bs]._int_workspace_buffer = self.int_workspace_buffer
+        self.prepare_metadata(batch)
+        metadata = batch.attn_metadata
+        assert isinstance(metadata, FIMetadata)
+        # Copy metadata tensors to capture buffers and update metadata to use capture buffers
+        # This ensures CUDA graph uses fixed memory addresses
+        indices_len = metadata.indices.shape[0]
+        capture.indices[:indices_len].copy_(metadata.indices)
+        capture.cu_seqlens_k[: bs + 1].copy_(metadata.cu_seqlens_k_cpu.to(capture.indices.device))
+        capture.seq_lens[:bs].copy_(metadata.seq_lens_cpu.to(capture.indices.device))
+        # Update metadata to point to capture buffers (critical for CUDA graph)
+        metadata.indices = capture.indices[:indices_len]
+        metadata.wrapper = self.graph_wrappers[bs]
+        metadata.positions = capture.positions[:bs]
+        batch.input_ids = capture.input_ids[:bs]
+        batch.out_loc = capture.out_loc[:bs]
+        self._initialize_metadata_once(metadata)
+
+    def prepare_for_replay(self, batch: TTSBatch) -> None:
+        metadata, bs = batch.attn_metadata, batch.padded_size
+        assert isinstance(metadata, FIMetadata) and not metadata.initialized
+        assert self.capture is not None and bs in self.capture_bs
+        capture = self.capture
+
+        if logger.isEnabledFor(logging.DEBUG) and _fi_debug_counter < 5:
+            logger.debug("prepare_for_replay: bs=%d, indices_len=%d", bs, metadata.indices.shape[0])
+            logger.debug("  cu_seqlens_k_cpu=%s", metadata.cu_seqlens_k_cpu.tolist())
+
+        # Copy all dynamic tensors to capture buffers (fixed memory addresses for CUDA graph)
+        capture.input_ids[:bs].copy_(batch.input_ids)
+        capture.out_loc[:bs].copy_(batch.out_loc)
+        capture.positions[:bs].copy_(metadata.positions)
+        # Copy indices (page table entries) - critical for KV cache lookups
+        indices_len = metadata.indices.shape[0]
+        capture.indices[:indices_len].copy_(metadata.indices)
+        # Copy cu_seqlens_k for the wrapper's indptr
+        capture.cu_seqlens_k[: bs + 1].copy_(metadata.cu_seqlens_k_cpu.to(capture.indices.device))
+        # Copy seq_lens for the wrapper
+        capture.seq_lens[:bs].copy_(metadata.seq_lens_cpu.to(capture.indices.device))
+        # Update metadata to point to capture buffers (CRITICAL: must match prepare_for_capture)
+        # This ensures CUDA graph uses fixed memory addresses for all inputs
+        metadata.indices = capture.indices[:indices_len]
+        metadata.positions = capture.positions[:bs]
+        metadata.wrapper = self.graph_wrappers[bs]
+        # Update batch input_ids and out_loc to point to capture buffers
+        batch.input_ids = capture.input_ids[:bs]
+        batch.out_loc = capture.out_loc[:bs]
+        self._initialize_metadata_once(metadata)

zonos2/attention/utils.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List
+
+import torch
+
+if TYPE_CHECKING:
+    from zonos2.core import TTSReq
+
+
+@dataclass
+class BaseCaptureData:
+    input_ids: torch.Tensor
+    seq_lens: torch.Tensor
+    positions: torch.Tensor
+    cu_seqlens_k: torch.Tensor
+    cu_seqlens_q: torch.Tensor
+    page_table: torch.Tensor
+    out_loc: torch.Tensor
+
+    @classmethod
+    def create(
+        cls, max_bs: int, max_seq_len: int, device: torch.device, frame_width: int = 1, **kwargs
+    ):
+        # Audio token frames are 2D (batch, frame_width); frame_width == 1 keeps a 1D layout
+        if frame_width > 1:
+            input_ids = torch.zeros((max_bs, frame_width), dtype=torch.int32, device=device)
+        else:
+            input_ids = torch.zeros((max_bs,), dtype=torch.int32, device=device)
+        return cls(
+            input_ids=input_ids,
+            seq_lens=torch.ones((max_bs,), dtype=torch.int32, device=device),
+            positions=torch.zeros((max_bs,), dtype=torch.int32, device=device),
+            cu_seqlens_k=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device),
+            cu_seqlens_q=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device),
+            page_table=torch.zeros((max_bs, max_seq_len), dtype=torch.int32, device=device),
+            out_loc=torch.zeros((max_bs,), dtype=torch.int32, device=device),
+            **kwargs,
+        )
+
+
+def make_positions(device: torch.device, reqs: List[TTSReq]) -> torch.Tensor:
+    needed_size = sum(req.extend_len for req in reqs)
+    indices_host = torch.empty(needed_size, dtype=torch.int32, pin_memory=True)
+    offset = 0
+    for req in reqs:
+        length = req.extend_len
+        torch.arange(
+            req.cached_len,
+            req.device_len,
+            dtype=torch.int32,
+            out=indices_host[offset : offset + length],
+        )
+        offset += length
+    return indices_host.to(device, non_blocking=True)

zonos2/benchmark/client.py

@@ -0,0 +1,501 @@
+from __future__ import annotations
+
+import asyncio
+import random
+import time
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Tuple, overload
+
+from openai import AsyncOpenAI as OpenAI
+from pydantic import BaseModel
+from tqdm.asyncio import tqdm
+from zonos2.utils import UNSET, Unset, init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass(frozen=True)
+class BenchmarkTrace:
+    timestamp: float
+    message: str  # unit (second)
+    output_length: int  # output length in tokens
+    input_length: int | None = None  # input length in tokens, optional
+
+
+@dataclass(frozen=True)
+class BenchOneResult:
+    tics: List[float]
+    input_len: int
+    output_len: int
+
+    def as_json(self) -> List[float]:
+        return [self.input_len, self.output_len] + self.tics
+
+    @staticmethod
+    def from_json(raw: List[float]) -> BenchOneResult:
+        # check raw[0] and raw[1] are integers
+        assert raw[0].is_integer() and raw[1].is_integer()
+        return BenchOneResult(tics=raw[2:], input_len=int(raw[0]), output_len=int(raw[1]))
+
+
+@dataclass(frozen=True)
+class RawResult:
+    input_len: int | None
+    output_len: int
+    message: str
+    tics: List[float]
+
+
+@dataclass
+class Counter:
+    current: int = 0
+    history_max: int = 0
+
+    def inc(self, n=1):
+        self.current += n
+        self.history_max = max(self.history_max, self.current)
+
+    def dec(self, n=1):
+        self.current -= n
+        assert self.current >= 0
+
+
+@dataclass
+class Console:
+    input_pbar: tqdm
+    output_pbar: tqdm
+    prefill_pbar: tqdm
+    decode_pbar: tqdm
+    disabled: bool
+    inflight_counter: Counter = field(default_factory=Counter)
+    queue_counter: Counter = field(default_factory=Counter)
+
+    def update_input(self, n=1):
+        self.input_pbar.update(n)
+        self.input_pbar.refresh()
+        self.inflight_counter.inc(n)
+        self.queue_counter.inc(n)
+
+    def update_output(self, n=1):
+        self.output_pbar.update(n)
+        self.output_pbar.refresh()
+        self.inflight_counter.dec(n)
+
+    def update_prefill(self, n=1):
+        self.prefill_pbar.update(n)
+        self.prefill_pbar.refresh()
+        self.queue_counter.dec(n)
+
+    def update_decode(self, n=1):
+        self.decode_pbar.update(n)
+
+    @contextmanager
+    def inflight(self, n=1):
+        self.update_input(n)
+        yield
+        self.update_output(n)
+
+    @contextmanager
+    def log_stats(self):
+        yield
+        self.input_pbar.close()
+        self.output_pbar.close()
+        self.prefill_pbar.close()
+        self.decode_pbar.close()
+        if not self.disabled:
+            max_inflight = self.inflight_counter.history_max
+            max_queue = self.queue_counter.history_max
+            logger.info(f"Max inflight requests: {max_inflight}, Max queued requests: {max_queue}")
+
+
+@dataclass(frozen=True)
+class BenchmarkResult:
+    raw_data: List[BenchOneResult]
+
+    def as_json(self) -> List[List[float]]:
+        return [r.as_json() for r in self.raw_data]
+
+    @staticmethod
+    def from_json(raw: List[List[float]]) -> BenchmarkResult:
+        return BenchmarkResult(raw_data=[BenchOneResult.from_json(r) for r in raw])
+
+
+def make_console(num_requests: int, sum_output_length: int, use_pbar: bool = True) -> Console:
+    BAR_FORMAT_0 = (
+        "{desc:<10} {percentage:3.0f}%|{bar}|"
+        " {n_fmt:>5}/{total_fmt} "
+        "[{rate_fmt:>12} {elapsed:>8}/{remaining:<8}]"
+    )
+    BAR_FORMAT_1 = BAR_FORMAT_0
+    n_fmt_align = 5
+    prefill_tokens = num_requests
+    decode_tokens = sum_output_length - prefill_tokens
+
+    if len(str(decode_tokens)) > n_fmt_align:
+        n_fmt_align = len(str(decode_tokens))
+        BAR_FORMAT_0 = BAR_FORMAT_0.replace("{n_fmt:>5}", "{n_fmt:>" + str(n_fmt_align) + "}")
+        BAR_FORMAT_1 = BAR_FORMAT_0
+
+    if len(str(prefill_tokens)) < len(str(decode_tokens)):
+        old_align_str = "{n_fmt:>" + str(n_fmt_align) + "}"
+        n_fmt_align += len(str(decode_tokens)) - len(str(prefill_tokens))
+        BAR_FORMAT_0 = BAR_FORMAT_0.replace(old_align_str, "{n_fmt:>" + str(n_fmt_align) + "}")
+
+    disabled = not use_pbar
+    input_pbar = tqdm(
+        total=num_requests,
+        desc="Requests sent",
+        position=0,
+        bar_format=BAR_FORMAT_0,
+        disable=disabled,
+    )
+    output_pbar = tqdm(
+        total=num_requests,
+        desc="Requests done",
+        position=1,
+        bar_format=BAR_FORMAT_0,
+        disable=disabled,
+    )
+    prefill_pbar = tqdm(
+        total=prefill_tokens,
+        desc="Prefill token",
+        position=2,
+        bar_format=BAR_FORMAT_0,
+        disable=disabled,
+    )
+    decode_pbar = tqdm(
+        total=decode_tokens,
+        desc="Decode token ",
+        position=3,
+        bar_format=BAR_FORMAT_1,
+        disable=disabled,
+    )
+    return Console(
+        input_pbar=input_pbar,
+        output_pbar=output_pbar,
+        prefill_pbar=prefill_pbar,
+        decode_pbar=decode_pbar,
+        disabled=disabled,
+    )
+
+
+def generate_prompt(tokenizer: Any, n: int) -> str:
+    """Generate a prompt of approximately `n` tokens using the provided tokenizer."""
+    vocab_size = tokenizer.vocab_size // 2
+    token_ids = [random.randint(0, vocab_size) for _ in range(n - 1)]
+
+    for _ in range(64):
+        prompt = tokenizer.decode(token_ids)
+        token_ids = tokenizer.encode(prompt, add_special_tokens=False)
+        if len(token_ids) == n:
+            return prompt
+        if len(token_ids) < n:
+            need = n - len(token_ids)
+            token_ids.extend([random.randint(0, vocab_size) for _ in range(need)])
+        else:
+            token_ids = token_ids[:n]
+
+    raise ValueError("Failed to generate a message of the desired length.")
+
+
+async def benchmark_one(
+    client: OpenAI,
+    prompt: str,
+    output_length: int,
+    model: str,
+    *,
+    pbar: Console | bool = True,
+    extra_body: Dict[str, Any] | None = None,
+    input_length: int | None = None,  # a hack to force input length
+) -> RawResult:
+    if isinstance(pbar, bool):
+        pbar = make_console(1, output_length, use_pbar=pbar)
+    with pbar.inflight(1):
+        kwargs = {
+            "ignore_eos": True,
+            "top_k": 1,
+        }
+        # this is an internal kwargs that might work for our system
+        if input_length is not None:
+            kwargs["input_length_override"] = input_length
+        kwargs.update(extra_body or {})  # can override kwargs
+        response = await client.chat.completions.create(
+            model=model,
+            stream=True,
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt,
+                },
+            ],
+            max_tokens=output_length,
+            temperature=0.0,
+            extra_body=kwargs,
+        )
+        tics = [time.perf_counter()]
+        async for _ in response:
+            tics.append(time.perf_counter())
+            if len(tics) == 2:
+                pbar.update_prefill()
+            elif len(tics) <= output_length + 1:
+                pbar.update_decode()
+        return RawResult(
+            input_len=input_length,
+            output_len=output_length,
+            message=prompt,
+            tics=tics,
+        )
+
+
+async def benchmark_one_batch(
+    client: OpenAI,
+    prompts: List[str],
+    output_lengths: List[int] | int,
+    model: str,
+    *,
+    extra_body: Dict[str, Any] | None = None,
+    input_lengths: List[int | None] | None = None,
+    pbar: Console | bool = True,
+) -> List[RawResult]:
+    if isinstance(output_lengths, int):
+        output_lengths = [output_lengths] * len(prompts)
+    if isinstance(pbar, bool):
+        pbar = make_console(len(prompts), sum(output_lengths), use_pbar=pbar)
+    if input_lengths is None:
+        l: List[int | None] = [None] * len(prompts)
+        input_lengths = l  # work-around for typing bug
+
+    tasks = [
+        benchmark_one(
+            client=client,
+            prompt=prompt,
+            output_length=output_length,
+            model=model,
+            pbar=pbar,
+            extra_body=extra_body,
+            input_length=input_length,
+        )
+        for prompt, output_length, input_length in zip(
+            prompts, output_lengths, input_lengths, strict=True
+        )
+    ]
+    with pbar.log_stats():
+        return await asyncio.gather(*tasks)
+
+
+async def benchmark_trace(
+    client: OpenAI,
+    msgs: List[BenchmarkTrace],
+    model: str,
+    *,
+    pbar: Console | bool = True,
+) -> List[RawResult]:
+    if isinstance(pbar, bool):
+        sum_output_len = sum(msg.output_length for msg in msgs)
+        pbar = make_console(len(msgs), sum_output_len, use_pbar=pbar)
+    start = time.perf_counter()
+    offset = min(msg.timestamp for msg in msgs) - 1
+
+    async def benchmark_timed(msg: BenchmarkTrace):
+        target = start + msg.timestamp - offset
+        await asyncio.sleep(max(0, target - time.perf_counter()))
+        return await benchmark_one(
+            client, msg.message, msg.output_length, model, pbar=pbar, input_length=msg.input_length
+        )
+
+    tasks = [benchmark_timed(msg) for msg in msgs]
+    with pbar.log_stats():
+        return await asyncio.gather(*tasks)
+
+
+@overload
+def process_benchmark_results(raw_data: List[RawResult], tokenizer: Any) -> BenchmarkResult: ...
+
+
+@overload
+def process_benchmark_results(raw_data: List[RawResult]) -> None: ...
+
+
+def process_benchmark_results(
+    raw_data: List[RawResult],
+    tokenizer: Any = UNSET,
+) -> BenchmarkResult | None:
+    accum_times: List[float] = []
+    first_times: List[float] = []
+    results = [r.tics for r in raw_data]
+    for tics in results:
+        deltas: List[float] = []
+        for i in range(len(tics) - 1):
+            diff = tics[i + 1] - tics[i]
+            deltas.append(diff)
+        first_times.append(deltas[0])
+        accum_times.extend(deltas[1:])
+
+    e2e_times = [tics[-1] - tics[0] for tics in results]
+    first_times.sort()
+    accum_times.sort()
+    e2e_times.sort()
+
+    def _print_stats(times: List[float], scale: float = 1.0) -> Tuple[float, ...]:
+        assert len(times) > 0
+        return (
+            scale * sum(times) / len(times),  # avg
+            scale * times[int(len(times) * 0.5)],  # p50
+            scale * times[int(len(times) * 0.9)],  # p90
+            scale * times[int(len(times) * 0.99)],  # p99
+            scale * max(times),  # max
+        )
+
+    def _fmt(x: float) -> str:
+        if x >= 1000:
+            return f"{int(x):>6}"
+        elif x >= 10:
+            return f"{x:>6.2f}"
+        else:
+            return f"{x:>6.4f}"
+
+    avg_ttft, p50_ttft, p90_ttft, p99_ttft, max_ttft = _print_stats(first_times, 1000)
+    avg_tpot, p50_tpot, p90_tpot, p99_tpot, max_tpot = _print_stats(accum_times, 1000)
+    avg_e2e, p50_e2e, p90_e2e, p99_e2e, max_e2e = _print_stats(e2e_times)
+
+    min_time = min(min(r) for r in results)
+    max_time = max(max(r) for r in results)
+    dur = max_time - min_time
+    assert dur > 0, "Duration must be positive"
+
+    num_tokens = sum(len(tic) for tic in results)
+    num_requests = len(results)
+
+    logger.info(f"Num requests: #{num_requests}, Num tokens: #{num_tokens}")
+    logger.info(
+        f"TTFT: {_fmt(avg_ttft)} ms (p50: {_fmt(p50_ttft)} ms, p90: {_fmt(p90_ttft)} ms,"
+        f" p99: {_fmt(p99_ttft)} ms, max: {_fmt(max_ttft)} ms)"
+    )
+    logger.info(
+        f"TPOT: {_fmt(avg_tpot)} ms (p50: {_fmt(p50_tpot)} ms, p90: {_fmt(p90_tpot)} ms,"
+        f" p99: {_fmt(p99_tpot)} ms, max: {_fmt(max_tpot)} ms)"
+    )
+    logger.info(
+        f"E2E:  {_fmt(avg_e2e) }  s (p50: {_fmt(p50_e2e) }  s, p90: {_fmt(p90_e2e) }  s,"
+        f" p99: {_fmt(p99_e2e) }  s, max: {_fmt(max_e2e) }  s)"
+    )
+    logger.info(f"Duration: {_fmt(dur)} s")
+    logger.info(f"Throughput: {_fmt(num_tokens / dur)} token/s, {_fmt(num_requests / dur)} req/s")
+
+    # normalize the time to start from zero
+    results = [[r - min_time for r in tics] for tics in results]
+    if isinstance(tokenizer, Unset):
+        return None
+
+    return BenchmarkResult(
+        raw_data=[
+            BenchOneResult(
+                tics=r.tics,
+                input_len=(
+                    r.input_len
+                    if r.input_len is not None
+                    else len(tokenizer.encode(r.message, add_special_tokens=False))
+                ),
+                output_len=r.output_len,
+            )
+            for r in raw_data
+        ]
+    )
+
+
+def read_qwen_trace(
+    file_path: str,
+    tokenizer: Any,
+    n: int | None = None,
+    dummy: bool = False,
+) -> List[BenchmarkTrace]:
+    class JSONInput(BaseModel):
+        chat_id: int
+        parent_chat_id: int
+        timestamp: float
+        input_length: int
+        output_length: int
+        type: str  # unused
+        turn: int  # unused
+        hash_ids: List[int]  # unused
+
+    with open(file_path, "r") as f:
+        lines = f.readlines()
+        if n is not None:
+            lines = lines[:n]
+    objs = [JSONInput.model_validate_json(line) for line in lines]
+    if dummy:
+        prompt = generate_prompt(tokenizer, max(obj.input_length for obj in objs))
+        ids = tokenizer.encode(prompt, add_special_tokens=False)
+        _get_prompt = lambda obj: tokenizer.decode(ids[: obj.input_length])
+    else:
+        _get_prompt = lambda obj: generate_prompt(tokenizer, obj.input_length)
+    return [
+        BenchmarkTrace(
+            timestamp=obj.timestamp,
+            message=_get_prompt(obj),
+            input_length=obj.input_length,
+            output_length=obj.output_length,
+        )
+        for obj in objs
+    ]
+
+
+def read_mooncake_trace(
+    file_path: str,
+    tokenizer: Any,
+    n: int | None = None,
+    dummy: bool = False,
+) -> List[BenchmarkTrace]:
+    class JSONInput(BaseModel):
+        timestamp: int
+        input_length: int
+        output_length: int
+        hash_ids: List[int]  # unused for now
+
+    with open(file_path, "r") as f:
+        lines = f.readlines()
+        if n is not None:
+            lines = lines[:n]
+    objs = [JSONInput.model_validate_json(line) for line in lines]
+    if dummy:
+        prompt = generate_prompt(tokenizer, max(obj.input_length for obj in objs))
+        ids = tokenizer.encode(prompt, add_special_tokens=False)
+        _get_prompt = lambda obj: tokenizer.decode(ids[: obj.input_length])
+    else:
+        _get_prompt = lambda obj: generate_prompt(tokenizer, obj.input_length)
+    return [
+        BenchmarkTrace(
+            timestamp=obj.timestamp / 1000,
+            message=_get_prompt(obj),
+            input_length=obj.input_length,
+            output_length=obj.output_length,
+        )
+        for obj in objs
+    ]
+
+
+def scale_traces(
+    traces: List[BenchmarkTrace],
+    scale: float,
+) -> List[BenchmarkTrace]:
+    min_tic = min(trace.timestamp for trace in traces)
+    return sorted(
+        [
+            BenchmarkTrace(
+                timestamp=(trace.timestamp - min_tic) * scale,
+                message=trace.message,
+                input_length=trace.input_length,
+                output_length=trace.output_length,
+            )
+            for trace in traces
+        ],
+        key=lambda x: x.timestamp,
+    )
+
+
+async def get_model_name(client: OpenAI) -> str:
+    async for model in client.models.list():
+        return model.id
+    raise ValueError("No models available")

zonos2/benchmark/perf.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, Tuple
+
+from zonos2.utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def perf_cuda(
+    f: Callable[[], Any],
+    *,
+    init_stream: bool = True,
+    repetitions: int = 10,
+    cuda_graph_repetitions: int | None = 10,
+) -> float:
+    import torch
+
+    assert repetitions > 0
+    tic = torch.cuda.Event(enable_timing=True)
+    toc = torch.cuda.Event(enable_timing=True)
+    stream = torch.cuda.Stream()
+    torch.cuda.synchronize()
+    if init_stream:
+        stream = torch.cuda.Stream()
+    else:
+        stream = torch.cuda.current_stream()
+
+    with torch.cuda.stream(stream):
+        f()
+        if N := cuda_graph_repetitions:
+            g = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(g):
+                for _ in range(N):
+                    f()
+            replay = g.replay
+            del g
+        else:
+            replay = f
+            N = 1
+
+        torch.cuda.synchronize()
+
+        replay()
+        tic.record()
+        for _ in range(repetitions):
+            replay()
+        toc.record()
+        toc.synchronize()
+        dur = tic.elapsed_time(toc)
+        return dur / (N * repetitions)
+
+
+def compare_memory_kernel_perf(
+    *,
+    baseline: Callable[[], Any],
+    our_impl: Callable[[], Any],
+    memory_footprint: int,  # in bytes
+    description: str = " ",
+    extra_kwargs: Dict[str, Any] | None = None,
+    need_latency: bool = True,
+) -> Tuple[float, float]:
+    extra_kwargs = extra_kwargs or {}
+
+    dur = perf_cuda(baseline, **extra_kwargs)
+    bandwidth_0 = memory_footprint / (dur * 1e6)  # GB/s
+    latency_msg = f"{dur:8.3f} ms | " if need_latency else ""
+    message_0 = f"Baseline: {latency_msg}{bandwidth_0:8.3f} GB/s"
+
+    dur = perf_cuda(our_impl, **extra_kwargs)
+    bandwidth_1 = memory_footprint / (dur * 1e6)  # GB/s
+    latency_msg = f"{dur:8.3f} ms | " if need_latency else ""
+    logger.info(f"{description}{message_0} | Our Impl: {latency_msg}{bandwidth_1:8.3f} GB/s")
+    return bandwidth_0, bandwidth_1

zonos2/core.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, List, Literal
+
+import torch
+
+if TYPE_CHECKING:
+    from zonos2.attention import BaseAttnBackend, BaseAttnMetadata
+    from zonos2.kvcache import BaseCacheHandle
+
+
+@dataclass
+class Context:
+    page_size: int
+    attn_backend: BaseAttnBackend
+    _batch: TTSBatch | None = field(default=None, init=False)
+
+    @property
+    def batch(self) -> TTSBatch:
+        assert self._batch is not None, "No active batch in context"
+        return self._batch
+
+    @contextmanager
+    def forward_batch(self, batch: TTSBatch):
+        assert self._batch is None, "Nested forward_batch is not allowed"
+        try:
+            self._batch = batch
+            yield
+        finally:
+            self._batch = None
+
+
+_GLOBAL_CTX: Context | None = None
+
+
+def set_global_ctx(ctx: Context):
+    global _GLOBAL_CTX
+    assert _GLOBAL_CTX is None, "Global context is already set"
+    _GLOBAL_CTX = ctx
+
+
+def get_global_ctx() -> Context:
+    assert _GLOBAL_CTX is not None, "Global context is not set"
+    return _GLOBAL_CTX
+
+
+# =============================================================================
+# TTS-specific data structures
+# =============================================================================
+
+
+@dataclass
+class TTSSamplingParams:
+    """Sampling parameters for TTS generation."""
+
+    temperature: float = 1.15
+    topk: int = 106
+    top_p: float = 0.0
+    min_p: float = 0.18
+    max_tokens: int = 1024
+    ignore_eos: bool = False
+    repetition_window: int = 50
+    repetition_penalty: float = 1.2
+    repetition_codebooks: int = 8
+    seed: int | None = None
+
+
+@dataclass(eq=False)
+class TTSReq:
+    """Request class for TTS generation with 2D token format.
+
+    Tokens are in unpacked format: [cb0, cb1, ..., cb8, text_token] per frame.
+    """
+
+    input_ids: torch.Tensor  # 2D CPU tensor (seq_len, frame_width)
+    table_idx: int
+    cached_len: int
+    output_len: int
+    uid: int
+    sampling_params: TTSSamplingParams
+    cache_handle: BaseCacheHandle
+    n_codebooks: int = 9
+    eoa_id: int = 1024
+    eos_frame: int = -1  # Aligned frame where EOS first appeared (-1 = not seen)
+    eos_countdown: int = -1  # Steps remaining after EOS (-1 = not in countdown)
+    total_generated: int = 0  # Total frames generated (for logging)
+    rng: torch.Generator | None = None  # Per-request RNG for deterministic sampling
+    speaker_embedding: torch.Tensor | None = None  # 1D CPU float32 tensor
+    speaker_token_position: int = -1  # Injection position within the prompt sequence
+
+    def __post_init__(self) -> None:
+        assert self.input_ids.is_cpu
+        assert self.input_ids.dim() == 2, "TTS input_ids must be 2D (seq_len, frame_width)"
+        self.device_len = len(self.input_ids)
+        self.max_device_len = len(self.input_ids) + self.output_len
+        assert 0 <= self.cached_len < self.device_len <= self.max_device_len
+
+        if self.speaker_embedding is not None:
+            emb = self.speaker_embedding
+            if emb.dim() == 2 and emb.shape[0] == 1:
+                emb = emb.squeeze(0)
+            if emb.dim() != 1:
+                raise ValueError(
+                    f"speaker_embedding must be 1D or (1, D), got shape {tuple(emb.shape)}"
+                )
+            self.speaker_embedding = emb.to(dtype=torch.float32, device="cpu")
+
+        if self.speaker_token_position < 0:
+            # Training convention: reserved speaker slot is at prompt position 0.
+            self.speaker_token_position = 0
+        if self.speaker_token_position >= self.device_len:
+            self.speaker_token_position = 0
+
+    @property
+    def frame_width(self) -> int:
+        """Number of elements per frame (n_codebooks + extras)."""
+        return self.input_ids.shape[-1]
+
+    @property
+    def remain_len(self) -> int:
+        return self.max_device_len - self.device_len
+
+    @property
+    def extend_len(self) -> int:
+        return self.device_len - self.cached_len
+
+    @property
+    def num_completion_tokens(self) -> int:
+        """Number of generated tokens (frames)."""
+        return self.device_len - self.cached_len
+
+    def complete_one(self) -> None:
+        self.cached_len = self.device_len
+        self.device_len += 1
+        self.total_generated += 1
+
+    def append_host(self, next_token: torch.Tensor) -> None:
+        """Append a single frame (unpacked token) to input_ids."""
+        assert next_token.dim() == 1, "next_token must be 1D (frame_width,)"
+        self.input_ids = torch.cat([self.input_ids, next_token.unsqueeze(0)], dim=0)
+
+    def can_decode(self) -> bool:
+        return self.remain_len > 0 and self.eos_countdown != 0
+
+    def check_eos(self, audio_codes: List[int]) -> bool:
+        """Check for EOS and update countdown state.
+
+        Args:
+            audio_codes: List of audio codebook values for one frame
+
+        Returns:
+            True if sequence is finished (countdown reached 0)
+        """
+        if self.sampling_params.ignore_eos:
+            return False
+
+        # Match Zonos2 reference inference: any sampled EOA codebook starts the
+        # delayed stop countdown. The aligned frame is shifted back by the
+        # highest EOA codebook index and clamped at zero.
+        # Use total_generated because this request only sees one decode frame at a time.
+        if self.eos_frame < 0:
+            step = self.total_generated - 1
+            eos_cols = [c == self.eoa_id for c in audio_codes[: self.n_codebooks]]
+            if any(eos_cols):
+                # First EOS: compute aligned frame
+                max_eos_cb = max(i for i, is_eos in enumerate(eos_cols) if is_eos)
+                self.eos_frame = max(0, step - max_eos_cb)
+                self.eos_countdown = self.n_codebooks + 1
+
+        # Decrement countdown
+        if self.eos_countdown > 0:
+            self.eos_countdown -= 1
+            if self.eos_countdown == 0:
+                return True
+
+        return False
+
+    def __repr__(self) -> str:
+        return (
+            f"{type(self).__name__}(table_idx={self.table_idx}, "
+            f"cached_len={self.cached_len}, device_len={self.device_len}, "
+            f"max_device_len={self.max_device_len}, eos_frame={self.eos_frame})"
+        )
+
+
+@dataclass
+class TTSBatch:
+    """Batch of TTS requests with 2D token format."""
+
+    reqs: List[TTSReq]
+    phase: Literal["prefill", "decode"]
+    # these fields should be set by scheduler
+    input_ids: torch.Tensor = field(init=False)  # (total_tokens, frame_width)
+    out_loc: torch.Tensor = field(init=False)
+    padded_reqs: List[TTSReq] = field(init=False)
+    # this field should be set by attention backend
+    attn_metadata: BaseAttnMetadata = field(init=False)
+    # Optional per-batch speaker conditioning data (set by TTS scheduler).
+    speaker_emb_values: torch.Tensor | None = field(default=None, init=False)
+    speaker_token_positions: torch.Tensor | None = field(default=None, init=False)
+
+    @property
+    def is_prefill(self) -> bool:
+        return self.phase == "prefill"
+
+    @property
+    def is_decode(self) -> bool:
+        return self.phase == "decode"
+
+    @property
+    def size(self) -> int:
+        return len(self.reqs)
+
+    @property
+    def padded_size(self) -> int:
+        return len(self.padded_reqs)

zonos2/distributed/__init__.py

@@ -0,0 +1,12 @@
+from .impl import DistributedCommunicator, destroy_distributed, enable_pynccl_distributed
+from .info import DistributedInfo, get_tp_info, set_tp_info, try_get_tp_info
+
+__all__ = [
+    "DistributedInfo",
+    "get_tp_info",
+    "set_tp_info",
+    "enable_pynccl_distributed",
+    "DistributedCommunicator",
+    "try_get_tp_info",
+    "destroy_distributed",
+]

zonos2/distributed/impl.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List
+
+import torch
+import torch.distributed as dist
+
+if TYPE_CHECKING:
+    from zonos2.distributed import DistributedInfo
+    from zonos2.kernel import PyNCCLCommunicator
+
+
+@dataclass
+class DistributedImpl(ABC):
+    @abstractmethod
+    def all_reduce(self, x: torch.Tensor) -> torch.Tensor: ...
+
+    @abstractmethod
+    def all_gather(self, x: torch.Tensor) -> torch.Tensor: ...
+
+
+@dataclass
+class TorchDistributedImpl(DistributedImpl):
+    def all_reduce(self, x: torch.Tensor) -> torch.Tensor:
+        tp_size = dist.get_world_size()
+        if tp_size == 1:
+            return x
+        dist.all_reduce(x, op=dist.ReduceOp.SUM)
+        return x
+
+    def all_gather(self, x: torch.Tensor) -> torch.Tensor:
+        tp_size = dist.get_world_size()
+        if tp_size == 1:
+            return x
+        shape = list(x.shape)
+        shape[0] = shape[0] * tp_size
+        out = torch.empty(shape, dtype=x.dtype, device=x.device)
+        dist.all_gather_into_tensor(out, x)
+        return out
+
+
+@dataclass
+class PyNCCLDistributedImpl(DistributedImpl):
+    comm: PyNCCLCommunicator
+
+    def all_reduce(self, x: torch.Tensor) -> torch.Tensor:
+        self.comm.all_reduce(x, "sum")
+        return x
+
+    def all_gather(self, x: torch.Tensor) -> torch.Tensor:
+        from .info import get_tp_info
+
+        world_size = get_tp_info().size
+        output_shape = list(x.shape)
+        output_shape[0] *= world_size
+        result = x.new_empty(output_shape)
+        self.comm.all_gather(result, x)
+        return result
+
+
+class DistributedCommunicator:
+    plugins: List[DistributedImpl] = [TorchDistributedImpl()]
+
+    def all_reduce(self, x: torch.Tensor) -> torch.Tensor:
+        return self.plugins[-1].all_reduce(x)
+
+    def all_gather(self, x: torch.Tensor) -> torch.Tensor:
+        return self.plugins[-1].all_gather(x)
+
+
+def enable_pynccl_distributed(
+    tp_info: DistributedInfo, tp_cpu_group: torch.distributed.ProcessGroup, max_bytes: int
+) -> None:
+    """
+    Enable PyNCCL-based distributed communication for tensor parallelism.
+    """
+    if tp_info.size == 1:
+        return
+    from zonos2.kernel import init_pynccl
+
+    comm = init_pynccl(
+        tp_rank=tp_info.rank,
+        tp_size=tp_info.size,
+        tp_cpu_group=tp_cpu_group,
+        max_size_bytes=max_bytes,
+    )
+
+    DistributedCommunicator.plugins.append(PyNCCLDistributedImpl(comm))
+
+
+def destroy_distributed() -> None:
+    """
+    Destroy all the distributed communication plugins.
+    """
+    DistributedCommunicator.plugins = []

zonos2/distributed/info.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class DistributedInfo:  # should not export from here
+    rank: int
+    size: int
+
+    def __post_init__(self):
+        assert 0 <= self.rank < self.size
+
+    def is_primary(self) -> bool:
+        return self.rank == 0
+
+
+_TP_INFO: DistributedInfo | None = None
+
+
+def set_tp_info(rank: int, size: int) -> None:
+    global _TP_INFO
+    if _TP_INFO is not None:
+        raise RuntimeError("TP info has been set")
+    _TP_INFO = DistributedInfo(rank, size)
+
+
+def get_tp_info() -> DistributedInfo:
+    if _TP_INFO is None:
+        raise RuntimeError("TP info has not been set")
+    return _TP_INFO
+
+
+def try_get_tp_info() -> DistributedInfo | None:
+    return _TP_INFO
+
+
+__all__ = ["DistributedInfo", "set_tp_info", "get_tp_info", "try_get_tp_info"]

zonos2/engine/__init__.py

@@ -0,0 +1,4 @@
+from .config import EngineConfig
+from .engine import Engine
+
+__all__ = ["Engine", "EngineConfig"]

zonos2/engine/config.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING, List
+
+import torch
+from zonos2.distributed import DistributedInfo
+from zonos2.utils import cached_load_checkpoint_config
+
+if TYPE_CHECKING:
+    from zonos2.models import ModelConfig
+
+
+@dataclass(frozen=True)
+class EngineConfig:
+    model_path: str
+    tp_info: DistributedInfo
+    dtype: torch.dtype
+    max_running_req: int = 256
+    attention_backend: str = "auto"
+    moe_backend: str = "fused_moe"
+    cuda_graph_bs: List[int] | None = None
+    cuda_graph_max_bs: int | None = None
+    page_size: int = 1
+    memory_ratio: float = 0.9
+    distributed_timeout: float = 60.0
+    use_dummy_weight: bool = False
+    use_pynccl: bool = True
+    max_seq_len_override: int | None = None
+    num_page_override: int | None = None  # if not None, will override the number of pages
+
+    @cached_property
+    def checkpoint_config(self):
+        return cached_load_checkpoint_config(self.model_path)
+
+    @cached_property
+    def model_config(self) -> ModelConfig:
+        from zonos2.models import ModelConfig
+
+        return ModelConfig.from_checkpoint_config(self.checkpoint_config)
+
+    @property
+    def max_seq_len(self) -> int:
+        if self.max_seq_len_override is not None:
+            return self.max_seq_len_override
+        return self.model_config.rotary_config.max_position
+
+    @property
+    def max_forward_len(self) -> int:
+        return self.max_seq_len
+
+    @property
+    def distributed_addr(self) -> str:
+        return "tcp://127.0.0.1:23333"

zonos2/engine/engine.py

@@ -0,0 +1,336 @@
+from __future__ import annotations
+
+import logging
+from datetime import timedelta
+from typing import Dict, Tuple
+
+import torch
+from zonos2.attention import create_attention_backend
+from zonos2.core import (
+    Context,
+    TTSBatch,
+    TTSReq,
+    TTSSamplingParams,
+    set_global_ctx,
+)
+from zonos2.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info
+from zonos2.kvcache import create_kvcache
+from zonos2.layers import set_rope_device
+from zonos2.models import create_model, load_checkpoint_weight
+from zonos2.utils import divide_even, init_logger, torch_dtype
+
+from .config import EngineConfig
+from .graph import GraphRunner, get_free_memory, mem_GB
+
+logger = init_logger(__name__)
+
+
+def create_page_table(shape: Tuple[int, int], device: torch.device) -> torch.Tensor:
+    return torch.zeros(shape, dtype=torch.int32, device=device)
+
+
+def _align_up_32(num: int) -> int:
+    return (num + 31) // 32 * 32
+
+
+def _get_frame_width(model_config) -> int:
+    """Get frame width for multi-codebook models.
+
+    Frame width = n_codebooks + (1 if text_vocab).
+    """
+    width = model_config.n_codebooks
+    if model_config.text_vocab is not None:
+        width += 1
+    return width
+
+
+class Engine:
+    def __init__(self, config: EngineConfig):
+        self.model_config = config.model_config
+        set_tp_info(rank=config.tp_info.rank, size=config.tp_info.size)
+
+        self.device = torch.device("cuda")
+        # Importing the model eagerly pulls in sgl_kernel, which creates a CUDA
+        # context on the current device. The scheduler entry point binds this
+        # process to its rank's GPU before that import, so an existing context is
+        # expected; only one on the wrong device indicates a real ordering bug.
+        if config.tp_info.size > 1 and (
+            torch.cuda.is_initialized()
+            and torch.cuda.current_device() != config.tp_info.rank
+        ):
+            raise RuntimeError(
+                f"CUDA was initialized on cuda:{torch.cuda.current_device()} before "
+                f"the engine could bind TP rank {config.tp_info.rank}; "
+                "set the device before importing model code."
+            )
+        if config.tp_info.size > 1:
+            torch.cuda.set_device(config.tp_info.rank)
+        self.stream = torch.cuda.Stream()
+        torch.cuda.set_stream(self.stream)
+        self.dtype = config.dtype
+
+        self.tp_cpu_group = self._init_communication(config)
+        init_free_memory = self._sync_get_memory()[1]
+        logger.info_rank0(f"Free memory before loading model: {mem_GB(init_free_memory)}")
+
+        # load model and determine number of pages
+        set_rope_device(self.device)
+        with torch.device("meta"), torch_dtype(config.dtype):
+            self.model = create_model(config)
+
+        # Log expected shapes before loading (for debugging dimension mismatches)
+        logger.info_rank0(f"Model config: num_kv_heads={config.model_config.num_kv_heads}, head_dim={config.model_config.head_dim}")
+        expected_kv_dim = config.model_config.num_kv_heads * config.model_config.head_dim
+        logger.info_rank0(f"Expected kv_dim per TP rank: {expected_kv_dim}")
+
+        state_dict = self._load_weight_state_dict(config)
+        self._check_speaker_lda_weights(state_dict)
+
+        # Debug: Check for temp keys in state_dict
+        if logger.isEnabledFor(logging.DEBUG):
+            temp_keys = [k for k in state_dict.keys() if '.temp' in k]
+            logger.debug("Found %d temp keys in state_dict: %s", len(temp_keys), temp_keys[:5])
+            if temp_keys:
+                first_temp = state_dict[temp_keys[0]]
+                logger.debug(
+                    "First temp tensor: shape=%s, dtype=%s, values=%s",
+                    first_temp.shape, first_temp.dtype, first_temp.flatten()[:4].tolist(),
+                )
+
+        # Validate wkv weight shapes match config (catch n_kv_heads mismatch early)
+        for key, tensor in state_dict.items():
+            if ".wkv." in key and "weight" in key:
+                # wkv.weight should be (2, kv_dim, hidden_size) or (2*kv_dim, hidden_size)
+                if tensor.dim() == 3:
+                    actual_kv_dim = tensor.shape[1]
+                else:
+                    actual_kv_dim = tensor.shape[0] // 2
+                if actual_kv_dim != expected_kv_dim:
+                    inferred_kv_heads = actual_kv_dim // config.model_config.head_dim
+                    raise ValueError(
+                        f"KV dimension mismatch for {key}: checkpoint has kv_dim={actual_kv_dim} "
+                        f"(implying n_kv_heads={inferred_kv_heads}), but config expects "
+                        f"kv_dim={expected_kv_dim} (n_kv_heads={config.model_config.num_kv_heads}). "
+                        f"Please ensure your checkpoint's params.json has the correct n_kv_heads value "
+                        f"matching the trained model weights."
+                    )
+                break  # Only need to check one layer
+
+        self.model.load_state_dict(state_dict)
+
+        # Debug: print attention temp values to verify they loaded correctly
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.debug("Model type: %s", type(self.model).__name__)
+            logger.debug("Model has layers: %s", hasattr(self.model, 'layers'))
+            if hasattr(self.model, 'layers'):
+                logger.debug("layers type: %s", type(self.model.layers))
+                if hasattr(self.model.layers, 'op_list'):
+                    logger.debug("num layers: %d", len(self.model.layers.op_list))
+                    for i, layer in enumerate(self.model.layers.op_list[:2]):
+                        logger.debug("Layer %d type: %s", i, type(layer).__name__)
+                        if hasattr(layer, 'attention'):
+                            attn = layer.attention
+                            logger.debug("Layer %d attention type: %s", i, type(attn).__name__)
+                            logger.debug("Layer %d has_qk_norm: %s", i, getattr(attn, 'has_qk_norm', 'N/A'))
+                            if hasattr(attn, 'temp'):
+                                temp = attn.temp
+                                if temp is not None:
+                                    logger.debug("Layer %d attention.temp: shape=%s, values=%s", i, temp.shape, temp.flatten()[:8].tolist())
+                                else:
+                                    logger.debug("Layer %d attention.temp is None", i)
+
+        self.num_pages = self.dummy_page = self._determine_num_pages(init_free_memory, config)
+        self.kv_cache = create_kvcache(
+            model_config=config.model_config,
+            num_pages=self.num_pages + 1,  # +1 for dummy page
+            device=self.device,
+            dtype=self.dtype,
+        )
+        # NOTE: make page table 128 aligned (32 * sizeof(int32) == 128 bytes)
+        self.max_seq_len = _align_up_32(min(config.max_seq_len, self.num_pages))
+        self.page_table = create_page_table(  # + 1 for dummy request
+            (config.max_running_req + 1, self.max_seq_len),
+            device=self.device,
+        )
+        self.attn_backend = create_attention_backend(
+            config.attention_backend,
+            config.model_config,
+            self.kv_cache,
+            self.page_table,
+        )
+        self.ctx = Context(page_size=1, attn_backend=self.attn_backend)
+        set_global_ctx(self.ctx)
+
+        post_free_memory = self._sync_get_memory()[0]
+        logger.info_rank0(f"Free memory after initialization: {mem_GB(post_free_memory)}")
+
+        frame_width = _get_frame_width(self.model_config)
+        self.dummy_req = TTSReq(
+            input_ids=torch.zeros((1, frame_width), dtype=torch.int32, device="cpu"),
+            table_idx=config.max_running_req,
+            cached_len=0,
+            output_len=1,
+            uid=-1,
+            sampling_params=TTSSamplingParams(),
+            cache_handle=None,  # type: ignore
+            n_codebooks=self.model_config.n_codebooks,
+            eoa_id=self.model_config.eoa_id,
+        )
+        self.page_table[self.dummy_req.table_idx].fill_(self.dummy_page)
+
+        self.graph_runner = GraphRunner(
+            stream=self.stream,
+            device=self.device,
+            model=self.model,
+            attn_backend=self.attn_backend,
+            cuda_graph_bs=config.cuda_graph_bs,
+            cuda_graph_max_bs=config.cuda_graph_max_bs,
+            free_memory=init_free_memory,
+            max_seq_len=self.max_seq_len,
+            vocab_size=self.model_config.codebook_size + 2,
+            dummy_req=self.dummy_req,
+        )
+
+    def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup:
+        if config.tp_info.size == 1 or config.use_pynccl:
+            torch.distributed.init_process_group(
+                backend="gloo",
+                rank=config.tp_info.rank,
+                world_size=config.tp_info.size,
+                timeout=timedelta(seconds=config.distributed_timeout),
+                init_method=config.distributed_addr,
+            )
+            tp_cpu_group = torch.distributed.group.WORLD
+            assert tp_cpu_group is not None
+            max_bytes = (
+                config.max_forward_len * config.model_config.hidden_size * self.dtype.itemsize
+            )
+            enable_pynccl_distributed(config.tp_info, tp_cpu_group, max_bytes)
+        else:
+            torch.distributed.init_process_group(
+                backend="nccl",
+                rank=config.tp_info.rank,
+                world_size=config.tp_info.size,
+                timeout=timedelta(seconds=config.distributed_timeout),
+                init_method=config.distributed_addr,
+            )
+            tp_cpu_group = torch.distributed.new_group(backend="gloo")
+            assert tp_cpu_group is not None
+        return tp_cpu_group
+
+    def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tensor]:
+        if config.use_dummy_weight:
+            return {
+                k: torch.randn_like(v, device=self.device)
+                for k, v in self.model.state_dict().items()
+            }
+        else:
+            return {
+                k: v.to(self.dtype)
+                for k, v in load_checkpoint_weight(config.model_path, self.device).items()
+            }
+
+    def _check_speaker_lda_weights(self, state_dict: Dict[str, torch.Tensor]) -> None:
+        lda_keys = [key for key in state_dict if key.startswith("speaker_lda_projection.")]
+        if getattr(self.model, "speaker_lda_projection", None) is None:
+            for key in lda_keys:
+                state_dict.pop(key)
+            if lda_keys:
+                logger.warning_rank0(
+                    "Dropped %d speaker_lda_projection weights because speaker LDA is not enabled.",
+                    len(lda_keys),
+                )
+            return
+
+        missing = {
+            "speaker_lda_projection.weight",
+            "speaker_lda_projection.bias",
+        }.difference(lda_keys)
+        if missing:
+            raise ValueError(
+                "speaker_lda_dim is configured but the checkpoint is missing "
+                f"{sorted(missing)}. Use a checkpoint with the LDA projection merged in."
+            )
+
+    def _determine_num_pages(self, old_free_memory: int, config: EngineConfig) -> int:
+        new_free_memory = self._sync_get_memory()[1]
+        cache_per_page = (
+            2  # key + value
+            * self.model_config.head_dim
+            * divide_even(self.model_config.num_kv_heads, config.tp_info.size)
+            * config.page_size
+            * self.dtype.itemsize
+            * self.model_config.num_layers
+        )
+        num_pages = config.num_page_override
+        if num_pages is None:
+            model_memory = old_free_memory - new_free_memory
+            available_memory = int(config.memory_ratio * old_free_memory) - model_memory
+            num_pages = available_memory // cache_per_page
+
+        assert num_pages > 1, "Not enough memory for KV cache, try reducing --num-tokens"
+        real_kv_size = num_pages * cache_per_page
+        logger.info(f"Allocating {num_pages} pages for KV cache, K + V = {mem_GB(real_kv_size)}")
+        return num_pages
+
+    def _sync_get_memory(self) -> Tuple[int, int]:
+        """Get the min and max free memory across TP ranks."""
+        torch.cuda.synchronize(self.device)
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats(self.device)
+        free_memory = get_free_memory(self.device)
+        free_mem_tensor = torch.tensor([free_memory, -free_memory], device="cpu", dtype=torch.int64)
+        torch.distributed.all_reduce(
+            free_mem_tensor, op=torch.distributed.ReduceOp.MIN, group=self.tp_cpu_group
+        )
+        min_free_memory = int(free_mem_tensor[0].item())
+        max_free_memory = -int(free_mem_tensor[1].item())
+        if max_free_memory - min_free_memory > 2 * 1024 * 1024 * 1024:
+            logger.error(
+                f"Memory across TP ranks are imbalanced:"
+                f" min {mem_GB(min_free_memory)}, max {mem_GB(max_free_memory)}"
+            )
+            raise RuntimeError("Memory across TP ranks are imbalanced")
+
+        return min_free_memory, max_free_memory
+
+    def forward_batch_tts(self, batch: TTSBatch) -> torch.Tensor:
+        """Forward pass for TTS batch, returning multi-codebook logits.
+
+        Args:
+            batch: TTS batch with 2D input tokens
+
+        Returns:
+            Logits tensor of shape (batch_size, n_codebooks, vocab_size)
+        """
+        assert torch.cuda.current_stream() == self.stream
+        with self.ctx.forward_batch(batch):
+            # For TTS, the model outputs multi-codebook logits directly
+            # The zonos2 model has a MultiParallelLMHead that outputs
+            # (batch, seq, n_codebooks, vocab_size)
+            if self.graph_runner.can_use_cuda_graph(batch):
+                logits = self.graph_runner.replay(batch)
+            else:
+                logits = self.model.forward()
+
+        # logits shape: (batch, n_codebooks, vocab) for decode
+        # or (total_tokens, n_codebooks, vocab) for prefill
+        # We only care about the last token per sequence
+        if batch.is_decode:
+            return logits[: batch.size]
+        else:
+            # For prefill, get the LAST token logits for each sequence
+            # logits shape: (total_tokens, n_codebooks, vocab)
+            # Need to extract logits at the last position of each sequence
+            last_indices = []
+            cumsum = 0
+            for req in batch.reqs:
+                cumsum += req.extend_len  # extend_len = device_len - cached_len
+                last_indices.append(cumsum - 1)
+            last_indices = torch.tensor(last_indices, device=logits.device)
+            return logits[last_indices]
+
+    def shutdown(self) -> None:
+        self.graph_runner.destroy_cuda_graphs()
+        torch.distributed.destroy_process_group()
+        destroy_distributed()

zonos2/engine/graph.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+import gc
+from typing import TYPE_CHECKING, Dict, List
+
+import torch
+from tqdm import tqdm
+from zonos2.core import TTSBatch, TTSReq, get_global_ctx
+from zonos2.distributed import get_tp_info
+from zonos2.utils import init_logger
+
+if TYPE_CHECKING:
+    from zonos2.attention import BaseAttnBackend
+    from zonos2.models import Zonos2ForCausalLM
+
+logger = init_logger(__name__)
+
+
+def _determine_cuda_graph_bs(
+    cuda_graph_bs: List[int] | None,
+    cuda_graph_max_bs: int | None,
+    free_memory: int,
+) -> List[int]:
+    if cuda_graph_bs is not None:
+        return cuda_graph_bs
+
+    free_memory_gb = free_memory / (1 << 30)
+    if cuda_graph_max_bs is None:
+        if free_memory_gb > 80:  # H200
+            cuda_graph_max_bs = 256
+        else:
+            cuda_graph_max_bs = 160
+
+    if cuda_graph_max_bs < 1:
+        return []
+
+    return [1, 2, 4] + list(range(8, cuda_graph_max_bs + 1, 8))
+
+
+def mem_GB(size: int) -> str:
+    return f"{size / (1024**3):.2f} GiB"
+
+
+def get_free_memory(device: torch.device) -> int:
+    return torch.cuda.mem_get_info(device)[0]
+
+
+class GraphRunner:
+    def __init__(
+        self,
+        stream: torch.cuda.Stream,
+        device: torch.device,
+        model: Zonos2ForCausalLM,
+        attn_backend: BaseAttnBackend,
+        cuda_graph_bs: List[int] | None,
+        cuda_graph_max_bs: int | None,
+        free_memory: int,
+        max_seq_len: int,
+        vocab_size: int,
+        dummy_req: TTSReq,
+    ) -> None:
+        cuda_graph_bs = _determine_cuda_graph_bs(
+            cuda_graph_bs=cuda_graph_bs,
+            cuda_graph_max_bs=cuda_graph_max_bs,
+            free_memory=free_memory,
+        )
+        self.attn_backend = attn_backend
+        if not cuda_graph_bs:
+            logger.info_rank0("CUDA graph is disabled.")
+            self.max_graph_bs = 0
+            self.graph_bs_list = []
+            self.dummy_req = dummy_req
+            self.stream = stream
+            self.device = device
+            self.graph_map = {}
+            return
+
+        self.max_graph_bs = max(cuda_graph_bs)
+        self.graph_bs_list = sorted(cuda_graph_bs)
+        self.dummy_req = dummy_req
+        self.stream = stream
+        self.device = device
+        self.graph_map = self._capture_graphs(max_seq_len, vocab_size, model)
+
+    def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: Zonos2ForCausalLM):
+        graph_map: Dict[int, torch.cuda.CUDAGraph] = {}
+        if self.max_graph_bs == 0:
+            logger.info_rank0("CUDA graph is disabled.")
+            return graph_map
+
+        self.logits = torch.empty(
+            (self.max_graph_bs, self.dummy_req.n_codebooks, vocab_size),
+            dtype=torch.float32,
+            device=self.device,
+        )
+        frame_width = self.dummy_req.frame_width
+        self.attn_backend.init_capture_graph(
+            max_seq_len=max_seq_len, bs_list=self.graph_bs_list, frame_width=frame_width
+        )
+
+        torch.cuda.synchronize(self.device)
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats(self.device)
+
+        logger.info_rank0(f"Start capturing CUDA graphs with sizes: {self.graph_bs_list}")
+        free_memory = get_free_memory(self.device)
+        logger.info_rank0(f"Free GPU memory before capturing CUDA graphs: {mem_GB(free_memory)}")
+
+        pbar = tqdm(
+            sorted(self.graph_bs_list, reverse=True),
+            desc="Preparing for capturing CUDA graphs...",
+            unit="batch",
+            disable=not get_tp_info().is_primary(),  # disable for non-primary ranks
+        )
+        pool = None
+        for bs in pbar:
+            free_memory = get_free_memory(self.device)
+            pbar.desc = f"Capturing graphs: bs = {bs:<3} | avail_mem = {mem_GB(free_memory)}"
+            pbar.refresh()
+            graph = torch.cuda.CUDAGraph()
+            batch = TTSBatch(reqs=[self.dummy_req] * bs, phase="decode")
+            self.attn_backend.prepare_for_capture(batch)
+            with get_global_ctx().forward_batch(batch):
+                self.logits[:bs] = model.forward()
+                with torch.cuda.graph(graph, pool=pool, stream=self.stream):
+                    self.logits[:bs] = model.forward()
+            if pool is None:
+                pool = graph.pool()
+            graph_map[bs] = graph
+
+        free_memory = get_free_memory(self.device)
+        logger.info_rank0(f"Free GPU memory after capturing CUDA graphs: {mem_GB(free_memory)}")
+        return graph_map
+
+    def can_use_cuda_graph(self, batch: TTSBatch) -> bool:
+        return batch.is_decode and batch.size <= self.max_graph_bs
+
+    def replay(self, batch: TTSBatch) -> torch.Tensor:
+        assert self.can_use_cuda_graph(batch)
+        g = self.graph_map[batch.padded_size]
+        self.attn_backend.prepare_for_replay(batch)
+        g.replay()
+        return self.logits[: batch.size]
+
+    def pad_batch(self, batch: TTSBatch) -> int:
+        padded_size = (  # choose the first available batch size
+            next(bs for bs in self.graph_bs_list if bs >= batch.size)
+            if self.can_use_cuda_graph(batch)
+            else batch.size
+        )
+        batch.padded_reqs = batch.reqs + [self.dummy_req] * (padded_size - batch.size)
+        return batch.padded_size - batch.size
+
+    # NOTE: This must be called before freeing NCCL resources to prevent program hang
+    def destroy_cuda_graphs(self) -> None:
+        del self.graph_map
+        gc.collect()

zonos2/engine/sample.py

@@ -0,0 +1,170 @@
+"""TTS sampler for multi-codebook audio generation."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, List
+
+import torch
+from zonos2.tts.sampler import sample_tts
+
+if TYPE_CHECKING:
+    from zonos2.core import TTSBatch
+
+
+def make_device_tensor(
+    data: List, dtype: torch.dtype, device: torch.device
+) -> torch.Tensor:
+    """Create a tensor on device from a list."""
+    return torch.tensor(data, dtype=dtype, pin_memory=True).to(
+        device, non_blocking=True
+    )
+
+
+@dataclass
+class TTSBatchSamplingArgs:
+    """Sampling arguments for a TTS batch."""
+
+    temperatures: torch.Tensor
+    top_ks: torch.Tensor
+    top_ps: torch.Tensor
+    min_ps: torch.Tensor
+    text_vocab: int
+    repetition_token_ids: torch.Tensor | None = None
+    repetition_penalties: torch.Tensor | None = None
+    generators: list[torch.Generator | None] = field(default_factory=list)
+
+
+@dataclass
+class TTSSampler:
+    """Sampler for multi-codebook TTS output."""
+
+    device: torch.device
+    n_codebooks: int
+    codebook_size: int
+    text_vocab: int
+
+    def prepare(
+        self, batch: TTSBatch, token_pool: torch.Tensor | None = None
+    ) -> TTSBatchSamplingArgs:
+        """Prepare sampling arguments from a TTS batch.
+
+        Args:
+            batch: TTS batch with requests
+            token_pool: Optional device-side token pool for recent generated history
+
+        Returns:
+            TTSBatchSamplingArgs with per-sequence sampling parameters
+        """
+        params = [r.sampling_params for r in batch.reqs]
+
+        MIN_T = 1e-6
+        temps = [max(p.temperature, MIN_T) for p in params]
+        top_ks = [p.topk if p.topk >= 1 else self.codebook_size for p in params]
+        top_ps = [min(max(p.top_p, 0.0), 1.0) for p in params]
+        min_ps = [max(p.min_p, 0.0) for p in params]
+        repetition_windows = [max(int(p.repetition_window), 0) for p in params]
+        repetition_penalties = [max(float(p.repetition_penalty), 1.0) for p in params]
+        repetition_codebooks = [
+            self.n_codebooks
+            if int(p.repetition_codebooks) < 0
+            else min(max(int(p.repetition_codebooks), 0), self.n_codebooks)
+            for p in params
+        ]
+        generators = [r.rng for r in batch.reqs]
+        repetition_token_ids = None
+        repetition_penalties_tensor = None
+
+        active_windows = [
+            min(window, r.total_generated, r.device_len)
+            if window > 0 and penalty > 1.0 and codebooks > 0
+            else 0
+            for r, window, penalty, codebooks in zip(
+                batch.reqs,
+                repetition_windows,
+                repetition_penalties,
+                repetition_codebooks,
+            )
+        ]
+        max_window = max(active_windows, default=0)
+        if max_window > 0:
+            repetition_token_ids = torch.full(
+                (len(batch.reqs), self.n_codebooks, max_window),
+                -1,
+                dtype=torch.int64,
+                device=self.device,
+            )
+            for i, (req, window) in enumerate(zip(batch.reqs, active_windows)):
+                if window <= 0:
+                    continue
+
+                start = req.device_len - window
+                if token_pool is None:
+                    history = req.input_ids[start : req.device_len, : self.n_codebooks].to(
+                        device=self.device, dtype=torch.int64, non_blocking=True
+                    )
+                else:
+                    history = token_pool[
+                        req.table_idx, start : req.device_len, : self.n_codebooks
+                    ].to(dtype=torch.int64)
+
+                history = history.transpose(0, 1).contiguous()
+                valid = (history >= 0) & (history < self.codebook_size)
+                valid[repetition_codebooks[i] :] = False
+                repetition_token_ids[i, :, -window:] = torch.where(
+                    valid, history, torch.full_like(history, -1)
+                )
+
+            repetition_penalties_tensor = make_device_tensor(
+                repetition_penalties, torch.float32, self.device
+            )
+
+        return TTSBatchSamplingArgs(
+            temperatures=make_device_tensor(temps, torch.float32, self.device),
+            top_ks=make_device_tensor(top_ks, torch.int32, self.device),
+            top_ps=make_device_tensor(top_ps, torch.float32, self.device),
+            min_ps=make_device_tensor(min_ps, torch.float32, self.device),
+            text_vocab=self.text_vocab,
+            repetition_token_ids=repetition_token_ids,
+            repetition_penalties=repetition_penalties_tensor,
+            generators=generators,
+        )
+
+    def sample(
+        self, logits: torch.Tensor, args: TTSBatchSamplingArgs
+    ) -> List[List[int]]:
+        """Sample from multi-codebook logits.
+
+        Args:
+            logits: Shape (B, n_codebooks, vocab_size)
+            args: Sampling arguments
+
+        Returns:
+            List of unpacked tokens [cb0, cb1, ..., cb8, text_placeholder]
+        """
+        return sample_tts(
+            logits=logits,
+            temperatures=args.temperatures,
+            top_ks=args.top_ks,
+            top_ps=args.top_ps,
+            min_ps=args.min_ps,
+            repetition_token_ids=args.repetition_token_ids,
+            repetition_penalties=args.repetition_penalties,
+            text_vocab=args.text_vocab,
+            generators=args.generators,
+        )
+
+    def sample_to_tensor(
+        self, logits: torch.Tensor, args: TTSBatchSamplingArgs
+    ) -> torch.Tensor:
+        """Sample and return as tensor instead of list.
+
+        Args:
+            logits: Shape (B, n_codebooks, vocab_size)
+            args: Sampling arguments
+
+        Returns:
+            Tensor of shape (B, n_codebooks + 1) with sampled tokens
+        """
+        tokens_list = self.sample(logits, args)
+        return torch.tensor(tokens_list, dtype=torch.int32, device=self.device)

zonos2/env.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import os
+from functools import partial
+from typing import Callable, Generic, TypeVar
+
+
+class BaseEnv:
+    def _init(self, name: str) -> None:
+        raise NotImplementedError
+
+
+T = TypeVar("T")
+
+
+class EnvVar(BaseEnv, Generic[T]):
+    def __init__(self, default_value: T, fn: Callable[[str], T]):
+        self.value = default_value
+        self.fn = fn
+        super().__init__()
+
+    def _init(self, name: str) -> None:
+        env_value = os.getenv(name)
+        if env_value is not None:
+            try:
+                self.value = self.fn(env_value)
+            except Exception:
+                pass
+
+    def __bool__(self):
+        return self.value
+
+    def __str__(self):
+        return str(self.value)
+
+
+_TO_BOOL = lambda x: x.lower() in ("1", "true", "yes")
+
+
+def _PARSE_MEM_BYTES(mem: str) -> int:
+    mem = mem.strip().upper()
+    if not mem[-1].isalpha():
+        return int(mem)
+    if mem.endswith("B"):
+        mem = mem[:-1]
+    UNIT_MAP = {"K": 1024, "M": 1024**2, "G": 1024**3}
+    return int(float(mem[:-1]) * UNIT_MAP[mem[-1]])
+
+
+ZONOS2_ENV_PREFIX = "ZONOS2_"
+EnvInt = partial(EnvVar[int], fn=int)
+EnvFloat = partial(EnvVar[float], fn=float)
+EnvBool = partial(EnvVar[bool], fn=_TO_BOOL)
+EnvOption = partial(EnvVar[bool | None], fn=_TO_BOOL, default_value=None)
+EnvMem = partial(EnvVar[int], fn=_PARSE_MEM_BYTES)
+
+
+class EnvClassSingleton:
+    _instance: EnvClassSingleton | None = None
+
+    # shell
+    SHELL_MAX_TOKENS = EnvInt(2048)
+    SHELL_TOP_K = EnvInt(-1)
+    SHELL_TOP_P = EnvFloat(1.0)
+    SHELL_TEMPERATURE = EnvFloat(0.6)
+
+    # backend runtime
+    FLASHINFER_USE_TENSOR_CORES = EnvOption()
+    DISABLE_OVERLAP_SCHEDULING = EnvBool(False)
+    PYNCCL_MAX_BUFFER_SIZE = EnvMem(1024**3)
+
+    def __new__(cls):
+        # single instance
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self) -> None:
+        for attr_name in dir(self):
+            if attr_name.startswith("_"):
+                continue
+            attr_value = getattr(self, attr_name)
+            assert isinstance(attr_value, BaseEnv)
+            attr_value._init(f"{ZONOS2_ENV_PREFIX}{attr_name}")
+
+
+ENV = EnvClassSingleton()

zonos2/kernel/__init__.py

@@ -0,0 +1,14 @@
+from .index import indexing
+from .pynccl import PyNCCLCommunicator, init_pynccl
+from .radix import fast_compare_key
+from .store import store_cache
+from .tensor import test_tensor
+
+__all__ = [
+    "indexing",
+    "fast_compare_key",
+    "store_cache",
+    "test_tensor",
+    "init_pynccl",
+    "PyNCCLCommunicator",
+]

zonos2/kernel/__main__.py

@@ -0,0 +1,47 @@
+assert __name__ == "__main__"
+
+
+def generate_clangd():
+    import os
+    import subprocess
+
+    from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path
+    from zonos2.kernel.utils import DEFAULT_INCLUDE
+    from zonos2.utils import init_logger
+
+    logger = init_logger(__name__)
+    logger.info("Generating .clangd file...")
+    include_paths = [find_include_path(), find_dlpack_include_path()] + DEFAULT_INCLUDE
+    status = subprocess.run(
+        args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
+        capture_output=True,
+        check=True,
+    )
+    compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0]
+    major, minor = compute_cap.split(".")
+    compile_flags = ",\n    ".join(
+        [
+            "-xcuda",
+            f"--cuda-gpu-arch=sm_{major}{minor}",
+            "-std=c++20",
+            "-Wall",
+            "-Wextra",
+        ]
+        + [f"-isystem{path}" for path in include_paths]
+    )
+    clangd_content = f"""
+CompileFlags:
+  Add: [
+    {compile_flags}
+  ]
+"""
+    if os.path.exists(".clangd"):
+        logger.warning(".clangd file already exists, nothing done.")
+        logger.warning(f"suggested content: {clangd_content}")
+    else:
+        with open(".clangd", "w") as f:
+            f.write(clangd_content)
+        logger.info(".clangd file generated.")
+
+
+generate_clangd()

zonos2/kernel/csrc/include/zonos2/nccl227.h

@@ -0,0 +1,571 @@
+/*************************************************************************
+ * Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#ifndef NCCL_H_
+#define NCCL_H_
+
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#if CUDART_VERSION >= 11000
+#include <cuda_bf16.h>
+#endif
+#if __cplusplus && CUDART_VERSION >= 11080
+#include <cuda_fp8.h>
+#endif
+
+#define NCCL_MAJOR 2
+#define NCCL_MINOR 27
+#define NCCL_PATCH 1
+#define NCCL_SUFFIX ""
+
+#define NCCL_VERSION_CODE 22701
+#define NCCL_VERSION(X,Y,Z) (((X) <= 2 && (Y) <= 8) ? (X) * 1000 + (Y) * 100 + (Z) : (X) * 10000 + (Y) * 100 + (Z))
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include <limits.h>
+
+/* Opaque handle to communicator */
+typedef struct ncclComm* ncclComm_t;
+typedef struct ncclWindow_vidmem* ncclWindow_t;
+#define NCCL_COMM_NULL NULL
+
+#define NCCL_UNIQUE_ID_BYTES 128
+typedef struct { char internal[NCCL_UNIQUE_ID_BYTES]; } ncclUniqueId;
+
+/* Error type */
+typedef enum { ncclSuccess                 =  0,
+               ncclUnhandledCudaError      =  1,
+               ncclSystemError             =  2,
+               ncclInternalError           =  3,
+               ncclInvalidArgument         =  4,
+               ncclInvalidUsage            =  5,
+               ncclRemoteError             =  6,
+               ncclInProgress              =  7,
+               ncclNumResults              =  8 } ncclResult_t;
+
+#define NCCL_CONFIG_UNDEF_INT INT_MIN
+#define NCCL_CONFIG_UNDEF_PTR NULL
+#define NCCL_SPLIT_NOCOLOR -1
+#define NCCL_UNDEF_FLOAT -1.0f
+
+/* Window Registration flags */
+#define NCCL_WIN_DEFAULT 0x00
+#define NCCL_WIN_COLL_SYMMETRIC 0x01
+
+#define NCCL_WIN_REQUIRED_ALIGNMENT 4096
+
+/* NCCL performance policy */
+#define NCCL_CTA_POLICY_DEFAULT 0x00
+#define NCCL_CTA_POLICY_EFFICIENCY 0x01
+#define NCCL_CTA_POLICY_ZERO 0x02
+
+/* ncclCommShrink flags*/
+#define NCCL_SHRINK_DEFAULT 0x00 /* shrink the parent communicator */
+#define NCCL_SHRINK_ABORT 0x01   /* First, terminate ongoing parent operations, and then shrink the parent communicator */
+
+/* Communicator configuration. Users can assign value to attributes to specify the
+ * behavior of a communicator. */
+typedef struct ncclConfig_v22800 {
+  /* attributes that users should never touch. */
+  size_t size;
+  unsigned int magic;
+  unsigned int version;
+  /* attributes that users are able to customize. */
+  int blocking;
+  int cgaClusterSize;
+  int minCTAs;
+  int maxCTAs;
+  const char *netName;
+  int splitShare;
+  int trafficClass;
+  const char *commName;
+  int collnetEnable;
+  int CTAPolicy;
+  int shrinkShare;
+  int nvlsCTAs;
+  int nChannelsPerNetPeer;
+  int nvlinkCentricSched;
+} ncclConfig_t;
+
+/* Config initializer must be assigned to initialize config structure when it is created.
+ * Not initialized config will result in NCCL error. */
+#define NCCL_CONFIG_INITIALIZER {                                       \
+  sizeof(ncclConfig_t), /* size */                                      \
+  0xcafebeef,           /* magic */                                     \
+  NCCL_VERSION(NCCL_MAJOR, NCCL_MINOR, NCCL_PATCH), /* version */       \
+  NCCL_CONFIG_UNDEF_INT,                    /* blocking */              \
+  NCCL_CONFIG_UNDEF_INT,                    /* cgaClusterSize */        \
+  NCCL_CONFIG_UNDEF_INT,                    /* minCTAs */               \
+  NCCL_CONFIG_UNDEF_INT,                    /* maxCTAs */               \
+  NCCL_CONFIG_UNDEF_PTR,                    /* netName */               \
+  NCCL_CONFIG_UNDEF_INT,                    /* splitShare */            \
+  NCCL_CONFIG_UNDEF_INT,                    /* trafficClass */          \
+  NCCL_CONFIG_UNDEF_PTR,                    /* commName */              \
+  NCCL_CONFIG_UNDEF_INT,                    /* collnetEnable */         \
+  NCCL_CONFIG_UNDEF_INT,                    /* CTAPolicy */             \
+  NCCL_CONFIG_UNDEF_INT,                    /* shrinkShare */           \
+  NCCL_CONFIG_UNDEF_INT,                    /* nvlsCTAs */              \
+  NCCL_CONFIG_UNDEF_INT,                    /* nChannelsPerNetPeer */   \
+  NCCL_CONFIG_UNDEF_INT,                    /* nvlinkCentricSched */    \
+}
+
+/* This struct will be used by ncclGroupSimulateEnd() API to query information about simulation. */
+typedef struct ncclSimInfo_v22200 {
+    size_t size;
+    unsigned int magic;
+    unsigned int version;
+    float estimatedTime;
+} ncclSimInfo_t;
+
+/* NCCL_SIM_INFO_INITIALIZER must be assigned to initialize simInfo structure when it is created.
+ * Not initialized simInfo will result in NCCL error. */
+#define NCCL_SIM_INFO_INITIALIZER {                                         \
+  sizeof(ncclSimInfo_t),                            /* size */              \
+  0x74685283,                                       /* magic */             \
+  NCCL_VERSION(NCCL_MAJOR, NCCL_MINOR, NCCL_PATCH), /* version */           \
+  NCCL_UNDEF_FLOAT                                  /* estimated time */    \
+}
+
+/* NCCL malloc and free function for all types of NCCL optimizations
+ * (e.g. user buffer registration). The actual allocated size might
+ * be larger than requested due to granularity requirement. */
+ncclResult_t  ncclMemAlloc(void** ptr, size_t size);
+ncclResult_t pncclMemAlloc(void** ptr, size_t size);
+
+ncclResult_t  ncclMemFree(void *ptr);
+ncclResult_t pncclMemFree(void *ptr);
+
+/* Return the NCCL_VERSION_CODE of the NCCL library in the supplied integer.
+ * This integer is coded with the MAJOR, MINOR and PATCH level of the
+ * NCCL library
+ */
+ncclResult_t  ncclGetVersion(int *version);
+ncclResult_t pncclGetVersion(int *version);
+
+/* Generates an Id to be used in ncclCommInitRank. ncclGetUniqueId should be
+ * called once and the Id should be distributed to all ranks in the
+ * communicator before calling ncclCommInitRank. */
+ncclResult_t  ncclGetUniqueId(ncclUniqueId* uniqueId);
+ncclResult_t pncclGetUniqueId(ncclUniqueId* uniqueId);
+
+/* Create a new communicator (multi thread/process version) with a configuration
+ * set by users. */
+ncclResult_t  ncclCommInitRankConfig(ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank, ncclConfig_t* config);
+ncclResult_t pncclCommInitRankConfig(ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank, ncclConfig_t* config);
+
+/* Creates a new communicator (multi thread/process version).
+ * rank must be between 0 and nranks-1 and unique within a communicator clique.
+ * Each rank is associated to a CUDA device, which has to be set before calling
+ * ncclCommInitRank.
+ * ncclCommInitRank implicitly syncronizes with other ranks, so it must be
+ * called by different threads/processes or use ncclGroupStart/ncclGroupEnd. */
+ncclResult_t  ncclCommInitRank(ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank);
+ncclResult_t pncclCommInitRank(ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank);
+
+/* Creates a clique of communicators (single process version).
+ * This is a convenience function to create a single-process communicator clique.
+ * Returns an array of ndev newly initialized communicators in comm.
+ * comm should be pre-allocated with size at least ndev*sizeof(ncclComm_t).
+ * If devlist is NULL, the first ndev CUDA devices are used.
+ * Order of devlist defines user-order of processors within the communicator. */
+ncclResult_t  ncclCommInitAll(ncclComm_t* comm, int ndev, const int* devlist);
+ncclResult_t pncclCommInitAll(ncclComm_t* comm, int ndev, const int* devlist);
+
+/* Finalize a communicator. ncclCommFinalize flushes all issued communications,
+ * and marks communicator state as ncclInProgress. The state will change to ncclSuccess
+ * when the communicator is globally quiescent and related resources are freed; then,
+ * calling ncclCommDestroy can locally free the rest of the resources (e.g. communicator
+ * itself) without blocking. */
+ncclResult_t  ncclCommFinalize(ncclComm_t comm);
+ncclResult_t pncclCommFinalize(ncclComm_t comm);
+
+/* Frees local resources associated with communicator object. */
+ncclResult_t  ncclCommDestroy(ncclComm_t comm);
+ncclResult_t pncclCommDestroy(ncclComm_t comm);
+
+/* Frees resources associated with communicator object and aborts any operations
+ * that might still be running on the device. */
+ncclResult_t  ncclCommAbort(ncclComm_t comm);
+ncclResult_t pncclCommAbort(ncclComm_t comm);
+
+/* Creates one or more communicators from an existing one.
+ * Ranks with the same color will end up in the same communicator.
+ * Within the new communicator, key will be used to order ranks.
+ * NCCL_SPLIT_NOCOLOR as color will indicate the rank will not be part of any group
+ * and will therefore return a NULL communicator.
+ * If config is NULL, the new communicator will inherit the original communicator's
+ * configuration*/
+ncclResult_t  ncclCommSplit(ncclComm_t comm, int color, int key, ncclComm_t *newcomm, ncclConfig_t* config);
+ncclResult_t pncclCommSplit(ncclComm_t comm, int color, int key, ncclComm_t *newcomm, ncclConfig_t* config);
+
+/* Shrink existing communicator.
+ * Ranks in excludeRanksList will be removed form the existing communicator.
+ * Within the new communicator, ranks will be re-ordered to fill the gap of removed ones.
+ * If config is NULL, the new communicator will inherit the original communicator's configuration
+ * The flag enables NCCL to adapt to various states of the parent communicator, see NCCL_SHRINK flags.*/
+ncclResult_t  ncclCommShrink(ncclComm_t comm, int* excludeRanksList, int excludeRanksCount, ncclComm_t* newcomm, ncclConfig_t* config, int shrinkFlags);
+ncclResult_t pncclCommShrink(ncclComm_t comm, int* excludeRanksList, int excludeRanksCount, ncclComm_t* newcomm, ncclConfig_t* config, int shrinkFlags);
+
+/* Creates a new communicator (multi thread/process version), similar to ncclCommInitRankConfig.
+ * Allows to use more than one ncclUniqueId (up to one per rank), indicated by nId, to accelerate the init operation.
+ * The number of ncclUniqueIds and their order must be the same for every rank.
+ */
+ncclResult_t ncclCommInitRankScalable(ncclComm_t* newcomm, int nranks, int myrank, int nId, ncclUniqueId* commIds, ncclConfig_t* config);
+ncclResult_t pncclCommInitRankScalable(ncclComm_t* newcomm, int nranks, int myrank, int nId, ncclUniqueId* commIds, ncclConfig_t* config);
+
+/* Returns a string for each error code. */
+const char*  ncclGetErrorString(ncclResult_t result);
+const char* pncclGetErrorString(ncclResult_t result);
+
+/* Returns a human-readable message of the last error that occurred. */
+const char*  ncclGetLastError(ncclComm_t comm);
+const char* pncclGetLastError(ncclComm_t comm);
+
+/* Reload environment variables that determine logging. */
+__attribute__ ((deprecated("ncclResetDebugInit is not supported as part of the NCCL API and will be removed in the future")))
+void  ncclResetDebugInit();
+__attribute__ ((deprecated("pncclResetDebugInit is not supported as part of the NCCL API and will be removed in the future")))
+void pncclResetDebugInit();
+
+/* Checks whether the comm has encountered any asynchronous errors */
+ncclResult_t  ncclCommGetAsyncError(ncclComm_t comm, ncclResult_t *asyncError);
+ncclResult_t pncclCommGetAsyncError(ncclComm_t comm, ncclResult_t *asyncError);
+
+/* Gets the number of ranks in the communicator clique. */
+ncclResult_t  ncclCommCount(const ncclComm_t comm, int* count);
+ncclResult_t pncclCommCount(const ncclComm_t comm, int* count);
+
+/* Returns the cuda device number associated with the communicator. */
+ncclResult_t  ncclCommCuDevice(const ncclComm_t comm, int* device);
+ncclResult_t pncclCommCuDevice(const ncclComm_t comm, int* device);
+
+/* Returns the user-ordered "rank" associated with the communicator. */
+ncclResult_t  ncclCommUserRank(const ncclComm_t comm, int* rank);
+ncclResult_t pncclCommUserRank(const ncclComm_t comm, int* rank);
+
+/* Register CUDA buffer for zero-copy operation */
+ncclResult_t  ncclCommRegister(const ncclComm_t comm, void* buff, size_t size, void** handle);
+ncclResult_t pncclCommRegister(const ncclComm_t comm, void* buff, size_t size, void** handle);
+
+/* Deregister CUDA buffer */
+ncclResult_t  ncclCommDeregister(const ncclComm_t comm, void* handle);
+ncclResult_t pncclCommDeregister(const ncclComm_t comm, void* handle);
+
+/* Register memory window  */
+ncclResult_t  ncclCommWindowRegister(ncclComm_t comm, void* buff, size_t size, ncclWindow_t* win, int winFlags);
+ncclResult_t pncclCommWindowRegister(ncclComm_t comm, void* buff, size_t size, ncclWindow_t* win, int winFlags);
+
+/* Deregister symmetric memory */
+ncclResult_t  ncclCommWindowDeregister(ncclComm_t comm, ncclWindow_t win);
+ncclResult_t pncclCommWindowDeregister(ncclComm_t comm, ncclWindow_t win);
+
+/* Reduction operation selector */
+typedef enum { ncclNumOps_dummy = 5 } ncclRedOp_dummy_t;
+typedef enum { ncclSum        = 0,
+               ncclProd       = 1,
+               ncclMax        = 2,
+               ncclMin        = 3,
+               ncclAvg        = 4,
+               /* ncclNumOps: The number of built-in ncclRedOp_t values. Also
+                * serves as the least possible value for dynamic ncclRedOp_t's
+                * as constructed by ncclRedOpCreate*** functions. */
+               ncclNumOps     = 5,
+               /* ncclMaxRedOp: The largest valid value for ncclRedOp_t.
+                * It is defined to be the largest signed value (since compilers
+                * are permitted to use signed enums) that won't grow
+                * sizeof(ncclRedOp_t) when compared to previous NCCL versions to
+                * maintain ABI compatibility. */
+               ncclMaxRedOp   = 0x7fffffff>>(32-8*sizeof(ncclRedOp_dummy_t))
+             } ncclRedOp_t;
+
+/* Data types */
+typedef enum { ncclInt8       = 0, ncclChar       = 0,
+               ncclUint8      = 1,
+               ncclInt32      = 2, ncclInt        = 2,
+               ncclUint32     = 3,
+               ncclInt64      = 4,
+               ncclUint64     = 5,
+               ncclFloat16    = 6, ncclHalf       = 6,
+               ncclFloat32    = 7, ncclFloat      = 7,
+               ncclFloat64    = 8, ncclDouble     = 8,
+               ncclBfloat16   = 9,
+               ncclFloat8e4m3 = 10,
+               ncclFloat8e5m2 = 11,
+               ncclNumTypes   = 12
+} ncclDataType_t;
+
+/* ncclScalarResidence_t: Location and dereferencing logic for scalar arguments. */
+typedef enum {
+  /* ncclScalarDevice: The scalar is in device-visible memory and will be
+   * dereferenced while the collective is running. */
+  ncclScalarDevice = 0,
+
+  /* ncclScalarHostImmediate: The scalar is in host-visible memory and will be
+   * dereferenced before the ncclRedOpCreate***() function returns. */
+  ncclScalarHostImmediate = 1
+} ncclScalarResidence_t;
+
+/*
+ * ncclRedOpCreatePreMulSum
+ *
+ * Creates a new reduction operator which pre-multiplies input values by a given
+ * scalar locally before reducing them with peer values via summation. For use
+ * only with collectives launched against *comm* and *datatype*. The
+ * *residence* argument indicates how/when the memory pointed to by *scalar*
+ * will be dereferenced. Upon return, the newly created operator's handle
+ * is stored in *op*.
+ */
+ncclResult_t  ncclRedOpCreatePreMulSum(ncclRedOp_t *op, void *scalar, ncclDataType_t datatype, ncclScalarResidence_t residence, ncclComm_t comm);
+ncclResult_t pncclRedOpCreatePreMulSum(ncclRedOp_t *op, void *scalar, ncclDataType_t datatype, ncclScalarResidence_t residence, ncclComm_t comm);
+
+/*
+ * ncclRedOpDestroy
+ *
+ * Destroys the reduction operator *op*. The operator must have been created by
+ * ncclRedOpCreatePreMul with the matching communicator *comm*. An operator may be
+ * destroyed as soon as the last NCCL function which is given that operator returns.
+ */
+ncclResult_t ncclRedOpDestroy(ncclRedOp_t op, ncclComm_t comm);
+ncclResult_t pncclRedOpDestroy(ncclRedOp_t op, ncclComm_t comm);
+
+/*
+ * Collective communication operations
+ *
+ * Collective communication operations must be called separately for each
+ * communicator in a communicator clique.
+ *
+ * They return when operations have been enqueued on the CUDA stream.
+ *
+ * Since they may perform inter-CPU synchronization, each call has to be done
+ * from a different thread or process, or need to use Group Semantics (see
+ * below).
+ */
+
+/*
+ * Reduce
+ *
+ * Reduces data arrays of length count in sendbuff into recvbuff using op
+ * operation.
+ * recvbuff may be NULL on all calls except for root device.
+ * root is the rank (not the CUDA device) where data will reside after the
+ * operation is complete.
+ *
+ * In-place operation will happen if sendbuff == recvbuff.
+ */
+ncclResult_t  ncclReduce(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype,
+    ncclRedOp_t op, int root, ncclComm_t comm, cudaStream_t stream);
+ncclResult_t pncclReduce(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype,
+    ncclRedOp_t op, int root, ncclComm_t comm, cudaStream_t stream);
+
+/*
+ * (deprecated) Broadcast (in-place)
+ *
+ * Copies count values from root to all other devices.
+ * root is the rank (not the CUDA device) where data resides before the
+ * operation is started.
+ *
+ * This operation is implicitely in place.
+ */
+ncclResult_t  ncclBcast(void* buff, size_t count, ncclDataType_t datatype, int root,
+    ncclComm_t comm, cudaStream_t stream);
+ncclResult_t pncclBcast(void* buff, size_t count, ncclDataType_t datatype, int root,
+    ncclComm_t comm, cudaStream_t stream);
+
+/*
+ * Broadcast
+ *
+ * Copies count values from root to all other devices.
+ * root is the rank (not the CUDA device) where data resides before the
+ * operation is started.
+ *
+ * In-place operation will happen if sendbuff == recvbuff.
+ */
+ncclResult_t  ncclBroadcast(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype, int root,
+    ncclComm_t comm, cudaStream_t stream);
+ncclResult_t pncclBroadcast(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype, int root,
+    ncclComm_t comm, cudaStream_t stream);
+
+/*
+ * All-Reduce
+ *
+ * Reduces data arrays of length count in sendbuff using op operation, and
+ * leaves identical copies of result on each recvbuff.
+ *
+ * In-place operation will happen if sendbuff == recvbuff.
+ */
+ncclResult_t  ncclAllReduce(const void* sendbuff, void* recvbuff, size_t count,
+    ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, cudaStream_t stream);
+ncclResult_t pncclAllReduce(const void* sendbuff, void* recvbuff, size_t count,
+    ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, cudaStream_t stream);
+
+/*
+ * Reduce-Scatter
+ *
+ * Reduces data in sendbuff using op operation and leaves reduced result
+ * scattered over the devices so that recvbuff on rank i will contain the i-th
+ * block of the result.
+ * Assumes sendcount is equal to nranks*recvcount, which means that sendbuff
+ * should have a size of at least nranks*recvcount elements.
+ *
+ * In-place operations will happen if recvbuff == sendbuff + rank * recvcount.
+ */
+ncclResult_t  ncclReduceScatter(const void* sendbuff, void* recvbuff,
+    size_t recvcount, ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
+    cudaStream_t stream);
+ncclResult_t pncclReduceScatter(const void* sendbuff, void* recvbuff,
+    size_t recvcount, ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
+    cudaStream_t stream);
+
+/*
+ * All-Gather
+ *
+ * Each device gathers sendcount values from other GPUs into recvbuff,
+ * receiving data from rank i at offset i*sendcount.
+ * Assumes recvcount is equal to nranks*sendcount, which means that recvbuff
+ * should have a size of at least nranks*sendcount elements.
+ *
+ * In-place operations will happen if sendbuff == recvbuff + rank * sendcount.
+ */
+ncclResult_t  ncclAllGather(const void* sendbuff, void* recvbuff, size_t sendcount,
+    ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream);
+ncclResult_t pncclAllGather(const void* sendbuff, void* recvbuff, size_t sendcount,
+    ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream);
+
+/*
+ * All-to-All
+ *
+ * Each device sends count values to all other devices and receives count values
+ * from all other devices. Data to send to destination rank j is taken from
+ * sendbuff+j*count and data received from source rank i is placed at
+ * recvbuff+i*count.
+ */
+ncclResult_t  ncclAlltoAll(const void* sendbuff, void* recvbuff, size_t count,
+    ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream);
+ncclResult_t pncclAlltoAll(const void* sendbuff, void* recvbuff, size_t count,
+    ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream);
+
+/*
+ * Gather
+ *
+ * Each rank sends count elements from sendbuff to the root rank.
+ * On the root rank, data from rank i is placed at recvbuff + i*count.
+ * On non-root ranks, recvbuff is not used.
+ * root is the rank where data will be gathered.
+ *
+ * In-place operations will happen if sendbuff == recvbuff + root * count.
+ */
+ncclResult_t  ncclGather(const void* sendbuff, void* recvbuff, size_t count,
+    ncclDataType_t datatype, int root, ncclComm_t comm, cudaStream_t stream);
+ncclResult_t pncclGather(const void* sendbuff, void* recvbuff, size_t count,
+    ncclDataType_t datatype, int root, ncclComm_t comm, cudaStream_t stream);
+
+/*
+ * Scatter
+ *
+ * On the root rank, count elements from sendbuff+i*count are sent to rank i.
+ * On non-root ranks, sendbuff is not used.
+ * Each rank receives count elements into recvbuff.
+ * root is the rank that will distribute the data.
+ *
+ * In-place operations will happen if recvbuff == sendbuff + root * count.
+ */
+ncclResult_t  ncclScatter(const void* sendbuff, void* recvbuff, size_t count,
+    ncclDataType_t datatype, int root, ncclComm_t comm, cudaStream_t stream);
+ncclResult_t pncclScatter(const void* sendbuff, void* recvbuff, size_t count,
+    ncclDataType_t datatype, int root, ncclComm_t comm, cudaStream_t stream);
+
+/*
+ * Send
+ *
+ * Send data from sendbuff to rank peer.
+ *
+ * Rank peer needs to call ncclRecv with the same datatype and the same count from this
+ * rank.
+ *
+ * This operation is blocking for the GPU. If multiple ncclSend and ncclRecv operations
+ * need to progress concurrently to complete, they must be fused within a ncclGroupStart/
+ * ncclGroupEnd section.
+ */
+ncclResult_t  ncclSend(const void* sendbuff, size_t count, ncclDataType_t datatype, int peer,
+    ncclComm_t comm, cudaStream_t stream);
+ncclResult_t pncclSend(const void* sendbuff, size_t count, ncclDataType_t datatype, int peer,
+    ncclComm_t comm, cudaStream_t stream);
+
+/*
+ * Receive
+ *
+ * Receive data from rank peer into recvbuff.
+ *
+ * Rank peer needs to call ncclSend with the same datatype and the same count to this
+ * rank.
+ *
+ * This operation is blocking for the GPU. If multiple ncclSend and ncclRecv operations
+ * need to progress concurrently to complete, they must be fused within a ncclGroupStart/
+ * ncclGroupEnd section.
+ */
+ncclResult_t pncclRecv(void* recvbuff, size_t count, ncclDataType_t datatype, int peer,
+    ncclComm_t comm, cudaStream_t stream);
+ncclResult_t  ncclRecv(void* recvbuff, size_t count, ncclDataType_t datatype, int peer,
+    ncclComm_t comm, cudaStream_t stream);
+
+/*
+ * Group semantics
+ *
+ * When managing multiple GPUs from a single thread, and since NCCL collective
+ * calls may perform inter-CPU synchronization, we need to "group" calls for
+ * different ranks/devices into a single call.
+ *
+ * Grouping NCCL calls as being part of the same collective operation is done
+ * using ncclGroupStart and ncclGroupEnd. ncclGroupStart will enqueue all
+ * collective calls until the ncclGroupEnd call, which will wait for all calls
+ * to be complete. Note that for collective communication, ncclGroupEnd only
+ * guarantees that the operations are enqueued on the streams, not that
+ * the operation is effectively done.
+ *
+ * Both collective communication and ncclCommInitRank can be used in conjunction
+ * of ncclGroupStart/ncclGroupEnd, but not together.
+ *
+ * Group semantics also allow to fuse multiple operations on the same device
+ * to improve performance (for aggregated collective calls), or to permit
+ * concurrent progress of multiple send/receive operations.
+ */
+
+/*
+ * Group Start
+ *
+ * Start a group call. All calls to NCCL until ncclGroupEnd will be fused into
+ * a single NCCL operation. Nothing will be started on the CUDA stream until
+ * ncclGroupEnd.
+ */
+ncclResult_t  ncclGroupStart();
+ncclResult_t pncclGroupStart();
+
+/*
+ * Group End
+ *
+ * End a group call. Start a fused NCCL operation consisting of all calls since
+ * ncclGroupStart. Operations on the CUDA stream depending on the NCCL operations
+ * need to be called after ncclGroupEnd.
+ */
+ncclResult_t  ncclGroupEnd();
+ncclResult_t pncclGroupEnd();
+
+/*
+ * Group Simulate End
+ *
+ * Simulate a ncclGroupEnd() call and return NCCL's simulation info in a struct.
+ */
+ncclResult_t  ncclGroupSimulateEnd(ncclSimInfo_t* simInfo);
+ncclResult_t pncclGroupSimulateEnd(ncclSimInfo_t* simInfo);
+
+#ifdef __cplusplus
+} // end extern "C"
+#endif
+
+#endif // end include guard

zonos2/kernel/csrc/include/zonos2/tensor.h

@@ -0,0 +1,496 @@
+#pragma once
+#include <zonos2/utils.h>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/dtype.h>
+
+#include <algorithm>
+#include <array>
+#include <concepts>
+#include <cstddef>
+#include <cstdint>
+#include <initializer_list>
+#include <optional>
+#include <ranges>
+#include <source_location>
+#include <span>
+#include <sstream>
+#include <string>
+#include <string_view>
+#include <type_traits>
+#include <utility>
+
+namespace host {
+
+namespace stdr = std::ranges;
+namespace stdv = std::views;
+
+namespace details {
+
+struct SizeRef;
+struct DTypeRef;
+struct DeviceRef;
+
+inline constexpr auto kAnyDeviceID = -1;
+inline constexpr auto kAnySize = static_cast<int64_t>(-1);
+inline constexpr auto kNullSize = static_cast<int64_t>(0);
+inline constexpr auto kNullDType = static_cast<DLDataTypeCode>(18u);
+inline constexpr auto kNullDevice = static_cast<DLDeviceType>(-1);
+
+template <typename T> struct dtype_trait {};
+
+template <std::integral T> struct dtype_trait<T> {
+  inline static constexpr auto value =
+      DLDataType{.code = std::is_signed_v<T> ? DLDataTypeCode::kDLInt
+                                             : DLDataTypeCode::kDLUInt,
+                 .bits = static_cast<std::uint8_t>(sizeof(T) * 8),
+                 .lanes = 1};
+};
+
+template <std::floating_point T> struct dtype_trait<T> {
+  inline static constexpr auto value =
+      DLDataType{.code = DLDataTypeCode::kDLFloat,
+                 .bits = static_cast<std::uint8_t>(sizeof(T) * 8),
+                 .lanes = 1};
+};
+
+template <DLDeviceType Code> struct device_trait {
+  inline static constexpr auto value =
+      DLDevice{.device_type = Code, .device_id = kAnyDeviceID};
+};
+
+template <typename... Ts>
+inline constexpr auto kDTypeList =
+    std::array<DLDataType, sizeof...(Ts)>{dtype_trait<Ts>::value...};
+
+template <DLDeviceType... Codes>
+inline constexpr auto kDeviceList =
+    std::array<DLDevice, sizeof...(Codes)>{device_trait<Codes>::value...};
+
+template <typename T> struct PrintAbleSpan {
+  explicit PrintAbleSpan(std::span<const T> data) : data(data) {}
+  std::span<const T> data;
+};
+
+// define DLDataType comparison and printing in root namespace
+inline constexpr auto kDeviceStringMap = [] {
+  constexpr auto map =
+      std::array<std::pair<DLDeviceType, std::string_view>, 16>{
+          std::pair{DLDeviceType::kDLCPU, "cpu"},
+          std::pair{DLDeviceType::kDLCUDA, "cuda"},
+          std::pair{DLDeviceType::kDLCUDAHost, "cuda_host"},
+          std::pair{DLDeviceType::kDLOpenCL, "opencl"},
+          std::pair{DLDeviceType::kDLVulkan, "vulkan"},
+          std::pair{DLDeviceType::kDLMetal, "metal"},
+          std::pair{DLDeviceType::kDLVPI, "vpi"},
+          std::pair{DLDeviceType::kDLROCM, "rocm"},
+          std::pair{DLDeviceType::kDLROCMHost, "rocm_host"},
+          std::pair{DLDeviceType::kDLExtDev, "ext_dev"},
+          std::pair{DLDeviceType::kDLCUDAManaged, "cuda_managed"},
+          std::pair{DLDeviceType::kDLOneAPI, "oneapi"},
+          std::pair{DLDeviceType::kDLWebGPU, "webgpu"},
+          std::pair{DLDeviceType::kDLHexagon, "hexagon"},
+          std::pair{DLDeviceType::kDLMAIA, "maia"},
+          std::pair{DLDeviceType::kDLTrn, "trn"},
+      };
+  constexpr auto max_type = stdr::max(map | stdv::keys);
+  auto result = std::array<std::string_view, max_type + 1>{};
+  for (const auto &[code, name] : map) {
+    result[static_cast<std::size_t>(code)] = name;
+  }
+  return result;
+}();
+
+struct PrintableDevice {
+  DLDevice device;
+};
+
+inline auto &operator<<(std::ostream &os, DLDevice device) {
+  const auto &mapping = kDeviceStringMap;
+  const auto entry = static_cast<std::size_t>(device.device_type);
+  RuntimeCheck(entry < mapping.size());
+  const auto name = mapping[entry];
+  RuntimeCheck(!name.empty(), "Unknown device: ", int(device.device_type));
+  os << name;
+  if (device.device_id != kAnyDeviceID)
+    os << "[" << device.device_id << "]";
+  return os;
+}
+
+inline auto &operator<<(std::ostream &os, PrintableDevice pd) {
+  return os << pd.device;
+}
+
+template <typename T>
+inline auto &operator<<(std::ostream &os, PrintAbleSpan<T> span) {
+  os << "[";
+  for (const auto i : stdv::iota(std::size_t{0}, span.data.size())) {
+    if (i > 0) {
+      os << ", ";
+    }
+    os << span.data[i];
+  }
+  os << "]";
+  return os;
+}
+
+} // namespace details
+
+struct SymbolicSize {
+public:
+  SymbolicSize(std::string_view annotation = {})
+      : m_value(details::kNullSize), m_annotation(annotation) {}
+
+  auto get_name() const -> std::string_view { return m_annotation; }
+  auto set_value(int64_t value) -> void {
+    RuntimeCheck(!this->has_value(), "Size value already set");
+    m_value = value;
+  }
+  auto has_value() const -> bool { return m_value != 0; }
+  auto get_value() const -> std::optional<int64_t> {
+    return this->has_value() ? std::optional{m_value} : std::nullopt;
+  }
+  auto unwrap() const -> int64_t {
+    RuntimeCheck(this->has_value(), "Size value is not set");
+    return m_value;
+  }
+
+  SymbolicSize(const SymbolicSize &) = delete;
+  SymbolicSize &operator=(const SymbolicSize &) = delete;
+
+  auto verify(int64_t size, const char *prefix, int64_t dim) -> void {
+    if (this->has_value()) {
+      if (m_value != size) {
+        [[unlikely]];
+        Panic("Size mismatch for ", m_name(prefix, dim), ": expected ", m_value,
+              " but got ", size);
+      }
+    } else {
+      this->set_value(size);
+    }
+  }
+
+  auto value_or_name(const char *prefix, int64_t dim) const -> std::string {
+    if (const auto value = this->get_value()) {
+      return std::to_string(*value);
+    } else {
+      return m_name(prefix, dim);
+    }
+  }
+
+private:
+  auto m_name(const char *prefix, int64_t dim) const -> std::string {
+    const auto annotation = this->get_name();
+    std::ostringstream os;
+    if (annotation.empty()) {
+      os << prefix << '#' << dim;
+    } else {
+      os << annotation << '(' << prefix << '#' << dim << ')';
+    }
+    return std::move(os).str();
+  }
+
+  std::int64_t m_value;
+  std::string_view m_annotation;
+};
+
+inline auto operator==(DLDevice lhs, DLDevice rhs) -> bool {
+  return lhs.device_type == rhs.device_type && lhs.device_id == rhs.device_id;
+}
+
+struct SymbolicDType {
+public:
+  SymbolicDType() : m_value({details::kNullDType, 0, 0}) {}
+
+  auto set_value(DLDataType value) -> void {
+    RuntimeCheck(!this->has_value(), "Dtype value already set");
+    RuntimeCheck(
+        m_check(value), "Dtype value [", value,
+        "] not in the allowed options: ", details::PrintAbleSpan{m_options});
+    m_value = value;
+  }
+  auto has_value() const -> bool { return m_value.code != details::kNullDType; }
+  auto get_value() const -> std::optional<DLDataType> {
+    return this->has_value() ? std::optional{m_value} : std::nullopt;
+  }
+  auto unwrap() const -> DLDataType {
+    RuntimeCheck(this->has_value(), "Dtype value is not set");
+    return m_value;
+  }
+
+  auto set_options(std::span<const DLDataType> options) -> void {
+    m_options = options;
+  }
+  template <typename... Ts> auto set_options() -> void {
+    m_options = details::kDTypeList<Ts...>;
+  }
+
+  auto verify(DLDataType dtype) -> void {
+    if (this->has_value()) {
+      RuntimeCheck(m_value == dtype, "DType mismatch: expected ", m_value,
+                   " but got ", dtype);
+    } else {
+      this->set_value(dtype);
+    }
+  }
+
+private:
+  auto m_check(DLDataType value) const -> bool {
+    return stdr::empty(m_options) ||
+           (stdr::find(m_options, value) != stdr::end(m_options));
+  }
+
+  std::span<const DLDataType> m_options;
+  DLDataType m_value;
+};
+
+struct SymbolicDevice {
+public:
+  SymbolicDevice() : m_value({details::kNullDevice, details::kAnyDeviceID}) {}
+
+  auto set_value(DLDevice value) -> void {
+    RuntimeCheck(!this->has_value(), "Device value already set");
+    RuntimeCheck(
+        m_check(value), "Device value [", details::PrintableDevice{value},
+        "] not in the allowed options: ", details::PrintAbleSpan{m_options});
+    m_value = value;
+  }
+  auto has_value() const -> bool {
+    return m_value.device_type != details::kNullDevice;
+  }
+  auto get_value() const -> std::optional<DLDevice> {
+    return this->has_value() ? std::optional{m_value} : std::nullopt;
+  }
+  auto unwrap() const -> DLDevice {
+    RuntimeCheck(this->has_value(), "Device value is not set");
+    return m_value;
+  }
+
+  auto set_options(std::span<const DLDevice> options) -> void {
+    m_options = options;
+  }
+  template <DLDeviceType... Codes> auto set_options() -> void {
+    m_options = details::kDeviceList<Codes...>;
+  }
+
+  auto verify(DLDevice device) -> void {
+    if (this->has_value()) {
+      RuntimeCheck(m_value == device, "Device mismatch: expected ",
+                   details::PrintableDevice{m_value}, " but got ",
+                   details::PrintableDevice{device});
+    } else {
+      this->set_value(device);
+    }
+  }
+
+private:
+  auto m_check(DLDevice value) const -> bool {
+    return stdr::empty(m_options) ||
+           (stdr::any_of(m_options, [value](const DLDevice &opt) {
+             // device type must exactly match
+             if (opt.device_type != value.device_type)
+               return false;
+             // device id can be wildcarded
+             return opt.device_id == details::kAnyDeviceID ||
+                    opt.device_id == value.device_id;
+           }));
+  }
+
+  std::span<const DLDevice> m_options;
+  DLDevice m_value;
+};
+
+struct TensorMatcher;
+
+namespace details {
+
+template <typename T> struct BaseRef {
+public:
+  BaseRef(const BaseRef &) = delete;
+  BaseRef &operator=(const BaseRef &) = delete;
+
+  auto operator->() const -> T * { return m_ref; }
+  auto operator*() const -> T & { return *m_ref; }
+  auto rebind(T &other) -> void { m_ref = &other; }
+
+  explicit BaseRef() : m_ref(&m_cache), m_cache() {}
+  BaseRef(T &size) : m_ref(&size), m_cache() {}
+
+private:
+  T *m_ref;
+  T m_cache;
+};
+
+struct SizeRef : BaseRef<SymbolicSize> {
+  using BaseRef::BaseRef;
+  SizeRef(int64_t value) {
+    if (value != kAnySize) {
+      (**this).set_value(value);
+    } else {
+      // otherwise, we can match any size
+    }
+  }
+};
+
+struct DTypeRef : BaseRef<SymbolicDType> {
+  using BaseRef::BaseRef;
+  DTypeRef(DLDataType options) { (**this).set_value(options); }
+  DTypeRef(std::initializer_list<DLDataType> options) {
+    (**this).set_options(options);
+  }
+  DTypeRef(std::span<const DLDataType> options) {
+    (**this).set_options(options);
+  }
+};
+
+struct DeviceRef : BaseRef<SymbolicDevice> {
+  using BaseRef::BaseRef;
+  DeviceRef(DLDevice options) { (**this).set_value(options); }
+  DeviceRef(std::initializer_list<DLDevice> options) {
+    (**this).set_options(options);
+  }
+  DeviceRef(std::span<const DLDevice> options) {
+    (**this).set_options(options);
+  }
+};
+
+} // namespace details
+
+struct TensorMatcher {
+private:
+  using SizeRef = details::SizeRef;
+  using DTypeRef = details::DTypeRef;
+  using DeviceRef = details::DeviceRef;
+  using Loc_t = std::source_location;
+
+public:
+  TensorMatcher(const TensorMatcher &) = delete;
+  TensorMatcher &operator=(const TensorMatcher &) = delete;
+
+  explicit TensorMatcher(std::initializer_list<SizeRef> shape)
+      : m_shape(shape), m_strides(), m_dtype() {}
+
+  auto with_strides(std::initializer_list<SizeRef> strides) && //
+      -> TensorMatcher && {
+    // no partial update allowed
+    RuntimeCheck(m_strides.size() == 0, "Strides already specified");
+    RuntimeCheck(m_shape.size() == strides.size(),
+                 "Strides size must match shape size");
+    m_strides = strides;
+    return std::move(*this);
+  }
+
+  template <typename... Ts>
+  auto with_dtype(DTypeRef &&dtype) && -> TensorMatcher && {
+    m_init_dtype();
+    m_dtype.rebind(*dtype);
+    return std::move(*this);
+  }
+
+  template <typename... Ts>
+    requires(sizeof...(Ts) > 0)
+  auto with_dtype() && -> TensorMatcher && {
+    m_init_dtype();
+    m_dtype->set_options<Ts...>();
+    return std::move(*this);
+  }
+
+  template <DLDeviceType... Codes>
+  auto with_device(DeviceRef &&device) && -> TensorMatcher && {
+    m_init_device();
+    m_device.rebind(*device);
+    return std::move(*this);
+  }
+
+  template <DLDeviceType... Codes>
+    requires(sizeof...(Codes) > 0)
+  auto with_device() && -> TensorMatcher && {
+    m_init_device();
+    m_device->set_options<Codes...>();
+    return std::move(*this);
+  }
+
+  // once we start verification, we cannot modify anymore
+  auto verify(tvm::ffi::TensorView view,
+              Loc_t loc = Loc_t::current()) const && -> const TensorMatcher && {
+    try {
+      this->m_verify_impl(view);
+    } catch (PanicError &e) {
+      auto oss = std::ostringstream{};
+      oss << "Tensor match failed for " << this->debug_str() << " at "
+          << loc.file_name() << ":" << loc.line()
+          << "\n- Root cause: " << e.detail();
+      throw PanicError(std::move(oss).str());
+    }
+    return std::move(*this);
+  }
+
+  auto debug_str() const -> std::string {
+    auto oss = std::ostringstream{};
+    oss << "Tensor<";
+    std::size_t dim = 0;
+    for (const auto &size_ref : m_shape) {
+      if (dim > 0) {
+        oss << ", ";
+      }
+      oss << size_ref->value_or_name("shape", dim++);
+    }
+    oss << ">";
+    if (m_strides.size() > 0) {
+      oss << " [strides=<";
+      dim = 0;
+      for (const auto &stride_ref : m_strides) {
+        if (dim > 0) {
+          oss << ", ";
+        }
+        oss << stride_ref->value_or_name("stride", dim++);
+      }
+      oss << ">]";
+    }
+    return std::move(oss).str();
+  }
+
+private:
+  auto m_verify_impl(tvm::ffi::TensorView view) const -> void {
+    const auto dim = static_cast<std::size_t>(view.dim());
+    RuntimeCheck(dim == m_shape.size(), "Tensor dimension mismatch: expected ",
+                 m_shape.size(), " but got ", dim);
+    for (const auto i : stdv::iota(std::size_t{0}, dim)) {
+      m_shape[i]->verify(view.size(i), "shape", i);
+    }
+    if (this->m_has_strides()) {
+      for (const auto i : stdv::iota(std::size_t{0}, dim)) {
+        if (view.size(i) != 1 || !m_strides[i]->has_value()) {
+          // skip stride check for size 1 dimension
+          m_strides[i]->verify(view.stride(i), "stride", i);
+        }
+      }
+    } else {
+      RuntimeCheck(view.is_contiguous(),
+                   "Tensor is not contiguous as expected");
+    }
+    // since we may double verify, we will force to check
+    m_dtype->verify(view.dtype());
+    m_device->verify(view.device());
+  }
+
+  auto m_init_dtype() -> void {
+    RuntimeCheck(!m_has_dtype, "DType already specified");
+    m_has_dtype = true;
+  }
+  auto m_init_device() -> void {
+    RuntimeCheck(!m_has_device, "Device already specified");
+    m_has_device = true;
+  }
+  auto m_has_strides() const -> bool { return !m_strides.empty(); }
+
+  std::span<const SizeRef> m_shape;
+  std::span<const SizeRef> m_strides;
+  DTypeRef m_dtype;
+  DeviceRef m_device;
+  bool m_has_dtype = false;
+  bool m_has_device = false;
+};
+
+} // namespace host

zonos2/kernel/csrc/include/zonos2/utils.cuh

@@ -0,0 +1,144 @@
+#pragma once
+
+#include <zonos2/utils.h>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/extra/c_env_api.h>
+
+#include <concepts>
+#include <cstddef>
+#include <source_location>
+#include <type_traits>
+
+namespace device {
+
+inline constexpr auto kWarpThreads = 32u;
+
+template <std::integral T, std::integral U>
+__always_inline __device__ constexpr auto div_ceil(T a, U b) {
+  return (a + b - 1) / b;
+}
+
+namespace pointer {
+
+// we only allow void * pointer arithmetic for safety
+
+template <typename T, std::integral... U>
+__always_inline __device__ auto offset(T *ptr, U... offset) -> void * {
+  static_assert(std::is_same_v<T, void>,
+                "Pointer arithmetic is only allowed for void* pointers");
+  return static_cast<char *>(ptr) + (... + offset);
+}
+
+template <typename T, std::integral... U>
+__always_inline __device__ auto offset(const T *ptr, U... offset) -> const
+    void * {
+  static_assert(std::is_same_v<T, void>,
+                "Pointer arithmetic is only allowed for void* pointers");
+  return static_cast<const char *>(ptr) + (... + offset);
+}
+
+} // namespace pointer
+
+namespace PDL {
+
+template <bool kUsePDL> __always_inline __device__ void wait() {
+  if constexpr (kUsePDL) {
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+  }
+}
+
+template <bool kUsePDL> __always_inline __device__ void launch() {
+  if constexpr (kUsePDL) {
+    asm volatile("griddepcontrol.launch_dependents;" :::);
+  }
+}
+
+} // namespace PDL
+
+} // namespace device
+
+namespace host {
+
+inline auto
+CUDA_CHECK(::cudaError_t error,
+           std::source_location location = std::source_location::current())
+    -> void {
+  if (error != ::cudaSuccess) {
+    [[unlikely]];
+    ::host::panic(location, "CUDA error: ", ::cudaGetErrorString(error));
+  }
+}
+
+inline auto
+CUDA_CHECK(std::source_location location = std::source_location::current())
+    -> void {
+  return CUDA_CHECK(::cudaGetLastError(), location);
+}
+
+template <auto F> inline void set_smem_once(std::size_t smem_size) {
+  static const auto last_smem_size = [&] {
+    CUDA_CHECK(::cudaFuncSetAttribute(
+        F, ::cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+    return smem_size;
+  }();
+  RuntimeCheck(
+      smem_size <= last_smem_size,
+      "Dynamic shared memory size exceeds the previously set maximum size: ",
+      last_smem_size, " bytes");
+}
+
+struct LaunchKernel {
+public:
+  explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, DLDevice device,
+                        std::size_t dynamic_shared_mem_bytes = 0) noexcept
+      : m_config(s_make_config(grid_dim, block_dim, resolve_device(device),
+                               dynamic_shared_mem_bytes)) {}
+
+  explicit LaunchKernel(dim3 grid_dim, dim3 block_dim, cudaStream_t stream,
+                        std::size_t dynamic_shared_mem_bytes = 0) noexcept
+      : m_config(s_make_config(grid_dim, block_dim, stream,
+                               dynamic_shared_mem_bytes)) {}
+
+  static auto resolve_device(DLDevice device) -> cudaStream_t {
+    return static_cast<cudaStream_t>(
+        ::TVMFFIEnvGetStream(device.device_type, device.device_id));
+  }
+
+  LaunchKernel(const LaunchKernel &) = delete;
+  LaunchKernel &operator=(const LaunchKernel &) = delete;
+
+  template <typename T, typename... Args>
+  auto operator()(T &&kernel, Args &&...args) const -> void {
+    CUDA_CHECK(
+        ::cudaLaunchKernelEx(&m_config, kernel, std::forward<Args>(args)...));
+  }
+
+  auto with_attr(bool use_pdl) -> LaunchKernel & {
+    if (use_pdl) {
+      m_attr_cache.id = ::cudaLaunchAttributeProgrammaticStreamSerialization;
+      m_attr_cache.val.programmaticStreamSerializationAllowed = 1;
+      m_config.attrs = &m_attr_cache;
+      m_config.numAttrs = 1;
+    } else {
+      m_config.numAttrs = 0;
+    }
+    return *this;
+  }
+
+private:
+  static auto s_make_config(dim3 grid_dim, dim3 block_dim, cudaStream_t stream,
+                            std::size_t smem) -> cudaLaunchConfig_t {
+    auto config = ::cudaLaunchConfig_t{};
+    config.gridDim = grid_dim;
+    config.blockDim = block_dim;
+    config.dynamicSmemBytes = smem;
+    config.stream = stream;
+    config.numAttrs = 0;
+    return config;
+  }
+  cudaLaunchConfig_t m_config;
+  cudaLaunchAttribute m_attr_cache;
+};
+
+} // namespace host

zonos2/kernel/csrc/include/zonos2/utils.h

@@ -0,0 +1,122 @@
+#pragma once
+
+// ref:
+// https://forums.developer.nvidia.com/t/c-20s-source-location-compilation-error-when-using-nvcc-12-1/258026/3
+#ifdef __CUDACC__
+#pragma push_macro("__cpp_consteval")
+#pragma push_macro("_NODISCARD")
+#pragma push_macro("__builtin_LINE")
+
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wbuiltin-macro-redefined"
+#define __cpp_consteval 201811L
+#pragma clang diagnostic pop
+
+#ifdef _NODISCARD
+#undef _NODISCARD
+#define _NODISCARD
+#endif
+
+#define consteval constexpr
+
+#include <source_location>
+
+#undef consteval
+#pragma pop_macro("__cpp_consteval")
+#pragma pop_macro("_NODISCARD")
+#else
+#include <source_location>
+#endif
+
+#include <dlpack/dlpack.h>
+
+#include <concepts>
+#include <ostream>
+#include <sstream>
+#include <utility>
+
+namespace host {
+
+struct PanicError : public std::runtime_error {
+public:
+  // copy and move constructors
+  PanicError(std::string msg) : runtime_error(msg), m_message(std::move(msg)) {}
+  auto detail() const -> std::string_view {
+    const auto sv = std::string_view{m_message};
+    const auto pos = sv.find(": ");
+    return pos == std::string_view::npos ? sv : sv.substr(pos + 2);
+  }
+
+private:
+  std::string m_message;
+};
+
+template <typename... Args>
+[[noreturn]]
+inline auto panic(std::source_location location, Args &&...args) -> void {
+  std::ostringstream os;
+  os << "Runtime check failed at " << location.file_name() << ":"
+     << location.line();
+  if constexpr (sizeof...(args) > 0) {
+    os << ": ";
+    (os << ... << std::forward<Args>(args));
+  } else {
+    os << " in " << location.function_name();
+  }
+  throw PanicError(std::move(os).str());
+}
+
+template <typename... Args> struct Panic {
+  explicit Panic(Args &&...args, std::source_location location =
+                                     std::source_location::current()) {
+    [[unlikely]];
+    ::host::panic(location, std::forward<Args>(args)...);
+  }
+  [[noreturn]] ~Panic() { std::terminate(); }
+};
+
+template <typename... Args> struct RuntimeCheck {
+  template <typename T>
+  explicit RuntimeCheck(
+      T &&condition, Args &&...args,
+      std::source_location location = std::source_location::current()) {
+    if (!condition) {
+      [[unlikely]];
+      ::host::panic(location, std::forward<Args>(args)...);
+    }
+  }
+};
+
+template <typename T, typename... Args>
+explicit RuntimeCheck(T &&, Args &&...) -> RuntimeCheck<Args...>;
+
+template <typename... Args> explicit Panic(Args &&...) -> Panic<Args...>;
+
+template <std::integral T, std::integral U>
+inline constexpr auto div_ceil(T a, U b) {
+  return (a + b - 1) / b;
+}
+
+inline auto dtype_bytes(DLDataType dtype) -> std::size_t {
+  return static_cast<std::size_t>(dtype.bits / 8);
+}
+
+namespace pointer {
+
+template <typename T, std::integral... U>
+inline auto offset(T *ptr, U... offset) -> void * {
+  static_assert(std::is_same_v<T, void>,
+                "Pointer arithmetic is only allowed for void* pointers");
+  return static_cast<char *>(ptr) + (... + offset);
+}
+
+template <typename T, std::integral... U>
+inline auto offset(const T *ptr, U... offset) -> const void * {
+  static_assert(std::is_same_v<T, void>,
+                "Pointer arithmetic is only allowed for void* pointers");
+  return static_cast<const char *>(ptr) + (... + offset);
+}
+
+} // namespace pointer
+
+} // namespace host

zonos2/kernel/csrc/include/zonos2/warp.cuh

@@ -0,0 +1,80 @@
+#pragma once
+#include <zonos2/utils.cuh>
+
+#include <sys/cdefs.h>
+
+#include <cstddef>
+
+namespace device::warp {
+
+namespace details {
+
+template <std::size_t kUnit> inline constexpr auto get_mem_package() {
+  if constexpr (kUnit == 16) {
+    return uint4{};
+  } else if constexpr (kUnit == 8) {
+    return uint2{};
+  } else if constexpr (kUnit == 4) {
+    return uint1{};
+  } else {
+    static_assert(kUnit == 16 || kUnit == 8 || kUnit == 4,
+                  "Unsupported memory package size");
+  }
+}
+
+inline constexpr auto resolve_unit_size(std::size_t x) -> std::size_t {
+  if (x % (16 * kWarpThreads) == 0)
+    return 16;
+  if (x % (8 * kWarpThreads) == 0)
+    return 8;
+  if (x % (4 * kWarpThreads) == 0)
+    return 4;
+  return 0; // trigger static assert in _get_mem_package
+}
+
+template <std::size_t kBytes, std::size_t kUnit>
+using mem_package_t = decltype(get_mem_package<kUnit>());
+
+} // namespace details
+
+template <std::size_t kBytes,
+          std::size_t kUnit = details::resolve_unit_size(kBytes)>
+__always_inline __device__ void copy(void *__restrict__ dst,
+                                     const void *__restrict__ src) {
+  using Package = details::mem_package_t<kBytes, kUnit>;
+  constexpr auto kBytesPerLoop = sizeof(Package) * kWarpThreads;
+  constexpr auto kLoopCount = kBytes / kBytesPerLoop;
+  static_assert(kBytes % kBytesPerLoop == 0,
+                "kBytes must be multiple of 128 bytes");
+
+  const auto dst_packed = static_cast<Package *>(dst);
+  const auto src_packed = static_cast<const Package *>(src);
+  const auto lane_id = threadIdx.x % kWarpThreads;
+
+#pragma unroll kLoopCount
+  for (std::size_t i = 0; i < kLoopCount; ++i) {
+    const auto j = i * kWarpThreads + lane_id;
+    dst_packed[j] = src_packed[j];
+  }
+}
+
+template <std::size_t kBytes,
+          std::size_t kUnit = details::resolve_unit_size(kBytes)>
+__always_inline __device__ void reset(void *__restrict__ dst) {
+  using Package = details::mem_package_t<kBytes, kUnit>;
+  constexpr auto kBytesPerLoop = sizeof(Package) * kWarpThreads;
+  constexpr auto kLoopCount = kBytes / kBytesPerLoop;
+  static_assert(kBytes % kBytesPerLoop == 0,
+                "warp_copy: kBytes must be multiple of 128 bytes");
+
+  const auto dst_ = static_cast<Package *>(dst);
+  const auto lane_id = threadIdx.x % kWarpThreads;
+  const auto zero_value = Package{};
+
+#pragma unroll kLoopCount
+  for (std::size_t i = 0; i < kLoopCount; ++i) {
+    dst_[i * kWarpThreads + lane_id] = zero_value;
+  }
+}
+
+} // namespace device::warp

zonos2/kernel/csrc/jit/index.cu

@@ -0,0 +1,172 @@
+#include <zonos2/tensor.h>
+#include <zonos2/utils.cuh>
+#include <zonos2/utils.h>
+#include <zonos2/warp.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/container/tuple.h>
+
+#include <bit>
+#include <concepts>
+#include <cstddef>
+#include <cstdint>
+
+namespace {
+
+struct IndexKernelParams {
+  void *__restrict__ output;
+  const void *__restrict__ weight;
+  const void *__restrict__ indice;
+  std::size_t num_warps;
+};
+
+struct MaskedKernelParams {
+  IndexKernelParams params;
+  std::size_t start;
+  std::size_t length;
+};
+
+template <std::size_t kNumThreads, std::size_t kMaxOccupancy, bool kUsePDL,
+          std::size_t kElementSize, std::size_t kNumSplits, std::integral T>
+__global__ __launch_bounds__(kNumThreads, kMaxOccupancy) void //
+    index_kernel(const __grid_constant__ IndexKernelParams params) {
+  using namespace device;
+  constexpr auto kSize = kElementSize;
+  constexpr auto kSizePerWarp = kSize / kNumSplits;
+  constexpr auto kWarpPerBlock = static_cast<unsigned>(kNumThreads / 32);
+
+  static_assert(kNumThreads % 32 == 0);
+  static_assert(std::has_single_bit(kNumSplits));
+  static_assert(kElementSize % kNumSplits == 0);
+
+  const auto &[output, weight, indices_, num_warps] = params;
+  const auto indices = static_cast<const T *>(indices_);
+  const auto warp_id =
+      (threadIdx.x / kWarpThreads) + blockIdx.x * kWarpPerBlock;
+  PDL::wait<kUsePDL>();
+
+  if (warp_id < num_warps) {
+    const auto pos = indices[warp_id / kNumSplits];
+    const auto dst = pointer::offset(output, warp_id * kSizePerWarp);
+    const auto src = pointer::offset(weight, pos * kSize,
+                                     (warp_id % kNumSplits) * kSizePerWarp);
+    warp::copy<kSizePerWarp>(dst, src);
+  }
+
+  PDL::launch<kUsePDL>();
+}
+
+template <std::size_t kNumThreads, std::size_t kMaxOccupancy, bool kUsePDL,
+          std::size_t kElementSize, std::size_t kNumSplits, std::integral T>
+__global__ __launch_bounds__(kNumThreads, kMaxOccupancy) void //
+    masked_index_kernel(
+        const __grid_constant__ MaskedKernelParams mask_params) {
+  using namespace device;
+  constexpr auto kSize = kElementSize;
+  constexpr auto kSizePerWarp = kSize / kNumSplits;
+  constexpr auto kWarpPerBlock = static_cast<unsigned>(kNumThreads / 32);
+
+  static_assert(kNumThreads % 32 == 0);
+  static_assert(std::has_single_bit(kNumSplits));
+  static_assert(kElementSize % kNumSplits == 0);
+
+  const auto &[params, start, length] = mask_params;
+  const auto &[output, weight, indices_, num_warps] = params;
+  const auto indices = static_cast<const T *>(indices_);
+  const auto warp_id =
+      (threadIdx.x / kWarpThreads) + blockIdx.x * kWarpPerBlock;
+
+  PDL::wait<kUsePDL>();
+
+  if (warp_id < num_warps) {
+    const auto pos = indices[warp_id / kNumSplits] - start;
+    const auto dst = pointer::offset(output, warp_id * kSizePerWarp);
+    if (pos < length) {
+      const auto src = pointer::offset(weight, pos * kSize,
+                                       (warp_id % kNumSplits) * kSizePerWarp);
+      warp::copy<kSizePerWarp>(dst, src);
+    } else {
+      warp::reset<kSizePerWarp>(dst);
+    }
+  }
+
+  PDL::launch<kUsePDL>();
+}
+
+template <std::size_t element_size,   // depends on data type and embedding dim
+          std::size_t num_splits = 1, // how many warps handles one element
+          std::size_t num_threads = 128,   // number of threads per block
+          std::size_t max_concurrency = 1, // max blocks per SM
+          bool use_pdl = false>
+struct IndexKernel {
+  static void run(const tvm::ffi::TensorView weights,
+                  const tvm::ffi::TensorView indices,
+                  const tvm::ffi::TensorView output,
+                  tvm::ffi::Optional<tvm::ffi::Tuple<int, int>> mask_opts) {
+    using namespace host;
+    auto D = SymbolicSize{"D"}; // embedding size
+    auto L = SymbolicSize{"L"}; // num indices
+    auto device_ = SymbolicDevice{};
+    auto weights_dtype_ = SymbolicDType{};
+    auto indices_dtype_ = SymbolicDType{};
+
+    TensorMatcher({-1, D}) //
+        .with_dtype(weights_dtype_)
+        .with_device<kDLCUDA>(device_)
+        .verify(weights);
+    TensorMatcher({L, D}) //
+        .with_dtype(weights_dtype_)
+        .with_device<kDLCUDA>(device_)
+        .verify(output);
+    TensorMatcher({L}) //
+        .with_dtype<int32_t, int64_t>(indices_dtype_)
+        .with_device<kDLCUDA>(device_)
+        .verify(indices);
+
+    const auto device = device_.unwrap();
+    const auto use_int32 = indices_dtype_.unwrap().bits == 32;
+    const auto num_indices = L.unwrap();
+    const auto entry_size = dtype_bytes(weights_dtype_.unwrap()) * D.unwrap();
+    RuntimeCheck(entry_size == element_size,
+                 "IndexKernel: element_size mismatch.");
+
+    constexpr auto kWarpPerBlock = num_threads / 32;
+    const auto num_warps = num_splits * num_indices;
+    const auto num_blocks = div_ceil(num_warps, kWarpPerBlock);
+    const auto params = IndexKernelParams{
+        .output = static_cast<char *>(output.data_ptr()),
+        .weight = static_cast<const char *>(weights.data_ptr()),
+        .indice = indices.data_ptr(),
+        .num_warps = num_warps,
+    };
+
+    if (mask_opts.has_value()) {
+      const auto &obj = mask_opts.value();
+      const auto [start, length] = obj;
+      const auto m_params = MaskedKernelParams{
+          .params = params,
+          .start = static_cast<std::size_t>(start),
+          .length = static_cast<std::size_t>(length),
+      };
+      const auto kernel =
+          use_int32 ? masked_index_kernel<num_threads, max_concurrency, use_pdl,
+                                          element_size, num_splits, int32_t>
+                    : masked_index_kernel<num_threads, max_concurrency, use_pdl,
+                                          element_size, num_splits, int64_t>;
+      LaunchKernel(num_blocks, num_threads, device)
+          .with_attr(use_pdl)(kernel, m_params);
+    } else {
+      const auto kernel =
+          use_int32 ? index_kernel<num_threads, max_concurrency, use_pdl,
+                                   element_size, num_splits, int32_t>
+                    : index_kernel<num_threads, max_concurrency, use_pdl,
+                                   element_size, num_splits, int64_t>;
+      LaunchKernel(num_blocks, num_threads, device)
+          .with_attr(use_pdl)(kernel, params);
+    }
+  }
+};
+
+} // namespace

zonos2/kernel/csrc/jit/store.cu

@@ -0,0 +1,123 @@
+#include <zonos2/tensor.h>
+#include <zonos2/utils.cuh>
+#include <zonos2/utils.h>
+#include <zonos2/warp.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <concepts>
+#include <cstddef>
+#include <cstdint>
+
+namespace {
+
+struct StoreKernelParams {
+  void *__restrict__ k_cache;
+  void *__restrict__ v_cache;
+  const void *__restrict__ indices;
+  const void *__restrict__ k;
+  const void *__restrict__ v;
+  std::size_t kv_cache_stride;
+  std::size_t kv_input_stride;
+  std::size_t length;
+};
+
+template <std::size_t kNumThreads, std::size_t kMaxOccupancy, bool kUsePDL,
+          std::size_t kElementSize, std::integral T>
+__global__ __launch_bounds__(kNumThreads, kMaxOccupancy) void //
+    store_kv_cache(const __grid_constant__ StoreKernelParams params) {
+  using namespace device;
+
+  constexpr auto kWarpPerBlock =
+      static_cast<unsigned>(kNumThreads / kWarpThreads);
+  static_assert(kNumThreads % kWarpThreads == 0);
+
+  const auto &[k_cache, v_cache, indices, k, v, kv_cache_stride,
+               kv_input_stride, length] = params;
+  const auto warp_id =
+      (threadIdx.x / kWarpThreads) + blockIdx.x * kWarpPerBlock;
+  PDL::wait<kUsePDL>();
+
+  // each warp handles one element
+  if (warp_id < length) {
+    const auto pos = static_cast<const T *>(indices)[warp_id];
+    const auto dst_k = pointer::offset(k_cache, pos * kv_cache_stride);
+    const auto src_k = pointer::offset(k, warp_id * kv_input_stride);
+    warp::copy<kElementSize>(dst_k, src_k);
+    const auto dst_v = pointer::offset(v_cache, pos * kv_cache_stride);
+    const auto src_v = pointer::offset(v, warp_id * kv_input_stride);
+    warp::copy<kElementSize>(dst_v, src_v);
+  }
+
+  PDL::launch<kUsePDL>();
+}
+
+template <std::size_t element_size, // depends on data type and embedding dim
+          std::size_t num_threads = 128,   // number of threads per block
+          std::size_t max_concurrency = 1, // max blocks per SM
+          bool use_pdl = false>
+struct StoreKernel {
+  static void run(const tvm::ffi::TensorView k_cache,
+                  const tvm::ffi::TensorView v_cache,
+                  const tvm::ffi::TensorView indices,
+                  const tvm::ffi::TensorView k, const tvm::ffi::TensorView v) {
+    using namespace host;
+    auto D = SymbolicSize{"D"}; // element size
+    auto L = SymbolicSize{"L"}; // length
+    auto X = SymbolicSize{"X"}; // stride kv cache
+    auto Y = SymbolicSize{"Y"}; // stride kv input
+    auto indices_dtype_ = SymbolicDType{};
+    auto dtype_ = SymbolicDType{};
+    auto device_ = SymbolicDevice{};
+
+    TensorMatcher({-1, D}) //
+        .with_strides({X, 1})
+        .with_device<kDLCUDA>(device_)
+        .with_dtype(dtype_)
+        .verify(k_cache)
+        .verify(v_cache);
+    TensorMatcher({L, D}) //
+        .with_strides({Y, 1})
+        .with_device<kDLCUDA>(device_)
+        .with_dtype(dtype_)
+        .verify(k)
+        .verify(v);
+    TensorMatcher({L}) //
+        .with_device<kDLCUDA>(device_)
+        .with_dtype<int32_t, int64_t>(indices_dtype_)
+        .verify(indices);
+
+    const auto dtype_size = dtype_bytes(dtype_.unwrap());
+    RuntimeCheck(element_size == dtype_size * D.unwrap());
+
+    const auto device = device_.unwrap();
+    const auto use_int32 = indices_dtype_.unwrap().bits == 32;
+    const auto length = static_cast<std::size_t>(L.unwrap());
+    const auto kv_cache_stride = X.unwrap() * dtype_size;
+    const auto kv_input_stride = Y.unwrap() * dtype_size;
+
+    const auto params = StoreKernelParams{
+        .k_cache = k_cache.data_ptr(),
+        .v_cache = v_cache.data_ptr(),
+        .indices = indices.data_ptr(),
+        .k = k.data_ptr(),
+        .v = v.data_ptr(),
+        .kv_cache_stride = kv_cache_stride,
+        .kv_input_stride = kv_input_stride,
+        .length = length,
+    };
+
+    constexpr auto kWarpPerBlock = num_threads / 32;
+    static_assert(num_threads % 32 == 0);
+    const auto num_blocks = div_ceil(length, kWarpPerBlock);
+    const auto kernel = use_int32
+                            ? store_kv_cache<num_threads, max_concurrency,
+                                             use_pdl, element_size, int32_t>
+                            : store_kv_cache<num_threads, max_concurrency,
+                                             use_pdl, element_size, int64_t>;
+    LaunchKernel(num_blocks, num_threads, device)
+        .with_attr(use_pdl)(kernel, params);
+  }
+};
+
+} // namespace

zonos2/kernel/csrc/src/pynccl.cu

@@ -0,0 +1,188 @@
+#include <zonos2/nccl227.h>
+#include <zonos2/tensor.h>
+#include <zonos2/utils.cuh>
+#include <zonos2/utils.h>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/registry.h>
+
+#include <bit>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <string_view>
+#include <unordered_map>
+
+namespace {
+
+using NCCLIDList = tvm::ffi::Array<char>;
+
+auto NCCL_CHECK(::ncclResult_t result) -> void {
+  if (result != ::ncclSuccess) {
+    host::RuntimeCheck(false, ::ncclGetErrorString(result));
+  }
+}
+
+auto get_uid(const NCCLIDList &wrapper) -> ncclUniqueId {
+  host::RuntimeCheck(wrapper.size() == NCCL_UNIQUE_ID_BYTES,
+                     "Invalid NCCL ID wrapper size");
+  ncclUniqueId id;
+  std::copy(wrapper.begin(), wrapper.end(), id.internal);
+  return id;
+}
+
+auto create_uid() -> NCCLIDList {
+  ncclUniqueId id;
+  NCCL_CHECK(::ncclGetUniqueId(&id));
+  return NCCLIDList(id.internal, id.internal + NCCL_UNIQUE_ID_BYTES);
+}
+
+const auto kNCCLReduceOPMap = std::unordered_map<std::string_view, ncclRedOp_t>{
+    {"sum", ::ncclSum}, {"prod", ::ncclProd}, {"max", ::ncclMax},
+    {"min", ::ncclMin}, {"avg", ::ncclAvg},
+};
+
+struct DLDataTypeHash {
+  auto operator()(const DLDataType &dtype) const noexcept -> std::size_t {
+    return std::bit_cast<std::uint32_t>(dtype);
+  }
+};
+
+template <typename = void>
+auto operator==(const DLDataType &a, const DLDataType &b) -> bool {
+  return a.code == b.code && a.bits == b.bits && a.lanes == b.lanes;
+}
+
+const auto kNCCLDtypeMap =
+    std::unordered_map<DLDataType, ncclDataType_t, DLDataTypeHash>{
+        {{DLDataTypeCode::kDLFloat, 16, 1}, ncclFloat16},
+        {{DLDataTypeCode::kDLBfloat, 16, 1}, ncclBfloat16},
+    };
+
+using std::shared_ptr;
+
+template <typename T> using shared_obj = shared_ptr<std::remove_pointer_t<T>>;
+template <auto Fn>
+inline constexpr auto template_fn =
+    [](auto &&...args) { return Fn(std::forward<decltype(args)>(args)...); };
+
+struct NCCLWrapper : public tvm::ffi::Object {
+public:
+  NCCLWrapper(int rank, int world_size, const size_t max_bytes, NCCLIDList uid)
+      : m_rank(rank), m_world_size(world_size), m_max_bytes(max_bytes) {
+    ncclUniqueId id = get_uid(uid);
+    ncclComm_t comm;
+    NCCL_CHECK(::ncclCommInitRank(&comm, m_world_size, id, m_rank));
+    m_comm = {comm, template_fn<::ncclCommDestroy>};
+
+    void *buf;
+    NCCL_CHECK(::ncclMemAlloc(&buf, max_bytes));
+    m_sym_mem = {buf, template_fn<::ncclMemFree>};
+
+    ncclWindow_t win;
+    NCCL_CHECK(::ncclCommWindowRegister(comm, buf, max_bytes, &win,
+                                        NCCL_WIN_COLL_SYMMETRIC));
+    m_win = {win, [comm = m_comm](ncclWindow_t w) {
+               return NCCL_CHECK(::ncclCommWindowDeregister(comm.get(), w));
+             }};
+  }
+
+  auto all_reduce(tvm::ffi::TensorView t, std::string op) const -> void {
+    using namespace host;
+    RuntimeCheck(t.device().device_type == kDLCUDA,
+                 "Tensor must be on CUDA device");
+    RuntimeCheck(t.is_contiguous(), "Tensor must be contiguous");
+    const auto size_dim = static_cast<size_t>(t.shape().Product());
+    const auto dtype = kNCCLDtypeMap.at(t.dtype());
+    const auto size_bytes = size_dim * (t.dtype().bits / 8);
+    const auto data_ptr = t.data_ptr();
+    const auto reduce_op = kNCCLReduceOPMap.at(op);
+    const auto stream = LaunchKernel::resolve_device(t.device());
+
+    if (size_bytes <= m_max_bytes) { // use internal buffer
+      const auto buf_ptr = m_sym_mem.get();
+      const auto need_memcpy = (buf_ptr != data_ptr);
+      if (need_memcpy) {
+        CUDA_CHECK(::cudaMemcpyAsync(buf_ptr, data_ptr, size_bytes,
+                                     ::cudaMemcpyDeviceToDevice, stream));
+      }
+      NCCL_CHECK(::ncclAllReduce(
+          /*sendbuff=*/buf_ptr,
+          /*recvbuff=*/buf_ptr,
+          /*count=*/size_dim,
+          /*datatype=*/dtype,
+          /*op=*/reduce_op,
+          /*comm=*/m_comm.get(),
+          /*stream=*/stream));
+      if (need_memcpy) {
+        CUDA_CHECK(::cudaMemcpyAsync(data_ptr, buf_ptr, size_bytes,
+                                     ::cudaMemcpyDeviceToDevice, stream));
+      }
+    } else {
+      NCCL_CHECK(::ncclAllReduce(
+          /*sendbuff=*/data_ptr,
+          /*recvbuff=*/data_ptr,
+          /*count=*/size_dim,
+          /*datatype=*/dtype,
+          /*op=*/reduce_op,
+          /*comm=*/m_comm.get(),
+          /*stream=*/stream));
+    }
+  }
+
+  auto all_gather(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) const
+      -> void {
+    using namespace host;
+    RuntimeCheck(src.device().device_type == kDLCUDA,
+                 "Tensor must be on CUDA device");
+    RuntimeCheck(src.is_contiguous(), "Tensor must be contiguous");
+    RuntimeCheck(dst.device().device_type == kDLCUDA,
+                 "Tensor must be on CUDA device");
+    RuntimeCheck(dst.is_contiguous(), "Tensor must be contiguous");
+    RuntimeCheck(dst.size(0) == src.size(0) * m_world_size,
+                 "Destination tensor has incorrect size");
+    const auto size_dim = static_cast<size_t>(src.shape().Product());
+    const auto dtype = kNCCLDtypeMap.at(src.dtype());
+    const auto src_ptr = src.data_ptr();
+    const auto dst_ptr = dst.data_ptr();
+    const auto stream = LaunchKernel::resolve_device(src.device());
+    // do not use internal buffer for all_gather, directly gather to output
+    // tensor
+    NCCL_CHECK(::ncclAllGather(
+        /*sendbuff=*/src_ptr,
+        /*recvbuff=*/dst_ptr,
+        /*sendcount=*/size_dim,
+        /*datatype=*/dtype,
+        /*comm=*/m_comm.get(),
+        /*stream=*/stream));
+  }
+
+  auto get_buffer() const -> void * { return m_sym_mem.get(); }
+
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("zonos2.NCCLWrapper", NCCLWrapper,
+                                    tvm::ffi::Object);
+
+private:
+  int m_rank;
+  int m_world_size;
+  size_t m_max_bytes;
+  shared_obj<ncclComm_t> m_comm;
+  shared_ptr<void> m_sym_mem;
+  shared_obj<ncclWindow_t> m_win;
+};
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+  namespace refl = tvm::ffi::reflection;
+  refl::ObjectDef<NCCLWrapper>()
+      .def(refl::init<int, int, size_t, NCCLIDList>(), "__init__")
+      .def("all_reduce", &NCCLWrapper::all_reduce)
+      .def("all_gather", &NCCLWrapper::all_gather)
+      .def("get_buffer", &NCCLWrapper::get_buffer);
+}
+
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(create_nccl_uid, &create_uid);
+
+} // namespace

zonos2/kernel/csrc/src/radix.cpp

@@ -0,0 +1,44 @@
+#include <zonos2/utils.h>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/extra/c_env_api.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/object.h>
+
+namespace {
+
+auto _is_1d_cpu_int_tensor(const tvm::ffi::TensorView tensor) -> bool {
+  return tensor.ndim() == 1 && tensor.is_contiguous() &&
+         tensor.device().device_type == kDLCPU &&
+         (tensor.dtype().code == kDLInt) &&
+         (tensor.dtype().bits == 32 || tensor.dtype().bits == 64);
+}
+
+auto fast_compare_key(const tvm::ffi::TensorView a,
+                      const tvm::ffi::TensorView b) -> size_t {
+  host::RuntimeCheck(_is_1d_cpu_int_tensor(a) && _is_1d_cpu_int_tensor(b),
+                     "Both tensors must be 1D CPU int tensors.");
+  host::RuntimeCheck(a.dtype() == b.dtype());
+  const auto a_ptr = a.data_ptr();
+  const auto b_ptr = b.data_ptr();
+  const auto common_len = std::min(a.size(0), b.size(0));
+  if (a.dtype().bits == 64) {
+    const auto a_ptr_64 = static_cast<const int64_t *>(a_ptr);
+    const auto b_ptr_64 = static_cast<const int64_t *>(b_ptr);
+    const auto diff_pos =
+        std::mismatch(a_ptr_64, a_ptr_64 + common_len, b_ptr_64);
+    return static_cast<size_t>(diff_pos.first - a_ptr_64);
+  } else {
+    const auto a_ptr_32 = static_cast<const int32_t *>(a_ptr);
+    const auto b_ptr_32 = static_cast<const int32_t *>(b_ptr);
+    const auto diff_pos =
+        std::mismatch(a_ptr_32, a_ptr_32 + common_len, b_ptr_32);
+    return static_cast<size_t>(diff_pos.first - a_ptr_32);
+  }
+}
+
+} // namespace
+
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(fast_compare_key, fast_compare_key);

zonos2/kernel/csrc/src/tensor.cpp

@@ -0,0 +1,31 @@
+#include <zonos2/tensor.h>
+#include <zonos2/utils.h>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/extra/c_env_api.h>
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/object.h>
+
+namespace {
+
+auto test(tvm::ffi::TensorView x, tvm::ffi::TensorView y) -> void {
+  auto N = host::SymbolicSize{"N"};
+  const auto M = 1024;
+  host::TensorMatcher({N, M})
+      .with_strides({-1, 1}) // -1 means any
+      .with_dtype<int, float>()
+      .with_device<kDLCPU>()
+      .verify(x);
+  host::TensorMatcher({N, M}) // default contiguous
+      .with_dtype({{kDLInt, 32, 1}, {kDLInt, 64, 1}})
+      .with_device({{kDLCUDA, 1}})
+      .verify(y);
+  host::RuntimeCheck(N.unwrap() % 4 == 0);
+}
+
+} // namespace
+
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(test, test);

zonos2/kernel/index.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import TYPE_CHECKING, Tuple
+
+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)
+
+
+@lru_cache(maxsize=None)
+def _jit_index_module(
+    element_size: int,
+    *,
+    num_splits: int = 1,
+    config: KernelConfig = DEFAULT_INDEX_KERNEL_CONFIG,
+) -> Module:
+    args = make_cpp_args(element_size, num_splits, *config)
+    return load_jit(
+        "index",
+        *args,
+        cuda_files=["index.cu"],
+        cuda_wrappers=[("launch", f"IndexKernel<{args}>::run")],
+    )
+
+
+def indexing(
+    weights: torch.Tensor,
+    indices: torch.Tensor,
+    *,
+    output: torch.Tensor | None = None,
+    vocab_range: Tuple[int, int] | None = None,  # (start, length)
+) -> torch.Tensor:
+    if output is None:
+        output = weights.new_empty(indices.shape[0], weights.shape[1])
+
+    element_size = weights.shape[1] * weights.element_size()
+    if element_size % 2048 == 0:
+        num_splits = 4
+    elif element_size % 1024 == 0:
+        num_splits = 2
+    else:
+        num_splits = 1
+    module = _jit_index_module(element_size, num_splits=num_splits)
+    module.launch(weights, indices, output, vocab_range)
+    return output

zonos2/kernel/moe_impl.py

@@ -0,0 +1,59 @@
+from typing import Any, Dict
+
+import torch
+import triton
+import triton.language as tl
+from zonos2.kernel.triton.fused_moe import fused_moe_kernel
+
+
+def fused_moe_kernel_triton(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    mul_routed_weight: bool,
+    top_k: int,
+    config: Dict[str, Any],
+    compute_type: tl.dtype,
+) -> None:
+    assert topk_weights.stride(1) == 1
+    assert sorted_token_ids.stride(0) == 1
+    padded_size = 0
+    grid = lambda META: (
+        triton.cdiv(sorted_token_ids.shape[0], META["BLOCK_SIZE_M"])
+        * triton.cdiv(B.shape[1], META["BLOCK_SIZE_N"]),
+    )
+    K = B.shape[2] - padded_size
+    if K % config["BLOCK_SIZE_K"] == 0:
+        even_Ks = True
+    else:
+        even_Ks = False
+    fused_moe_kernel[grid](
+        A,
+        B,
+        C,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        B.shape[1],
+        B.shape[2] - padded_size,
+        sorted_token_ids.shape[0],
+        topk_ids.numel(),
+        A.stride(0),
+        A.stride(1),
+        B.stride(0),
+        B.stride(2),
+        B.stride(1),
+        C.stride(1),
+        C.stride(2),
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        top_k=top_k,
+        compute_type=compute_type,
+        even_Ks=even_Ks,
+        **config,
+    )

zonos2/kernel/pynccl.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import TYPE_CHECKING, Any, Literal
+
+from zonos2.env import ENV
+
+from .utils import load_aot
+
+if TYPE_CHECKING:
+    from abc import abstractmethod
+
+    import torch
+    from tvm_ffi import Module
+
+    class PyNCCLCommunicator:
+        @abstractmethod
+        def all_reduce(self, input: torch.Tensor, op: Literal["sum"]) -> None: ...
+        @abstractmethod
+        def all_gather(self, output: torch.Tensor, input: torch.Tensor) -> None: ...
+        @abstractmethod
+        def get_buffer(self) -> int: ...
+
+else:
+    PyNCCLCommunicator = Any
+
+
+@lru_cache(maxsize=None)
+def _load_nccl_module() -> Module:
+    return load_aot("pynccl", cuda_files=["pynccl.cu"], extra_ldflags=["-lnccl"])
+
+
+@lru_cache(maxsize=None)
+def _get_pynccl_wrapper_cls():
+    import tvm_ffi
+
+    @tvm_ffi.register_object("zonos2.NCCLWrapper")
+    class PyNCCLImpl(tvm_ffi.Object):
+        def __init__(self, *args):
+            self.__ffi_init__(*args)
+
+    return PyNCCLImpl
+
+
+def init_pynccl(
+    *,
+    tp_rank: int,
+    tp_size: int,
+    tp_cpu_group: torch.distributed.ProcessGroup,
+    max_size_bytes: int = 0,
+) -> PyNCCLCommunicator:
+    import torch
+
+    max_size_bytes = min(max_size_bytes, ENV.PYNCCL_MAX_BUFFER_SIZE.value)
+
+    module = _load_nccl_module()
+    cls = _get_pynccl_wrapper_cls()
+
+    if tp_rank == 0:
+        id_list = [module.create_nccl_uid()]
+        torch.distributed.broadcast_object_list(
+            id_list,
+            src=0,
+            group=tp_cpu_group,
+        )
+    else:
+        id_list = [None]
+        torch.distributed.broadcast_object_list(
+            id_list,
+            src=0,
+            group=tp_cpu_group,
+        )
+
+    nccl_id = id_list[0]
+    assert not nccl_id is None, f"Failed to get NCCL unique ID on {tp_rank = }"
+
+    # bypass type checking for the FFI object
+    return cls(tp_rank, tp_size, max_size_bytes, nccl_id)  # type: ignore

zonos2/kernel/radix.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import TYPE_CHECKING
+
+from .utils import load_aot
+
+if TYPE_CHECKING:
+    import torch
+    from tvm_ffi import Module
+
+
+@lru_cache(maxsize=None)
+def _load_radix_module() -> Module:
+    return load_aot("radix", cpp_files=["radix.cpp"])
+
+
+def fast_compare_key(x: torch.Tensor, y: torch.Tensor) -> int:
+    # compare 2 1-D int cpu tensors for equality
+    return _load_radix_module().fast_compare_key(x, y)

zonos2/kernel/store.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import TYPE_CHECKING
+
+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)
+
+
+@lru_cache(maxsize=None)
+def _jit_store_module(
+    element_size: int,
+    *,
+    config: KernelConfig = DEFAULT_INDEX_KERNEL_CONFIG,
+) -> Module:
+    args = make_cpp_args(element_size, *config)
+    return load_jit(
+        "store",
+        *args,
+        cuda_files=["store.cu"],
+        cuda_wrappers=[("launch", f"StoreKernel<{args}>::run")],
+    )
+
+
+def store_cache(
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    indices: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+) -> None:
+    num_tokens = k_cache.shape[0]
+    k_cache = k_cache.view(num_tokens, -1)
+    v_cache = v_cache.view(num_tokens, -1)
+    element_size = k_cache.shape[1] * k_cache.element_size()
+    module = _jit_store_module(element_size)
+    module.launch(k_cache, v_cache, indices, k, v)

zonos2/kernel/tensor.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import TYPE_CHECKING
+
+from .utils import load_aot
+
+if TYPE_CHECKING:
+    import torch
+    from tvm_ffi import Module
+
+
+@lru_cache(maxsize=None)
+def _load_test_tensor_module() -> Module:
+    return load_aot("test_tensor", cpp_files=["tensor.cpp"])
+
+
+def test_tensor(x: torch.Tensor, y: torch.Tensor) -> int:
+    return _load_test_tensor_module().test(x, y)

zonos2/kernel/triton/fused_moe.py

@@ -0,0 +1,230 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _moe_sum_reduce_kernel(
+    input_ptr,
+    input_stride_0,
+    input_stride_1,
+    input_stride_2,
+    output_ptr,
+    output_stride_0,
+    output_stride_1,
+    token_num: int,
+    topk_num: int,
+    hidden_dim: int,
+    routed_scaling_factor: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DIM: tl.constexpr,
+    NUM_STAGE: tl.constexpr,
+):
+    input_stride_0 = tl.cast(input_stride_0, dtype=tl.int64)
+    input_stride_1 = tl.cast(input_stride_1, dtype=tl.int64)
+    output_stride_0 = tl.cast(output_stride_0, dtype=tl.int64)
+
+    token_block_id = tl.program_id(0)
+    dim_block_id = tl.program_id(1)
+
+    token_start = token_block_id * BLOCK_M
+    token_end = min((token_block_id + 1) * BLOCK_M, token_num)
+
+    dim_start = dim_block_id * BLOCK_DIM
+    dim_end = min((dim_block_id + 1) * BLOCK_DIM, hidden_dim)
+
+    offs_dim = dim_start + tl.arange(0, BLOCK_DIM)
+
+    for token_index in range(token_start, token_end):
+        accumulator = tl.zeros((BLOCK_DIM,), dtype=tl.float32)
+        input_t_ptr = input_ptr + token_index * input_stride_0 + offs_dim
+        for i in tl.range(0, topk_num, num_stages=NUM_STAGE):
+            tmp = tl.load(input_t_ptr + i * input_stride_1, mask=offs_dim < dim_end, other=0.0)
+            accumulator += tmp
+        accumulator = accumulator * routed_scaling_factor
+        store_t_ptr = output_ptr + token_index * output_stride_0 + offs_dim
+        tl.store(
+            store_t_ptr,
+            accumulator.to(input_ptr.dtype.element_ty),
+            mask=offs_dim < dim_end,
+        )
+
+
+@triton.jit
+def fused_moe_kernel(
+    # Pointers to matrices
+    a_ptr,
+    b_ptr,
+    c_ptr,
+    topk_weights_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    # Matrix dimensions
+    N,
+    K,
+    EM,
+    num_valid_tokens,
+    # The stride variables represent how much to increase the ptr by when
+    # moving by 1 element in a particular dimension. E.g. `stride_am` is
+    # how much to increase `a_ptr` by to get the element one row down
+    # (A has M rows).
+    stride_am,
+    stride_ak,
+    stride_be,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    MUL_ROUTED_WEIGHT: tl.constexpr,
+    top_k: tl.constexpr,
+    compute_type: tl.constexpr,
+    even_Ks: tl.constexpr,
+):
+    """
+    Implements the fused computation for a Mixture of Experts (MOE) using
+    token and expert matrices.
+
+    Key Parameters:
+    - A: The input tensor representing tokens with shape (*, K), where '*' can
+        be any shape representing batches and K is the feature dimension of
+        each token.
+    - B: The stacked MOE weight tensor with shape (E, N, K), where E is
+        the number of experts, K is the input feature dimension, and N is
+        the output feature dimension.
+    - C: The output cache tensor with shape (M, topk, N), where M is the
+        total number of tokens post padding, topk is the number of times
+        each token is repeated, and N is the output feature dimension.
+    - sorted_token_ids: A tensor containing the sorted indices of tokens,
+        repeated topk times and arranged by the expert index they are
+        assigned to.
+    - expert_ids: A tensor containing the indices of the expert for each
+        block. It determines which expert matrix from B should be used for
+        each block in A.
+
+    This kernel performs the multiplication of a token by its corresponding
+    expert matrix as determined by `expert_ids`. The sorting of
+    `sorted_token_ids` by expert index and padding ensures divisibility by
+    BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix
+    multiplication across different blocks processed by the same expert.
+    """
+    # -----------------------------------------------------------
+    # Map program ids `pid` to the block of C it should compute.
+    # This is done in a grouped ordering to promote L2 data reuse.
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    # ----------------------------------------------------------
+    # Create pointers for the first blocks of A and B.
+    # We will advance this pointer as we move in the K direction
+    # and accumulate
+    # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers
+    # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
+        return
+    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
+    offs_token = offs_token.to(tl.int64)
+    token_mask = offs_token < num_valid_tokens
+
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = a_ptr + (offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak)
+
+    off_experts = tl.load(expert_ids_ptr + pid_m)
+    b_ptrs = (
+        b_ptr
+        + off_experts * stride_be
+        + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+    )
+
+    # -----------------------------------------------------------
+    # Iterate to compute a block of the C matrix.
+    # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
+    # of fp32 values for higher accuracy.
+    # `accumulator` will be converted back to fp16 after the loop.
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        # Load the next block of A and B, generate a mask by checking the
+        # K dimension.
+        if even_Ks:
+            a = tl.load(
+                a_ptrs,
+                mask=token_mask[:, None],
+                other=0.0,
+            )
+            b = tl.load(b_ptrs)
+        else:
+            a = tl.load(
+                a_ptrs,
+                mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
+                other=0.0,
+            )
+            b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+
+        # We accumulate along the K dimension.
+
+        accumulator += tl.dot(a, b)
+        # Advance the ptrs to the next K block.
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * stride_bk
+
+    if MUL_ROUTED_WEIGHT:
+        moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
+        accumulator = accumulator * moe_weight[:, None]
+
+    accumulator = accumulator.to(compute_type)
+    # -----------------------------------------------------------
+    # Write back the block of the output
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
+    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
+    tl.store(c_ptrs, accumulator, mask=c_mask)
+
+
+def moe_sum_reduce_triton(input: torch.Tensor, output: torch.Tensor, routed_scaling_factor: float):
+    assert input.is_contiguous()
+    assert output.is_contiguous()
+
+    token_num, topk_num, hidden_dim = input.shape
+    assert output.shape[0] == token_num and output.shape[1] == hidden_dim
+
+    BLOCK_M = 1
+    BLOCK_DIM = 2048
+    NUM_STAGE = 1
+    num_warps = 8
+
+    grid = (
+        triton.cdiv(token_num, BLOCK_M),
+        triton.cdiv(hidden_dim, BLOCK_DIM),
+    )
+
+    _moe_sum_reduce_kernel[grid](
+        input,
+        *input.stride(),
+        output,
+        *output.stride(),
+        token_num=token_num,
+        topk_num=topk_num,
+        hidden_dim=hidden_dim,
+        routed_scaling_factor=routed_scaling_factor,
+        BLOCK_M=BLOCK_M,
+        BLOCK_DIM=BLOCK_DIM,
+        NUM_STAGE=NUM_STAGE,
+        num_warps=num_warps,
+    )
+    return

zonos2/kernel/utils.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import pathlib
+from typing import TYPE_CHECKING, List, NamedTuple, Tuple, TypeAlias, Union
+
+if TYPE_CHECKING:
+    from tvm_ffi import Module
+
+KERNEL_PATH = pathlib.Path(__file__).parent / "csrc"
+DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")]
+DEFAULT_CFLAGS = ["-std=c++20", "-O3"]
+DEFAULT_CUDA_CFLAGS = ["-std=c++20", "-O3", "--expt-relaxed-constexpr"]
+DEFAULT_LDFLAGS = []
+CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool]
+
+
+class CppArgList(list[str]):
+    def __str__(self) -> str:
+        return ", ".join(self)
+
+
+class KernelConfig(NamedTuple):
+    num_threads: int
+    max_occupancy: int
+    use_pdl: bool
+
+    @property
+    def template_args(self) -> str:
+        pdl = "true" if self.use_pdl else "false"
+        return f"{self.num_threads},{self.max_occupancy},{pdl}"
+
+
+def _make_name(*args: str) -> str:
+    return "zonos2__" + "_".join(str(arg) for arg in args)
+
+
+def _make_wrapper(tup: Tuple[str, str]) -> str:
+    export_name, kernel_name = tup
+    return f"TVM_FFI_DLL_EXPORT_TYPED_FUNC({export_name}, ({kernel_name}));"
+
+
+def make_cpp_args(*args: CPP_TEMPLATE_TYPE) -> CppArgList:
+    def _convert(arg: CPP_TEMPLATE_TYPE) -> str:
+        if isinstance(arg, bool):
+            return "true" if arg else "false"
+        if isinstance(arg, (int, float)):
+            return str(arg)
+        raise TypeError(f"Unsupported argument type for cpp template: {type(arg)}")
+
+    return CppArgList(_convert(arg) for arg in args)
+
+
+def load_aot(
+    *args: str,
+    cpp_files: List[str] | None = None,
+    cuda_files: List[str] | None = None,
+    extra_cflags: List[str] | None = None,
+    extra_cuda_cflags: List[str] | None = None,
+    extra_ldflags: List[str] | None = None,
+    extra_include_paths: List[str] | None = None,
+    build_directory: str | None = None,
+) -> Module:
+    from tvm_ffi.cpp import load
+
+    cpp_files = cpp_files or []
+    cuda_files = cuda_files or []
+    extra_cflags = extra_cflags or []
+    extra_cuda_cflags = extra_cuda_cflags or []
+    extra_ldflags = extra_ldflags or []
+    extra_include_paths = extra_include_paths or []
+
+    cpp_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cpp_files]
+    cuda_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cuda_files]
+
+    return load(
+        _make_name(*args),
+        cpp_files=cpp_files,
+        cuda_files=cuda_files,
+        extra_cflags=DEFAULT_CFLAGS + extra_cflags,
+        extra_cuda_cflags=DEFAULT_CUDA_CFLAGS + extra_cuda_cflags,
+        extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags,
+        extra_include_paths=DEFAULT_INCLUDE + extra_include_paths,
+        build_directory=build_directory,
+    )
+
+
+def load_jit(
+    *args: str,
+    cpp_files: List[str] | None = None,
+    cuda_files: List[str] | None = None,
+    cpp_wrappers: List[Tuple[str, str]] | None = None,
+    cuda_wrappers: List[Tuple[str, str]] | None = None,
+    extra_cflags: List[str] | None = None,
+    extra_cuda_cflags: List[str] | None = None,
+    extra_ldflags: List[str] | None = None,
+    extra_include_paths: List[str] | None = None,
+    build_directory: str | None = None,
+) -> Module:
+    from tvm_ffi.cpp import load_inline
+
+    cpp_files = cpp_files or []
+    cuda_files = cuda_files or []
+    cpp_wrappers = cpp_wrappers or []
+    cuda_wrappers = cuda_wrappers or []
+    extra_cflags = extra_cflags or []
+    extra_cuda_cflags = extra_cuda_cflags or []
+    extra_ldflags = extra_ldflags or []
+    extra_include_paths = extra_include_paths or []
+
+    # include cpp files
+    cpp_paths = [(KERNEL_PATH / "jit" / f).resolve() for f in cpp_files]
+    cpp_sources = [f'#include "{path}"' for path in cpp_paths]
+    cpp_sources += [_make_wrapper(tup) for tup in cpp_wrappers]
+
+    # include cuda files
+    cuda_paths = [(KERNEL_PATH / "jit" / f).resolve() for f in cuda_files]
+    cuda_sources = [f'#include "{path}"' for path in cuda_paths]
+    cuda_sources += [_make_wrapper(tup) for tup in cuda_wrappers]
+
+    return load_inline(
+        _make_name(*args),
+        cpp_sources=cpp_sources,
+        cuda_sources=cuda_sources,
+        extra_cflags=DEFAULT_CFLAGS + extra_cflags,
+        extra_cuda_cflags=DEFAULT_CUDA_CFLAGS + extra_cuda_cflags,
+        extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags,
+        extra_include_paths=DEFAULT_INCLUDE + extra_include_paths,
+        build_directory=build_directory,
+    )

zonos2/kvcache/__init__.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol
+
+from zonos2.utils import Registry
+
+if TYPE_CHECKING:
+    import torch
+    from zonos2.models import ModelConfig
+
+from .base import (
+    BaseCacheHandle,
+    BaseCacheManager,
+    BaseKVCache,
+    KVCacheLayout,
+    SizeInfo,
+)
+
+
+class CacheManagerCreator(Protocol):
+    def __call__(self, device: torch.device) -> BaseCacheManager: ...
+
+
+SUPPORTED_CACHE_MANAGER = Registry[CacheManagerCreator]("Cache Manager")
+
+
+def create_kvcache(
+    model_config: ModelConfig,
+    num_pages: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    cache_layout: KVCacheLayout = KVCacheLayout.LayerFirst,
+) -> BaseKVCache:
+    from .mha_pool import MHAKVCache  # TODO: support other variants (e.g. MLA)
+
+    return MHAKVCache(
+        num_kv_heads=model_config.num_kv_heads,
+        num_pages=num_pages,
+        kv_layout=cache_layout,
+        num_layers=model_config.num_layers,
+        head_dim=model_config.head_dim,
+        device=device,
+        dtype=dtype,
+    )
+
+
+@SUPPORTED_CACHE_MANAGER.register("naive")
+def create_naive_cache_manager(device: torch.device):
+    from .naive_manager import NaiveCacheManager
+
+    return NaiveCacheManager(device=device)
+
+
+@SUPPORTED_CACHE_MANAGER.register("radix")
+def create_radix_cache_manager(device: torch.device):
+    from .radix_manager import RadixCacheManager
+
+    return RadixCacheManager(device=device)
+
+
+def create_cache_manager(device: torch.device, type: str) -> BaseCacheManager:
+    return SUPPORTED_CACHE_MANAGER[type](device)
+
+
+__all__ = [
+    "create_kvcache",
+    "create_cache_manager",
+    "BaseKVCache",
+    "KVCacheLayout",
+    "BaseCacheHandle",
+    "BaseCacheManager",
+    "SizeInfo",
+    "SUPPORTED_CACHE_MANAGER",
+]

zonos2/kvcache/base.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+import enum
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import NamedTuple, Tuple
+
+import torch
+
+
+class BaseKVCache(ABC):
+    """
+    Base class for key-value caches.
+    This class defines the interface for key-value caches used.
+    """
+
+    @abstractmethod
+    def k_cache(self, index: int) -> torch.Tensor: ...
+
+    @abstractmethod
+    def v_cache(self, index: int) -> torch.Tensor: ...
+
+    @abstractmethod
+    def store_kv(
+        self, k: torch.Tensor, v: torch.Tensor, out_loc: torch.Tensor, layer_id: int
+    ) -> None: ...
+
+    @property
+    @abstractmethod
+    def device(self) -> torch.device: ...
+
+    @property
+    @abstractmethod
+    def dtype(self) -> torch.dtype: ...
+
+    @property
+    @abstractmethod
+    def num_layers(self) -> int: ...
+
+
+class KVCacheLayout(enum.Enum):
+    LayerFirst = enum.auto()
+    PageFirst = enum.auto()
+
+
+@dataclass(frozen=True)
+class BaseCacheHandle(ABC):
+    cached_len: int
+
+
+class SizeInfo(NamedTuple):
+    evictable_size: int
+    protected_size: int
+
+    @property
+    def total_size(self) -> int:
+        return self.evictable_size + self.protected_size
+
+
+class BaseCacheManager(ABC):
+    @abstractmethod
+    def match_prefix(self, input_ids: torch.Tensor) -> Tuple[BaseCacheHandle, torch.Tensor]:
+        """
+        Match prefix and return the indices of the matched prefix in the cache.
+        This operation will not modify the cache.
+        The returned indices is only safe to use when the handle is locked.
+
+        Args:
+            input_ids (torch.Tensor): The input ids to match. Shape: (seq_len,)
+        Returns:
+            handle (BaseCacheHandle): The handle to the matched prefix.
+            indices (torch.Tensor): The indices of the longest-matched prefix in the cache.
+        """
+
+    @abstractmethod
+    def lock_handle(self, handle: BaseCacheHandle, unlock: bool = False) -> None:
+        """
+        Lock or unlock a cache handle.
+        This operation will not modify the cache, but change the size info only.
+        When a handle is locked, it cannot be evicted.
+        Handles must be locked before the previously-returned tensor of `match_prefix` is used.
+        Otherwise it may be evicted by calling evict.
+
+        Args:
+            handle (BaseCacheHandle): The cache handle to lock or unlock.
+            unlock (bool): Whether to unlock the handle. Defaults to False.
+        """
+
+    @abstractmethod
+    def insert_prefix(self, input_ids: torch.Tensor, indices: torch.Tensor) -> int:
+        """
+        Insert a new prefix into the cache.
+        This operation will modify the cache.
+        Args:
+            input_ids (torch.Tensor): The input ids to insert. Shape: (seq_len,)
+            indices (torch.Tensor): The indices to store the new prefix. Shape: (seq_len,)
+
+        Returns:
+            int: The length of prefix that is already in the cache. This part is not
+                 inserted, so the caller should free these indices.
+        """
+
+    @abstractmethod
+    def evict(self, size: int) -> torch.Tensor:
+        """
+        Evict some prefixes from the cache to free up space.
+        This operation will modify the cache.
+        Note that evict 0 is always safe and does nothing.
+        Note that the actual evict size may be larger than the requested size.
+        Args:
+            size (int): The size to evict.
+
+        Returns:
+            torch.Tensor: The indices evicted. Shape: (evict_size,)
+        Raises:
+            RuntimeError: If the requested size is larger than the evictable size.
+        """
+
+    @abstractmethod
+    def reset(self) -> None:
+        """Reset the cache manager and the underlying cache."""
+
+    @property
+    @abstractmethod
+    def size_info(self) -> SizeInfo:
+        """Get the size information of the cache."""
+
+    @abstractmethod
+    def check_integrity(self) -> None:
+        """Check the integrity of the cache. Raise an error if the cache is corrupted."""

zonos2/kvcache/mha_pool.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import torch
+from zonos2.distributed import get_tp_info
+from zonos2.utils import divide_even
+
+from .base import BaseKVCache, KVCacheLayout
+
+
+class MHAKVCache(BaseKVCache):
+    """
+    Base class for key-value caches.
+    This class defines the interface for key-value caches used in LLMs.
+    """
+
+    def __init__(
+        self,
+        num_kv_heads: int,
+        num_layers: int,
+        head_dim: int,
+        num_pages: int,
+        dtype: torch.dtype,
+        kv_layout: KVCacheLayout,
+        device: torch.device,
+    ):
+        tp_info = get_tp_info()
+        local_kv_heads = divide_even(num_kv_heads, tp_info.size)
+        match kv_layout:
+            case KVCacheLayout.PageFirst:
+                kv_buffer = torch.empty(
+                    (2, num_pages, num_layers, local_kv_heads, head_dim),
+                    device=device,
+                    dtype=dtype,
+                ).permute(0, 2, 1, 3, 4)
+            case KVCacheLayout.LayerFirst:
+                kv_buffer = torch.empty(
+                    (2, num_layers, num_pages, local_kv_heads, head_dim),
+                    device=device,
+                    dtype=dtype,
+                )
+            case _:
+                raise ValueError(f"Unsupported kv_layout: {kv_layout}")
+        self._kv_buffer = kv_buffer.view(2, num_layers, num_pages, 1, local_kv_heads, head_dim)
+        self._num_layers = num_layers
+        self._k_buffer = self._kv_buffer[0]
+        self._v_buffer = self._kv_buffer[1]
+        self._device = device
+        self._storage_shape = (num_pages, local_kv_heads, head_dim)
+
+    def k_cache(self, index: int) -> torch.Tensor:
+        return self._k_buffer[index]
+
+    def v_cache(self, index: int) -> torch.Tensor:
+        return self._v_buffer[index]
+
+    def store_kv(
+        self, k: torch.Tensor, v: torch.Tensor, out_loc: torch.Tensor, layer_id: int
+    ) -> None:
+        from zonos2.kernel import store_cache
+
+        store_cache(
+            k_cache=self._k_buffer[layer_id].view(self._storage_shape),
+            v_cache=self._v_buffer[layer_id].view(self._storage_shape),
+            indices=out_loc,
+            k=k,
+            v=v,
+        )
+
+    @property
+    def device(self) -> torch.device:
+        return self._device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self._kv_buffer.dtype
+
+    @property
+    def num_layers(self) -> int:
+        return self._num_layers

zonos2/kvcache/naive_manager.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from typing import Tuple
+
+import torch
+
+from .base import BaseCacheHandle, BaseCacheManager, SizeInfo
+
+
+class NaiveCacheHandle(BaseCacheHandle):
+    pass
+
+
+class NaiveCacheManager(BaseCacheManager):
+    def __init__(self, device: torch.device):
+        self.device = device
+        self.empty_tensor = torch.empty(0, dtype=torch.int32, device=device)
+        super().__init__()
+
+    def match_prefix(self, input_ids: torch.Tensor) -> Tuple[NaiveCacheHandle, torch.Tensor]:
+        _ = input_ids  # unused
+        return NaiveCacheHandle(0), self.empty_tensor
+
+    def lock_handle(self, handle: BaseCacheHandle, unlock: bool = False) -> None:
+        _ = handle, unlock  # unused
+
+    def insert_prefix(self, input_ids: torch.Tensor, indices: torch.Tensor) -> int:
+        assert len(indices) == len(input_ids)
+        return len(indices)
+
+    def evict(self, size: int) -> torch.Tensor:
+        if size == 0:
+            return self.empty_tensor
+        raise NotImplementedError("NaiveCacheManager does not support eviction.")
+
+    def reset(self) -> None:
+        pass
+
+    @property
+    def size_info(self) -> SizeInfo:
+        return SizeInfo(evictable_size=0, protected_size=0)
+
+    def check_integrity(self) -> None:
+        pass

zonos2/kvcache/radix_manager.py

@@ -0,0 +1,221 @@
+from __future__ import annotations
+
+import heapq
+import time
+from dataclasses import dataclass
+from typing import Dict, List, Tuple
+
+import torch
+
+from .base import BaseCacheHandle, BaseCacheManager, SizeInfo
+
+
+class RadixTreeNode:
+    counter: int = 0
+
+    def __init__(self, tic: int | None = None) -> None:
+        self.children: Dict[int, RadixTreeNode] = {}
+        self._parent: RadixTreeNode | None = None
+        self.ref_count: int = 0
+        self.uuid = RadixTreeNode.counter
+        RadixTreeNode.counter += 1
+        self.timestamp = tic or time.monotonic_ns()
+
+        # these fields should be updated later
+        self._key: torch.Tensor
+        self._value: torch.Tensor
+        self._length: int
+
+    def set_key_value(self, key: torch.Tensor, value: torch.Tensor) -> None:
+        assert len(key) == len(value)
+        self._key = key
+        self._value = value
+        self._length = len(key)
+
+    def set_parent(self, parent: RadixTreeNode) -> None:
+        self._parent = parent
+        parent.children[int(self._key[0].item())] = self
+
+    @property
+    def length(self) -> int:
+        return self._length
+
+    @property
+    def parent(self) -> RadixTreeNode:
+        assert self._parent is not None
+        return self._parent
+
+    @property
+    def value(self) -> torch.Tensor:
+        return self._value
+
+    def is_root(self) -> bool:
+        return self._parent is None
+
+    def is_leaf(self) -> bool:
+        return len(self.children) == 0
+
+    def get_match_len(self, input_ids: torch.Tensor) -> int:
+        from zonos2.kernel import fast_compare_key
+
+        # compare key and input_ids, find the first diff
+        return fast_compare_key(self._key, input_ids)
+
+    def _split_at(self, pos: int) -> RadixTreeNode:
+        assert 0 < pos < self.length
+        parent = self.parent
+
+        new_node = RadixTreeNode(self.timestamp)
+        new_node.set_key_value(self._key[:pos], self._value[:pos])
+        new_node.set_parent(parent)
+        new_node.ref_count = self.ref_count
+
+        self.set_key_value(self._key[pos:], self._value[pos:])
+        self.set_parent(new_node)
+
+        return new_node
+
+    def __lt__(self, other: RadixTreeNode) -> bool:
+        return self.timestamp < other.timestamp
+
+
+@dataclass(frozen=True)
+class RadixCacheHandle(BaseCacheHandle):
+    node: RadixTreeNode
+
+
+class RadixCacheManager(BaseCacheManager):
+    def __init__(self, device: torch.device):
+        self.device = device
+        self.empty_tensor = torch.empty(0, dtype=torch.int32, device=device)
+        super().__init__()
+        self.root_node = RadixTreeNode()
+        self.root_node.ref_count = 1  # root is always protected
+        self.evictable_size = 0
+        self.protected_size = 0
+
+    def lock_handle(self, handle: BaseCacheHandle, unlock: bool = False) -> None:
+        assert isinstance(handle, RadixCacheHandle)
+        node = handle.node
+        if unlock:
+            while not node.is_root():
+                node.ref_count -= 1
+                assert node.ref_count >= 0
+                if node.ref_count == 0:
+                    self.evictable_size += node.length
+                    self.protected_size -= node.length
+                node = node.parent
+        else:
+            while not node.is_root():
+                if node.ref_count == 0:
+                    self.evictable_size -= node.length
+                    self.protected_size += node.length
+                node.ref_count += 1
+                node = node.parent
+
+    def match_prefix(self, input_ids: torch.Tensor) -> Tuple[RadixCacheHandle, torch.Tensor]:
+        node, prefix_len = self._walk(input_ids)
+        if prefix_len == 0:
+            assert node.is_root() and node is self.root_node and prefix_len == 0
+            return RadixCacheHandle(prefix_len, node), self.empty_tensor
+        value_list: List[torch.Tensor] = []
+        matched_node = node
+        while not node.is_root():
+            value_list.append(node.value)
+            node = node.parent
+        value_list.reverse()
+        return RadixCacheHandle(prefix_len, matched_node), torch.cat(value_list)
+
+    def insert_prefix(self, input_ids: torch.Tensor, indices: torch.Tensor) -> int:
+        node, prefix_len = self._walk(input_ids)
+        assert prefix_len <= len(input_ids)
+        if prefix_len < len(input_ids):
+            new_node = RadixTreeNode()
+            new_node.set_key_value(input_ids[prefix_len:], indices[prefix_len:])
+            new_node.set_parent(node)
+            self.evictable_size += new_node.length
+        return prefix_len
+
+    def _walk(self, input_ids: torch.Tensor) -> Tuple[RadixTreeNode, int]:
+        prefix_len = 0
+        indice_len = len(input_ids)
+        node = self.root_node
+        tic = time.monotonic_ns()
+
+        while prefix_len < indice_len:
+            this_id = int(input_ids[prefix_len].item())
+            if this_id not in node.children:
+                return node, prefix_len
+
+            node = node.children[this_id]
+
+            # NOTE: at least 1 char is matched, so match_len >= 1
+            match_len = node.get_match_len(input_ids[prefix_len:])
+            prefix_len += match_len
+
+            # need to split the node if not fully matched
+            if match_len != node.length:
+                node = node._split_at(match_len)
+                return node, prefix_len
+
+            # update timestamp for accessed node
+            node.timestamp = tic
+
+        return node, prefix_len
+
+    def evict(self, size: int) -> torch.Tensor:
+        if size == 0:
+            return self.empty_tensor
+        assert (
+            size <= self.evictable_size
+        ), f"Cannot evict {size}, only {self.evictable_size} is evictable"
+
+        leave_nodes = self._collect_leave_nodes_for_evict()
+        heapq.heapify(leave_nodes)
+        evicted_indices: List[torch.Tensor] = []
+        evicted_size = 0
+
+        while evicted_size < size:
+            assert (
+                leave_nodes
+            ), f"Cannot evict enough cache, need {size}, only {evicted_size} evicted"
+            node = heapq.heappop(leave_nodes)
+            assert node.ref_count == 0 and node.is_leaf() and not node.is_root()
+            evicted_size += node.length
+            evicted_indices.append(node.value)
+            self.evictable_size -= node.length
+            parent = node.parent
+            del parent.children[int(node._key[0].item())]
+            # NOTE: root is always protected, so won't be evicted
+            if parent.is_leaf() and parent.ref_count == 0:
+                heapq.heappush(leave_nodes, parent)
+
+        return torch.cat(evicted_indices)
+
+    def _collect_leave_nodes_for_evict(self) -> List[RadixTreeNode]:
+        nodes: List[RadixTreeNode] = [self.root_node]
+        leave_nodes: List[RadixTreeNode] = []
+
+        while len(nodes) > 0:
+            node = nodes.pop()
+            if node.is_leaf():
+                if node.ref_count == 0:
+                    leave_nodes.append(node)
+            else:
+                for child in node.children.values():
+                    nodes.append(child)
+
+        return leave_nodes
+
+    def reset(self) -> None:
+        raise NotImplementedError("RadixManager.reset is not implemented")
+
+    @property
+    def size_info(self) -> SizeInfo:
+        return SizeInfo(
+            evictable_size=self.evictable_size,
+            protected_size=self.protected_size,
+        )
+
+    def check_integrity(self) -> None:
+        pass

zonos2/layers/__init__.py

@@ -0,0 +1,32 @@
+from .activation import silu_and_mul
+from .attention import AttentionLayer
+from .base import BaseOP, OPList, StateLessOP
+from .embedding import ParallelLMHead, VocabParallelEmbedding
+from .linear import (
+    ChunkedLinear,
+    LinearColParallelMerged,
+    LinearOProj,
+    LinearQKVMerged,
+    LinearRowParallel,
+)
+from .norm import RMSNorm, RMSNormFused
+from .rotary import get_rope, set_rope_device
+
+__all__ = [
+    "silu_and_mul",
+    "AttentionLayer",
+    "BaseOP",
+    "StateLessOP",
+    "OPList",
+    "VocabParallelEmbedding",
+    "ParallelLMHead",
+    "ChunkedLinear",
+    "LinearColParallelMerged",
+    "LinearRowParallel",
+    "LinearOProj",
+    "LinearQKVMerged",
+    "RMSNorm",
+    "RMSNormFused",
+    "get_rope",
+    "set_rope_device",
+]

zonos2/layers/activation.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch
+
+
+def silu_and_mul(x: torch.Tensor) -> torch.Tensor:
+    from flashinfer import silu_and_mul
+
+    return silu_and_mul(x)
+
+
+__all__ = ["silu_and_mul"]