diff --git a/.flake8 b/.flake8 index 63396a2..a247806 100644 --- a/.flake8 +++ b/.flake8 @@ -35,3 +35,7 @@ per-file-ignores = __init__.py:F401 tests/*:D100,D101,D102,D103 setup.py:D100,D101,D102,D103,B009 + # pd_vllm implements vLLM connector / abstract-profile interfaces, so many + # method params are unused by design (U100); black collapses the ``...`` + # interface stubs onto one line (E704). + tilert/pd_vllm/*:U100,E704 diff --git a/README.md b/README.md index 2fb1ee2..4242f04 100644 --- a/README.md +++ b/README.md @@ -20,19 +20,21 @@ ______________________________________________________________________ ## πŸ“° News +- πŸ”€ **2026-07-14 Β· [v0.1.5](https://github.com/tile-ai/TileRT/releases/tag/v0.1.5) Released**. Introduce **PD (prefill–decode) disaggregation** β€” vLLM prefill + TileRT decode, behind an OpenAI-compatible endpoint. Supported on GLM-5/5.1 and DeepSeek-V3.2. + - πŸ’₯ **2026-06-08 Β· [Breaking 1000 TPS on a 1T Model](https://www.tilert.ai/blog/breaking-1000-tps.html)**. In collaboration with [Xiaomi MiMo](https://mimo.xiaomi.com/blog/mimo-tilert-1000tps), TileRT pushes [**MiMo-V2.5-Pro-UltraSpeed**](https://platform.xiaomimimo.com/docs/en-US/model-intro/mimo-v2.5-pro-ultraspeed) past **1000 tokens/s** on a **1-trillion-parameter** model through extreme model–system co-design β€” a first without custom silicon, all on a single 8-GPU node. - πŸš€ **2026-06-01 Β· [v0.1.4](https://github.com/tile-ai/TileRT/releases/tag/v0.1.4) Released**. A major performance upgrade for both DeepSeek-V3.2 and GLM-5, with model quality unchanged. See the benchmark charts for details. - 🏭 **2026-05-22 Β· [TileRT in Production](https://www.tilert.ai/blog/speed-as-the-next-scaling-law-zh.html)**. [**GLM-5.1-highspeed**](https://docs.bigmodel.cn/cn/guide/models/text/glm-5.1-highspeed) is now live on Z.ai, powered by TileRT β€” from experimental prototype to real production. +
+ Key Milestones + - :fire: **2026-02-14 Β· [Try the Online Demo](https://www.tilert.ai/)**. Our online demo is now live! Experience ultra-low-latency inference with **GLM-5** and **DeepSeek-V3.2**. [Try it now !](https://www.tilert.ai) - πŸŽ‰ **2026-02-14 Β· [v0.1.3](https://github.com/tile-ai/TileRT/releases/tag/v0.1.3) Released**. The v0.1.3 release introduces full support for the latest GLM-5 model, achieving up to 500 tokens/s on GLM-5-FP8 and up to 600 tokens/s on DeepSeek-V3.2. -
- Key Milestones - - πŸš€ **2026-01-26 Β· [v0.1.2-alpha.1](https://github.com/tile-ai/TileRT/releases/tag/v0.1.2-alpha.1)**. **Multi-Token Prediction (MTP)** is now available in TileRT! With mtp=3, we achieve decoding rates of up to **590 tokens/s** under synthetic workloads. - ⚑ **2025-12-23 Β· [v0.1.1](https://github.com/tile-ai/TileRT/releases/tag/v0.1.1)**. Achieved ~**35% further reduction** (3 ~ 4x speedup over baseline) in end-to-end token generation latency on a single node with **8Γ— NVIDIA B200**. @@ -54,9 +56,9 @@ To achieve this, TileRT introduces a **tile-level runtime engine**. Leveraging a The project is actively evolving, and the underlying compiler techniques will be gradually shared with the community as they are integrated into **TileLang** and **TileScale**.

- GLM-5.1-FP8 token generation speed on 8Γ— B200 with TileRT v0.1.4 + GLM-5.1-FP8 token generation speed on 8Γ— B200 with TileRT v0.1.5
- GLM-5.1-FP8 token generation speed on 8Γ— NVIDIA B200 with TileRT v0.1.4. Output length 1K, input length 1K–192K. Bars compare TileRT without MTP, with MTP at average acceptance length 3.2, and the peak under best-case MTP acceptance. + GLM-5.1-FP8 token generation speed on 8Γ— NVIDIA B200 with TileRT v0.1.5. Output length 1K, input length 1K–192K. Bars compare TileRT without MTP, with MTP at average acceptance length 3.2, and the peak under best-case MTP acceptance (4.0).

______________________________________________________________________ @@ -64,11 +66,11 @@ ______________________________________________________________________ ## Installation > \[!IMPORTANT\] -> TileRT v0.1.4 is distributed as a **pre-built binary wheel**. The wheel is linked against the exact ABI of the versions listed below. Other combinations of Python, CUDA, or PyTorch versions are **untested and not guaranteed to work** β€” please reproduce this environment for a supported setup. +> TileRT v0.1.5 is distributed as a **pre-built binary wheel**. The wheel is linked against the exact ABI of the versions listed below. Other combinations of Python, CUDA, or PyTorch versions are **untested and not guaranteed to work** β€” please reproduce this environment for a supported setup. -### Build environment of the v0.1.4 wheel +### Build environment of the v0.1.5 wheel -The official `tilert==0.1.4` wheel on PyPI was compiled against the following stack. Treat these as **hard requirements**, not lower bounds. +The official `tilert==0.1.5` wheel on PyPI was compiled against the following stack. Treat these as **hard requirements**, not lower bounds. | Component | Pinned version | | ---------------- | --------------------------------------------------- | @@ -83,7 +85,7 @@ The official `tilert==0.1.4` wheel on PyPI was compiled against the following st ### Recommended: pre-built Docker image The pinned build environment above is preinstalled in our official image -β€” this is the **recommended way to run v0.1.4** and avoids any version +β€” this is the **recommended way to run v0.1.5** and avoids any version drift on the host. The image is mirrored to two registries; pull from whichever is reachable: @@ -104,18 +106,18 @@ docker run --rm -it --gpus all --ipc=host \ ghcr.io/tile-ai/tilert:cu132-latest # Inside the container β€” install from PyPI: -pip install tilert==0.1.4 +pip install tilert==0.1.5 # Or pin the exact wheel from the GitHub Release page directly # (same artifact, useful when PyPI is unreachable): -pip install https://github.com/tile-ai/TileRT/releases/download/v0.1.4/tilert-0.1.4-cp312-cp312-manylinux_2_28_x86_64.whl +pip install https://github.com/tile-ai/TileRT/releases/download/v0.1.5/tilert-0.1.5-cp312-cp312-manylinux_2_28_x86_64.whl ``` Verify the install: ```bash python -c "import tilert, torch; print('tilert', tilert.__version__, '/ torch', torch.__version__, '/ cuda', torch.version.cuda)" -# Expected: tilert 0.1.4 / torch 2.11.0+cu130 / cuda 13.0 +# Expected: tilert 0.1.5 / torch 2.11.0+cu130 / cuda 13.0 ``` Proceed to [Getting Started](#getting-started) to download and convert model weights. @@ -143,7 +145,7 @@ python -m tilert.models.preprocess.weight_converter \ --save_dir "/path/to/DeepSeek-V3.2-TileRT" ``` -For **GLM-5**: +For **GLM-5/5.1**: ```bash python -m tilert.models.preprocess.weight_converter \ @@ -178,7 +180,7 @@ python -m tilert.generate --model deepseek_v3_2 --max-new-tokens 1000 ``` > \[!NOTE\] -> v0.1.4 ships **two independent backend libraries** (`libtilert_dsv32.so` +> v0.1.5 ships **two independent backend libraries** (`libtilert_dsv32.so` > and `libtilert_glm5.so`) and loads exactly one per Python process via > `tilert.load_backend(model_type)`. Run DeepSeek-V3.2 and GLM-5 in > separate processes β€” they cannot coexist in a single interpreter. @@ -306,6 +308,110 @@ This example highlights how MTP enables TileRT to efficiently generate longer ou For the full list of CLI flags (sampling, batching, benchmark modes, …), run `python -m tilert.generate --help`. +## Disaggregated Serving: vLLM Prefill + TileRT Decode + +TileRT can run as the **decode engine behind a vLLM prefill**, integrated through vLLM's V1 `KVConnector` interface. The connector, decode server, and router all ship inside the `tilert` wheel under `tilert.pd_vllm` β€” no vLLM fork or patch is needed (the connector loads via vLLM's standard `kv_connector_module_path`). Latency-critical requests are routed to the TileRT decode pool; other traffic can stay on native vLLM decode. + +**Prerequisites** + +- Convert the model weights for TileRT decode (see [Step 2](#step-2-shard-weights-with-weight_converter)). +- On the **prefill** node, a vLLM build with V1 disaggregation and support for the GLM-5/5.1 / DeepSeek-V3.2 (DSA) model and the `fp8_ds_mla` KV-cache dtype. Install `tilert` in the same environment so the connector plugin is importable. +- **The KV-cache dtype must match on both ends.** These examples use fp8: `--kv-cache-dtype fp8_ds_mla` on the vLLM prefill and `--kv-cache-dtype fp8` on the TileRT decode (a mismatch is rejected at the connector handshake). +- The examples use the **NIXL** transfer engine. On multi-NIC hosts, pin NIXL to the RDMA NICs via `UCX_NET_DEVICES` (otherwise UCX may pick the wrong interface). Mooncake is also supported (`--transport mooncake` on the decode, `"tilert_transport": "mooncake"` on the prefill). + +Commands below use GLM-5/5.1. For DeepSeek-V3.2, use `--model deepseek_v3_2`, the DeepSeek-V3.2-TileRT weights, and `--parser none`. + +### Topology A: vLLM prefill β†’ TileRT decode + +Three processes β€” a TileRT decode server, a stock vLLM prefill, and an OpenAI-compatible router: + +```bash +# 1) TileRT decode node +python -m tilert.pd_vllm.decode_server \ + --engine tilert --model glm5 \ + --model-weights-dir /path/to/GLM-5.1-FP8-TileRT \ + --with-mtp --max-seq-len 202752 \ + --kv-cache-dtype fp8 --transport nixl \ + --ctrl-port 5556 --http-port 5557 + +# 2) vLLM prefill (stock vLLM; the TileRT connector loads as a plugin). +# The MTP speculative config is required: the prefill populates the +# draft-layer KV that decode-side speculation resumes from. +export UCX_NET_DEVICES=mlx5_1:1,mlx5_2:1,... # pin NIXL to the RDMA NICs (multi-NIC hosts) +vllm serve /path/to/GLM-5.1-FP8 \ + --served-model-name glm5 --port 8000 \ + --tensor-parallel-size 8 --enforce-eager --trust-remote-code \ + --return-tokens-as-token-ids --gpu-memory-utilization 0.75 \ + --kv-cache-dtype fp8_ds_mla \ + --speculative-config '{"method": "mtp", "num_speculative_tokens": 1}' \ + --kv-transfer-config '{ + "kv_connector": "TileRTConnector", + "kv_connector_module_path": "tilert.pd_vllm.prefill_connector", + "kv_role": "kv_producer", + "kv_connector_extra_config": { + "tilert_host": "", "tilert_ctrl_port": 5556, + "tilert_model": "glm5", "tilert_max_seq_len": 202752, + "tilert_transport": "nixl"}}' + +# 3) Router β€” OpenAI-compatible ingress for the TileRT pool +python -m tilert.pd_vllm.pd_router \ + --vllm-url http://:8000 \ + --decode :5556:5557 \ + --model-path /path/to/GLM-5.1-FP8 \ + --parser glm47 --port 23333 +``` + +Send OpenAI requests to `http://:23333/v1/chat/completions`. The router runs the prefill on vLLM (first token), hands the attention state to the TileRT decode node over RDMA, and streams the completion back. + +### Topology B: shared prefill β†’ TileRT decode **and** native vLLM decode + +One prefill pool feeds two decode pools side by side, composed under vLLM's `MultiConnector`. Each request is claimed by exactly one connector β€” the TileRT connector claims requests marked with `tilert_host`, and vLLM's native connector handles the rest β€” so latency-critical traffic goes to TileRT while general traffic stays on native vLLM decode, behind the same OpenAI surface. + +```bash +# 1) TileRT decode node (identical to Topology A) +python -m tilert.pd_vllm.decode_server --engine tilert --model glm5 \ + --model-weights-dir /path/to/GLM-5.1-FP8-TileRT --with-mtp \ + --max-seq-len 202752 --kv-cache-dtype fp8 --transport nixl \ + --ctrl-port 5556 --http-port 5557 + +# 2) Native vLLM decode node β€” vLLM's standard disaggregation (NixlConnector consumer) +export UCX_NET_DEVICES=mlx5_1:1,mlx5_2:1,... +vllm serve /path/to/GLM-5.1-FP8 --served-model-name glm5 --port 8001 \ + --tensor-parallel-size 8 --enforce-eager --trust-remote-code \ + --return-tokens-as-token-ids --kv-cache-dtype fp8_ds_mla \ + --speculative-config '{"method": "mtp", "num_speculative_tokens": 1}' \ + --kv-transfer-config '{"kv_connector": "NixlConnector", "kv_role": "kv_consumer"}' + +# 3) Shared vLLM prefill β€” MultiConnector[ NixlConnector + TileRTConnector ] +export UCX_NET_DEVICES=mlx5_1:1,mlx5_2:1,... +vllm serve /path/to/GLM-5.1-FP8 --served-model-name glm5 --port 8000 \ + --tensor-parallel-size 8 --enforce-eager --trust-remote-code \ + --return-tokens-as-token-ids --gpu-memory-utilization 0.75 \ + --kv-cache-dtype fp8_ds_mla \ + --speculative-config '{"method": "mtp", "num_speculative_tokens": 1}' \ + --kv-transfer-config '{ + "kv_connector": "MultiConnector", "kv_role": "kv_producer", + "kv_connector_extra_config": {"connectors": [ + {"kv_connector": "NixlConnector", "kv_role": "kv_producer"}, + {"kv_connector": "TileRTConnector", + "kv_connector_module_path": "tilert.pd_vllm.prefill_connector", + "kv_role": "kv_producer", + "kv_connector_extra_config": { + "tilert_host": "", "tilert_ctrl_port": 5556, + "tilert_model": "glm5", "tilert_max_seq_len": 202752, + "tilert_transport": "nixl"}}]}}' + +# 4a) TileRT router β€” latency-critical traffic β†’ TileRT pool +python -m tilert.pd_vllm.pd_router --vllm-url http://:8000 \ + --decode :5556:5557 --model-path /path/to/GLM-5.1-FP8 \ + --parser glm47 --port 23333 + +# 4b) General traffic β†’ native vLLM decode pool, via vLLM's standard NixlConnector +# disaggregation proxy, pointing prefill :8000 β†’ native decode :8001. +``` + +**Note.** Running NIXL end to end (both the native and TileRT connectors in NIXL mode) lets the shared prefill use a single transfer library. Only the prefill's `--kv-transfer-config` differs from Topology A; the TileRT decode node is unchanged, and the native decode instance plus its proxy follow vLLM's usual `NixlConnector` disaggregation setup. + ## Status & Future Work TileRT is currently offered as a preview release, and we’re just getting started. diff --git a/assets/glm5_tilert_mtp.png b/assets/glm5_tilert_mtp.png index c4db83f..f0aeb99 100644 Binary files a/assets/glm5_tilert_mtp.png and b/assets/glm5_tilert_mtp.png differ diff --git a/pyproject.toml b/pyproject.toml index 13aa491..407c497 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,6 +6,7 @@ description = "TileRT" readme = "README.md" requires-python = ">=3.11" license = {text = "MIT"} +authors = [{name = "TileRT-team", email = "contact@tilert.ai"}] classifiers = [ "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", @@ -14,7 +15,7 @@ classifiers = [ ] dependencies = [ - # Pinned to the exact ABI the v0.1.4 wheel was built against. ``torch`` must + # Pinned to the exact ABI the v0.1.5 wheel was built against. ``torch`` must # come from PyTorch's cu130 index (``--index-url # https://download.pytorch.org/whl/cu130``); installing from PyPI yields a # CUDA build that does not match the cu130-linked tilert binary. @@ -64,7 +65,7 @@ dev = [ Homepage = "https://github.com/tile-ai/TileRT" Issues = "https://github.com/tile-ai/TileRT/issues" -# Note: this repository ships the public sources that match the v0.1.4 wheel. +# Note: this repository ships the public sources that match the v0.1.5 wheel. # The wheel itself is built in the development repo (TileRT-dev/TileRT) with # scikit-build-core; no [build-system] block is declared here on purpose so # nobody accidentally runs ``pip wheel .`` against this presentation copy. @@ -110,6 +111,16 @@ ignore_missing_imports = true exclude = ["3rd-party/"] explicit_package_bases = true +# pd_vllm is glue over vLLM's connector API; its interface methods take +# ``Any``-typed params by contract, so full def-level annotation is low value. +# Relax annotation-completeness there, but keep real type checks (arg-type, +# union-attr, etc.) on so genuine mistakes still surface. +[[tool.mypy.overrides]] +module = "tilert.pd_vllm.*" +disallow_untyped_defs = false +disallow_incomplete_defs = false +warn_return_any = false + [tool.bandit] exclude_dirs = ["tests", "3rd-party"] skips = ["B101", "B311", "B404", "B603", "B607"] diff --git a/requirements.txt b/requirements.txt index fd4a9ba..c22551d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -# Runtime dependencies for the v0.1.4 wheel, pinned to the exact ABI the +# Runtime dependencies for the v0.1.5 wheel, pinned to the exact ABI the # wheel was built against. ``torch`` must be installed from PyTorch's cu130 # index β€” PyPI's default ``torch`` is a different CUDA build and will not load # the cu130-linked tilert binary: diff --git a/tilert/models/glm_5/__init__.py b/tilert/models/glm_5/__init__.py index e69de29..8fddd5f 100644 --- a/tilert/models/glm_5/__init__.py +++ b/tilert/models/glm_5/__init__.py @@ -0,0 +1 @@ +"""GLM-5 model package.""" diff --git a/tilert/models/glm_5/_dsa_v32/generator.py b/tilert/models/glm_5/_dsa_v32/generator.py deleted file mode 100644 index 26ee685..0000000 --- a/tilert/models/glm_5/_dsa_v32/generator.py +++ /dev/null @@ -1,531 +0,0 @@ -"""DSA show hands for deepseek v3.2.""" - -import math -import time - -import torch -from transformers import AutoTokenizer - -from tilert import logger -from tilert.models.glm_5._dsa_v32.model_args import ModelArgs -from tilert.models.glm_5._dsa_v32.modules.end2end import ShowHandsDSALayer -from tilert.models.glm_5._dsa_v32.temp_var_indices import Idx -from tilert.tilert_init import tilert_init - -__all__ = [ - "DSAv32Generator", - "stats_time", -] - - -def stats_time(time_list: list[float], title: str) -> None: - if len(time_list) > 0: - avg_time = sum(time_list) / len(time_list) - std_dev = math.sqrt(sum((x - avg_time) ** 2 for x in time_list) / len(time_list)) - logger.info(title) - logger.info(f"--Average time taken to generate token: {avg_time * 1000:.4f} ms") - logger.info(f"--Standard deviation of time: {std_dev * 1000:.4f} ms") - logger.info(f"--Effective tokens per second: {1 / avg_time:.4f}") - - -class DSAv32Generator: - def __init__( - self, - model_args: ModelArgs, - max_new_tokens: int = 100, - temperature: float = 1.0, - model_weights_dir: str = "", - with_mtp: bool = False, - use_topp: bool = False, - top_p: float = 0.9, - top_k: int = 256, - sampling_seed: int = 42, - ): - """Initialize the DSAv32Generator. - - Args: - max_new_tokens: Maximum number of new tokens to generate. Defaults to 100. - temperature: Temperature for sampling. Defaults to 1.0. - model_weights_dir: Path of the model weights directory. - with_mtp: Whether to use MTP (Multi-Token Prediction) for speculative decoding. - use_topp: Whether to use top-p (nucleus) sampling instead of top-1 (argmax). - top_p: Top-p threshold for nucleus sampling. Defaults to 0.9. - top_k: Number of top-k candidates for top-p sampling. Defaults to 256. - sampling_seed: Sampling seed for top-p (fixed per request). Defaults to 42. - """ - torch.set_num_threads(64) - self.model_weights_dir = model_weights_dir - - self.max_new_tokens = max_new_tokens - self.temperature = temperature - self.with_mtp = with_mtp - self.use_topp = use_topp - self.top_p = top_p - self.top_k = top_k - self.sampling_seed = sampling_seed - - self.config = model_args - self.tokenizer = AutoTokenizer.from_pretrained( - self.model_weights_dir, trust_remote_code=True - ) # nosec B615 - self.eos_id = self.tokenizer.eos_token_id - self.batch_size = 1 - - self.default_device = torch.device("cuda:0") - - self.decode_layer = ShowHandsDSALayer( - model_args=self.config, - model_path=self.model_weights_dir, - with_mtp=with_mtp, - use_topp=use_topp, - top_p=top_p, - top_k=top_k, - ) - - self.mtp_seq_len = 4 if with_mtp else 1 - - def init(self) -> None: - """Initialize the ShowHandsGenerator.""" - tilert_init() - - def cleanup(self) -> None: - """Cleanup the ShowHandsGenerator.""" - self.decode_layer.cleanup() - - def init_random_weights(self) -> None: - """Random initialize the weights.""" - self.decode_layer.init_random_weights() - - def from_pretrained(self) -> None: - """Load the model weights from the given path.""" - self.decode_layer.from_pretrained(self.model_weights_dir) - - def extract_ffn_cache(self) -> tuple[dict[int, list], dict[int, set[str]]]: - """Extract MOE/MLP op objects and skip keys from current loaded weights. - - Returns: - Tuple of (cached_ffn_ops_per_device, skip_keys_per_device). - """ - from tilert.models.glm_5._dsa_v32.modules.end2end import ( - _extract_ffn_ops, - _get_moe_weight_keys, - ) - - cached_ffn_ops: dict[int, list] = {} - skip_keys: dict[int, set[str]] = {} - for device_id in range(self.decode_layer.num_devices): - dsa = self.decode_layer._dsa_objects[device_id] - if dsa is None: - raise RuntimeError(f"Device {device_id} Dsa not available for cache extraction") - cached_ffn_ops[device_id] = _extract_ffn_ops(dsa) - skip_keys[device_id] = _get_moe_weight_keys(dsa) - return cached_ffn_ops, skip_keys - - def from_pretrained_with_cache( - self, - cached_ffn_ops_per_device: dict[int, list], - skip_keys_per_device: dict[int, set[str]], - ) -> None: - """Load weights reusing cached MOE/MLP ops.""" - self.decode_layer.from_pretrained_with_cache( - self.model_weights_dir, cached_ffn_ops_per_device, skip_keys_per_device - ) - - def update_sampling_params( - self, - temperature: float = 1.0, - top_p: float = 0.95, - top_k: int = 256, - use_topp: bool = True, - ) -> None: - """Update sampling parameters for the next generation.""" - self.temperature = temperature - self.use_topp = use_topp - self.top_p = top_p - self.top_k = top_k - self.decode_layer.update_sampling_config( - temperature=temperature, top_p=top_p, top_k=top_k, use_topp=use_topp - ) - - @torch.inference_mode() - def generate( - self, - prompt: str, - print_log: bool = True, - with_mtp: bool | None = None, - prompt_tokens: list[int] | None = None, - ) -> tuple[str, list[float], list[int], int]: - """Main function to load the model and perform single sequence generation. - - Args: - prompt: The input prompt string. - print_log: Whether to print generation logs. - with_mtp: Override MTP mode for this call. None uses self.with_mtp. - Requires MTP weights to have been loaded (self.with_mtp=True). - prompt_tokens: Pre-tokenized prompt tokens. If provided, skip tokenization - and use these tokens directly (useful for exact-length benchmarking). - - Returns: - Tuple of (result_text, time_list, accepted_counts, prompt_len). - accepted_counts is empty for non-MTP mode. - """ - active_mtp = with_mtp if with_mtp is not None else self.with_mtp - if active_mtp and not self.with_mtp: - raise ValueError("Cannot use MTP mode: MTP weights were not loaded") - self.decode_layer.set_sampling_seed(self.sampling_seed, with_mtp=active_mtp) - if active_mtp: - return self._generate_with_mtp(prompt, print_log, prompt_tokens=prompt_tokens) - result, time_list, prompt_len = self._generate_without_mtp( - prompt, print_log, with_mtp=active_mtp, prompt_tokens=prompt_tokens - ) - return result, time_list, [], prompt_len - - def _generate_without_mtp( - self, - prompt: str, - print_log: bool = True, - with_mtp: bool = False, - prompt_tokens: list[int] | None = None, - ) -> tuple[str, list[float], int]: - """Standard generation without MTP.""" - if prompt_tokens is None: - prompt_tokens = self.tokenizer.apply_chat_template( - [{"role": "user", "content": prompt}], add_generation_prompt=True - ) - - max_seq_len = self.config.max_seq_len - prompt_len = len(prompt_tokens) - total_len = min(max_seq_len, self.max_new_tokens + prompt_len) - - tokens = torch.full( - (self.batch_size, total_len), -1, dtype=torch.long, device=self.default_device - ) - tokens[0, :prompt_len] = torch.tensor( - prompt_tokens, dtype=torch.long, device=self.default_device - ) - prompt_mask = tokens != -1 - - prev_pos = 0 - finished = torch.tensor( - [False] * self.batch_size, dtype=torch.bool, device=self.default_device - ) - - time_list = [] - for cur_pos_val in range(1, total_len): - start_time = time.time() - multi_devices_results = self.decode_layer.forward( - tokens[0, prev_pos], with_mtp=with_mtp - ) - end_time = time.time() - time_list.append(end_time - start_time) - - intermediates, *_ = multi_devices_results[0] - next_token = intermediates[Idx.TOKEN_OUT][0][0] - - next_token = torch.where( - prompt_mask[0, cur_pos_val], tokens[0, cur_pos_val], next_token - ) - tokens[0, cur_pos_val] = next_token - finished |= torch.logical_and(~prompt_mask[0, cur_pos_val], next_token == self.eos_id) - prev_pos = cur_pos_val - if cur_pos_val >= prompt_len: - decoded_tokens = self.tokenizer.decode( - [next_token.item()], skip_special_tokens=True - ) - if print_log: - print(decoded_tokens, end="", flush=True) - - if finished.all(): - break - - if print_log: - print("\n") - logger.info(f"--Number of tokens generated: {len(time_list)}") - - stats_time(time_list, "==== Performance ====") - print("\n") - - self.decode_layer.reset_sequence() - - completion_tokens = [] - for _, toks in enumerate(tokens.tolist()): - toks = toks[prompt_len : prompt_len + self.max_new_tokens] - if self.eos_id in toks: - toks = toks[: toks.index(self.eos_id)] - completion_tokens.append(toks) - - decoded_tokens = self.tokenizer.batch_decode(completion_tokens, skip_special_tokens=True) - - return f"{decoded_tokens[0]}\n" if decoded_tokens else "", time_list, prompt_len - - def _generate_with_mtp( - self, - prompt: str, - print_log: bool = True, - prompt_tokens: list[int] | None = None, - ) -> tuple[str, list[float], list[int], int]: - """Generation with MTP (Multi-Token Prediction) speculative decoding.""" - if prompt_tokens is None: - prompt_tokens = self.tokenizer.apply_chat_template( - [{"role": "user", "content": prompt}], add_generation_prompt=True - ) - - max_seq_len = self.config.max_seq_len - prompt_len = len(prompt_tokens) - total_len = min(max_seq_len, self.max_new_tokens + prompt_len) - - tokens = torch.full( - (self.batch_size, total_len), -1, dtype=torch.long, device=self.default_device - ) - tokens[0, :prompt_len] = torch.tensor( - prompt_tokens, dtype=torch.long, device=self.default_device - ) - - prefill_time_list = [] - decode_time_list = [] - decode_accepted_counts = [] - cur_pos = 0 - - while cur_pos < prompt_len - 1: - draft_end = min(cur_pos + self.mtp_seq_len, prompt_len) - draft_tokens = tokens[0, cur_pos:draft_end].clone() - actual_token_count = draft_tokens.shape[0] - - if actual_token_count < self.mtp_seq_len: - pad_token = draft_tokens[-1].item() - padding = torch.full( - (self.mtp_seq_len - actual_token_count,), - pad_token, - dtype=torch.long, - device=self.default_device, - ) - draft_tokens = torch.cat([draft_tokens, padding]) - - draft_tokens = draft_tokens.reshape(1, self.mtp_seq_len).to(torch.int32) - - mtp_extra_pos = cur_pos + self.mtp_seq_len - if mtp_extra_pos < prompt_len: - mtp_extra_token = int(tokens[0, mtp_extra_pos].item()) - else: - mtp_extra_token = int(tokens[0, draft_end - 1].item()) - self.decode_layer.set_prefill_mtp_extra_token(mtp_extra_token) - - self.decode_layer.set_prefill_valid_tokens(actual_token_count) - - start_time = time.time() - self.decode_layer.forward(draft_tokens, with_mtp=True) - end_time = time.time() - prefill_time_list.append(end_time - start_time) - - cur_pos += actual_token_count - - cur_pos = prompt_len - 1 - self.set_cur_pos(prompt_len - 1) - - self.decode_layer.set_prefill_valid_tokens(0) - - finished = False - while cur_pos < total_len - 1 and not finished: - if cur_pos == prompt_len - 1: - last_token = tokens[0, prompt_len - 1].item() - draft_tokens = torch.full( - (self.mtp_seq_len,), - last_token, - dtype=torch.long, - device=self.default_device, - ) - draft_tokens = draft_tokens.reshape(1, self.mtp_seq_len).to(torch.int32) - else: - draft_tokens = self.decode_layer.get_next_draft_tokens(0).reshape( - 1, self.mtp_seq_len - ) - - start_time = time.time() - self.decode_layer.forward(draft_tokens, with_mtp=True) - end_time = time.time() - decode_time_list.append(end_time - start_time) - - num_accepted = self.decode_layer.get_num_accepted(0) - predicted_tokens = self.decode_layer.get_predicted_tokens(0).flatten() - decode_accepted_counts.append(num_accepted) - - num_output_tokens = num_accepted - for i in range(num_output_tokens): - if cur_pos + 1 + i >= total_len: - break - new_token = int(predicted_tokens[i].item()) - tokens[0, cur_pos + 1 + i] = new_token - - if cur_pos + 1 + i >= prompt_len and print_log: - decoded_text = self.tokenizer.decode([new_token], skip_special_tokens=True) - print(decoded_text, end="", flush=True) - - if new_token == self.eos_id: - finished = True - break - - cur_pos += num_accepted - - if print_log: - print("\n") - total_tokens = sum(decode_accepted_counts) - logger.info(f"--Number of forward calls (decode): {len(decode_accepted_counts)}") - logger.info(f"--Total tokens generated: {total_tokens}") - if len(decode_accepted_counts) > 0: - avg_accepted = sum(decode_accepted_counts) / len(decode_accepted_counts) - min_accepted = min(decode_accepted_counts) - max_accepted = max(decode_accepted_counts) - logger.info( - f"--Accepted tokens per call: mean={avg_accepted:.2f}, " - f"min={min_accepted}, max={max_accepted}" - ) - - if decode_time_list: - total_decode_time = sum(decode_time_list) - effective_tps = total_tokens / total_decode_time if total_decode_time > 0 else 0 - avg_time_ms = total_decode_time / len(decode_time_list) * 1000 - logger.info(f"--Avg forward time: {avg_time_ms:.2f}ms") - logger.info(f"--Effective TPS (with MTP): {effective_tps:.2f} tokens/s") - - print("\n") - - self.decode_layer.reset_sequence() - - completion_tokens = [] - for _, toks in enumerate(tokens.tolist()): - toks = toks[prompt_len : prompt_len + self.max_new_tokens] - toks = [t for t in toks if t != -1] - if self.eos_id in toks: - toks = toks[: toks.index(self.eos_id)] - completion_tokens.append(toks) - - decoded_tokens = self.tokenizer.batch_decode(completion_tokens, skip_special_tokens=True) - - return ( - f"{decoded_tokens[0]}\n" if decoded_tokens else "", - decode_time_list, - decode_accepted_counts, - prompt_len, - ) - - def inject_cache( - self, - layer_caches: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]], - start_pos: int = 0, - end_pos: int | None = None, - ) -> None: - """Inject external cache data into TileRT. - - This API allows injecting pre-computed KI/KV/PE cache data from an external - prefill system, enabling prefill-decode disaggregation. - - Args: - layer_caches: List of (ki, kv, pe) tuples for each layer (0 to NUM_LAYERS-1). - Each tensor should be BF16 with shape [seqlen, dim] where: - - ki: [seqlen, 128] - compressed key - - kv: [seqlen, 512] - compressed key-value - - pe: [seqlen, 64] - position encoding cache - start_pos: Start position in cache to write (0-indexed). Defaults to 0. - end_pos: End position in cache (exclusive). If None, uses seqlen from tensors. - - Example: - >>> # Load cache from external prefill system - >>> layer_caches = [] # List of 61 (ki, kv, pe) tuples - >>> for layer_id in range(61): - ... ki = load_ki_for_layer(layer_id) # [seqlen, 128] bf16 - ... kv = load_kv_for_layer(layer_id) # [seqlen, 512] bf16 - ... pe = load_pe_for_layer(layer_id) # [seqlen, 64] bf16 - ... layer_caches.append((ki, kv, pe)) - >>> generator.inject_cache(layer_caches, start_pos=0) - >>> generator.set_cur_pos(seqlen) # Set RoPE position - >>> # Continue generation from cache - """ - num_layers = len(layer_caches) - if num_layers == 0: - logger.warning("inject_cache called with empty layer_caches") - return - - first_ki, _, _ = layer_caches[0] - seqlen = first_ki.size(0) - if end_pos is None: - end_pos = start_pos + seqlen - - cache_len = end_pos - start_pos - logger.info(f"Injecting cache: {num_layers} layers, positions [{start_pos}, {end_pos})") - - num_devices = self.decode_layer.num_devices - - for device_id in range(num_devices): - _, caches, _, _ = self.decode_layer._get_device_result(device_id) - - for layer_id, (ki, kv, pe) in enumerate(layer_caches): - if layer_id >= num_layers: - logger.warning(f"Layer index {layer_id} is out of bounds, skipping.") - break - - base_idx = layer_id * 3 - - ki_src = ki[:cache_len].to(f"cuda:{device_id}") - kv_src = kv[:cache_len].to(f"cuda:{device_id}") - pe_src = pe[:cache_len].to(f"cuda:{device_id}") - - caches[base_idx + 0][0, start_pos:end_pos, :].copy_(ki_src) - caches[base_idx + 1][0, start_pos:end_pos, :].copy_(kv_src) - caches[base_idx + 2][0, start_pos:end_pos, :].copy_(pe_src) - - logger.info(f"Cache injection completed for {num_devices} devices") - - def set_cur_pos(self, cur_pos: int) -> None: - """Set the current position for RoPE. - - This should be called after inject_cache() to ensure the runtime position - matches the injected cache length, for correct RoPE position encoding - during continued generation. - - Args: - cur_pos: The current sequence position (typically the length of prefilled tokens). - - Example: - >>> generator.inject_cache(layer_caches, start_pos=0) - >>> generator.set_cur_pos(prefill_len) # Set position to prefill length - >>> # Now generate continues from the correct position - """ - if self.with_mtp: - num_devices = self.decode_layer.num_devices - for device_id in range(num_devices): - intermediates, _, _, _ = self.decode_layer._get_device_result(device_id) - cur_pos_tensor = intermediates[Idx.CUR_POS] - cur_pos_tensor.fill_(cur_pos) - else: - torch.ops.tilert.dsa_show_hands_set_cur_pos(cur_pos) - - def inject_last_hidden_state(self, last_hidden_state: torch.Tensor) -> None: - """Inject the last hidden state for MTP mode. - - For MTP (Multi-Token Prediction), the MTP preprocess layer needs the - last hidden state from the main model's last token. - - Args: - last_hidden_state: [hidden_size] or [1, hidden_size] BF16 tensor. - The hidden state of the last token from prefill. - - Example: - >>> # After inject_cache, inject the last hidden state for MTP - >>> generator.inject_last_hidden_state(last_hidden_state) - >>> # Then set cur_pos and start generation - """ - if not self.with_mtp: - logger.warning("inject_last_hidden_state called but with_mtp is False, skipping") - return - - if last_hidden_state.dim() == 1: - last_hidden_state = last_hidden_state.unsqueeze(0) - - num_devices = self.decode_layer.num_devices - for device_id in range(num_devices): - intermediates, _, _, _ = self.decode_layer._get_device_result(device_id) - lhs_tensor = intermediates[Idx.LAST_HIDDEN_STATES] - lhs_src = last_hidden_state.to(f"cuda:{device_id}") - lhs_tensor[0, 0, :].copy_(lhs_src.squeeze(0)) - - logger.info(f"Injected last_hidden_state to {num_devices} devices") diff --git a/tilert/models/glm_5/_dsa_v32/model_args.py b/tilert/models/glm_5/_dsa_v32/model_args.py index 441b684..143fbd7 100644 --- a/tilert/models/glm_5/_dsa_v32/model_args.py +++ b/tilert/models/glm_5/_dsa_v32/model_args.py @@ -10,43 +10,7 @@ @dataclass class ModelArgs: - """ - Data class for defining model arguments and hyperparameters. - - Attributes: - arch_name (str): Architecture name. - max_batch_size (int): Maximum batch size. - max_seq_len (int): Maximum sequence length. - dtype (Literal["bf16", "fp8"]): Data type for computations. - scale_fmt (Optional[str]): Format for quantization scale. - vocab_size (int): Vocabulary size. - dim (int): Model dimension. - inter_dim (int): Intermediate dimension for MLP layers. - moe_inter_dim (int): Intermediate dimension for MoE layers. - n_layers (int): Number of transformer layers. - n_dense_layers (int): Number of dense layers in the model. - n_heads (int): Number of attention heads. - n_routed_experts (int): Number of routed experts for MoE layers. - n_shared_experts (int): Number of shared experts for MoE layers. - n_activated_experts (int): Number of activated experts in MoE layers. - n_expert_groups (int): Number of expert groups. - n_limited_groups (int): Number of limited groups for MoE routing. - score_func (Literal["softmax", "sigmoid"]): Scoring function for MoE routing. - route_scale (float): Scaling factor for routing scores. - q_lora_rank (int): LoRA rank for query projections. - kv_lora_rank (int): LoRA rank for key-value projections. - qk_nope_head_dim (int): Dimension for query-key projections without positional embeddings. - qk_rope_head_dim (int): Dimension for query-key projections with rotary embeddings. - v_head_dim (int): Dimension for value projections. - original_seq_len (Optional[int]): Original sequence length. - rope_theta (float): Base for rotary positional encoding. - rope_factor (Optional[float]): Scaling factor for extended sequence lengths. - beta_fast (Optional[int]): Fast beta correction factor. - beta_slow (Optional[int]): Slow beta correction factor. - mscale (float): Scaling factor for extended attention. - index_head_dim (int): Dimension for index head. - index_topk (int): Top-k for index head. - """ + """Data class for defining model arguments and hyperparameters.""" arch_name = "deepseek_v3_2" @@ -54,6 +18,7 @@ class ModelArgs: max_seq_len: int = 160 * 1024 dtype: Literal["bf16", "fp8"] = "fp8" scale_fmt: str | None = None + fp8_kv_cache: bool = False vocab_size: int = 129280 dim: int = 7168 diff --git a/tilert/models/glm_5/_dsa_v32/modules/__init__.py b/tilert/models/glm_5/_dsa_v32/modules/__init__.py deleted file mode 100644 index 937085b..0000000 --- a/tilert/models/glm_5/_dsa_v32/modules/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -"""DeepSeek v3.2 high-level Python modules (MLA, MLP, MTP, etc.).""" - -__all__ = [ - "dsa", - "end2end", - "mla", - "mlp", - "moe", - "mtp", - "mtp_preprocess", -] diff --git a/tilert/models/glm_5/_dsa_v32/ops/__init__.py b/tilert/models/glm_5/_dsa_v32/ops/__init__.py index a58dab8..40fc65b 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/__init__.py +++ b/tilert/models/glm_5/_dsa_v32/ops/__init__.py @@ -8,11 +8,6 @@ DownAllReduceAlgorithm, down_allreduce, ) -from tilert.models.glm_5._dsa_v32.ops.eh_proj_allreduce import ( - EHProjAllReduce, - EHProjAllReduceAlgorithm, - eh_proj_allreduce, -) from tilert.models.glm_5._dsa_v32.ops.expert_down_allreduce import ( ExpertDownAllReduce, ExpertDownAllReduceAlgorithm, @@ -22,29 +17,13 @@ ExpertSelectUpGateSiLU, ExpertSelectUpGateSiLUAlgorithm, ) -from tilert.models.glm_5._dsa_v32.ops.flash_sparse_mla import ( - FlashSparseMLACombineAlgorithm, - flash_sparse_mla, -) from tilert.models.glm_5._dsa_v32.ops.layernorm_rope_rotate import ( LayerNormRoPERotateAlgorithm, layernorm_rope_rotate, ) -from tilert.models.glm_5._dsa_v32.ops.padded_allreduce_add import ( - PaddedAllReduceAdd, - PaddedAllReduceAddAlgorithm, - padded_allreduce_add, -) from tilert.models.glm_5._dsa_v32.ops.projo_wkvb import ProjoWKVbAlgorithm, projo_wkvb from tilert.models.glm_5._dsa_v32.ops.projq_wqb import ProjqWqbAlgorithm, projq_wqb from tilert.models.glm_5._dsa_v32.ops.projx_wis import ProjxWisAlgorithm, projx_wis -from tilert.models.glm_5._dsa_v32.ops.qkv_rope import ( - QKVRoPE, - QKVRoPEAlgorithm, - QKVRoPERefWeightsAlias, - QKVRoPETilertWeightsAlias, - qkv_rope, -) from tilert.models.glm_5._dsa_v32.ops.receive_selected_token_ids import ( receive_selected_token_ids, ) @@ -88,13 +67,6 @@ rotate, rotate_activation, ) -from tilert.models.glm_5._dsa_v32.ops.sparse_index import sparse_index, sparse_index_topk -from tilert.models.glm_5._dsa_v32.ops.topk import TopK, topk_accurate, topk_approximate -from tilert.models.glm_5._dsa_v32.ops.unproj_o_allreduce import ( - UnProjOAllReduce, - UnProjOAllReduceAlgorithm, - unproj_o_allreduce, -) __all__ = [ "down_allreduce", @@ -118,9 +90,6 @@ "RotateTilertWeightsAlias", "layernorm_rope_rotate", "LayerNormRoPERotateAlgorithm", - "TopK", - "topk_approximate", - "topk_accurate", "sparse_index", "sparse_index_topk", "flash_sparse_mla", @@ -132,8 +101,6 @@ "QKVRoPEAlgorithm", "QKVRoPERefWeightsAlias", "QKVRoPETilertWeightsAlias", - "eh_proj_allreduce", - "EHProjAllReduceAlgorithm", "rmsnorm_quant", "RmsnormProjqWqi", "RmsnormProjqWqiAlgorithm", diff --git a/tilert/models/glm_5/_dsa_v32/ops/broadcast_selected_token_ids.py b/tilert/models/glm_5/_dsa_v32/ops/broadcast_selected_token_ids.py index f6bf2a8..1621a14 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/broadcast_selected_token_ids.py +++ b/tilert/models/glm_5/_dsa_v32/ops/broadcast_selected_token_ids.py @@ -15,17 +15,7 @@ def broadcast_selected_token_ids( model_arch: str, compute_kernel_type: str = "bf16", ) -> None: - """Broadcast idx_selects [1,S,2048] int32 from GPU 0 to peer GPUs. - - Args: - idx_selects: Source tensor [1, S, 2048] int32 on GPU 0. - peer_bufs: Device pointer array [N] int64 β€” each entry is a peer - buffer address. - flag_val: Synchronization flag value. - profile_logs: Profile logs tensor. - model_arch: Model architecture ("deepseek_v3_2" or "glm_5"). - compute_kernel_type: Compute kernel type ("bf16"). - """ + """Broadcast idx_selects [1,S,2048] int32 from GPU 0 to peer GPUs.""" torch.ops.tilert.broadcast_selected_token_ids_op( idx_selects, peer_bufs, diff --git a/tilert/models/glm_5/_dsa_v32/ops/down_allreduce.py b/tilert/models/glm_5/_dsa_v32/ops/down_allreduce.py index 38b305c..8ddabf0 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/down_allreduce.py +++ b/tilert/models/glm_5/_dsa_v32/ops/down_allreduce.py @@ -32,20 +32,7 @@ def down_allreduce( model_arch: str, compute_kernel_type: str = "bf16", ) -> None: - """ - Fused operation of down and allreduce. - - Args: - vec_in: Input tensor. - mat_in: Input tensor. - mat_scale: Input tensor. - x_in: Input tensor. - flag: Input flag. - vec_out: Output tensor. - profile_logs: Profile logs tensor (1D). - model_arch: Model architecture ("deepseek_v3_2" or "glm_5"). - compute_kernel_type: Compute kernel type ("bf16"). - """ + """Fused operation of down and allreduce.""" torch.ops.tilert.down_allreduce_op( vec_in, mat_in, @@ -63,6 +50,8 @@ class DownAllReduceAlgorithm(Enum): """DownAllReduce algorithm""" GENERAL = "general" + BF16MMA = "bf16mma" + BF16MMA_V2 = "bf16mma_v2" DownAllReduceWeightsConverter = ExpertDownAllReduceWeightsConverter @@ -88,7 +77,11 @@ class DownAllReduce(TileRTModule): _SUPPORTED_ALGORITHMS = { "deepseek_v3_2": [DownAllReduceAlgorithm.GENERAL], - "glm_5": [DownAllReduceAlgorithm.GENERAL], + "glm_5": [ + DownAllReduceAlgorithm.GENERAL, + DownAllReduceAlgorithm.BF16MMA, + DownAllReduceAlgorithm.BF16MMA_V2, + ], } def __init__( @@ -119,10 +112,14 @@ def __init__( self.moe_inter_scale_dim_per_device = self.moe_inter_dim_per_device // self.block_size self.algorithm = algorithm - if self.arch_name in ("deepseek_v3_2", "glm_5"): - self.compute_kernel_type = "bf16" - else: + if self.arch_name not in ("deepseek_v3_2", "glm_5"): raise ValueError(f"Unsupported architecture: {self.arch_name}") + if self.algorithm == DownAllReduceAlgorithm.BF16MMA: + self.compute_kernel_type = "bf16mma" + elif self.algorithm == DownAllReduceAlgorithm.BF16MMA_V2: + self.compute_kernel_type = "bf16mma_v2" + else: + self.compute_kernel_type = "bf16" self.model_arch = self.arch_name @@ -153,12 +150,7 @@ def tilert_tensor_alias(self) -> list[str]: return self.tilert_weights_alias.tilert_tensor_alias def get_weights_list(self) -> list[torch.Tensor]: - """ - Get the weights list. - - Returns: - List of weights. - """ + """Get the weights list.""" return [self.tilert_weights, self.tilert_scales] def device_sharding( @@ -166,15 +158,7 @@ def device_sharding( weights_dict: dict[str, torch.Tensor], key_prefix: str, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Device sharding. - - Args: - weights_dict: Dictionary of weights. - key_prefix: Key prefix. - Returns: - Tuple of weights. - """ + """Device sharding.""" down_proj_weight_key = f"{key_prefix}.down_proj.weight" down_proj_scale_key = f"{key_prefix}.down_proj.weight_scale_inv" down_proj_weight = weights_dict[down_proj_weight_key] @@ -216,13 +200,7 @@ def init_reference_weights( key_prefix: str, device_id: int = 0, ) -> None: - """ - Initialize the reference weights. - - Args: - state_dict: State dictionary. - device_id: Device ID. - """ + """Initialize the reference weights.""" sharded_list = self.device_sharding(state_dict, key_prefix) down_weights = sharded_list[0][device_id] @@ -235,25 +213,19 @@ def init_reference_weights( self.ref_down = torch.stack(down_list, dim=0) def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - """ - Initialize the tilert weights. - - Args: - state_dict: State dictionary. - """ + """Initialize the tilert weights.""" assert self.algorithm is not None, "Algorithm is not set" + converter_algorithm = ( + DownAllReduceAlgorithm.BF16MMA + if self.algorithm == DownAllReduceAlgorithm.BF16MMA_V2 + else self.algorithm + ) self.tilert_weights, self.tilert_scales = DownAllReduceWeightsConverter( self.model_args, self.num_devices - ).dispatch(self.algorithm, [state_dict[alias] for alias in self.tensor_alias]) + ).dispatch(converter_algorithm, [state_dict[alias] for alias in self.tensor_alias]) def init_tilert_vars(self, batch_size: int, seq_len: int, device_id: int = 0) -> None: - """ - Initialize the tilert variables. - - Args: - batch_size: Batch size. - seq_len: Sequence length. - """ + """Initialize the tilert variables.""" self.hidden_out = torch.zeros( (batch_size, seq_len, self.dim), dtype=torch.bfloat16, @@ -262,8 +234,10 @@ def init_tilert_vars(self, batch_size: int, seq_len: int, device_id: int = 0) -> self.profile_logs = get_profile_log_tensor(device=f"cuda:{device_id}") self.is_init = True - def init_random_weights(self, device_id: int = 0) -> None: + def init_random_weights(self, device_id: int | None = None) -> None: """Initialize the random weights.""" + if device_id is None: + device_id = self.device_id scale_dtype = torch.float32 if self.arch_name == "glm_5" else torch.bfloat16 down_weights = torch.randn( self.dim, self.inter_dim, dtype=torch.bfloat16, device=f"cuda:{device_id}" @@ -292,15 +266,7 @@ def golden_forward( self, vec_in: torch.Tensor, ) -> torch.Tensor: - """ - Forward pass for the down-project module. - - Args: - vec_in: Input vector. - - Returns: - Output tensor. - """ + """Forward pass for the down-project module.""" assert self.ref_down is not None bsz = vec_in.shape[0] assert bsz == 1 diff --git a/tilert/models/glm_5/_dsa_v32/ops/eh_proj_allreduce.py b/tilert/models/glm_5/_dsa_v32/ops/eh_proj_allreduce.py deleted file mode 100644 index fe0b71f..0000000 --- a/tilert/models/glm_5/_dsa_v32/ops/eh_proj_allreduce.py +++ /dev/null @@ -1,293 +0,0 @@ -"""EHProjAllReduce operation module.""" - -from dataclasses import dataclass -from enum import Enum - -import torch - -from tilert.models.base import TileRTModule, TilertWeightsConverter -from tilert.models.glm_5._dsa_v32.model_args import ModelArgs -from tilert.utils import get_profile_log_tensor - -__all__ = [ - "eh_proj_allreduce", - "EHProjAllReduceTilertWeightsAlias", -] - - -def eh_proj_allreduce( - vec_in_enorm: torch.Tensor, - vec_in_hnorm: torch.Tensor, - w_eh: torch.Tensor, - flag: int, - vec_out: torch.Tensor, - profile_logs: torch.Tensor, - model_arch: str, -) -> None: - """ - Fused operation of EHProj and allreduce. - - Args: - vec_in_enorm: Input tensor of shape (1, seq_len, 7168). - vec_in_hnorm: Input tensor of shape (1, seq_len, 7168). - w_eh: Input tensor of shape (7168, 1792) or (128, 7, 56, 256). - flag: Input tensor. - vec_out: Output tensor of shape (1, seq_len, 7168). - profile_logs: Profile logs tensor (1D). - model_arch: Model architecture string. - """ - compute_kernel_type = "bf16" - torch.ops.tilert.eh_proj_allreduce_op( - vec_in_enorm, - vec_in_hnorm, - w_eh, - flag, - vec_out, - profile_logs, - model_arch, - compute_kernel_type, - torch.empty(0, dtype=torch.int64, device=vec_in_enorm.device), - ) - - -class EHProjAllReduceAlgorithm(Enum): - """EHProjAllReduce algorithm""" - - GENERAL = "general" - - -class EHProjAllReduceWeightsConverter(TilertWeightsConverter): - """EHProj weights converter""" - - def convert_to_general(self, weights_list: list[torch.Tensor]) -> tuple[torch.Tensor]: - """ - Convert the weights to general format. - - Args: - weights_list: List of weights. - - Returns: - Tuple of weights. - """ - args = self.model_args - assert args.arch_name == "deepseek_v3_2" or args.arch_name == "glm_5" - dim = args.dim - num_sms = 128 - dim_per_sm = dim // num_sms - in_dim = dim * 2 - in_dim_per_device = in_dim // self.num_devices - stages = in_dim_per_device // 256 - - with torch.inference_mode(): - (proj_weights,) = weights_list - proj_weights = proj_weights.reshape(num_sms, dim_per_sm, stages, 256) - proj_weights = proj_weights.transpose(1, 2) - return (proj_weights.contiguous(),) - - -@dataclass -class EHProjAllReduceTilertWeightsAlias: - """TileRT weights alias for EHProjAllReduce.""" - - eh_proj_weights = "eh_proj_weights" - - @property - def tilert_tensor_alias(self) -> list[str]: - return [self.eh_proj_weights] - - def __call__(self) -> list[str]: - return self.tilert_tensor_alias - - -class EHProjAllReduce(TileRTModule): - """EHProjAllReduce module""" - - _SUPPORTED_ALGORITHMS = { - "deepseek_v3_2": [EHProjAllReduceAlgorithm.GENERAL], - "glm_5": [EHProjAllReduceAlgorithm.GENERAL], - } - - def __init__( - self, - model_args: ModelArgs, - num_devices: int, - algorithm: EHProjAllReduceAlgorithm = EHProjAllReduceAlgorithm.GENERAL, - ): - super().__init__( - self.__class__.__name__, - model_args=model_args, - num_devices=num_devices, - ) - - self.arch_name = self.model_args.arch_name - self.dim = self.model_args.dim - - self.algorithm = algorithm - - self.ref_proj: torch.Tensor | None = None - - self.tilert_proj: torch.Tensor | None = None - - self.hidden_out: torch.Tensor | None = None - - self.profile_logs: torch.Tensor | None = None - self.is_init = False - - self.tilert_weights_alias = EHProjAllReduceTilertWeightsAlias() - - self.tensor_alias: list[str] = [ - "eh_proj_weights", - ] - - self.ref_tensor_alias: list[str] = [ - "eh_proj.weight", - ] - - @property - def tilert_tensor_alias(self) -> list[str]: - return self.tilert_weights_alias.tilert_tensor_alias - - def get_weights_list(self) -> list[torch.Tensor]: - """ - Get the weights list. - - Returns: - List of weights. - """ - return [self.tilert_proj] - - def device_sharding( - self, - weights_dict: dict[str, torch.Tensor], - key_prefix: str | None = None, - ) -> tuple[torch.Tensor]: - """ - Device sharding. - - Args: - weights_dict: Dictionary of weights. - key_prefix: Key prefix. - Returns: - Tuple of weights. - """ - eh_proj_key = "eh_proj.weight" - if key_prefix is not None: - eh_proj_key = f"{key_prefix}.eh_proj.weight" - - eh_proj_weight = weights_dict[eh_proj_key] - in_dim = eh_proj_weight.shape[1] - out_dim = eh_proj_weight.shape[0] - in_dim_per_device = in_dim // self.num_devices - eh_proj_weight = eh_proj_weight.reshape(out_dim, self.num_devices, in_dim_per_device) - eh_proj_weight = eh_proj_weight.transpose(0, 1) - return (eh_proj_weight.contiguous(),) - - def init_reference_weights( - self, - state_dict: dict[str, torch.Tensor], - key_prefix: str | None = None, - device_id: int = 0, - ) -> None: - """ - Initialize the reference weights. - - Args: - state_dict: State dictionary. - device_id: Device ID. - """ - sharded_list = self.device_sharding(state_dict, key_prefix) - - eh_proj_weight = sharded_list[0][device_id] - - self.ref_proj = eh_proj_weight - - def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - """ - Initialize the tilert weights. - - Args: - state_dict: State dictionary. - """ - assert self.algorithm is not None - (self.tilert_proj,) = EHProjAllReduceWeightsConverter( - self.model_args, self.num_devices - ).dispatch(self.algorithm, [state_dict[alias] for alias in self.tensor_alias]) - - def init_tilert_vars(self, batch_size: int, seq_len: int, device_id: int = 0) -> None: - """ - Initialize the tilert variables. - - Args: - batch_size: Batch size. - seq_len: Sequence length. - """ - self.hidden_out = torch.zeros( - (batch_size, seq_len, self.dim), - dtype=torch.bfloat16, - device=f"cuda:{device_id}", - ) - self.profile_logs = get_profile_log_tensor(device=f"cuda:{device_id}") - self.is_init = True - - def init_random_weights(self, device_id: int = 0) -> None: - """Initialize the random weights.""" - proj_weights = torch.randn( - self.dim, self.dim * 2, dtype=torch.bfloat16, device=f"cuda:{device_id}" - ) - - tensor_list = [ - proj_weights, - ] - state_dict = dict(zip(self.ref_tensor_alias, tensor_list)) - - self.init_reference_weights(state_dict, None, device_id) - sharded_list = self.device_sharding(state_dict, None) - sharded_state_dict = { - alias: sharded_list[i][device_id] for i, alias in enumerate(self.tensor_alias) - } - self.init_tilert_weights(sharded_state_dict) - - def golden_forward( - self, - vec_in_enorm: torch.Tensor, - vec_in_hnorm: torch.Tensor, - device_id: int = 0, - ) -> torch.Tensor: - """ - Forward pass for the down-project module. - - Args: - vec_in_enorm: Input vector of shape (1, seq_len, 7168). - vec_in_hnorm: Input vector of shape (1, seq_len, 7168). - - Returns: - Output tensor. - """ - assert self.ref_proj is not None - bsz = vec_in_enorm.shape[0] - assert bsz == 1 - - vec_in_concat = torch.cat([vec_in_enorm, vec_in_hnorm], dim=-1) - dim_per_device = (self.dim * 2) // self.num_devices - vec_in_slice = vec_in_concat[ - ..., dim_per_device * device_id : dim_per_device * device_id + dim_per_device - ] - return vec_in_slice @ self.ref_proj.T - - def tilert_forward( - self, - vec_in_enorm: torch.Tensor, - vec_in_hnorm: torch.Tensor, - flag: int, - ) -> torch.Tensor: - assert self.hidden_out is not None - eh_proj_allreduce( - vec_in_enorm, - vec_in_hnorm, - self.tilert_proj, - flag, - self.hidden_out, - self.profile_logs, - model_arch=self.model_args.arch_name, - ) - return self.hidden_out diff --git a/tilert/models/glm_5/_dsa_v32/ops/expert_down_allreduce.py b/tilert/models/glm_5/_dsa_v32/ops/expert_down_allreduce.py index b0e6b24..19e98e6 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/expert_down_allreduce.py +++ b/tilert/models/glm_5/_dsa_v32/ops/expert_down_allreduce.py @@ -18,9 +18,6 @@ ] -VALID_SEQ_LENS = (1, 2, 4) - - def expert_down_allreduce( vec_in: torch.Tensor, mat_in: torch.Tensor, @@ -30,25 +27,11 @@ def expert_down_allreduce( x_in: torch.Tensor, flag: int, vec_out: torch.Tensor, - profile_logs: torch.Tensor, model_arch: str, compute_kernel_type: str = "bf16", + profile_logs: torch.Tensor | None = None, ) -> None: - """ - Fused expert down + allreduce (unified for DSv32 and GLM5). - - Args: - vec_in: [1, seq_len, n_experts, 256], bfloat16. - mat_in: [n_experts, dim, 256], float8_e4m3fn. - mat_scale: [n_experts, 1024, 2], bfloat16 (DSv32) or float32 (GLM5). - indices: [1, seq_len, 8], int32. - scores: [1, seq_len, 8], float32. - x_in: [1, seq_len, dim], bfloat16. - flag: User flag. - vec_out: [1, seq_len, dim], bfloat16 (output). - profile_logs: 1D tensor for profile logs. - compute_kernel_type: "bf16". - """ + """Fused expert down + allreduce (unified for DSv32 and GLM5).""" torch.ops.tilert.expert_down_allreduce_op( vec_in, mat_in, @@ -58,9 +41,9 @@ def expert_down_allreduce( x_in, flag, vec_out, - profile_logs, model_arch, compute_kernel_type, + profile_logs, ) @@ -68,6 +51,8 @@ class ExpertDownAllReduceAlgorithm(Enum): """ExpertDownAllReduce algorithm.""" GENERAL = "general" + BF16MMA = "bf16mma" + GLM5_FP4_HMMA = "glm5_fp4_hmma" class ExpertDownAllReduceWeightsConverter(TilertWeightsConverter): @@ -87,6 +72,26 @@ def _swizzle_qmma_8x32(mat_in: torch.Tensor) -> torch.Tensor: pre_shape = mat_in.shape[:-2] return mat_in.reshape(*pre_shape, 8, 2, 4, 4).transpose(-2, -3).contiguous() + @staticmethod + def _swizzle_bf16mma_full_16x32(mat_in: torch.Tensor) -> torch.Tensor: + assert mat_in.shape[-2] == 16 and mat_in.shape[-1] == 32 + assert mat_in.dtype == torch.float8_e4m3fn + pre = mat_in.shape[:-2] + mat = mat_in.reshape(*pre, 2, 8, 2, 2, 4, 2) + n = len(pre) + mat = mat.permute(*range(n), 1 + n, 4 + n, 2 + n, 3 + n, 0 + n, 5 + n) + return mat.reshape(*pre, 32, 16).contiguous() + + @staticmethod + def _swizzle_bf16mma_partial_8x32(mat_in: torch.Tensor) -> torch.Tensor: + assert mat_in.shape[-2] == 8 and mat_in.shape[-1] == 32 + assert mat_in.dtype == torch.float8_e4m3fn + pre = mat_in.shape[:-2] + mat = mat_in.reshape(*pre, 8, 2, 2, 4, 2) + n = len(pre) + mat = mat.permute(*range(n), 0 + n, 3 + n, 1 + n, 2 + n, 4 + n) + return mat.reshape(*pre, 32, 8).contiguous() + def convert_to_general( self, weights_list: list[torch.Tensor] ) -> tuple[torch.Tensor, torch.Tensor]: @@ -155,10 +160,124 @@ def convert_to_general( mat_scale_tilert = mat_scale_tilert.to(torch.bfloat16) return mat_in_swizzled.contiguous(), mat_scale_tilert.contiguous() + def convert_to_bf16mma( + self, weights_list: list[torch.Tensor] + ) -> tuple[torch.Tensor, torch.Tensor]: + args = self.model_args + assert args.arch_name in ( + "deepseek_v3_2", + "glm_5", + ), "BF16 MMA wired for DSv32 / GLM5 only." + dim = args.dim + num_sms = 128 + dim_per_sm = dim // num_sms + expert_dim = args.moe_inter_dim // 8 + k_chunks = expert_dim // 32 + scale_cols = expert_dim // args.block_size + n_full_tiles = dim_per_sm // 16 + remainder_rows = dim_per_sm % 16 + full_rows = n_full_tiles * 16 + + with torch.inference_mode(): + mat_in, scale_in = weights_list + exp_num = mat_in.shape[0] + mat_per_cta = mat_in.reshape(exp_num, num_sms, dim_per_sm, expert_dim) + + full_part = mat_per_cta[:, :, :full_rows, :] + full_tiles = full_part.reshape( + exp_num, num_sms, n_full_tiles, 16, k_chunks, 32 + ).transpose(3, 4) + full_swizzled = self._swizzle_bf16mma_full_16x32(full_tiles) + full_swizzled = full_swizzled.reshape( + exp_num, num_sms, n_full_tiles * k_chunks * 32 * 16 + ) + + mats = [full_swizzled] + if remainder_rows > 0: + partial_part = mat_per_cta[:, :, full_rows:, :] + partial_tiles = partial_part.reshape( + exp_num, num_sms, 1, remainder_rows, k_chunks, 32 + ).transpose(3, 4) + partial_swizzled = self._swizzle_bf16mma_partial_8x32(partial_tiles) + partial_swizzled = partial_swizzled.reshape( + exp_num, num_sms, k_chunks * 32 * remainder_rows + ) + mats.append(partial_swizzled) + + mat_swizzled = torch.cat(mats, dim=2) if len(mats) > 1 else mats[0] + mat_swizzled = mat_swizzled.reshape(exp_num, dim, expert_dim) + + mat_scale_tilert = ( + scale_in.reshape(exp_num, dim // args.block_size, 1, scale_cols) + .repeat(1, 1, 16, 1) + .reshape(exp_num, num_sms, -1) + ) + target_cols_per_sm = 1024 * scale_cols // num_sms + pad_amount = target_cols_per_sm - mat_scale_tilert.shape[-1] + if pad_amount > 0: + padding_zeros = torch.zeros( + (exp_num, num_sms, pad_amount), + dtype=scale_in.dtype, + device=scale_in.device, + ) + mat_scale_tilert = torch.cat([mat_scale_tilert, padding_zeros], dim=2) + mat_scale_tilert = mat_scale_tilert.reshape(exp_num, 1024, scale_cols) + if args.arch_name == "glm_5": + mat_scale_tilert = mat_scale_tilert.to(torch.float32) + else: + mat_scale_tilert = mat_scale_tilert.to(torch.bfloat16) + + return mat_swizzled.contiguous(), mat_scale_tilert.contiguous() + + def convert_to_glm5_fp4_hmma( + self, weights_list: list[torch.Tensor] + ) -> tuple[torch.Tensor, torch.Tensor]: + from tilert.models.common_mxfp4 import ( + _unpack_fp4_nibbles_last, + build_down_weights_mma_natural, + ) + + assert ( + len(weights_list) == 2 + ), f"convert_to_glm5_fp4_hmma expects 2 tensors, got {len(weights_list)}" + down_fp4, down_e8m0 = weights_list + arch = self.model_args.arch_name + assert arch == "glm_5", f"GLM5_FP4_HMMA down converter is GLM5-only, got arch={arch}" + + dim = self.model_args.dim + moe_inter_pd = self.model_args.moe_inter_dim // self.num_devices + + with torch.inference_mode(): + if down_fp4.shape[-1] == moe_inter_pd: + down_nib = down_fp4.to(torch.uint8).contiguous() + else: + assert down_fp4.shape[-1] == moe_inter_pd // 2, ( + "routed fp4 down last dim must be inter_pd or inter_pd/2; " + f"got {down_fp4.shape[-1]} (inter_pd={moe_inter_pd})" + ) + down_nib = _unpack_fp4_nibbles_last(down_fp4) + down_e8 = down_e8m0.to(torch.uint8).contiguous() + + n_routed = down_nib.shape[0] + assert down_nib.shape == (n_routed, dim, moe_inter_pd), ( + f"down_fp4 must be (n_routed, {dim}, {moe_inter_pd}); " + f"got {tuple(down_nib.shape)}" + ) + + device = down_nib.device + e_total = n_routed + 1 + u8 = {"dtype": torch.uint8, "device": device} + full_nib = torch.zeros(e_total, dim, moe_inter_pd, **u8) + full_e8 = torch.zeros(e_total, dim, moe_inter_pd // 32, **u8) + full_nib[1:] = down_nib + full_e8[1:] = down_e8 + down_packed = build_down_weights_mma_natural(full_nib, full_e8, dim, moe_inter_pd) + dummy = torch.zeros(1, dtype=torch.float32, device=device) + return down_packed, dummy + @dataclass class ExpertDownAllReduceTilertWeightsAlias: - """TileRT weights alias for ExpertDownAllReduce.""" exp_down_weights = "exp_down_weights" exp_down_scales = "exp_down_scales" @@ -167,6 +286,9 @@ class ExpertDownAllReduceTilertWeightsAlias: def tilert_tensor_alias(self) -> list[str]: return [self.exp_down_weights, self.exp_down_scales] + def glm5_fp4_tilert_tensor_alias(self) -> list[str]: + return [self.exp_down_weights, self.exp_down_scales] + def __call__(self) -> list[str]: return self.tilert_tensor_alias @@ -175,8 +297,15 @@ class ExpertDownAllReduce(TileRTModule): """ExpertDownAllReduce module.""" _SUPPORTED_ALGORITHMS = { - "deepseek_v3_2": [ExpertDownAllReduceAlgorithm.GENERAL], - "glm_5": [ExpertDownAllReduceAlgorithm.GENERAL], + "deepseek_v3_2": [ + ExpertDownAllReduceAlgorithm.GENERAL, + ExpertDownAllReduceAlgorithm.BF16MMA, + ], + "glm_5": [ + ExpertDownAllReduceAlgorithm.GENERAL, + ExpertDownAllReduceAlgorithm.BF16MMA, + ExpertDownAllReduceAlgorithm.GLM5_FP4_HMMA, + ], } def __init__( @@ -210,6 +339,8 @@ def __init__( if self.arch_name in ("deepseek_v3_2", "glm_5"): self.compute_kernel_type = "bf16" + if algorithm == ExpertDownAllReduceAlgorithm.BF16MMA: + self.compute_kernel_type = "bf16mma" else: raise ValueError(f"Unsupported architecture: {self.arch_name}") @@ -228,6 +359,13 @@ def __init__( def tilert_tensor_alias(self) -> list[str]: return self.tilert_weights_alias.tilert_tensor_alias + def set_algorithm(self, algorithm: Enum) -> None: + super().set_algorithm(algorithm) + if algorithm == ExpertDownAllReduceAlgorithm.BF16MMA: + self.compute_kernel_type = "bf16mma" + elif algorithm == ExpertDownAllReduceAlgorithm.GENERAL: + self.compute_kernel_type = "bf16" + def get_weights_list(self) -> list[torch.Tensor]: return [self.tilert_weights, self.tilert_scales] @@ -261,11 +399,50 @@ def process_down_weights( ) return down_proj_weight, down_proj_scale + @staticmethod + def _split_last_axis(t: torch.Tensor, num_devices: int) -> torch.Tensor: + d, inter = t.shape[-2], t.shape[-1] + assert ( + inter % num_devices == 0 + ), f"down last-axis {inter} not divisible by num_devices {num_devices}" + return ( + t.reshape(d, num_devices, inter // num_devices) + .transpose(0, 1) + .reshape(num_devices, 1, d, inter // num_devices) + ) + + def process_down_weights_fp4( + self, + key_prefix: str, + weights_hf: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + n_dev = self.num_devices + down_w = weights_hf[f"{key_prefix}.down_proj.weight"] + down_s = weights_hf[f"{key_prefix}.down_proj.weight_scale"] + return self._split_last_axis(down_w, n_dev), self._split_last_axis(down_s, n_dev) + + def device_sharding_fp4( + self, + weights_dict: dict[str, torch.Tensor], + key_prefix: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Shard routed-only down weight + scale across devices.""" + dw, ds = [], [] + for exp_id in range(self.n_routed_experts): + down_weights, down_scales = self.process_down_weights_fp4( + f"{key_prefix}.experts.{exp_id}", weights_dict + ) + dw.append(down_weights) + ds.append(down_scales) + return torch.cat(dw, dim=1).contiguous(), torch.cat(ds, dim=1).contiguous() + def device_sharding( self, weights_dict: dict[str, torch.Tensor], key_prefix: str, ) -> tuple[torch.Tensor, torch.Tensor]: + if self.algorithm == ExpertDownAllReduceAlgorithm.GLM5_FP4_HMMA: + return self.device_sharding_fp4(weights_dict, key_prefix) assert self.n_shared_experts == 1, "Only one shared expert is supported" down_weights_list = [] down_scales_list = [] @@ -300,13 +477,31 @@ def init_reference_weights( weight_dequant(down_weight, down_scale) for down_weight, down_scale in zip(down_weights, down_scales) ] - self.ref_down = torch.stack(down_list, dim=0) + self.ref_down = torch.stack([t.to(torch.bfloat16) for t in down_list], dim=0) + + def get_tilert_weights_alias(self) -> list[str]: + """Return the alias list keyed into ``state_dict`` for this op.""" + return list(self.tilert_weights_alias()) def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: assert self.algorithm is not None, "Algorithm is not set" - self.tilert_weights, self.tilert_scales = ExpertDownAllReduceWeightsConverter( - self.model_args, self.num_devices - ).dispatch(self.algorithm, [state_dict[alias] for alias in self.tensor_alias]) + if self.algorithm == ExpertDownAllReduceAlgorithm.GLM5_FP4_HMMA: + assert ( + self.arch_name == "glm_5" + ), f"GLM5_FP4_HMMA is GLM5-only, got arch={self.arch_name}" + converter = ExpertDownAllReduceWeightsConverter(self.model_args, self.num_devices) + self.tilert_weights, self.tilert_scales = converter.convert_to_glm5_fp4_hmma( + [state_dict[alias] for alias in self.tensor_alias] + ) + self.is_tilert_weights_init = True + return + aliases = [state_dict[alias] for alias in self.tensor_alias] + self.tilert_weights, self.tilert_scales = ( + torch.ops.tilert.expert_down_allreduce__convert_weights( + aliases, self.model_arch, self.compute_kernel_type + ) + ) + self.is_tilert_weights_init = True def init_tilert_vars(self, batch_size: int, seq_len: int, device_id: int = 0) -> None: self.hidden_out = torch.zeros( @@ -317,7 +512,9 @@ def init_tilert_vars(self, batch_size: int, seq_len: int, device_id: int = 0) -> self.profile_logs = get_profile_log_tensor(device=f"cuda:{device_id}") self.is_init = True - def init_random_weights(self, device_id: int = 0) -> None: + def init_random_weights(self, device_id: int | None = None) -> None: + if device_id is None: + device_id = self.device_id n = self.n_routed_experts + 1 dev = f"cuda:{device_id}" down_weights = list( @@ -389,9 +586,9 @@ def tilert_forward( x_in, flag, self.hidden_out, - self.profile_logs, self.model_arch, self.compute_kernel_type, + self.profile_logs, ) return self.hidden_out diff --git a/tilert/models/glm_5/_dsa_v32/ops/expert_sel_up_gate_silu.py b/tilert/models/glm_5/_dsa_v32/ops/expert_sel_up_gate_silu.py index e2d96eb..c5c3fbd 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/expert_sel_up_gate_silu.py +++ b/tilert/models/glm_5/_dsa_v32/ops/expert_sel_up_gate_silu.py @@ -29,10 +29,14 @@ def expert_select_up_gate_silu( hidden_out: torch.Tensor, expert_probs_out: torch.Tensor, expert_indices_out: torch.Tensor, - profile_logs: torch.Tensor, + profile_logs: torch.Tensor | None = None, algorithm: str = "fp8mma", *, model_arch: str, + tid2eid: torch.Tensor | None = None, + token_id: torch.Tensor | None = None, + experts_weights_32row: torch.Tensor | None = None, + shared_experts_weights: torch.Tensor | None = None, ) -> None: """Expert SelectUpGateSiLU operation.""" torch.ops.tilert.expert_select_up_gate_silu_op( @@ -46,6 +50,10 @@ def expert_select_up_gate_silu( profile_logs, model_arch, algorithm, + tid2eid, + token_id, + experts_weights_32row, + shared_experts_weights, ) @@ -71,6 +79,16 @@ def ref_tensor_alias(self) -> list[str]: + [f"{self.key_prefix}.experts.{i}.up_proj.weight_scale_inv" for i in range(n)] ) + def fp4_routed_tensor_alias(self) -> list[str]: + n = self.n_routed_experts + return ( + [f"{self.key_prefix}.gate.e_score_correction_bias"] + + [f"{self.key_prefix}.experts.{i}.gate_proj.weight" for i in range(n)] + + [f"{self.key_prefix}.experts.{i}.up_proj.weight" for i in range(n)] + + [f"{self.key_prefix}.experts.{i}.gate_proj.weight_scale" for i in range(n)] + + [f"{self.key_prefix}.experts.{i}.up_proj.weight_scale" for i in range(n)] + ) + def __call__(self) -> list[str]: return self.ref_tensor_alias @@ -104,6 +122,8 @@ class ExpertSelectUpGateSiLUAlgorithm(Enum): FP8MMA = "fp8mma" FP16MMA = "fp16mma" + BF16MMA = "bf16mma" + GLM5_FP4_HMMA = "glm5_fp4_hmma" class ExpertSelectUpGateSiLUWeightsConverter(TilertWeightsConverter): @@ -135,16 +155,6 @@ def _swizzle_mma_16x16(mat_in: torch.Tensor) -> torch.Tensor: def tilert_to_tilert_144sm( mat_in: torch.Tensor, mat_scale_in: torch.Tensor, mma_type: str | None = None ) -> torch.Tensor: - """ - Convert tilert weights and scales to tilert_144sm input format. - - Args: - mat_in: tilert weights - mat_scale_in: tilert scales - mma_type: MMA type, None,"16x32" or "16x16" - Returns: - tilert_144sm weights and scales - """ exp_num = mat_in.shape[0] assert mat_in.shape == (exp_num, 512, 7168) assert mat_scale_in.shape == (exp_num, 4, 64) @@ -198,15 +208,6 @@ def tilert_to_tilert_144sm( def tilert_to_tilert_144sm_mma( mat_in: torch.Tensor, mat_scale_in: torch.Tensor, mma_type: str = "16x32" ) -> torch.Tensor: - """ - Convert tilert weights and scales to tilert_144sm_mma input format. - - Args: - mat_in: tilert weights - mat_scale_in: tilert scales - Returns: - tilert_144sm weights and scales - """ return ExpertSelectUpGateSiLUWeightsConverter.tilert_to_tilert_144sm( mat_in, mat_scale_in, mma_type ) @@ -303,31 +304,74 @@ def convert_to_mma( def convert_to_fp8mma( self, weights_list: list[torch.Tensor] ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Convert the weights to fp8mma format. - - Args: - weights: List of weights. - - Returns: - Tuple of weights. - """ return self.convert_to_mma(weights_list, "fp8mma") def convert_to_fp16mma( self, weights_list: list[torch.Tensor] ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Convert the weights to fp16mma format. - - Args: - weights: List of weights. + return self.convert_to_mma(weights_list, "fp16mma") - Returns: - Tuple of weights. - """ + def convert_to_bf16mma( + self, weights_list: list[torch.Tensor] + ) -> tuple[torch.Tensor, torch.Tensor]: return self.convert_to_mma(weights_list, "fp16mma") + def convert_to_glm5_fp4_hmma( + self, weights_list: list[torch.Tensor] + ) -> tuple[torch.Tensor, torch.Tensor]: + from tilert.models.common_mxfp4 import ( + _unpack_fp4_nibbles_last, + build_ug_weights_mma_natural, + ) + + assert ( + len(weights_list) == 5 + ), f"convert_to_glm5_fp4_hmma expects 5 tensors, got {len(weights_list)}" + bias, gate_fp4, gate_e8m0, up_fp4, up_e8m0 = weights_list + arch = self.model_args.arch_name + assert arch == "glm_5", f"GLM5_FP4_HMMA converter is GLM5-only, got arch={arch}" + + dim = self.model_args.dim + moe_inter_pd = self.model_args.moe_inter_dim // self.num_devices + + with torch.inference_mode(): + + def _ensure_unpacked(t: torch.Tensor) -> torch.Tensor: + if t.shape[-1] == dim: + return t.to(torch.uint8).contiguous() + assert t.shape[-1] == dim // 2, ( + "routed fp4 weight last dim must be dim or dim/2; " + f"got {t.shape[-1]} (dim={dim})" + ) + return _unpack_fp4_nibbles_last(t) + + gate_nib = _ensure_unpacked(gate_fp4) + up_nib = _ensure_unpacked(up_fp4) + gate_e8 = gate_e8m0.to(torch.uint8).contiguous() + up_e8 = up_e8m0.to(torch.uint8).contiguous() + + n_routed = gate_nib.shape[0] + assert gate_nib.shape == (n_routed, moe_inter_pd, dim), ( + f"gate_fp4 must be (n_routed, {moe_inter_pd}, {dim}); " + f"got {tuple(gate_nib.shape)}" + ) + + device = gate_nib.device + e_total = n_routed + 1 + u8 = {"dtype": torch.uint8, "device": device} + + def _slot0(nib: torch.Tensor, e8: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + full_nib = torch.zeros(e_total, moe_inter_pd, dim, **u8) + full_e8 = torch.zeros(e_total, moe_inter_pd, dim // 32, **u8) + full_nib[1:] = nib + full_e8[1:] = e8 + return full_nib, full_e8 + + g_nib, g_e8 = _slot0(gate_nib, gate_e8) + u_nib, u_e8 = _slot0(up_nib, up_e8) + ug_packed = build_ug_weights_mma_natural(g_nib, g_e8, u_nib, u_e8, dim, moe_inter_pd) + return bias.float().contiguous(), ug_packed + class ExpertSelectUpGateSiLU(TileRTModule): """ExpertSelectUpGateSiLU module""" @@ -336,10 +380,12 @@ class ExpertSelectUpGateSiLU(TileRTModule): "deepseek_v3_2": [ ExpertSelectUpGateSiLUAlgorithm.FP8MMA, ExpertSelectUpGateSiLUAlgorithm.FP16MMA, + ExpertSelectUpGateSiLUAlgorithm.BF16MMA, ], "glm_5": [ ExpertSelectUpGateSiLUAlgorithm.FP8MMA, ExpertSelectUpGateSiLUAlgorithm.FP16MMA, + ExpertSelectUpGateSiLUAlgorithm.GLM5_FP4_HMMA, ], } @@ -391,7 +437,11 @@ def __init__( self.tilert_bias: torch.Tensor | None = None self.tilert_weights: torch.Tensor | None = None - self.tilert_scales = torch.zeros(1, dtype=torch.bfloat16, device=torch.device("cuda")) + self.tilert_scales = ( + torch.zeros(1, dtype=torch.bfloat16, device=torch.device("cuda")) + if torch.cuda.is_available() + else None + ) self.hidden_out: torch.Tensor | None = None self.expert_probs: torch.Tensor | None = None @@ -417,12 +467,7 @@ def tilert_tensor_alias(self) -> list[str]: return self._tilert_tensor_alias def get_weights_list(self) -> list[torch.Tensor]: - """ - Get the weights list. - - Returns: - List of weights. - """ + """Get the weights list.""" return [self.tilert_bias, self.tilert_weights, self.tilert_scales] @staticmethod @@ -454,16 +499,59 @@ def process_gate_up_weights( up_proj_scale = up_proj_scale.reshape(num_devices, 1, in_scale_dim_per_device, scale_dim) return gate_proj_weight, gate_proj_scale, up_proj_weight, up_proj_scale - def device_sharding(self, weights_map: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """ - Device sharding: ref state dict -> tilert sharded tensors (num_devices, ...). + @staticmethod + def _split_inter_axis(t: torch.Tensor, num_devices: int) -> torch.Tensor: + inter, k = t.shape[-2], t.shape[-1] + assert ( + inter % num_devices == 0 + ), f"moe-inter {inter} not divisible by num_devices {num_devices}" + return t.reshape(num_devices, 1, inter // num_devices, k) + + def process_gate_up_weights_fp4( + self, + key_prefix: str, + weights_hf: dict[str, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + n_dev = self.num_devices + gate_w = weights_hf[f"{key_prefix}.gate_proj.weight"] + gate_s = weights_hf[f"{key_prefix}.gate_proj.weight_scale"] + up_w = weights_hf[f"{key_prefix}.up_proj.weight"] + up_s = weights_hf[f"{key_prefix}.up_proj.weight_scale"] + return ( + self._split_inter_axis(gate_w, n_dev), + self._split_inter_axis(gate_s, n_dev), + self._split_inter_axis(up_w, n_dev), + self._split_inter_axis(up_s, n_dev), + ) + + def _device_sharding_fp4(self, weights_map: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + ref_alias = self.ref_weights_alias + key_prefix = ref_alias.key_prefix + bias = weights_map[f"{key_prefix}.gate.e_score_correction_bias"] + bias = bias[None, :].repeat(self.num_devices, 1) - Args: - weights_map: State dict keyed by ref_weights_alias(). + gw, gs, uw, us = [], [], [], [] + for exp_id in range(self.n_routed_experts): + g_w, g_s, u_w, u_s = self.process_gate_up_weights_fp4( + f"{key_prefix}.experts.{exp_id}", weights_map + ) + gw.append(g_w) + gs.append(g_s) + uw.append(u_w) + us.append(u_s) + tilert_alias = self.tilert_weights_alias + return { + tilert_alias.exp_bias: bias, + tilert_alias.exp_gate_weights: torch.cat(gw, dim=1), + tilert_alias.exp_gate_scales: torch.cat(gs, dim=1), + tilert_alias.exp_up_weights: torch.cat(uw, dim=1), + tilert_alias.exp_up_scales: torch.cat(us, dim=1), + } + + def device_sharding(self, weights_map: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + if self.algorithm == ExpertSelectUpGateSiLUAlgorithm.GLM5_FP4_HMMA: + return self._device_sharding_fp4(weights_map) - Returns: - Dict keyed by tilert_weights_alias() with (num_devices, ...) tensors. - """ ref_alias = self.ref_weights_alias key_prefix = ref_alias.key_prefix @@ -513,13 +601,6 @@ def init_reference_weights( state_dict: dict[str, torch.Tensor], device_id: int | None = None, ) -> None: - """ - Initialize the reference weights. - - Args: - state_dict: State dict keyed by ref_weights_alias(). - device_id: Device ID; defaults to self.device_id. - """ did = self.device_id if device_id is None else device_id sharded = self.device_sharding(state_dict) @@ -537,24 +618,36 @@ def init_reference_weights( ref_up_list = [ weight_dequant(up_weights[i], up_scales[i]) for i in range(up_weights.shape[0]) ] - self.ref_gate = torch.stack(ref_gate_list, dim=0) - self.ref_up = torch.stack(ref_up_list, dim=0) + self.ref_gate = torch.stack([t.to(torch.bfloat16) for t in ref_gate_list], dim=0) + self.ref_up = torch.stack([t.to(torch.bfloat16) for t in ref_up_list], dim=0) + + def get_tilert_weights_alias(self) -> list[str]: + return list(self.tilert_weights_alias()) def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - """Initialize the tilert weights.""" assert self.algorithm is not None, "Algorithm is not set" + if self.algorithm == ExpertSelectUpGateSiLUAlgorithm.GLM5_FP4_HMMA: + assert ( + self.arch_name == "glm_5" + ), f"GLM5_FP4_HMMA is GLM5-only, got arch={self.arch_name}" + a = self.tilert_weights_alias + converter = ExpertSelectUpGateSiLUWeightsConverter(self.model_args, self.num_devices) + self.tilert_bias, self.tilert_weights = converter.convert_to_glm5_fp4_hmma( + [ + state_dict[a.exp_bias], + state_dict[a.exp_gate_weights], + state_dict[a.exp_gate_scales], + state_dict[a.exp_up_weights], + state_dict[a.exp_up_scales], + ] + ) + return + weights_list = [state_dict[alias] for alias in self.tilert_weights_alias()] converter = ExpertSelectUpGateSiLUWeightsConverter(self.model_args, self.num_devices) self.tilert_bias, self.tilert_weights = converter.dispatch(self.algorithm, weights_list) def init_tilert_vars(self, batch_size: int, seq_len: int, device: str = "cuda") -> None: - """ - Initialize the tilert variables. - - Args: - batch_size: Batch size. - seq_len: Sequence length. - """ self.hidden_out = torch.zeros( ( batch_size, @@ -580,12 +673,6 @@ def init_tilert_vars(self, batch_size: int, seq_len: int, device: str = "cuda") self.is_init = True def init_random_weights(self, device: str = "cuda") -> None: - """ - Initialize the random weights. - - Returns: - None - """ n = self.n_routed_experts + 1 bias = torch.randn(self.n_routed_experts, dtype=torch.float32, device=device) gate_weights = list( @@ -705,6 +792,8 @@ def tilert_forward( self, x_in: torch.Tensor, scores: torch.Tensor, + tid2eid: torch.Tensor | None = None, + token_id: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: assert self.algorithm is not None, "Algorithm is not set" expert_select_up_gate_silu( @@ -718,5 +807,7 @@ def tilert_forward( self.profile_logs, self.algorithm.value, model_arch=self.model_args.arch_name, + tid2eid=tid2eid, + token_id=token_id, ) return self.hidden_out, self.expert_probs, self.expert_indices diff --git a/tilert/models/glm_5/_dsa_v32/ops/flash_sparse_mla.py b/tilert/models/glm_5/_dsa_v32/ops/flash_sparse_mla.py deleted file mode 100644 index 1d4cc00..0000000 --- a/tilert/models/glm_5/_dsa_v32/ops/flash_sparse_mla.py +++ /dev/null @@ -1,261 +0,0 @@ -"""Flash Sparse MLA operation module.""" - -import math -from enum import Enum - -import torch - -from tilert.models.base import TileRTModule -from tilert.models.glm_5._dsa_v32.model_args import ModelArgs -from tilert.utils import get_profile_log_tensor - -__all__ = [ - "flash_sparse_mla", - "FlashSparseMLACombine", -] - - -def flash_sparse_mla( - query: torch.Tensor, - query_pe: torch.Tensor, - key_value: torch.Tensor, - key_pe: torch.Tensor, - indices: torch.Tensor, - cur_pos: torch.Tensor, - output: torch.Tensor, - profile_logs: torch.Tensor, - split_size: int = 64, - compute_kernel_type: str = "bf16mma", - *, - model_arch: str, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Flash Sparse MLA operation for GLM5. - - Args: - query: Query tensor. (bs, seqlen, heads, dim) - query_pe: Query position embedding tensor. (bs, seqlen, heads, pe_dim) - key_value: Key-value tensor. (bs, seqlen_kv, dim) - key_pe: Key position embedding tensor. (bs, seqlen_kv, pe_dim) - indices: Indices tensor. (bs, seqlen, topk) - cur_pos: cur_pos tensor. (1) - output: Output tensor. - profile_logs: Profile logs tensor. - split_size: Number of splits. - """ - batch, seqlen, heads, hidden_dim = query.shape - if split_size != 64: - raise ValueError( - "The current implementation of flash_sparse_mla_op only supports split_size=64" - ) - if batch != 1: - raise ValueError("The current implementation of flash_sparse_mla_op only supports batch=1") - if seqlen > 4: - raise ValueError( - "The current implementation of flash_sparse_mla_op only supports seqlen<=4" - ) - - seqlen_kv = key_value.shape[1] - index_len = indices.shape[-1] - if index_len > seqlen_kv: - raise ValueError("index_len must be less than or equal to seqlen_kv") - - device = query.device - acc_type = torch.float32 - - dim = key_value.shape[-1] - max_num_splits = 32 - - lse = torch.empty((batch, seqlen, heads), device=device, dtype=acc_type) - lse_acc = torch.empty((batch, seqlen, heads, max_num_splits), device=device, dtype=acc_type) - output_acc = torch.empty( - batch, seqlen, heads, max_num_splits, dim, device=device, dtype=acc_type - ) - - if heads not in (8, 10, 16, 20): - raise ValueError(f"Unsupported heads: {heads}") - torch.ops.tilert.flash_sparse_mla_op( - query, - query_pe, - key_value, - key_pe, - indices, - cur_pos, - output, - output_acc, - lse, - lse_acc, - split_size, - model_arch, - compute_kernel_type, - profile_logs, - torch.empty(0, dtype=torch.int64, device=query.device), - ) - return lse, lse_acc, output_acc - - -class FlashSparseMLACombineAlgorithm(Enum): - """FlashSparseMLACombine algorithm.""" - - BF16MMA = "bf16mma" - - -class FlashSparseMLACombine(TileRTModule): - """Flash Sparse MLA combine module; no weights, uses model_args for scale and config.""" - - _SUPPORTED_ALGORITHMS = { - "deepseek_v3_2": [FlashSparseMLACombineAlgorithm.BF16MMA], - "glm_5": [FlashSparseMLACombineAlgorithm.BF16MMA], - } - - def __init__( - self, - model_args: ModelArgs, - num_devices: int, - layer_idx: int = 0, - ): - super().__init__( - type(self).__name__, - model_args=model_args, - num_devices=num_devices, - layer_idx=layer_idx, - ) - self.tilert_tensor_alias: list[str] = [] - self.ref_tensor_alias: list[str] = [] - - scale = (model_args.qk_nope_head_dim + model_args.qk_rope_head_dim) ** -0.5 - if model_args.rope_factor is None: - mscale = 1.0 - else: - mscale = 0.1 * math.log(model_args.rope_factor) + 1.0 - self.softmax_scale = scale * mscale * mscale - - self.profile_logs = get_profile_log_tensor() - - def init_reference_weights( - self, state_dict: dict[str, torch.Tensor], device_id: int = 0 - ) -> None: - del state_dict, device_id - self.is_ref_weights_init = True - - def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - del state_dict - self.is_tilert_weights_init = True - - def init_random_weights(self) -> None: - self.is_ref_weights_init = True - self.is_tilert_weights_init = True - - def init_tilert_vars(self, batch_size: int, seq_len: int) -> None: - del batch_size, seq_len - self.profile_logs = get_profile_log_tensor() - self.is_var_init = True - - def golden_forward( - self, - q_nope: torch.Tensor, - q_pe: torch.Tensor, - kv_cache: torch.Tensor, - pe_cache: torch.Tensor, - topk_indices: torch.Tensor, - cur_pos: torch.Tensor, - ) -> torch.Tensor: - """Flash Sparse MLA golden version. - - Args: - q_nope: Query tensor. (bs, seqlen, heads, dim) - q_pe: Query position embedding tensor. (bs, seqlen, heads, pe_dim) - kv_cache: Key-value tensor. (bs, seqlen_kv, dim) - pe_cache: Key position embedding tensor. (bs, seqlen_kv, pe_dim) - topk_indices: Indices tensor. (bs, seqlen, topk) - cur_pos: cur_pos tensor. (1) - """ - batch_size = q_nope.shape[0] - seqlen = q_nope.shape[1] - seqlen_kv = kv_cache.shape[1] - - start_pos = int(cur_pos.item()) - mask = ( - torch.full((seqlen, seqlen_kv), float("-inf")).triu_(start_pos + 1) - if seqlen > 1 - else None - ) - - scores = ( - torch.einsum("bshc,btc->bsht", q_nope.float(), kv_cache.float()) - + torch.einsum("bshr,btr->bsht", q_pe.float(), pe_cache.float()) - ) * self.softmax_scale - index_mask = torch.full( - (batch_size, seqlen, seqlen_kv), float("-inf"), device=q_nope.device - ).scatter_(-1, topk_indices, 0) - if mask is not None: - index_mask += mask - - scores += index_mask.unsqueeze(2) - scores = scores.softmax(dim=-1, dtype=torch.float32) - return torch.einsum("bsht,btc->bshc", scores.to(torch.bfloat16), kv_cache) - - def tilert_forward( - self, - q_nope: torch.Tensor, - q_pe: torch.Tensor, - kv_cache: torch.Tensor, - pe_cache: torch.Tensor, - topk_indices: torch.Tensor, - cur_pos: torch.Tensor, - ) -> torch.Tensor: - """Flash Sparse MLA tilert version. - - Args: - q_nope: Query tensor. (bs, seqlen, heads, dim) - q_pe: Query position embedding tensor. (bs, seqlen, heads, pe_dim) - kv_cache: Key-value tensor. (bs, seqlen_kv, dim) - pe_cache: Key position embedding tensor. (bs, seqlen_kv, pe_dim) - topk_indices: Indices tensor. (bs, seqlen, topk) - cur_pos: cur_pos tensor. (1) - """ - batch_size, seqlen, heads, dim = q_nope.shape - v_dim = kv_cache.shape[-1] - - topk_indices = topk_indices.to(torch.int32) - topk_indices = topk_indices[..., : kv_cache.shape[1]] - device = q_nope.device - if any(t.device != device for t in (q_pe, kv_cache, pe_cache, topk_indices, cur_pos)): - raise RuntimeError( - "flash_sparse_mla inputs must be on the same device: " - f"q_nope={device}, q_pe={q_pe.device}, kv_cache={kv_cache.device}, " - f"pe_cache={pe_cache.device}, topk_indices={topk_indices.device}, " - f"cur_pos={cur_pos.device}" - ) - if self.profile_logs is not None and self.profile_logs.device != device: - self.profile_logs = get_profile_log_tensor(device_index=device.index, device=device) - output = torch.zeros( - (batch_size, seqlen, heads, v_dim), dtype=torch.bfloat16, device=device - ) - flash_sparse_mla( - q_nope, - q_pe, - kv_cache, - pe_cache, - topk_indices, - cur_pos, - output, - self.profile_logs, - model_arch=self.model_args.arch_name, - ) - return output - - def to_tilert_weights(self) -> None: - raise NotImplementedError("to_tilert_weights not implemented") - - def __call__( - self, - q_nope: torch.Tensor, - q_pe: torch.Tensor, - kv_cache: torch.Tensor, - pe_cache: torch.Tensor, - topk_indices: torch.Tensor, - cur_pos: torch.Tensor, - ) -> torch.Tensor: - if self.flag_enable_tilert: - return self.tilert_forward(q_nope, q_pe, kv_cache, pe_cache, topk_indices, cur_pos) - return self.golden_forward(q_nope, q_pe, kv_cache, pe_cache, topk_indices, cur_pos) diff --git a/tilert/models/glm_5/_dsa_v32/ops/head_proj_w16a16_hmma.py b/tilert/models/glm_5/_dsa_v32/ops/head_proj_w16a16_hmma.py new file mode 100644 index 0000000..a8a2cc9 --- /dev/null +++ b/tilert/models/glm_5/_dsa_v32/ops/head_proj_w16a16_hmma.py @@ -0,0 +1,47 @@ +"""HeadProj BF16-MMA operation for DeepSeek-V3.2 / GLM5.""" + +from __future__ import annotations + +import torch + +__all__ = [ + "head_proj_w16a16_hmma", + "swizzle_head_proj_weight_bf16mma", +] + + +def head_proj_w16a16_hmma( + hidden_in: torch.Tensor, + weight_in: torch.Tensor, + logits_out: torch.Tensor, + profile_logs: torch.Tensor, + model_arch: str, + compute_kernel_type: str = "w16a16_hmma", +) -> None: + torch.ops.tilert.head_proj_op( + hidden_in, + weight_in, + logits_out, + model_arch, + compute_kernel_type, + profile_logs, + ) + + +def _swizzle_mma_16x16(mat_in: torch.Tensor) -> torch.Tensor: + assert mat_in.shape[-2] == 16 and mat_in.shape[-1] == 16 + pre = mat_in.shape[:-2] + x = mat_in.reshape(*pre, 2, 8, 2, 4, 2).transpose(-4, -3).transpose(-5, -4) + return x.reshape(*pre, 2 * 2, 8 * 4, 2).transpose(-3, -2) + + +def swizzle_head_proj_weight_bf16mma(weight: torch.Tensor) -> torch.Tensor: + n, k = weight.shape + assert n % 16 == 0 and k % 1024 == 0, "head_proj weight must be /16 in N and /1024 in K" + n_tiles = n // 16 + k_pages = k // 1024 + k_inner = 1024 // 16 + w = weight.reshape(n_tiles, 16, k_pages, k_inner, 16) + w = w.permute(0, 2, 3, 1, 4).contiguous() + w = _swizzle_mma_16x16(w) + return w.contiguous() diff --git a/tilert/models/glm_5/_dsa_v32/ops/layernorm_rope_rotate.py b/tilert/models/glm_5/_dsa_v32/ops/layernorm_rope_rotate.py index 4fc8c0d..00a3bc9 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/layernorm_rope_rotate.py +++ b/tilert/models/glm_5/_dsa_v32/ops/layernorm_rope_rotate.py @@ -31,23 +31,6 @@ def layernorm_rope_rotate( model_arch: str, compute_kernel_type: str = "general", ) -> None: - """ - Layernorm_rope_rotate operation. - - Layernorm_rope_rotate the input tensor `input_raw` and stores the result in `k_cache_raw`. - - Args: - input_raw (torch.Tensor): The input tensor. - cur_pos (torch.Tensor): The current position tensor. - k_cache_raw (torch.Tensor): The output tensor where the result will be stored. - weight (torch.Tensor): The weight tensor. - bias (torch.Tensor): The bias tensor. - freqs_cis (torch.Tensor): The frequency tensor. - profile_logs (torch.Tensor): Tensor for storing profiling logs. - - Returns: - None - """ if input_raw.dtype != torch.bfloat16: raise ValueError("input must be a bfloat16 tensor.") if cur_pos.dtype != torch.int32: @@ -162,15 +145,6 @@ def get_weights_list(self) -> list[torch.Tensor]: return [self.tilert_weight, self.tilert_bias] def device_sharding(self, weights_map: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """ - Device sharding: replicate weight and bias for each device. - - Args: - weights_map: Map from ref weight alias to tensor. - - Returns: - Map from tilert weight alias to (num_devices, ...) tensors. - """ k_weight = weights_map[self.ref_weights_alias.k_weight][None, ...].repeat( self.num_devices, 1 ) @@ -215,7 +189,7 @@ def golden_forward(self, idx_k: torch.Tensor, freqs_cis: torch.Tensor) -> torch. k_pe, k_nope = torch.split( k, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 ) - k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis).squeeze(2) + k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis, interleaved=False).squeeze(2) k = torch.cat([k_pe, k_nope], dim=-1) return rotate_activation(k) diff --git a/tilert/models/glm_5/_dsa_v32/ops/padded_allreduce_add.py b/tilert/models/glm_5/_dsa_v32/ops/padded_allreduce_add.py deleted file mode 100644 index a6490c9..0000000 --- a/tilert/models/glm_5/_dsa_v32/ops/padded_allreduce_add.py +++ /dev/null @@ -1,147 +0,0 @@ -"""PaddedAllReduceAdd operation module.""" - -from enum import Enum - -import torch - -from tilert.models.base import TileRTModule -from tilert.models.glm_5._dsa_v32.model_args import ModelArgs -from tilert.utils import get_profile_log_tensor - -__all__ = [ - "padded_allreduce_add", - "PaddedAllReduceAdd", -] - - -def padded_allreduce_add( - partial_buf: torch.Tensor, - x_in: torch.Tensor, - flag: int, - vec_out: torch.Tensor, - profile_logs: torch.Tensor, - model_arch: str, - compute_kernel_type: str = "bf16", -) -> None: - """Padded AllReduce + residual add for Device Group A (GPU 0). - - GPU 0 contributes zeros to the 8-GPU AllReduce, then adds the residual. - - Args: - partial_buf: Zero-filled partial buffer [1, L, hidden_dim] bf16. - x_in: Residual input [1, L, hidden_dim] bf16. - flag: AllReduce sync flag. - vec_out: Output tensor [1, L, hidden_dim] bf16. - profile_logs: Profile logs tensor. - model_arch: Model architecture ("deepseek_v3_2" or "glm_5"). - compute_kernel_type: Compute kernel type ("bf16"). - """ - torch.ops.tilert.padded_allreduce_add_op( - partial_buf, x_in, flag, vec_out, profile_logs, model_arch, compute_kernel_type - ) - - -class PaddedAllReduceAddAlgorithm(Enum): - """PaddedAllReduceAdd algorithm.""" - - BF16 = "bf16" - - -class PaddedAllReduceAdd(TileRTModule): - """PaddedAllReduceAdd module β€” zero-partial AllReduce + residual add.""" - - _SUPPORTED_ALGORITHMS = { - "deepseek_v3_2": [PaddedAllReduceAddAlgorithm.BF16], - "glm_5": [PaddedAllReduceAddAlgorithm.BF16], - } - - def __init__( - self, - model_args: ModelArgs, - num_devices: int, - device_id: int = 0, - ): - super().__init__( - self.__class__.__name__, - model_args=model_args, - num_devices=num_devices, - device_id=device_id, - ) - - self.dim = self.model_args.dim - - self.partial_buf: torch.Tensor | None = None - - self.hidden_out: torch.Tensor | None = None - - self.profile_logs: torch.Tensor | None = None - self.is_var_init = False - - def init_tilert_vars(self, batch_size: int, seq_len: int) -> None: - """Allocate output buffer and persistent zero-filled partial buffer. - - Args: - batch_size: Batch size. - seq_len: Sequence length. - """ - self.hidden_out = torch.zeros( - (batch_size, seq_len, self.dim), - dtype=torch.bfloat16, - device=f"cuda:{self.device_id}", - ) - self.partial_buf = torch.zeros( - (batch_size, seq_len, self.dim), - dtype=torch.bfloat16, - device=f"cuda:{self.device_id}", - ) - self.profile_logs = get_profile_log_tensor(device=f"cuda:{self.device_id}") - self.is_var_init = True - - def golden_forward( - self, - x_in: torch.Tensor, - ) -> torch.Tensor: - """Golden reference: allreduce(zeros) + x_in = x_in (single-GPU). - - On a single GPU, allreduce of zeros returns zeros, so output = x_in. - - Args: - x_in: Residual input [1, L, hidden_dim]. - - Returns: - Output tensor (copy of x_in). - """ - return x_in.clone() - - def tilert_forward( - self, - x_in: torch.Tensor, - flag: int, - ) -> torch.Tensor: - """Run TileRT kernel forward. - - Args: - x_in: Residual input [1, L, hidden_dim]. - flag: AllReduce sync flag. - - Returns: - Output tensor [1, L, hidden_dim]. - """ - assert self.hidden_out is not None - assert self.partial_buf is not None - assert self.profile_logs is not None - padded_allreduce_add( - self.partial_buf, - x_in, - flag, - self.hidden_out, - self.profile_logs, - model_arch=self.model_args.arch_name, - ) - return self.hidden_out - - def __call__( - self, - x_in: torch.Tensor, - ) -> torch.Tensor: - return self.golden_forward(x_in) diff --git a/tilert/models/glm_5/_dsa_v32/ops/projo_wkvb.py b/tilert/models/glm_5/_dsa_v32/ops/projo_wkvb.py index 3e99f0e..36fdf73 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/projo_wkvb.py +++ b/tilert/models/glm_5/_dsa_v32/ops/projo_wkvb.py @@ -30,18 +30,6 @@ def projo_wkvb( model_arch: str, compute_kernel_type: str = "fp16mma", ) -> None: - """ - Define the ProjOWkvb operation. - - Args: - o_in: Input tensor. - wkv_b_b: Weight tensor. - wkv_b_scales: Scale tensor. - output: Output tensor. - profile_logs: Profile logs tensor. - model_arch: Model architecture ("deepseek_v3_2" or "glm_5"). - compute_kernel_type: Kernel type ("fp16mma" for both DSv32 and GLM5). - """ torch.ops.tilert.projo_wkvb_op( o_in, wkv_b_b, @@ -68,7 +56,6 @@ def __init__(self, model_args: ModelArgs, num_devices: int): @staticmethod def _swizzle_mma_16x16(mat_in: torch.Tensor) -> torch.Tensor: - """Swizzle a [*, 16, 16] block for the packed weight layout.""" assert mat_in.shape[-2] == 16 and mat_in.shape[-1] == 16 pre_shape = mat_in.shape[:-2] mat_in = mat_in.reshape(*pre_shape, 2, 8, 2, 4, 2).transpose(-4, -3).transpose(-5, -4) @@ -76,7 +63,6 @@ def _swizzle_mma_16x16(mat_in: torch.Tensor) -> torch.Tensor: @staticmethod def _swizzle_mma_16x16_for_pages(mat_in: torch.Tensor, k_dim: int, pages: int) -> torch.Tensor: - """Swizzle a [*, 16, K] matrix for the paged weight layout.""" assert mat_in.shape[-2] == 16 and mat_in.shape[-1] == k_dim pre_shape = mat_in.shape[:-2] k_per_page = k_dim // pages @@ -87,17 +73,15 @@ def _swizzle_mma_16x16_for_pages(mat_in: torch.Tensor, k_dim: int, pages: int) - return mat_in.contiguous() def convert_to_fp16mma(self, weights: list[torch.Tensor]) -> torch.Tensor: - """Convert weights to the packed format expected by the kernel.""" with torch.inference_mode(): wkv_b_b, wkv_b_b_scales = self.convert_to_general(weights) n_heads = wkv_b_b.size(0) v_head_dim = wkv_b_b.size(1) kv_lora_rank = wkv_b_b.size(2) - num_ctas = 80 - rows_per_cta = (n_heads * v_head_dim) // num_ctas - is_glm5 = self.model_args.arch_name == "glm_5" + num_ctas = (n_heads * v_head_dim) // 32 if is_glm5 else 80 + rows_per_cta = (n_heads * v_head_dim) // num_ctas w_flat = wkv_b_b.reshape(num_ctas, rows_per_cta // 16, 16, kv_lora_rank) w_swizzled = ProjoWKVbWeightsConverter._swizzle_mma_16x16_for_pages( @@ -131,7 +115,6 @@ def convert_to_fp16mma(self, weights: list[torch.Tensor]) -> torch.Tensor: return torch.cat([w_bytes, scales_raw, padding], dim=-1).contiguous() def convert_to_bf16mma(self, weights: list[torch.Tensor]) -> torch.Tensor: - """Convert weights to the packed format expected by the BF16 kernel.""" with torch.inference_mode(): tilert_wkv_b_weights, tilert_wkv_b_scales = weights @@ -169,7 +152,8 @@ def convert_to_bf16mma(self, weights: list[torch.Tensor]) -> torch.Tensor: wkv_bf16 = (w.float() * s).to(torch.bfloat16) n_heads = n_local_heads - num_ctas = 80 + is_glm5 = self.model_args.arch_name == "glm_5" + num_ctas = (n_heads * v_head_dim) // 32 if is_glm5 else 80 rows_per_cta = (n_heads * v_head_dim) // num_ctas w_flat = wkv_bf16.reshape(num_ctas, rows_per_cta // 16, 16, kv_lora_rank) @@ -328,15 +312,6 @@ def get_weights_list(self) -> list[torch.Tensor]: return [self.tilert_wkv_b_b, self.tilert_wkv_b_b_scales] def device_sharding(self, weights_map: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """ - Device sharding: split weights and scales per device. - - Args: - weights_map: Map from ref weight alias to tensor. - - Returns: - Map from tilert weight alias to (num_devices, ...) tensors. - """ kv_b_proj_weight = weights_map[self.ref_weights_alias.wkv_b_weights] kv_b_proj_weight_scale = weights_map[self.ref_weights_alias.wkv_b_scales] @@ -403,7 +378,6 @@ def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: self.init_tilert_weights_hmma(state_dict) def init_tilert_weights_hmma(self, state_dict: dict[str, torch.Tensor]) -> None: - """Initialize with HMMA-packed weights.""" packed = ProjoWKVbWeightsConverter(self.model_args, self.num_devices).dispatch( ProjoWKVbAlgorithm.FP16MMA, [ @@ -416,7 +390,6 @@ def init_tilert_weights_hmma(self, state_dict: dict[str, torch.Tensor]) -> None: self.compute_kernel_type = "fp16mma" def init_tilert_weights_hmma_bf16(self, state_dict: dict[str, torch.Tensor]) -> None: - """Initialize with BF16 HMMA-packed weights (dequantized, no scales).""" packed = ProjoWKVbWeightsConverter(self.model_args, self.num_devices).dispatch( ProjoWKVbAlgorithm.BF16MMA, [ diff --git a/tilert/models/glm_5/_dsa_v32/ops/projq_wqb.py b/tilert/models/glm_5/_dsa_v32/ops/projq_wqb.py index c40ca51..7ad3d1d 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/projq_wqb.py +++ b/tilert/models/glm_5/_dsa_v32/ops/projq_wqb.py @@ -31,18 +31,7 @@ def projq_wqb( *, model_arch: str, ) -> None: - """ - Define the ProjqWqb operation. - - Args: - q_nope_in: Input tensor. - wkv_b_a: Weight tensor. - wkv_b_a_scales: Scale tensor. - output: Output tensor. - profile_logs: Profile logs tensor. - compute_kernel_type: Kernel type ("fp16mma"). - model_arch: Model architecture ("deepseek_v3_2" or "glm_5"). - """ + """Define the ProjqWqb operation.""" torch.ops.tilert.projq_wqb_op( q_nope_in, wkv_b_a, @@ -70,7 +59,7 @@ def __init__(self, model_args: ModelArgs, num_devices: int, head_dim_block_size: @staticmethod def _swizzle_mma_16x16(mat_in: torch.Tensor) -> torch.Tensor: - """Swizzle a [*, 16, 16] block for the packed weight layout.""" + """Swizzle a [*, 16, 16] sub-block for the MMA kernel.""" assert mat_in.shape[-2] == 16 and mat_in.shape[-1] == 16 pre_shape = mat_in.shape[:-2] mat_in = mat_in.reshape(*pre_shape, 2, 8, 2, 4, 2).transpose(-4, -3).transpose(-5, -4) @@ -78,7 +67,7 @@ def _swizzle_mma_16x16(mat_in: torch.Tensor) -> torch.Tensor: @staticmethod def _swizzle_mma_16x16_for_pages(mat_in: torch.Tensor, k_dim: int, pages: int) -> torch.Tensor: - """Swizzle a [*, 16, K] matrix for the paged weight layout.""" + """Swizzle [*, 16, K] matrix for paged MMA layout.""" assert mat_in.shape[-2] == 16 and mat_in.shape[-1] == k_dim pre_shape = mat_in.shape[:-2] k_per_page = k_dim // pages @@ -89,17 +78,16 @@ def _swizzle_mma_16x16_for_pages(mat_in: torch.Tensor, k_dim: int, pages: int) - return mat_in.contiguous() def convert_to_fp16mma(self, weights: list[torch.Tensor]) -> torch.Tensor: - """Convert weights to the packed format expected by the kernel.""" + """Convert weights to the FP16 MMA packed format.""" with torch.inference_mode(): wkv_b_a, wkv_b_a_scales = self.convert_to_general(weights) n_heads = wkv_b_a.size(0) head_dim = wkv_b_a.size(2) kv_lora_rank = wkv_b_a.size(1) - num_ctas = 80 - rows_per_cta = (n_heads * kv_lora_rank) // num_ctas - is_glm5 = self.model_args.arch_name == "glm_5" + num_ctas = (n_heads * kv_lora_rank) // 64 if is_glm5 else 80 + rows_per_cta = (n_heads * kv_lora_rank) // num_ctas w_flat = wkv_b_a.reshape(num_ctas, rows_per_cta // 16, 16, head_dim) w_swizzled = self._swizzle_mma_16x16_for_pages(w_flat, head_dim, pages=1) @@ -129,7 +117,7 @@ def convert_to_fp16mma(self, weights: list[torch.Tensor]) -> torch.Tensor: return torch.cat([w_bytes, scales_raw, padding], dim=-1).contiguous() def convert_to_bf16mma(self, weights: list[torch.Tensor]) -> torch.Tensor: - """Convert weights to the packed format expected by the BF16 kernel.""" + """Convert weights to the BF16 MMA packed format.""" with torch.inference_mode(): tilert_wkv_b_weights, tilert_wkv_b_scales = weights @@ -153,7 +141,8 @@ def convert_to_bf16mma(self, weights: list[torch.Tensor]) -> torch.Tensor: n_heads = n_local_heads head_dim = nope_head_dim - num_ctas = 80 + is_glm5 = self.model_args.arch_name == "glm_5" + num_ctas = (n_heads * kv_lora_rank) // 64 if is_glm5 else 80 rows_per_cta = (n_heads * kv_lora_rank) // num_ctas w_flat = wkv_bf16.reshape(num_ctas, rows_per_cta // 16, 16, head_dim) @@ -313,15 +302,6 @@ def get_weights_list(self) -> list[torch.Tensor]: return [self.tilert_wkv_b_a, self.tilert_wkv_b_a_scales] def device_sharding(self, weights_map: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """ - Device sharding: split weights and scales per device. - - Args: - weights_map: Map from ref weight alias to tensor. - - Returns: - Map from tilert weight alias to (num_devices, ...) tensors. - """ kv_b_proj_weight = weights_map[self.ref_weights_alias.wkv_b_weights] kv_b_proj_weight_scale = weights_map[self.ref_weights_alias.wkv_b_scales] @@ -388,33 +368,29 @@ def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: self.init_tilert_weights_hmma(state_dict) def init_tilert_weights_hmma(self, state_dict: dict[str, torch.Tensor]) -> None: - """Initialize with HMMA-packed weights.""" - packed = ProjqWqbWeightsConverter( - self.model_args, self.num_devices, self.head_dim_block_size - ).dispatch( - ProjqWqbAlgorithm.FP16MMA, + packed, dummy_scales = torch.ops.tilert.projq_wkvb__convert_weights( [ state_dict[self.tilert_weights_alias.wkv_b_weights], state_dict[self.tilert_weights_alias.wkv_b_scales], ], + self.model_args.arch_name, + "fp16mma", ) self.tilert_wkv_b_a = packed - self.tilert_wkv_b_a_scales = torch.empty(1, dtype=torch.float8_e4m3fn, device=packed.device) + self.tilert_wkv_b_a_scales = dummy_scales self.compute_kernel_type = "fp16mma" def init_tilert_weights_hmma_bf16(self, state_dict: dict[str, torch.Tensor]) -> None: - """Initialize with BF16 HMMA-packed weights (dequantized, no scales).""" - packed = ProjqWqbWeightsConverter( - self.model_args, self.num_devices, self.head_dim_block_size - ).dispatch( - ProjqWqbAlgorithm.BF16MMA, + packed, dummy_scales = torch.ops.tilert.projq_wkvb__convert_weights( [ state_dict[self.tilert_weights_alias.wkv_b_weights], state_dict[self.tilert_weights_alias.wkv_b_scales], ], + self.model_args.arch_name, + "bf16mma", ) self.tilert_wkv_b_a = packed - self.tilert_wkv_b_a_scales = torch.empty(1, dtype=torch.float8_e4m3fn, device=packed.device) + self.tilert_wkv_b_a_scales = dummy_scales self.compute_kernel_type = "bf16mma" def init_random_weights(self) -> None: diff --git a/tilert/models/glm_5/_dsa_v32/ops/projx_wis.py b/tilert/models/glm_5/_dsa_v32/ops/projx_wis.py index e13b4e0..1784629 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/projx_wis.py +++ b/tilert/models/glm_5/_dsa_v32/ops/projx_wis.py @@ -26,17 +26,6 @@ def projx_wis( profile_logs: torch.Tensor, model_arch: str, ) -> None: - """ - Define the ProjxWis operation. - - Args: - x_in: Input tensor. - w: Weight tensor. - output: Output tensor. - compute_kernel_type: Compute kernel type ("bf16" or "bf16mma"). - profile_logs: Profile logs tensor. - model_arch: Model architecture ("deepseek_v3_2" or "glm_5"). - """ torch.ops.tilert.proj_w_op(x_in, w, output, model_arch, compute_kernel_type, profile_logs) @@ -125,7 +114,7 @@ def __init__( @staticmethod def _swizzle_mma_16x16(mat_in: torch.Tensor) -> torch.Tensor: - """Swizzle each 16x16 BF16 tile for the packed weight layout.""" + """Swizzle each 16x16 BF16 tile for MMA loading.""" assert mat_in.shape[-2] == 16 and mat_in.shape[-1] == 16 pre_shape = mat_in.shape[:-2] mat_in = mat_in.reshape(*pre_shape, 2, 8, 2, 4, 2).transpose(-4, -3).transpose(-5, -4) @@ -135,7 +124,6 @@ def _swizzle_mma_16x16(mat_in: torch.Tensor) -> torch.Tensor: def _to_hmma_layout( w_orig: torch.Tensor, n_ctas: int, rows_per_cta: int, x_dim: int, num_pages: int ) -> torch.Tensor: - """Convert [output_dim, x_dim] BF16 weights to the packed kernel layout.""" cols_per_page = x_dim // num_pages n_k_tiles = cols_per_page // 16 w = w_orig.reshape(n_ctas, rows_per_cta, num_pages, cols_per_page) @@ -153,15 +141,6 @@ def get_weights_list(self) -> list[torch.Tensor]: return [self.tilert_w] def device_sharding(self, weights_map: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """ - Device sharding: replicate weight for each device. - - Args: - weights_map: Map from ref weight alias to tensor. - - Returns: - Map from tilert weight alias to (num_devices, ...) tensors. - """ w = weights_map[self.ref_weights_alias.w_weights] if self.compute_kernel_type == "bf16mma": n_ctas, rows_per_cta, num_pages = self._HMMA_CONFIGS[self.dim] @@ -177,7 +156,17 @@ def init_reference_weights(self, state_dict: dict[str, torch.Tensor]) -> None: self.is_ref_weights_init = True def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - self.tilert_w = state_dict[self.tilert_weights_alias.w_weights].detach().clone() + w = state_dict[self.tilert_weights_alias.w_weights].detach().clone() + if ( + self.compute_kernel_type == "bf16mma" + and w.dim() == 2 + and (w.shape[0] == self.index_n_heads and w.shape[1] == self.dim) + ): + n_ctas, rows_per_cta, num_pages = self._HMMA_CONFIGS[self.dim] + w = self._to_hmma_layout( + w.to(torch.bfloat16), n_ctas, rows_per_cta, self.dim, num_pages + ) + self.tilert_w = w self.is_tilert_weights_init = True def init_random_weights(self) -> None: diff --git a/tilert/models/glm_5/_dsa_v32/ops/projx_wqaki.py b/tilert/models/glm_5/_dsa_v32/ops/projx_wqaki.py index 367d5fe..6f8a1e3 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/projx_wqaki.py +++ b/tilert/models/glm_5/_dsa_v32/ops/projx_wqaki.py @@ -19,18 +19,6 @@ def projx_wqaki( *, model_arch: str, ) -> None: - """FP8 projection for q, ki. - - Args: - x_quant: FP8 quantized hidden states [1, seq_len, hidden_dim]. - x_scale: Scale factors for x_quant. - wqaki: Packed FP8 weights + scales for q, ki. - out_q: Output q tensor. - out_ki: Output ki tensor. - profile_logs: Profile logs tensor. - compute_kernel_type: Kernel type ("fp8mma", "fp8mma_68cta", "fp8mma_136cta"). - model_arch: Model architecture ("deepseek_v3_2" or "glm_5"). - """ torch.ops.tilert.projx_wqaki_op( x_quant, x_scale, @@ -62,7 +50,6 @@ def convert_dsv32( wki: torch.Tensor, wki_scale: torch.Tensor, ) -> torch.Tensor: - """Convert DSV3.2 weights to the packed format expected by the kernel.""" with torch.inference_mode(): wq_a_scale = wq_a_scale.to(torch.bfloat16) wki_scale = wki_scale.to(torch.bfloat16) @@ -153,7 +140,6 @@ def convert_glm5_68cta( wki: torch.Tensor, wki_scale: torch.Tensor, ) -> torch.Tensor: - """Convert GLM5 weights to the packed format expected by the kernel.""" with torch.inference_mode(): wq_a_scale = wq_a_scale.to(torch.float32) wki_scale = wki_scale.to(torch.float32) @@ -195,6 +181,23 @@ def convert_glm5_68cta( ) return torch.cat([wqaki_raw, wqaki_scales, wqaki_padding], dim=-1).contiguous() + @staticmethod + def convert_glm5_68cta_w8a16( + wq_a: torch.Tensor, + wq_a_scale: torch.Tensor, + wki: torch.Tensor, + wki_scale: torch.Tensor, + ) -> torch.Tensor: + from tilert.models.glm_5._dsa_v32.ops.rmsnorm_projx_wqkva import ( + RMSNormProjQKVAW8A16MMAWeightsConverter, + ) + + with torch.inference_mode(): + dim = 6144 + w_fp8 = torch.cat([wq_a.reshape(2048, dim), wki.reshape(128, dim)], dim=0).contiguous() + scales = torch.cat([wq_a_scale, wki_scale], dim=0).to(torch.float32).contiguous() + return RMSNormProjQKVAW8A16MMAWeightsConverter.pack_lane_major(w_fp8, scales, dim) + @staticmethod def convert_glm5_136cta( wq_a: torch.Tensor, @@ -202,7 +205,6 @@ def convert_glm5_136cta( wki: torch.Tensor, wki_scale: torch.Tensor, ) -> torch.Tensor: - """Convert GLM5 weights to the packed format expected by the kernel.""" with torch.inference_mode(): wq_a_scale = wq_a_scale.to(torch.float32) wki_scale = wki_scale.to(torch.float32) diff --git a/tilert/models/glm_5/_dsa_v32/ops/projx_wqkva.py b/tilert/models/glm_5/_dsa_v32/ops/projx_wqkva.py index 6ade7af..b94a36f 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/projx_wqkva.py +++ b/tilert/models/glm_5/_dsa_v32/ops/projx_wqkva.py @@ -10,6 +10,7 @@ from tilert.models.glm_5._dsa_v32.ops.rmsnorm_projx_wqkva import ( RMSNormProjQKVAFP8MMAWeightsConverter, RMSNormProjQKVAFP16MMAWeightsConverter, + RMSNormProjQKVAW8A16MMAWeightsConverter, ) from tilert.utils import get_profile_log_tensor @@ -32,7 +33,7 @@ def projx_wqkva( *, model_arch: str, ) -> None: - """FP8 MMA projection for q, kv, pe_cache (DSV3.2).""" + """Standalone FP8 QMMA projection for q, kv, pe_cache.""" torch.ops.tilert.projx_wqkva_op( x_quant, x_scale, @@ -101,14 +102,19 @@ class ProjXWqkvaAlgorithm(Enum): FP8MMA = "fp8mma" FP16MMA = "fp16mma" + W8A16HMMA = "w8a16_hmma" class ProjXWqkva(TileRTModule): - """FP8 MMA projection module for q, kv, pe_cache.""" + """Standalone FP8 QMMA GEMV for q, kv, pe_cache projections.""" _SUPPORTED_ALGORITHMS = { - "deepseek_v3_2": [ProjXWqkvaAlgorithm.FP8MMA], - "glm_5": [ProjXWqkvaAlgorithm.FP8MMA, ProjXWqkvaAlgorithm.FP16MMA], + "deepseek_v3_2": [ProjXWqkvaAlgorithm.FP8MMA, ProjXWqkvaAlgorithm.W8A16HMMA], + "glm_5": [ + ProjXWqkvaAlgorithm.FP8MMA, + ProjXWqkvaAlgorithm.FP16MMA, + ProjXWqkvaAlgorithm.W8A16HMMA, + ], } def __init__( @@ -156,11 +162,12 @@ def set_algorithm(self, algorithm: Enum) -> None: super().set_algorithm(algorithm) if algorithm == ProjXWqkvaAlgorithm.FP16MMA: self.compute_kernel_type = "fp16mma" + elif algorithm == ProjXWqkvaAlgorithm.W8A16HMMA: + self.compute_kernel_type = "w8a16_hmma" else: self.compute_kernel_type = "fp8mma" def device_sharding(self, weights_map: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Repeat weights for device sharding.""" q_a_proj_weight = weights_map[self.ref_weights_alias.q_a_weights][None, ...].repeat( self.num_devices, 1, 1 ) @@ -222,6 +229,20 @@ def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: hidden_dim=self.dim, q_lora_rank=self.q_lora_rank, ) + elif self.algorithm == ProjXWqkvaAlgorithm.W8A16HMMA: + self.tilert_wqkva, _ = ( + RMSNormProjQKVAW8A16MMAWeightsConverter.convert_to_w8a16_mma_gemv( + wq_a, + wq_a_scale, + wkv_a, + wkv_a_scale, + w_pe, + w_pe_scale, + dummy_gamma, + hidden_dim=self.dim, + q_lora_rank=self.q_lora_rank, + ) + ) else: self.tilert_wqkva, _ = RMSNormProjQKVAFP8MMAWeightsConverter.convert_to_fp8_mma_gemv( wq_a, @@ -256,9 +277,9 @@ def init_random_weights(self) -> None: tensor_list = [ torch.randn(self.dim, dtype=torch.float32), torch.randn(self.q_lora_rank, self.dim, dtype=torch.bfloat16).to(torch.float8_e4m3fn), - torch.randn(q_scale_dim, dim_scale_dim, dtype=scale_dtype), + torch.randn(q_scale_dim, dim_scale_dim, dtype=scale_dtype).abs(), torch.randn(kv_mqa_rows, self.dim, dtype=torch.bfloat16).to(torch.float8_e4m3fn), - torch.randn(kv_mqa_scale_dim, dim_scale_dim, dtype=scale_dtype), + torch.randn(kv_mqa_scale_dim, dim_scale_dim, dtype=scale_dtype).abs(), ] ref_state_dict = dict(zip(self.ref_weights_alias(), tensor_list)) self.init_reference_weights(ref_state_dict) @@ -277,7 +298,10 @@ def golden_forward( assert self.ref_wkv_a is not None assert self.ref_w_pe is not None - if self.algorithm == ProjXWqkvaAlgorithm.FP16MMA: + if self.algorithm in ( + ProjXWqkvaAlgorithm.FP16MMA, + ProjXWqkvaAlgorithm.W8A16HMMA, + ): x_float = x_quant.float() else: x_fp8 = x_quant.to(torch.float32) diff --git a/tilert/models/glm_5/_dsa_v32/ops/qkv_rope.py b/tilert/models/glm_5/_dsa_v32/ops/qkv_rope.py deleted file mode 100644 index 7f16a1c..0000000 --- a/tilert/models/glm_5/_dsa_v32/ops/qkv_rope.py +++ /dev/null @@ -1,192 +0,0 @@ -"""QKV Rope operation module.""" - -from dataclasses import dataclass -from enum import Enum - -import torch - -from tilert.models.base import TileRTModule -from tilert.models.glm_5._dsa_v32.model_args import ModelArgs -from tilert.models.utils import apply_rotary_emb -from tilert.utils import get_profile_log_tensor - -__all__ = [ - "qkv_rope", - "QKVRoPE", - "QKVRoPERefWeightsAlias", - "QKVRoPETilertWeightsAlias", -] - - -def qkv_rope( - pe_cache: torch.Tensor, - kv_cache: torch.Tensor, - rope_freqs: torch.Tensor, - cur_pos: torch.Tensor, - profile_logs: torch.Tensor, - model_arch: str, - compute_kernel_type: str = "general", -) -> None: - """ - Perform QKV Rope operation. - - Args: - pe_cache: Q PE tensor (bsz, seq, n_local_heads, qk_rope_head_dim). - kv_cache: K PE cache (bsz, seq, qk_rope_head_dim). - rope_freqs: Rope frequencies tensor. - cur_pos: Current position tensor. - profile_logs: Profile logs tensor. - model_arch: Model architecture string. - compute_kernel_type: Compute kernel type string. - """ - torch.ops.tilert.qkv_rope_op( - pe_cache, - kv_cache, - rope_freqs, - cur_pos, - model_arch, - compute_kernel_type, - profile_logs, - ) - - -@dataclass -class QKVRoPERefWeightsAlias: - """Reference weights alias for QKVRoPE (no weights).""" - - @property - def ref_tensor_alias(self) -> list[str]: - return [] - - def __call__(self) -> list[str]: - return self.ref_tensor_alias - - -@dataclass -class QKVRoPETilertWeightsAlias: - """TileRT weights alias for QKVRoPE (no weights).""" - - @property - def tilert_tensor_alias(self) -> list[str]: - return [] - - def __call__(self) -> list[str]: - return self.tilert_tensor_alias - - -class QKVRoPEAlgorithm(Enum): - """QKVRoPE algorithm.""" - - GENERAL = "general" - - -class QKVRoPE(TileRTModule): - """QKV RoPE module. Unified for deepseek_v3_2 and glm_5.""" - - _SUPPORTED_ALGORITHMS = { - "deepseek_v3_2": [QKVRoPEAlgorithm.GENERAL], - "glm_5": [QKVRoPEAlgorithm.GENERAL], - } - - def __init__( - self, - model_args: ModelArgs, - num_devices: int = 1, - device_id: int = 0, - layer_idx: int = 0, - ref_weights_alias: QKVRoPERefWeightsAlias | None = None, - ) -> None: - super().__init__( - self.__class__.__name__, - model_args=model_args, - num_devices=num_devices, - device_id=device_id, - layer_idx=layer_idx, - ) - self.tilert_weights_alias = QKVRoPETilertWeightsAlias() - self.ref_weights_alias = ( - ref_weights_alias if ref_weights_alias is not None else QKVRoPERefWeightsAlias() - ) - self.n_local_heads = model_args.n_heads // num_devices - self.qk_rope_head_dim = model_args.qk_rope_head_dim - self.profile_logs: torch.Tensor | None = None - - def get_weights_list(self) -> list[torch.Tensor]: - return [] - - def device_sharding(self, weights_map: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - del weights_map - return {} - - def init_reference_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - del state_dict - pass - - def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - del state_dict - pass - - def init_random_weights(self) -> None: - pass - - def init_tilert_vars(self, batch_size: int, seq_len: int) -> None: - del batch_size, seq_len - self.profile_logs = get_profile_log_tensor() - self.is_var_init = True - - def golden_forward( - self, - q_pe: torch.Tensor, - pe_cache: torch.Tensor, - start_pos: int, - freqs_cis: torch.Tensor, - bsz: int, - seqlen: int, - ) -> torch.Tensor: - end_pos = start_pos + seqlen - - k_pe = pe_cache[:bsz, start_pos:end_pos] - k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis) - pe_cache[:bsz, start_pos:end_pos] = k_pe.squeeze(2) - - return apply_rotary_emb(q_pe, freqs_cis) - - def tilert_forward( - self, - q_pe: torch.Tensor, - pe_cache: torch.Tensor, - start_pos: int, - freqs_cis: torch.Tensor, - bsz: int, - seqlen: int, - ) -> torch.Tensor: - assert self.profile_logs is not None - end_pos = start_pos + seqlen - - q_pe_rope = q_pe.clone() - rope_freqs = torch.view_as_real(freqs_cis).reshape(*freqs_cis.shape[:-1], -1) - cur_pos = torch.tensor([start_pos], dtype=torch.int32) - - qkv_rope( - q_pe_rope, - pe_cache[:bsz, start_pos:end_pos], - rope_freqs, - cur_pos, - self.profile_logs, - model_arch=self.model_args.arch_name, - ) - - return q_pe_rope - - def __call__( - self, - q_pe: torch.Tensor, - pe_cache: torch.Tensor, - start_pos: int, - freqs_cis: torch.Tensor, - bsz: int, - seqlen: int, - ) -> torch.Tensor: - if self.flag_enable_tilert: - return self.tilert_forward(q_pe, pe_cache, start_pos, freqs_cis, bsz, seqlen) - return self.golden_forward(q_pe, pe_cache, start_pos, freqs_cis, bsz, seqlen) diff --git a/tilert/models/glm_5/_dsa_v32/ops/receive_selected_token_ids.py b/tilert/models/glm_5/_dsa_v32/ops/receive_selected_token_ids.py index 508d13e..80bb1ee 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/receive_selected_token_ids.py +++ b/tilert/models/glm_5/_dsa_v32/ops/receive_selected_token_ids.py @@ -8,25 +8,15 @@ def receive_selected_token_ids( - ll_buf: torch.Tensor, + recv_buf: torch.Tensor, dst: torch.Tensor, expected_flag: int, profile_logs: torch.Tensor, model_arch: str, compute_kernel_type: str = "bf16", ) -> None: - """Receive idx_selects from GPU 0. - - Args: - ll_buf: Receive buffer on this GPU (written by GPU 0). - dst: Destination idx_selects tensor [1, S, 2048] int32. - expected_flag: Expected synchronization flag value. - profile_logs: Profile logs tensor. - model_arch: Model architecture ("deepseek_v3_2" or "glm_5"). - compute_kernel_type: Compute kernel type ("bf16"). - """ torch.ops.tilert.receive_selected_token_ids_op( - ll_buf, + recv_buf, dst, expected_flag, model_arch, diff --git a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_head_proj.py b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_head_proj.py index fa2086d..813e504 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_head_proj.py +++ b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_head_proj.py @@ -26,7 +26,6 @@ def rmsnorm_head_proj( model_arch: str, compute_kernel_type: str = "general", ) -> None: - """RMS Norm Head Projection operation.""" torch.ops.tilert.rmsnorm_head_proj_op( hidden_in, gamma_in, @@ -60,20 +59,18 @@ def tilert_to_tilert_native_bf16_warp_gemv( def convert_to_general( self, weights_list: list[torch.Tensor] ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Convert the weights to general format. - - Args: - weights_list: List of weights. - - Returns: - Tuple of weights. - """ args = self.model_args assert args.arch_name == "deepseek_v3_2" or args.arch_name == "glm_5" with torch.inference_mode(): rmsnorm_gamma, mat_in = weights_list + if args.arch_name == "glm_5": + from tilert.models.glm_5._dsa_v32.ops.head_proj_w16a16_hmma import ( + swizzle_head_proj_weight_bf16mma, + ) + + weights = swizzle_head_proj_weight_bf16mma(mat_in.contiguous()) + return rmsnorm_gamma.float(), weights logits_dim = mat_in.shape[-2] dim = mat_in.shape[-1] num_steps = dim // 1024 @@ -150,27 +147,13 @@ def tilert_tensor_alias(self) -> list[str]: return self.tilert_weights_alias() def get_weights_list(self) -> list[torch.Tensor]: - """ - Get the weights list. - - Returns: - List of weights. - """ + """Get the weights list.""" return [self.tilert_rmsnorm_gamma, self.tilert_head_proj] def device_sharding( self, weights_dict: dict[str, torch.Tensor], ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Device sharding. - - Args: - weights_dict: Dictionary of weights. - key_prefix: Key prefix. - Returns: - Tuple of weights. - """ rmsnorm_gamma_key = "model.norm.weight" head_proj_key = "lm_head.weight" rmsnorm_gamma = weights_dict[rmsnorm_gamma_key][None, ...] @@ -181,13 +164,6 @@ def device_sharding( return rmsnorm_gamma.contiguous(), head_proj.contiguous() def init_reference_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - """ - Initialize the reference weights. - - Args: - state_dict: State dictionary. - device_id: Device ID. - """ sharded_list = self.device_sharding(state_dict) gamma, head_proj = sharded_list[0][self.device_id], sharded_list[1][self.device_id] @@ -195,25 +171,12 @@ def init_reference_weights(self, state_dict: dict[str, torch.Tensor]) -> None: self.ref_head_proj = head_proj def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - """ - Initialize the tilert weights. - - Args: - state_dict: State dictionary. - """ assert self.algorithm is not None self.tilert_rmsnorm_gamma, self.tilert_head_proj = RMSNormHeadProjWeightsConverter( self.model_args, self.num_devices ).dispatch(self.algorithm, [state_dict[alias] for alias in self.tilert_weights_alias()]) def init_tilert_vars(self, batch_size: int, seq_len: int) -> None: - """ - Initialize the tilert variables. - - Args: - batch_size: Batch size. - seq_len: Sequence length. - """ self.hidden_rmsnorm_out = torch.zeros( (batch_size, seq_len, self.dim), dtype=torch.bfloat16, @@ -227,8 +190,9 @@ def init_tilert_vars(self, batch_size: int, seq_len: int) -> None: self.profile_logs = get_profile_log_tensor(device=f"cuda:{self.device_id}") self.is_init = True - def init_random_weights(self, device_id: int = 0) -> None: - """Initialize the random weights.""" + def init_random_weights(self, device_id: int | None = None) -> None: + if device_id is None: + device_id = self.device_id rmsnorm_gamma = torch.randn(self.dim, dtype=torch.float32, device=f"cuda:{device_id}") head_proj = torch.randn( self.logits_dim, self.dim, dtype=torch.bfloat16, device=f"cuda:{device_id}" @@ -252,15 +216,6 @@ def golden_forward( self, hidden_in: torch.Tensor, ) -> torch.Tensor: - """ - Forward pass for the down-project module. - - Args: - hidden_in: Input hidden. - - Returns: - Output tensor. - """ assert self.ref_rmsnorm_gamma is not None assert self.ref_head_proj is not None bsz = hidden_in.shape[0] diff --git a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_kv.py b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_kv.py index 81d161c..99cde46 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_kv.py +++ b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_kv.py @@ -26,18 +26,6 @@ def rmsnorm_kv( model_arch: str, compute_kernel_type: str = "general", ) -> None: - """ - Define the RMSNormKV operation. - - Args: - kv: Input tensor. - gamma: Weight tensor. - cur_pos: Current position tensor. - kv_cache: Output tensor. - profile_logs: Profile logs tensor. - model_arch: Model architecture string. - compute_kernel_type: Compute kernel type string. - """ torch.ops.tilert.rmsnorm_kv_op( kv, gamma, cur_pos, kv_cache, model_arch, compute_kernel_type, profile_logs ) @@ -75,14 +63,15 @@ class KVRMSNormAlgorithm(Enum): """KVRMSNorm algorithm.""" GENERAL = "general" + FP8 = "fp8" class KVRMSNorm(TileRTModule): """KVRMSNorm module: RMSNorm on KV tensor with in-place write to kv_cache.""" _SUPPORTED_ALGORITHMS = { - "deepseek_v3_2": [KVRMSNormAlgorithm.GENERAL], - "glm_5": [KVRMSNormAlgorithm.GENERAL], + "deepseek_v3_2": [KVRMSNormAlgorithm.GENERAL, KVRMSNormAlgorithm.FP8], + "glm_5": [KVRMSNormAlgorithm.GENERAL, KVRMSNormAlgorithm.FP8], } def __init__( @@ -126,40 +115,27 @@ def get_weights_list(self) -> list[torch.Tensor]: return [self.tilert_kv_norm_weight] def device_sharding(self, weights_map: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """ - Device sharding: replicate gamma for each device. - - Args: - weights_map: Map from ref weight alias to tensor. - - Returns: - Map from tilert weight alias to (num_devices, ...) tensors. - """ gamma = weights_map[self.ref_weights_alias.kv_norm_weight][None, ...].repeat( self.num_devices, 1 ) return {self.tilert_weights_alias.kv_norm_gamma: gamma} def init_reference_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - """Initialize reference weights from state dict.""" self.ref_norm_gamma = state_dict[self.ref_weights_alias.kv_norm_weight].contiguous() assert ( self.ref_norm_gamma.shape[-1] == self.kv_lora_rank ), f"kv_norm weight shape must be ({self.kv_lora_rank},), got {self.ref_norm_gamma.shape}" def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - """Initialize TileRT weights from state dict.""" gamma = state_dict[self.tilert_weights_alias.kv_norm_gamma] self.tilert_kv_norm_weight = gamma.float().detach().clone().contiguous() def init_tilert_vars(self, batch_size: int, seq_len: int) -> None: - """Allocate TileRT profiling buffer.""" del batch_size, seq_len self.profile_logs = get_profile_log_tensor() self.is_var_init = True def init_random_weights(self) -> None: - """Initialize random reference and TileRT weights for testing.""" ref_state_dict = { self.ref_weights_alias.kv_norm_weight: torch.randn( self.kv_lora_rank, dtype=torch.float32 @@ -172,7 +148,6 @@ def init_random_weights(self) -> None: def golden_forward( self, kv: torch.Tensor, kv_cache: torch.Tensor, start_pos: int, bsz: int, seqlen: int ) -> None: - """Reference forward: RMSNorm and write to kv_cache.""" assert self.ref_norm_gamma is not None end_pos = start_pos + seqlen out = torch.nn.functional.rms_norm( @@ -180,6 +155,10 @@ def golden_forward( ).to(kv.dtype) kv_cache[:bsz, start_pos:end_pos].copy_(out) + @property + def is_fp8(self) -> bool: + return self.algorithm == KVRMSNormAlgorithm.FP8 + def tilert_forward( self, kv: torch.Tensor, kv_cache: torch.Tensor, start_pos: int, bsz: int, seqlen: int ) -> None: @@ -194,6 +173,7 @@ def tilert_forward( kv_cache[:bsz], self.profile_logs, model_arch=self.model_args.arch_name, + compute_kernel_type="fp8" if self.is_fp8 else "general", ) def __call__( diff --git a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projq_wqb.py b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projq_wqb.py index 92d7a99..aab6c43 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projq_wqb.py +++ b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projq_wqb.py @@ -1,4 +1,4 @@ -"""RmsnormProjqWqb operation module.""" +"""RmsnormProjqWqb operation module (Device Group B, TP7).""" import math from dataclasses import dataclass @@ -47,15 +47,11 @@ class RmsnormProjqWqbAlgorithm(Enum): """RmsnormProjqWqb algorithm.""" FP16MMA = "fp16mma" + BF16MMA = "bf16mma" class RmsnormProjqWqbWeightsConverter(TilertWeightsConverter): - """Weights converter for RmsnormProjqWqb. - - Supports configurations where n_heads is not evenly divisible by - num_devices; in that case n_local_heads is padded and padded head - weight rows are zero-filled. - """ + """Weights converter for RmsnormProjqWqb.""" kBf16NumCtas = 80 kGemvPageSize = 8 @@ -83,9 +79,12 @@ def __init__(self, model_args: ModelArgs, num_devices: int): self.qk_dim = self.qk_head_dim * self.n_local_heads self.qk_qdim = self.qk_dim // self.block_size - assert self.qk_dim % (self.kBf16NumCtas * self.kGemvPageSize) == 0, ( - f"qk_dim ({self.qk_dim}) must be divisible by " - f"kBf16NumCtas * kGemvPageSize ({self.kBf16NumCtas * self.kGemvPageSize})" + kRowsPerCta = 32 + qk_nope_dim = self.qk_nope_head_dim * self.n_local_heads + qk_pe_dim = self.qk_rope_head_dim * self.n_local_heads + assert qk_nope_dim % kRowsPerCta == 0 and qk_pe_dim % kRowsPerCta == 0, ( + f"qk_nope_dim ({qk_nope_dim}) and qk_pe_dim ({qk_pe_dim}) must each " + f"be divisible by rows_per_cta ({kRowsPerCta})" ) assert self.qk_dim % self.block_size == 0, ( f"qk_dim ({self.qk_dim}) must be divisible by block_size ({self.block_size}) " @@ -94,7 +93,6 @@ def __init__(self, model_args: ModelArgs, num_devices: int): @classmethod def _compute_n_local_heads(cls, n_total_heads: int, num_devices: int, qk_head_dim: int) -> int: - """Compute padded n_local_heads per device.""" if n_total_heads % num_devices == 0: return n_total_heads // num_devices @@ -114,22 +112,6 @@ def _redistribute_heads( qk_head_dim: int, block_size: int, ) -> tuple[list[torch.Tensor], list[torch.Tensor]]: - """Redistribute heads across devices with padding. - - Args: - wq_b_full: [n_total_heads * qk_head_dim, q_lora_dim] full weight. - wq_b_scale_full: [n_total_heads * qk_head_dim // block_size, q_lora_qdim] full scale. - n_total_heads: Total number of heads (e.g. 128). - n_local_heads: Target heads per GPU (padded, e.g. 20). - num_devices: Number of devices (e.g. 7). - qk_head_dim: Head dimension (e.g. 192). - block_size: Quantization block size (e.g. 128). - - Returns: - Lists of per-device (wq_b, wq_b_scale) with shape - [n_local_heads * qk_head_dim, q_lora_dim] and - [n_local_heads * qk_head_dim // block_size, q_lora_qdim]. - """ total_rows = n_total_heads * qk_head_dim rows_per_dev = n_local_heads * qk_head_dim scale_rows_per_dev = rows_per_dev // block_size @@ -185,7 +167,6 @@ def _swizzle_mma_16x16(mat_in: torch.Tensor) -> torch.Tensor: def _swizzle_mma_16x16_for_pages( mat_in: torch.Tensor, q_lora_dim: int, pages: int ) -> torch.Tensor: - """Swizzle a 16xK matrix for the paged weight layout (K divisible by 16).""" assert mat_in.shape[-2] == 16 and mat_in.shape[-1] == q_lora_dim k_per_page = q_lora_dim // pages n_k_tiles = k_per_page // 16 @@ -201,7 +182,6 @@ def _common_to_tilert_fp16mma( wq_b_scales_raw: torch.Tensor, rmsnorm_gamma: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Convert common weights to the packed TileRT FP16 layout.""" pages = 2 rows_per_cta = 32 @@ -275,10 +255,14 @@ def _common_to_tilert_fp16mma( tilert_gamma = rmsnorm_gamma.float().detach().clone() return tilert_wqb, tilert_wqb_scales, tilert_gamma + def convert_to_bf16mma( + self, weights: list[torch.Tensor] + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return self.convert_to_fp16mma(weights) + def convert_to_fp16mma( self, weights: list[torch.Tensor] ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Convert common-format weights to TileRT FP16 MMA layout.""" with torch.inference_mode(): wq_b, wq_b_scale, q_norm_weight = weights return self._common_to_tilert_fp16mma(wq_b, wq_b_scale, q_norm_weight) @@ -325,11 +309,17 @@ def __call__(self) -> list[str]: class RmsnormProjqWqb(TileRTModule): - """RmsnormProjqWqb module: RMSNorm + Q projection (wq_b only).""" + """RmsnormProjqWqb module: RMSNorm + Q projection (wq_b only, TP7).""" _SUPPORTED_ALGORITHMS = { - "deepseek_v3_2": [RmsnormProjqWqbAlgorithm.FP16MMA], - "glm_5": [RmsnormProjqWqbAlgorithm.FP16MMA], + "deepseek_v3_2": [ + RmsnormProjqWqbAlgorithm.FP16MMA, + RmsnormProjqWqbAlgorithm.BF16MMA, + ], + "glm_5": [ + RmsnormProjqWqbAlgorithm.FP16MMA, + RmsnormProjqWqbAlgorithm.BF16MMA, + ], } def __init__( @@ -384,7 +374,7 @@ def get_weights_list(self) -> list[torch.Tensor]: return [self.tilert_q_norm_weight, self.tilert_wq_b, self.tilert_wq_b_scales] def device_sharding(self, weights_map: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Redistribute heads across devices with padding.""" + """Redistribute 128 heads into 7 GPUs Γ— 20 slots with padding.""" gamma = weights_map[self.ref_weights_alias.rmsnorm_gamma][None, ...].repeat( self.num_devices, 1 ) @@ -489,7 +479,7 @@ def golden_forward(self, q: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: assert self.ref_wq_b is not None bsz, seqlen, _ = q.shape - if bsz != 1 or seqlen not in [1, 2, 4]: + if bsz != 1 or seqlen not in [1, 2, 4, 8]: raise ValueError(f"Invalid batch size or sequence length: bsz={bsz}, seqlen={seqlen}") qr = torch.nn.functional.rms_norm(q.float(), [q.size(-1)], self.ref_q_norm, self.eps).to( @@ -510,7 +500,7 @@ def tilert_forward(self, q: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: assert self.profile_logs is not None bsz, seqlen, _ = q.shape - if bsz != 1 or seqlen not in [1, 2, 4]: + if bsz != 1 or seqlen not in [1, 2, 4, 8]: raise ValueError(f"Invalid batch size or sequence length: bsz={bsz}, seqlen={seqlen}") assert self.algorithm is not None, "Algorithm is not set" diff --git a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projq_wqi.py b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projq_wqi.py index 4f4d07f..31c0a8c 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projq_wqi.py +++ b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projq_wqi.py @@ -1,4 +1,4 @@ -"""RmsnormProjqWqi operation module (IQ-only projection).""" +"""RmsnormProjqWqi operation module (GLM5 v2, IQ-only projection).""" from dataclasses import dataclass from enum import Enum @@ -44,6 +44,7 @@ class RmsnormProjqWqiAlgorithm(Enum): """RmsnormProjqWqi algorithm.""" FP16MMA = "fp16mma" + BF16MMA = "bf16mma" class RmsnormProjqWqiWeightsConverter(TilertWeightsConverter): @@ -71,7 +72,6 @@ def _swizzle_mma_16x16(mat_in: torch.Tensor) -> torch.Tensor: def _swizzle_mma_16x16_for_pages( mat_in: torch.Tensor, q_lora_rank: int, pages: int ) -> torch.Tensor: - """Swizzle a 16xK matrix for the paged weight layout (K divisible by 16).""" assert mat_in.shape[-2] == 16 and mat_in.shape[-1] == q_lora_rank pre_shape = mat_in.shape[:-2] k_per_page = q_lora_rank // pages @@ -87,7 +87,6 @@ def _common_to_tilert_fp16mma( wqi_scales: torch.Tensor, rmsnorm_gamma: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Convert common weights to the packed TileRT FP16 layout (IQ only).""" sms = 128 k_per_page = 1024 if self.model_args.arch_name == "glm_5" else 512 pages = self.q_lora_dim // k_per_page @@ -128,14 +127,14 @@ def _common_to_tilert_fp16mma( tilert_gamma = rmsnorm_gamma.float().detach().clone() return tilert_wqi, tilert_wqi_scales, tilert_gamma - def convert_to_fp16mma( + def convert_to_bf16mma( self, weights: list[torch.Tensor] ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Convert common-format weights to TileRT FP16 MMA layout. + return self.convert_to_fp16mma(weights) - Args: - weights: [wqi, wqi_scale, q_norm_weight]. - """ + def convert_to_fp16mma( + self, weights: list[torch.Tensor] + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: with torch.inference_mode(): wqi, wqi_scale, q_norm_weight = weights return self._common_to_tilert_fp16mma(wqi, wqi_scale, q_norm_weight) @@ -177,8 +176,14 @@ class RmsnormProjqWqi(TileRTModule): """RmsnormProjqWqi module: RMSNorm + W_qi projection (IQ only, GLM5 v2).""" _SUPPORTED_ALGORITHMS = { - "deepseek_v3_2": [RmsnormProjqWqiAlgorithm.FP16MMA], - "glm_5": [RmsnormProjqWqiAlgorithm.FP16MMA], + "deepseek_v3_2": [ + RmsnormProjqWqiAlgorithm.FP16MMA, + RmsnormProjqWqiAlgorithm.BF16MMA, + ], + "glm_5": [ + RmsnormProjqWqiAlgorithm.FP16MMA, + RmsnormProjqWqiAlgorithm.BF16MMA, + ], } def __init__( @@ -240,7 +245,6 @@ def device_sharding(self, weights_map: dict[str, torch.Tensor]) -> dict[str, tor } def init_reference_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - """Initialize reference weights from common-format state dict.""" self.ref_q_norm = state_dict[self.tilert_weights_alias.rmsnorm_gamma] wqi = weight_dequant( state_dict[self.tilert_weights_alias.wqi_weights], @@ -249,21 +253,19 @@ def init_reference_weights(self, state_dict: dict[str, torch.Tensor]) -> None: self.ref_wqi = wqi.contiguous() def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - """Initialize TileRT weights from common-format state dict.""" + assert self.algorithm is not None, "Algorithm is not set" weights = [ + state_dict[self.tilert_weights_alias.rmsnorm_gamma], state_dict[self.tilert_weights_alias.wqi_weights], state_dict[self.tilert_weights_alias.wqi_scales], - state_dict[self.tilert_weights_alias.rmsnorm_gamma], ] - assert self.algorithm is not None, "Algorithm is not set" self.tilert_wqi, self.tilert_wqi_scales, self.tilert_q_norm_weight = ( - RmsnormProjqWqiWeightsConverter(self.model_args, self.num_devices).dispatch( - self.algorithm, weights + torch.ops.tilert.rmsnorm_projq_wqi__convert_weights( + weights, self.model_args.arch_name, self.algorithm.value ) ) def init_random_weights(self) -> None: - """Initialize random reference and TileRT weights for testing.""" q_norm = torch.randn(self.q_lora_rank, dtype=torch.float32) wqi = torch.randn(self.index_head_dim, self.q_lora_rank, dtype=torch.bfloat16).to( torch.float8_e4m3fn @@ -281,7 +283,6 @@ def init_random_weights(self) -> None: self.init_tilert_weights(ref_state) def init_tilert_vars(self, batch_size: int, seq_len: int) -> None: - """Allocate TileRT output buffers.""" self.iq = torch.zeros( batch_size, seq_len, self.index_n_heads, self.head_dim, dtype=torch.bfloat16 ) @@ -289,7 +290,6 @@ def init_tilert_vars(self, batch_size: int, seq_len: int) -> None: self.is_var_init = True def golden_forward(self, q: torch.Tensor) -> torch.Tensor: - """Reference forward: RMSNorm + W_qi_b linear projection.""" assert self.ref_q_norm is not None assert self.ref_wqi is not None diff --git a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projx_wqakis.py b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projx_wqakis.py index 8813d6a..2702b42 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projx_wqakis.py +++ b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projx_wqakis.py @@ -21,7 +21,7 @@ class RMSNormProjxWqakisWeightsConverter(TilertWeightsConverter): - """Weight converter for RMSNormProjxWqakis.""" + """Weight converter for RMSNormProjxWqakis (decoupled FP8 MMA).""" def __init__(self, model_args: ModelArgs, num_devices: int): super().__init__(model_args, num_devices) @@ -29,14 +29,6 @@ def __init__(self, model_args: ModelArgs, num_devices: int): def convert_to_decoupled( self, weights: list[torch.Tensor] ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Convert weights to decoupled FP8 MMA format. - - Args: - weights: [gamma, wq_a, wq_a_scale, wki, wki_scale, wis, wis_scale] - - Returns: - (wqaki_packed, wis_bf16, gamma) - """ arch_name = self.model_args.arch_name x_rmsnorm_gamma, wq_a, wq_a_scale, wki, wki_scale, wis, _wis_scale = weights @@ -111,14 +103,21 @@ class RMSNormProjxWqakisAlgorithm(Enum): """RMSNormProjxWqakis algorithm.""" FP8MMA = "fp8mma" + W8A16HMMA = "w8a16_hmma" class RMSNormProjxWqakis(TileRTModule): - """Decoupled RMSNorm + GEMV(W_q_a, W_ki, W_is).""" + """Decoupled RMSNorm + GEMV(W_q_a, W_ki, W_is) for Device Group A.""" _SUPPORTED_ALGORITHMS = { - "deepseek_v3_2": [RMSNormProjxWqakisAlgorithm.FP8MMA], - "glm_5": [RMSNormProjxWqakisAlgorithm.FP8MMA], + "deepseek_v3_2": [ + RMSNormProjxWqakisAlgorithm.FP8MMA, + RMSNormProjxWqakisAlgorithm.W8A16HMMA, + ], + "glm_5": [ + RMSNormProjxWqakisAlgorithm.FP8MMA, + RMSNormProjxWqakisAlgorithm.W8A16HMMA, + ], } def __init__( @@ -183,7 +182,6 @@ def get_weights_list(self) -> list[torch.Tensor]: return [self.tilert_norm_gamma, self.tilert_wqakis, self.tilert_wis] def device_sharding(self, weights_map: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Repeat weights for device sharding.""" input_layernorm_weight = ( weights_map[self.ref_weights_alias.x_rmsnorm_gamma][None, ...] .float() @@ -235,6 +233,31 @@ def init_reference_weights(self, state_dict: dict[str, torch.Tensor]) -> None: def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: tilert_aliases = self.tilert_weights_alias() weights_list = [state_dict[alias] for alias in tilert_aliases] + if self.algorithm == RMSNormProjxWqakisAlgorithm.W8A16HMMA: + gamma, wq_a, wq_a_scale, wki, wki_scale, wis, _wis_scale = weights_list + if self.arch_name == "glm_5": + self.tilert_wqakis = ProjxWqakiWeightsConverter.convert_glm5_68cta_w8a16( + wq_a, wq_a_scale, wki, wki_scale + ) + else: + from tilert.models.glm_5._dsa_v32.ops.rmsnorm_projx_wqkva import ( + RMSNormProjQKVAW8A16MMAWeightsConverter, + ) + + w_fp8 = torch.cat( + [ + wq_a.reshape(self.q_lora_rank, self.dim), + wki.reshape(self.idx_head_dim, self.dim), + ], + dim=0, + ).contiguous() + scales = torch.cat([wq_a_scale, wki_scale], dim=0).to(torch.float32).contiguous() + self.tilert_wqakis = RMSNormProjQKVAW8A16MMAWeightsConverter.pack_lane_major( + w_fp8, scales, self.dim + ) + self.tilert_wis = wis.to(torch.bfloat16) + self.tilert_norm_gamma = gamma.float() + return converter = RMSNormProjxWqakisWeightsConverter(self.model_args, self.num_devices) result = converter.convert_to_decoupled(weights_list) self.tilert_wqakis, self.tilert_wis, self.tilert_norm_gamma = result @@ -280,7 +303,6 @@ def golden_forward( self, x: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Pure PyTorch reference: RMSNorm -> q, ki, idx_scores.""" assert self.ref_norm_gamma is not None assert self.ref_wq_a is not None assert self.ref_wki is not None @@ -302,7 +324,6 @@ def tilert_forward( self, x: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Run RMSNorm + ProjXWqaki + ProjXWis via TileRT CUDA kernels.""" rmsnorm_quant( x.to(torch.bfloat16), self.tilert_norm_gamma, diff --git a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projx_wqkva.py b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projx_wqkva.py index 5343357..9e5bbc2 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projx_wqkva.py +++ b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_projx_wqkva.py @@ -16,7 +16,7 @@ class RMSNormProjQKVAFP8MMAWeightsConverter: - """Weight converter: pack FP8 weights into the kernel's packed layout.""" + """Weight converter: pack FP8 weights into WqkvaPagedShared layout for the FP8 MMA kernel.""" HIDDEN_DIM = 6144 Q_LORA_RANK = 2048 @@ -32,11 +32,9 @@ class RMSNormProjQKVAFP8MMAWeightsConverter: MAT_BYTES = ROWS_PER_CTA * COLS_PER_PAGE SCALE_OFFSET = MAT_BYTES - PAGE_BYTES = ((MAT_BYTES + 128 + 127) // 128) * 128 @staticmethod def _swizzle_mma_16x32(mat_in: torch.Tensor) -> torch.Tensor: - """Swizzle [*, 16, 32] tiles into the packed weight layout.""" assert mat_in.shape[-2] == 16 and mat_in.shape[-1] == 32 pre_shape = mat_in.shape[:-2] mat_in = mat_in.reshape(*pre_shape, 2, 8, 2, 4, 4).transpose(-4, -3).transpose(-5, -4) @@ -55,14 +53,7 @@ def convert_to_fp8_mma_gemv( hidden_dim: int = 6144, q_lora_rank: int = 2048, ) -> tuple[torch.Tensor, torch.Tensor]: - """Pack FP8 weights for the FP8 MMA kernel. - - Args: - hidden_dim: Model hidden dimension. - q_lora_rank: Q projection rank. - """ C = RMSNormProjQKVAFP8MMAWeightsConverter - block_size = C.BLOCK_SIZE kv_lora_rank = C.KV_LORA_RANK qk_rope_head_dim = C.QK_ROPE_HEAD_DIM @@ -73,54 +64,90 @@ def convert_to_fp8_mma_gemv( expected = qk_rope_head_dim * hidden_dim assert w_pe.numel() == expected, f"w_pe numel {w_pe.numel()} != expected {expected}" + return C._pack_per_row_no_requant( + wq_a, + wkv_a, + w_pe, + wq_a_scale, + wkv_a_scale, + w_pe_scale, + attn_norm_weight, + hidden_dim=hidden_dim, + q_lora_rank=q_lora_rank, + ) + + @staticmethod + def _pack_per_row_no_requant( + wq_a: torch.Tensor, + wkv_a: torch.Tensor, + w_pe: torch.Tensor, + wq_a_scale: torch.Tensor, + wkv_a_scale: torch.Tensor, + w_pe_scale: torch.Tensor, + attn_norm_weight: torch.Tensor, + *, + hidden_dim: int, + q_lora_rank: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + C = RMSNormProjQKVAFP8MMAWeightsConverter + kv_lora_rank = C.KV_LORA_RANK + qk_rope_head_dim = C.QK_ROPE_HEAD_DIM total_rows = q_lora_rank + kv_lora_rank + qk_rope_head_dim num_ctas = total_rows // C.ROWS_PER_CTA num_pages = hidden_dim // C.COLS_PER_PAGE - wq_a_f = weight_dequant(wq_a.reshape(q_lora_rank, hidden_dim), wq_a_scale) - wkv_a_f = weight_dequant(wkv_a.reshape(kv_lora_rank, hidden_dim), wkv_a_scale) - w_pe_f = weight_dequant(w_pe.reshape(qk_rope_head_dim, hidden_dim), w_pe_scale) - w_float = torch.cat([wq_a_f, wkv_a_f, w_pe_f], dim=0) - - w_blocks = w_float.reshape(total_rows, hidden_dim // block_size, block_size) - col_max = w_blocks.abs().amax(dim=(0, 2)) - fp8_max = torch.finfo(torch.float8_e4m3fn).max - w_scales = (col_max / fp8_max).clamp(min=1e-12) - - scales_expanded = w_scales.repeat_interleave(block_size) - w_scaled = w_float / scales_expanded.unsqueeze(0) - w_fp8 = w_scaled.to(torch.float8_e4m3fn) + num_tiles = C.COLS_PER_PAGE // 32 + blk = C.BLOCK_SIZE + num_blk_page = C.COLS_PER_PAGE // blk + num_blk_total = hidden_dim // blk + + w_fp8 = torch.cat( + [ + wq_a.reshape(q_lora_rank, hidden_dim), + wkv_a.reshape(kv_lora_rank, hidden_dim), + w_pe.reshape(qk_rope_head_dim, hidden_dim), + ], + dim=0, + ).contiguous() + + def _bcast_block_scale(scale: torch.Tensor, rows: int) -> torch.Tensor: + s = scale.to(torch.float32).reshape(-1, num_blk_total) + return s.repeat_interleave(blk, dim=0)[:rows] + + w_scales = torch.cat( + [ + _bcast_block_scale(wq_a_scale, q_lora_rank), + _bcast_block_scale(wkv_a_scale, kv_lora_rank), + _bcast_block_scale(w_pe_scale, qk_rope_head_dim), + ], + dim=0, + ).clamp(min=1e-12) assert C.MAT_BYTES == C.SCALE_OFFSET, "Layout mismatch: scales must follow mat" - assert block_size == C.COLS_PER_PAGE // C.SCALES_PER_PAGE, "Block size mismatch" - assert w_scales.numel() == num_pages * C.SCALES_PER_PAGE, "Scale count mismatch" w_bytes = w_fp8.view(torch.uint8) - num_tiles = C.COLS_PER_PAGE // 32 - mat = w_bytes.reshape(num_ctas, C.ROWS_PER_CTA, num_pages, C.COLS_PER_PAGE) mat = mat.transpose(1, 2) - mat = mat.reshape(num_ctas, num_pages, 2, 16, num_tiles, 32) mat = mat.transpose(3, 4) mat = C._swizzle_mma_16x32(mat) mat = mat.contiguous().reshape(num_ctas, num_pages, C.MAT_BYTES) - scales_f32 = w_scales.reshape(num_pages, C.SCALES_PER_PAGE).to(torch.float32).contiguous() - scales_bytes = scales_f32.view(torch.uint8) - scales_bytes = scales_bytes.unsqueeze(0).expand(num_ctas, -1, -1) - - pad_size = C.PAGE_BYTES - C.MAT_BYTES - C.SCALES_PER_PAGE * 4 - padding = torch.zeros(num_ctas, num_pages, pad_size, dtype=torch.uint8, device=w_fp8.device) - - packed = torch.cat([mat, scales_bytes, padding], dim=-1) + sc = w_scales.reshape(num_ctas, C.ROWS_PER_CTA, num_pages, num_blk_page) + sc = sc.permute(0, 2, 1, 3).contiguous() + scales_bytes = ( + sc.to(torch.float32) + .reshape(num_ctas, num_pages, C.ROWS_PER_CTA * num_blk_page) + .view(torch.uint8) + ) + packed = torch.cat([mat, scales_bytes], dim=-1) packed = packed.contiguous().reshape(-1) return packed.view(torch.float8_e4m3fn), attn_norm_weight.clone() class RMSNormProjQKVAFP16MMAWeightsConverter: - """Weight converter: pack FP16 weights for the kernel.""" + """Weight converter: pack FP16 weights for the MMA kernel.""" KV_LORA_RANK = 512 QK_ROPE_HEAD_DIM = 64 @@ -130,7 +157,6 @@ class RMSNormProjQKVAFP16MMAWeightsConverter: @staticmethod def _swizzle_mma_16x16(mat_in: torch.Tensor) -> torch.Tensor: - """Swizzle [*, 16, 16] tiles into the packed weight layout.""" assert mat_in.shape[-2] == 16 and mat_in.shape[-1] == 16 pre_shape = mat_in.shape[:-2] mat_in = mat_in.reshape(*pre_shape, 2, 8, 2, 4, 2).transpose(-4, -3).transpose(-5, -4) @@ -149,7 +175,6 @@ def convert_to_fp16_mma_gemv( hidden_dim: int = 6144, q_lora_rank: int = 2048, ) -> tuple[torch.Tensor, torch.Tensor]: - """Pack weights into the FP16 layout expected by the kernel.""" C = RMSNormProjQKVAFP16MMAWeightsConverter kv_lora_rank = C.KV_LORA_RANK qk_rope_head_dim = C.QK_ROPE_HEAD_DIM @@ -182,10 +207,125 @@ def convert_to_fp16_mma_gemv( return packed.view(torch.float16), attn_norm_weight.clone() +class RMSNormProjQKVAW8A16MMAWeightsConverter: + """Pack FP8 weight + block scale into packed format.""" + + KV_LORA_RANK = 512 + QK_ROPE_HEAD_DIM = 64 + ROWS_PER_CTA = 32 + COLS_PER_PAGE = 1024 + BLOCK_SIZE = 128 + NUM_WARPS = 8 + MMA_K = 16 + M_TILES_PER_CTA = ROWS_PER_CTA // 16 + K_TILES_PER_WARP = COLS_PER_PAGE // (NUM_WARPS * MMA_K) + SCALES_PER_PAGE = COLS_PER_PAGE // BLOCK_SIZE + PAGE_MAT_BYTES = M_TILES_PER_CTA * K_TILES_PER_WARP * NUM_WARPS * 32 * 8 + PAGE_BYTES = PAGE_MAT_BYTES + 128 + + @staticmethod + def _permute_mma_a_fragment_16x16(tile: torch.Tensor) -> torch.Tensor: + assert tile.shape[-2:] == (16, 16) + pre = tile.shape[:-2] + return ( + tile.reshape(*pre, 2, 8, 2, 4, 2) + .permute( + *range(len(pre)), + len(pre) + 1, + len(pre) + 3, + len(pre) + 2, + len(pre) + 0, + len(pre) + 4, + ) + .contiguous() + .reshape(*pre, 32, 8) + ) + + @staticmethod + def convert_to_w8a16_mma_gemv( + wq_a: torch.Tensor, + wq_a_scale: torch.Tensor, + wkv_a: torch.Tensor, + wkv_a_scale: torch.Tensor, + w_pe: torch.Tensor, + w_pe_scale: torch.Tensor, + attn_norm_weight: torch.Tensor, + *, + hidden_dim: int = 6144, + q_lora_rank: int = 2048, + ) -> tuple[torch.Tensor, torch.Tensor]: + C = RMSNormProjQKVAW8A16MMAWeightsConverter + kv_lora_rank = C.KV_LORA_RANK + qk_rope_head_dim = C.QK_ROPE_HEAD_DIM + rows_per_cta = C.ROWS_PER_CTA + cols_per_page = C.COLS_PER_PAGE + + w_fp8 = torch.cat( + [ + wq_a.reshape(q_lora_rank, hidden_dim), + wkv_a.reshape(kv_lora_rank, hidden_dim), + w_pe.reshape(qk_rope_head_dim, hidden_dim), + ], + dim=0, + ).contiguous() + assert w_fp8.dtype == torch.float8_e4m3fn, f"expected fp8 weight, got {w_fp8.dtype}" + + scales = ( + torch.cat([wq_a_scale, wkv_a_scale, w_pe_scale], dim=0).to(torch.float32).contiguous() + ) + + total_rows = q_lora_rank + kv_lora_rank + qk_rope_head_dim + num_ctas = total_rows // rows_per_cta + num_pages = hidden_dim // cols_per_page + expected_scale_rows = (total_rows + C.BLOCK_SIZE - 1) // C.BLOCK_SIZE + assert scales.shape == (expected_scale_rows, hidden_dim // C.BLOCK_SIZE), ( + f"scales {tuple(scales.shape)} != " + f"{(expected_scale_rows, hidden_dim // C.BLOCK_SIZE)}" + ) + + del num_ctas, num_pages + return C.pack_lane_major(w_fp8, scales, hidden_dim), attn_norm_weight.clone() + + @classmethod + def pack_lane_major( + cls, w_fp8: torch.Tensor, scales: torch.Tensor, hidden_dim: int + ) -> torch.Tensor: + assert w_fp8.dtype == torch.float8_e4m3fn, f"expected fp8 weight, got {w_fp8.dtype}" + rows_per_cta = cls.ROWS_PER_CTA + cols_per_page = cls.COLS_PER_PAGE + total_rows = w_fp8.shape[0] + num_ctas = total_rows // rows_per_cta + num_pages = hidden_dim // cols_per_page + scales = scales.to(torch.float32).contiguous() + device = w_fp8.device + w_bytes = w_fp8.view(torch.uint8) + + w = w_bytes.reshape(num_ctas, rows_per_cta, num_pages, cols_per_page) + w = w.reshape(num_ctas, cls.M_TILES_PER_CTA, 16, num_pages, cols_per_page) + w = w.permute(0, 3, 1, 2, 4).contiguous() + w = w.reshape( + num_ctas, num_pages, cls.M_TILES_PER_CTA, 16, cls.NUM_WARPS, cls.K_TILES_PER_WARP, 16 + ) + w = w.permute(0, 1, 2, 5, 4, 3, 6).contiguous() + w_lane = cls._permute_mma_a_fragment_16x16(w) + mat_blob = w_lane.contiguous().reshape(num_ctas, num_pages, cls.PAGE_MAT_BYTES) + + cta_idx = torch.arange(num_ctas, device=device) + scale_row = cta_idx // (cls.BLOCK_SIZE // rows_per_cta) + cta_scales = scales[scale_row].reshape(num_ctas, num_pages, cls.SCALES_PER_PAGE) + scale_bytes = cta_scales.contiguous().view(torch.uint8) + + out = torch.zeros(num_ctas, num_pages, cls.PAGE_BYTES, dtype=torch.uint8, device=device) + out[:, :, : cls.PAGE_MAT_BYTES] = mat_blob + out[:, :, cls.PAGE_MAT_BYTES : cls.PAGE_MAT_BYTES + cls.SCALES_PER_PAGE * 4] = scale_bytes + return out.reshape(-1).contiguous().view(torch.float8_e4m3fn) + + class RMSNormProjxWqkvaAlgorithm(Enum): """RMSNormProjxWqkva algorithm.""" DECOUPLED = "decoupled" + W8A16HMMA = "w8a16_hmma" class RMSNormProjxWqkvaWeightsConverter(TilertWeightsConverter): @@ -197,11 +337,6 @@ def __init__(self, model_args: ModelArgs, num_devices: int): def convert_to_fp8_mma_gemv( self, weights: list[torch.Tensor] ) -> tuple[torch.Tensor, torch.Tensor]: - """Convert tilert weights list to the FP8 kernel-ready format. - - Args: - weights: [gamma, wq_a, wq_a_scale, wkv_a, wkv_a_scale, w_pe, w_pe_scale] - """ gamma, wq_a, wq_a_scale, wkv_a, wkv_a_scale, w_pe, w_pe_scale = weights return RMSNormProjQKVAFP8MMAWeightsConverter.convert_to_fp8_mma_gemv( wq_a, @@ -218,11 +353,6 @@ def convert_to_fp8_mma_gemv( def convert_to_fp16_mma_gemv( self, weights: list[torch.Tensor] ) -> tuple[torch.Tensor, torch.Tensor]: - """Convert tilert weights list to the FP16 kernel-ready format. - - Args: - weights: [gamma, wq_a, wq_a_scale, wkv_a, wkv_a_scale, w_pe, w_pe_scale] - """ gamma, wq_a, wq_a_scale, wkv_a, wkv_a_scale, w_pe, w_pe_scale = weights return RMSNormProjQKVAFP16MMAWeightsConverter.convert_to_fp16_mma_gemv( wq_a, @@ -288,11 +418,17 @@ def __call__(self) -> list[str]: class RMSNormProjxWqkva(TileRTModule): - """Fused RMSNorm + GEMV(W_q_a, W_kv_a, W_pe).""" + """Fused RMSNorm + GEMV(W_q_a, W_kv_a, W_pe) for Device Group B.""" _SUPPORTED_ALGORITHMS = { - "deepseek_v3_2": [RMSNormProjxWqkvaAlgorithm.DECOUPLED], - "glm_5": [RMSNormProjxWqkvaAlgorithm.DECOUPLED], + "deepseek_v3_2": [ + RMSNormProjxWqkvaAlgorithm.DECOUPLED, + RMSNormProjxWqkvaAlgorithm.W8A16HMMA, + ], + "glm_5": [ + RMSNormProjxWqkvaAlgorithm.DECOUPLED, + RMSNormProjxWqkvaAlgorithm.W8A16HMMA, + ], } def __init__( @@ -354,7 +490,6 @@ def get_weights_list(self) -> list[torch.Tensor]: return [self.tilert_norm_gamma, self.tilert_wqkva, self.tilert_wqkva_scales] def device_sharding(self, weights_map: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Repeat weights for device sharding.""" input_layernorm_weight = ( weights_map[self.ref_weights_alias.x_rmsnorm_gamma][None, ...] .float() @@ -405,8 +540,27 @@ def init_reference_weights(self, state_dict: dict[str, torch.Tensor]) -> None: def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: tilert_aliases = self.tilert_weights_alias() weights_list = [state_dict[alias] for alias in tilert_aliases] - converter = RMSNormProjxWqkvaWeightsConverter(self.model_args, self.num_devices) - self.tilert_wqkva, self.tilert_norm_gamma = converter.convert_to_fp8_mma_gemv(weights_list) + if self.algorithm == RMSNormProjxWqkvaAlgorithm.W8A16HMMA: + gamma, wq_a, wq_a_scale, wkv_a, wkv_a_scale, w_pe, w_pe_scale = weights_list + self.tilert_wqkva, self.tilert_norm_gamma = ( + RMSNormProjQKVAW8A16MMAWeightsConverter.convert_to_w8a16_mma_gemv( + wq_a, + wq_a_scale, + wkv_a, + wkv_a_scale, + w_pe, + w_pe_scale, + gamma.float(), + hidden_dim=self.dim, + q_lora_rank=self.q_lora_rank, + ) + ) + self.tilert_norm_gamma = self.tilert_norm_gamma.float().contiguous() + else: + converter = RMSNormProjxWqkvaWeightsConverter(self.model_args, self.num_devices) + self.tilert_wqkva, self.tilert_norm_gamma = converter.convert_to_fp8_mma_gemv( + weights_list + ) self.tilert_wqkva_scales = torch.zeros((1,), dtype=torch.float32) def init_tilert_vars(self, batch_size: int, seq_len: int, max_len: int = 128) -> None: @@ -437,9 +591,9 @@ def init_random_weights(self) -> None: tensor_list = [ torch.randn(self.dim, dtype=torch.float32), torch.randn(self.q_lora_rank, self.dim, dtype=torch.bfloat16).to(torch.float8_e4m3fn), - torch.randn(q_scale_dim, dim_scale_dim, dtype=scale_dtype), + torch.randn(q_scale_dim, dim_scale_dim, dtype=scale_dtype).abs(), torch.randn(kv_mqa_rows, self.dim, dtype=torch.bfloat16).to(torch.float8_e4m3fn), - torch.randn(kv_mqa_scale_dim, dim_scale_dim, dtype=scale_dtype), + torch.randn(kv_mqa_scale_dim, dim_scale_dim, dtype=scale_dtype).abs(), ] ref_state_dict = dict(zip(self.ref_weights_alias(), tensor_list)) self.init_reference_weights(ref_state_dict) @@ -475,7 +629,6 @@ def tilert_forward( x: torch.Tensor, cur_pos: int = 0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Run RMSNorm + 3-way GEMV via the TileRT CUDA kernels.""" assert self.cur_pos is not None assert self.pe_cache_out is not None self.cur_pos.fill_(cur_pos) diff --git a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_quant.py b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_quant.py index 1d399c5..c977d08 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_quant.py +++ b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_quant.py @@ -27,17 +27,6 @@ def rmsnorm_quant( *, model_arch: str, ) -> None: - """ - Rmsnorm with optional activation quantization. - - Args: - hidden_in: Input tensor (..., dim). - gamma_in: RMSNorm gamma (dim,). - hidden_out: RMSNorm output (..., dim). - quant_hidden_out: Optional quantized output (..., dim). If None, no quant. - quant_hidden_scale_out: Optional quant scale (..., dim // block_size). If None, no quant. - profile_logs: Optional profile logs tensor. - """ if profile_logs is None: raise ValueError("profile_logs is required when calling rmsnorm_quant.") diff --git a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_up_gate_silu.py b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_up_gate_silu.py index 25adae9..1792aec 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_up_gate_silu.py +++ b/tilert/models/glm_5/_dsa_v32/ops/rmsnorm_up_gate_silu.py @@ -13,6 +13,9 @@ ExpertSelectUpGateSiLU, ExpertSelectUpGateSiLUWeightsConverter, ) +from tilert.models.glm_5._dsa_v32.ops.rmsnorm_projx_wqkva import ( + RMSNormProjQKVAW8A16MMAWeightsConverter, +) from tilert.utils import get_profile_log_tensor __all__ = [ @@ -49,6 +52,63 @@ class RMSNormUpGateSiLUAlgorithm(Enum): FP8MMA = "fp8mma" FP16MMA = "fp16mma" + BF16MMA = "bf16mma" + BF16MMA_V2 = "w8a16_hmma_v2" + + +def _pack_up_gate_w8a16_v2( + gate_w: torch.Tensor, + gate_s: torch.Tensor, + up_w: torch.Tensor, + up_s: torch.Tensor, + hidden_dim: int = 6144, +) -> torch.Tensor: + """Pack GLM5 dense gate/up FP8 weight + block scales into the V2 blob.""" + C = RMSNormProjQKVAW8A16MMAWeightsConverter + rows_per_cta = C.ROWS_PER_CTA + cols_per_page = C.COLS_PER_PAGE + block = C.BLOCK_SIZE + num_warps = C.NUM_WARPS + k_tiles = C.K_TILES_PER_WARP + scales_per_set = C.SCALES_PER_PAGE + page_mat_bytes = C.PAGE_MAT_BYTES + page_bytes = C.PAGE_BYTES + + gate_w = gate_w.reshape(-1, hidden_dim).contiguous() + up_w = up_w.reshape(-1, hidden_dim).contiguous() + assert gate_w.dtype == torch.float8_e4m3fn and up_w.dtype == torch.float8_e4m3fn + inter = gate_w.shape[0] + num_ctas = inter // 16 + num_pages = hidden_dim // cols_per_page + device = gate_w.device + + gate_s = gate_s.reshape(-1, hidden_dim // block).to(torch.float32).contiguous() + up_s = up_s.reshape(-1, hidden_dim // block).to(torch.float32).contiguous() + + gate_c = gate_w.reshape(num_ctas, 16, hidden_dim) + up_c = up_w.reshape(num_ctas, 16, hidden_dim) + w = torch.cat([gate_c, up_c], dim=1).contiguous() + w_bytes = w.view(torch.uint8) + + w = w_bytes.reshape(num_ctas, rows_per_cta, num_pages, cols_per_page) + w = w.reshape(num_ctas, C.M_TILES_PER_CTA, 16, num_pages, cols_per_page) + w = w.permute(0, 3, 1, 2, 4).contiguous() + w = w.reshape(num_ctas, num_pages, C.M_TILES_PER_CTA, 16, num_warps, k_tiles, 16) + w = w.permute(0, 1, 2, 5, 4, 3, 6).contiguous() + w_lane = C._permute_mma_a_fragment_16x16(w) + mat_blob = w_lane.contiguous().reshape(num_ctas, num_pages, page_mat_bytes) + + cta_idx = torch.arange(num_ctas, device=device) + scale_row = cta_idx // (block // 16) + gate_cta = gate_s[scale_row].reshape(num_ctas, num_pages, scales_per_set) + up_cta = up_s[scale_row].reshape(num_ctas, num_pages, scales_per_set) + scale16 = torch.cat([gate_cta, up_cta], dim=2).contiguous() + scale_bytes = scale16.view(torch.uint8) + + out = torch.zeros(num_ctas, num_pages, page_bytes, dtype=torch.uint8, device=device) + out[:, :, :page_mat_bytes] = mat_blob + out[:, :, page_mat_bytes : page_mat_bytes + 16 * 4] = scale_bytes + return out.reshape(-1).contiguous().view(torch.float8_e4m3fn) RMSNormUpGateSiLUWeightsConverter = ExpertSelectUpGateSiLUWeightsConverter @@ -83,8 +143,15 @@ class RMSNormUpGateSiLU(TileRTModule): """RMSNormUpGateSiLU module""" _SUPPORTED_ALGORITHMS = { - "deepseek_v3_2": [RMSNormUpGateSiLUAlgorithm.FP8MMA, RMSNormUpGateSiLUAlgorithm.FP16MMA], - "glm_5": [RMSNormUpGateSiLUAlgorithm.FP8MMA, RMSNormUpGateSiLUAlgorithm.FP16MMA], + "deepseek_v3_2": [ + RMSNormUpGateSiLUAlgorithm.FP8MMA, + RMSNormUpGateSiLUAlgorithm.FP16MMA, + RMSNormUpGateSiLUAlgorithm.BF16MMA, + ], + "glm_5": [ + RMSNormUpGateSiLUAlgorithm.BF16MMA, + RMSNormUpGateSiLUAlgorithm.BF16MMA_V2, + ], } def __init__( @@ -146,12 +213,6 @@ def tilert_tensor_alias(self) -> list[str]: return self.tilert_weights_alias() def get_weights_list(self) -> list[torch.Tensor]: - """ - Get the weights list. - - Returns: - List of weights. - """ return [self.tilert_norm_gamma, self.tilert_weights, self.tilert_scales] def device_sharding( @@ -159,15 +220,6 @@ def device_sharding( weights_dict: dict[str, torch.Tensor], key_prefix: str, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Device sharding. - - Args: - weights_dict: Dictionary of weights. - - Returns: - Tuple of weights. - """ rmsnorm_gamma_key = f"{key_prefix}.post_attention_layernorm.weight" if ".mlp" in key_prefix: key_prefix_without_mlp = key_prefix.replace(".mlp", "") @@ -210,13 +262,6 @@ def init_reference_weights( key_prefix: str, device_id: int = 0, ) -> None: - """ - Initialize the reference weights. - - Args: - state_dict: State dictionary. - device_id: Device ID. - """ sharded_list = self.device_sharding(state_dict, key_prefix) gamma = sharded_list[0][device_id] @@ -237,25 +282,23 @@ def init_reference_weights( self.ref_up = torch.stack(ref_up_list, dim=0) def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - """ - Initialize the tilert weights. - - Args: - state_dict: State dictionary. - """ assert self.algorithm is not None, "Algorithm is not set" + aliases = self.tilert_weights_alias() + if self.algorithm == RMSNormUpGateSiLUAlgorithm.BF16MMA_V2: + self.tilert_norm_gamma = state_dict[aliases[0]].float().contiguous() + self.tilert_weights = _pack_up_gate_w8a16_v2( + state_dict[aliases[1]], + state_dict[aliases[2]], + state_dict[aliases[3]], + state_dict[aliases[4]], + hidden_dim=self.dim, + ) + return self.tilert_norm_gamma, self.tilert_weights = RMSNormUpGateSiLUWeightsConverter( self.model_args, self.num_devices - ).dispatch(self.algorithm, [state_dict[alias] for alias in self.tilert_weights_alias()]) + ).dispatch(self.algorithm, [state_dict[alias] for alias in aliases]) def init_tilert_vars(self, batch_size: int, seq_len: int, dev_id: int = 0) -> None: - """ - Initialize the tilert variables. - - Args: - batch_size: Batch size. - seq_len: Sequence length. - """ self.hidden_out = torch.zeros( ( batch_size, @@ -270,13 +313,9 @@ def init_tilert_vars(self, batch_size: int, seq_len: int, dev_id: int = 0) -> No self.profile_logs = get_profile_log_tensor(device=f"cuda:{dev_id}") self.is_init = True - def init_random_weights(self, dev_id: int = 0) -> None: - """ - Initialize the random weights. - - Returns: - None - """ + def init_random_weights(self, dev_id: int | None = None) -> None: + if dev_id is None: + dev_id = self.device_id gamma = torch.randn(self.dim, dtype=torch.float32, device=f"cuda:{dev_id}") gate_weights = torch.randn( self.inter_dim, self.dim, dtype=torch.bfloat16, device=f"cuda:{dev_id}" diff --git a/tilert/models/glm_5/_dsa_v32/ops/rotate.py b/tilert/models/glm_5/_dsa_v32/ops/rotate.py index 10a46f1..9656acd 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/rotate.py +++ b/tilert/models/glm_5/_dsa_v32/ops/rotate.py @@ -65,21 +65,12 @@ def rotate( profile_logs: torch.Tensor, model_arch: str, compute_kernel_type: str = "general", + kv_cache: torch.Tensor | None = None, + cur_pos: torch.Tensor | None = None, + cache_base: int = 0, + cache_stride: int = 0, + cache_compressed: bool = False, ) -> None: - """ - Rotate (hadamard transform) operation. - - Args: - input_raw (torch.Tensor): The input tensor [..., head, 128]. - output_raw (torch.Tensor): The output tensor where the result will be stored. - freqs_cis_raw (torch.Tensor): The frequency tensor. - profile_logs (torch.Tensor): Tensor for storing profiling logs. - model_arch: Model architecture string. - compute_kernel_type: Compute kernel type string. - - Returns: - None - """ torch.ops.tilert.rotate_op( input_raw, output_raw, @@ -87,6 +78,11 @@ def rotate( model_arch, compute_kernel_type, profile_logs, + kv_cache, + cur_pos, + cache_base, + cache_stride, + cache_compressed, ) @@ -121,11 +117,7 @@ class RotateAlgorithm(Enum): class Rotate(TileRTModule): - """Rotate module: RoPE on first qk_rope_head_dim dims + hadamard transform. - - Unified for deepseek_v3_2 (index_n_heads=64) and glm_5 (index_n_heads=32). - No weights; uses model_args for dimensions. - """ + """Rotate module: RoPE on first qk_rope_head_dim dims + hadamard transform.""" _SUPPORTED_ALGORITHMS = { "deepseek_v3_2": [RotateAlgorithm.GENERAL], @@ -193,7 +185,7 @@ def golden_forward( [self.qk_rope_head_dim, self.index_head_dim - self.qk_rope_head_dim], dim=-1, ) - q_pe_idx = apply_rotary_emb(q_pe_idx, freqs_cis) + q_pe_idx = apply_rotary_emb(q_pe_idx, freqs_cis, interleaved=False) idx_q = torch.cat([q_pe_idx, q_nope_idx], dim=-1) return rotate_activation(idx_q) diff --git a/tilert/models/glm_5/_dsa_v32/ops/sparse_index.py b/tilert/models/glm_5/_dsa_v32/ops/sparse_index.py deleted file mode 100644 index ca69c49..0000000 --- a/tilert/models/glm_5/_dsa_v32/ops/sparse_index.py +++ /dev/null @@ -1,135 +0,0 @@ -"""Sparse index operation module.""" - -import torch - -__all__ = [ - "sparse_index", - "sparse_index_topk", -] - - -def sparse_index( - q: torch.Tensor, # noqa: VNE001 - kv: torch.Tensor, - weights: torch.Tensor, - logits: torch.Tensor, - cur_pos: int, - profile_logs: torch.Tensor, - compute_kernel_type: str = "bf16", - *, - model_arch: str, -) -> None: - """ - Sparse index operation. - - Calculate sparse index using q * kv * weights. - - Args: - q (torch.Tensor): The query tensor. - kv (torch.Tensor): The key-value tensor. - weights (torch.Tensor): The weights tensor. - logits (torch.Tensor): The logits tensor. - cur_pos (int): The position of the first token. - profile_logs (torch.Tensor): Tensor for storing profiling logs. - compute_kernel_type (str): Kernel type ("bf16"). - model_arch (str): Model architecture ("deepseek_v3_2"). - - Returns: - None - """ - if q.dtype != torch.bfloat16: - raise ValueError("input must be a bfloat16 tensor.") - if kv.dtype != torch.bfloat16: - raise ValueError("kv must be a bfloat16 tensor.") - if weights.dtype != torch.bfloat16: - raise ValueError("weights must be a bfloat16 tensor.") - if logits.dtype != torch.float32: - raise ValueError("logits must be a float32 tensor.") - - head = q.shape[-2] - dim = q.shape[-1] - - if head != 64 and head != 32: - raise ValueError( - f"Unsupported head size: {head}. Sparse index op currently only \ - supports a head number of 64 or 32." - ) - if dim != 128: - raise ValueError("dim must be 128, as we precompute scale inner kernel") - - device = q.device - if any(t.device != device for t in (kv, weights, logits, profile_logs)): - raise ValueError( - "sparse_index inputs must be on the same device: " - f"q={device}, kv={kv.device}, weights={weights.device}, " - f"logits={logits.device}, profile_logs={profile_logs.device}" - ) - if model_arch == "deepseek_v3_2" and head == 32: - model_arch = "glm_5" - torch.ops.tilert.sparse_index_op( - q, kv, weights, logits, cur_pos, model_arch, compute_kernel_type, profile_logs - ) - - -def sparse_index_topk( - q: torch.Tensor, # noqa: VNE001 - kv: torch.Tensor, - weights: torch.Tensor, - logits: torch.Tensor, - indices: torch.Tensor, - cur_pos: int, - profile_logs: torch.Tensor, -) -> None: - """ - Sparse index operation. - - Calculate sparse index using q * kv * weights. - - Args: - q (torch.Tensor): The query tensor. - kv (torch.Tensor): The key-value tensor. - weights (torch.Tensor): The weights tensor. - logits (torch.Tensor): The logits tensor. - cur_pos (int): The position of the first token. - profile_logs (torch.Tensor): Tensor for storing profiling logs. - - Returns: - None - """ - if q.dtype != torch.bfloat16: - raise ValueError("input must be a bfloat16 tensor.") - if kv.dtype != torch.bfloat16: - raise ValueError("kv must be a bfloat16 tensor.") - if weights.dtype != torch.bfloat16: - raise ValueError("weights must be a bfloat16 tensor.") - if logits.dtype != torch.float32: - raise ValueError("logits must be a float32 tensor.") - - seqlen = q.shape[-3] - head = q.shape[-2] - dim = q.shape[-1] - - if head not in (32, 64): - raise ValueError( - f"Unsupported head size: {head}. Sparse index topk fused op " - "supports head number of 32 (GLM5) or 64 (DSV3.2)." - ) - if dim != 128: - raise ValueError("dim must be 128, as we precompute scale inner kernel") - - device = q.device - if any(t.device != device for t in (kv, weights, logits, indices, profile_logs)): - raise ValueError( - "sparse_index inputs must be on the same device: " - f"q={device}, kv={kv.device}, weights={weights.device}, " - f"logits={logits.device}, profile_logs={profile_logs.device}" - ) - workspace = torch.zeros(seqlen, (200 * 1024 + 260), dtype=torch.int32, device=device) - if head == 64: - torch.ops.tilert.sparse_index_topk_dsv32_op( - q, kv, weights, logits, cur_pos, indices, workspace, profile_logs - ) - else: - torch.ops.tilert.sparse_index_topk_glm5_op( - q, kv, weights, logits, cur_pos, indices, workspace, profile_logs - ) diff --git a/tilert/models/glm_5/_dsa_v32/ops/topk.py b/tilert/models/glm_5/_dsa_v32/ops/topk.py deleted file mode 100644 index bb9dfbb..0000000 --- a/tilert/models/glm_5/_dsa_v32/ops/topk.py +++ /dev/null @@ -1,168 +0,0 @@ -"""topk operations module.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import torch -import torch.nn as nn - -from tilert.utils import get_profile_log_tensor - -if TYPE_CHECKING: - from tilert.models.glm_5._dsa_v32.model_args import ModelArgs - - -__all__ = [ - "TopK", - "topk_approximate", - "topk_accurate", -] - - -def topk_approximate( - logits: torch.Tensor, - seq_len: int, - topk: int, - profile_logs: torch.Tensor, - model_arch: str, - compute_kernel_type: str = "general", -) -> torch.Tensor: - """ - Topk approximate operation. - - Topk approximate the input tensor `logits` and stores the result in `output_raw`. - - Args: - logits (torch.Tensor): The input tensor. - seq_len (int): valid data of logits.shape[-1] - topk (int): The number of topk to approximate. - profile_logs (torch.Tensor): The profile logs tensor. - - Returns: - indices (torch.Tensor): The output tensor. - """ - if logits.dtype != torch.float32: - raise ValueError("logits must be a float32 tensor.") - - if topk != 2048: - raise ValueError("topk must be 2048.") - batch = logits.shape[0] - if batch != 1: - raise ValueError("batch must be 1 in this version") - - indices = torch.zeros(batch, topk, dtype=torch.int32, device=logits.device) - torch.ops.tilert.topk_approximate_op( - logits, indices, seq_len, model_arch, compute_kernel_type, profile_logs - ) - - return indices - - -def topk_accurate( - logits: torch.Tensor, - seq_len: int, - topk: int, - profile_logs: torch.Tensor, - model_arch: str, - compute_kernel_type: str = "general", -) -> torch.Tensor: - """ - Topk approximate operation. - - Topk approximate the input tensor `logits` and stores the result in `output_raw`. - - Args: - logits (torch.Tensor): The input tensor. - seq_len (int): length of last samples, - for k=logits.shape[1] samples, the length is - seq-k+1, seq-k+2, ..., seq-1, seq - topk (int): The number of topk to approximate. - profile_logs (torch.Tensor): The profile logs tensor. - Returns: - indices (torch.Tensor): The output tensor. - """ - if logits.dtype != torch.float32: - raise ValueError("logits must be a float32 tensor.") - - if topk not in (512, 2048): - raise ValueError("topk must be 512 or 2048.") - - assert logits.shape[0] == 1, "batch must be 1 in this version" - num_samples = logits.shape[1] - - indices = torch.zeros(num_samples, topk, dtype=torch.int32, device=logits.device) - indices_ws = torch.zeros(1, num_samples, 4, topk * 2, dtype=torch.int32, device=logits.device) - torch.ops.tilert.topk_accurate_op( - logits, - indices, - seq_len - num_samples, - indices_ws, - model_arch, - compute_kernel_type, - profile_logs, - ) - - return indices - - -class TopK(nn.Module): - """TopK operation with optional approximate kernel. - - Wraps topk_accurate / topk_approximate and provides golden_forward - (reference implementation) and tilert_forward (TileRT kernel). - """ - - def __init__(self, use_approximate: bool = False, model_args: ModelArgs | None = None) -> None: - super().__init__() - self.use_approximate = use_approximate - if model_args is None: - from tilert.models.glm_5._dsa_v32.model_args import ModelArgs - - model_args = ModelArgs() - self.model_args = model_args - - def golden_forward( - self, - logits: torch.Tensor, - topk: int, - ) -> torch.Tensor: - """Reference forward: torch.topk on the last dimension. - - Args: - logits: Scores tensor, shape (batch, ..., seq_len). - topk: Number of top indices to return. - - Returns: - Indices of top-k values along the last dimension. - """ - seq_len = logits.shape[-1] - return logits.topk(min(topk, seq_len), dim=-1)[1] - - def tilert_forward( - self, - logits: torch.Tensor, - topk: int, - ) -> torch.Tensor: - """Tilert forward: batch of samples with varying valid length. - - Args: - logits: Shape (batch, num_samples, cache_len). - topk: Number of top indices to return. - - Returns: - Indices tensor of shape (batch, num_samples, topk). - """ - profile_logs = get_profile_log_tensor(device=logits.device) - cache_len = logits.shape[-1] - if self.use_approximate: - indices = topk_approximate( - logits, cache_len, topk, profile_logs, model_arch=self.model_args.arch_name - ) - else: - indices = topk_accurate( - logits, cache_len, topk, profile_logs, model_arch=self.model_args.arch_name - ) - if indices.dim() == 2: - return indices.unsqueeze(0) - return indices diff --git a/tilert/models/glm_5/_dsa_v32/ops/unproj_o_allreduce.py b/tilert/models/glm_5/_dsa_v32/ops/unproj_o_allreduce.py index 257acf5..d1ed826 100644 --- a/tilert/models/glm_5/_dsa_v32/ops/unproj_o_allreduce.py +++ b/tilert/models/glm_5/_dsa_v32/ops/unproj_o_allreduce.py @@ -31,20 +31,6 @@ def unproj_o_allreduce( model_arch: str, compute_kernel_type: str = "bf16", ) -> None: - """ - Fused operation of unprojection and allreduce. - - Args: - vec_in: Input tensor. - mat_in: Input tensor. - mat_scale: Input tensor. - x_in: Input tensor. - flag: Input flag. - vec_out: Output tensor. - profile_logs: Profile logs tensor. - model_arch: Model architecture ("deepseek_v3_2" or "glm_5"). - compute_kernel_type: Compute kernel type ("bf16", "fp16mma"). - """ torch.ops.tilert.unproj_o_allreduce_op( vec_in, mat_in, @@ -62,6 +48,7 @@ class UnProjOAllReduceAlgorithm(Enum): """UnprojOAllReduce algorithm""" FP16MMA = "fp16mma" + BF16MMA = "bf16mma" @dataclass @@ -108,7 +95,6 @@ def convert_to_fp16mma_128cta( self, weights_list: list[torch.Tensor], ) -> tuple[torch.Tensor, torch.Tensor]: - """Convert weights to the packed kernel layout (GLM5 or DSV3.2).""" with torch.inference_mode(): mat, scales = weights_list if scales.dtype != torch.float32: @@ -186,6 +172,12 @@ def convert_to_fp16mma_128cta( dummy_scales = torch.zeros(1, dtype=torch.float32, device=mat.device) return mat_all, dummy_scales + def convert_to_bf16mma( + self, + weights_list: list[torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.convert_to_fp16mma(weights_list) + def convert_to_fp16mma( self, weights_list: list[torch.Tensor], @@ -254,9 +246,11 @@ class UnProjOAllReduce(TileRTModule): _SUPPORTED_ALGORITHMS = { "deepseek_v3_2": [ UnProjOAllReduceAlgorithm.FP16MMA, + UnProjOAllReduceAlgorithm.BF16MMA, ], "glm_5": [ UnProjOAllReduceAlgorithm.FP16MMA, + UnProjOAllReduceAlgorithm.BF16MMA, ], } @@ -314,24 +308,9 @@ def __init__( self.is_var_init = False def get_weights_list(self) -> list[torch.Tensor]: - """ - Get the weights list. - - Returns: - List of weights. - """ return [self.tilert_weights, self.tilert_scales] def device_sharding(self, weights_map: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """ - Device sharding. - - Args: - weights_map: Map from ref weight alias to tensor (full model). - - Returns: - Map from tilert weight alias to (num_devices, ...) tensors. - """ unproj_o_weight = weights_map[self.ref_weights_alias.o_proj_weight] unproj_o_scale = weights_map[self.ref_weights_alias.o_proj_scale_inv] @@ -393,13 +372,6 @@ def init_reference_weights( state_dict: dict[str, torch.Tensor], device_id: int | None = None, ) -> None: - """ - Initialize the reference weights. - - Args: - state_dict: State dictionary keyed by ref weight alias (full model). - device_id: Device ID for this shard; defaults to self.device_id. - """ did = self.device_id if device_id is None else device_id sharded = self.device_sharding(state_dict) weights = sharded[self.tilert_weights_alias.unproj_weights][did] @@ -407,12 +379,6 @@ def init_reference_weights( self.ref_unproj_o = weight_dequant(weights, scales) def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - """ - Initialize the tilert weights. - - Args: - state_dict: State dictionary keyed by tilert weight alias (per-device). - """ assert self.algorithm is not None, "Algorithm is not set" self.tilert_weights, self.tilert_scales = UnProjOAllReduceWeightsConverter( self.model_args, self.num_devices @@ -422,13 +388,6 @@ def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: ) def init_tilert_vars(self, batch_size: int, seq_len: int) -> None: - """ - Initialize the tilert variables. - - Args: - batch_size: Batch size. - seq_len: Sequence length. - """ self.hidden_out = torch.zeros( (batch_size, seq_len, self.dim), dtype=torch.bfloat16, @@ -438,7 +397,6 @@ def init_tilert_vars(self, batch_size: int, seq_len: int) -> None: self.is_var_init = True def init_random_weights(self) -> None: - """Initialize the random weights.""" unproj_o_weights = torch.randn( self.dim, self.n_heads * self.head_dim, @@ -469,15 +427,6 @@ def golden_forward( self, vec_in: torch.Tensor, ) -> torch.Tensor: - """ - Forward pass for the down-project module. - - Args: - vec_in: Input vector. - - Returns: - Output tensor. - """ assert self.ref_unproj_o is not None bsz = vec_in.shape[0] seq_len = vec_in.shape[1] diff --git a/tilert/models/glm_5/generator.py b/tilert/models/glm_5/generator.py index b3e8ddd..18c422a 100644 --- a/tilert/models/glm_5/generator.py +++ b/tilert/models/glm_5/generator.py @@ -4,13 +4,12 @@ import time import torch -from transformers import AutoTokenizer +from transformers import AutoTokenizer, PreTrainedTokenizerFast from tilert import logger -from tilert.models.glm_5._dsa_v32.generator import stats_time from tilert.models.glm_5._dsa_v32.model_args import ModelArgs -from tilert.models.glm_5._dsa_v32.modules.end2end import ShowHandsDSALayer -from tilert.models.glm_5._dsa_v32.temp_var_indices import Idx +from tilert.models.glm_5.modules.end2end import ShowHandsDSALayer +from tilert.models.glm_5.temp_var_indices import Idx from tilert.tilert_init import tilert_init __all__ = [ @@ -34,18 +33,7 @@ def __init__( enable_thinking: bool = False, sampling_seed: int = 42, ): - """Initialize the ShowHandsGeneratorGlm5. - - Args: - max_new_tokens: Maximum number of new tokens to generate. Defaults to 100. - temperature: Temperature for sampling. Defaults to 1.0. - model_weights_dir: Path of the model weights directory. - with_mtp: Whether to use MTP (Multi-Token Prediction) for speculative decoding. - top_p: Top-p (nucleus) sampling threshold. Defaults to 0.9. - top_k: Top-k sampling threshold. Defaults to 256. - use_topp: Whether to use top-p sampling. Defaults to False (top-1 argmax). - enable_thinking: Whether to enable thinking mode in chat template. - """ + """Initialize the ShowHandsGeneratorGlm5.""" torch.set_num_threads(64) self.model_weights_dir = model_weights_dir @@ -56,9 +44,14 @@ def __init__( self.sampling_seed = sampling_seed self.config = model_args - self.tokenizer = AutoTokenizer.from_pretrained( - self.model_weights_dir, trust_remote_code=True - ) # nosec B615 + try: + self.tokenizer = AutoTokenizer.from_pretrained( + self.model_weights_dir, trust_remote_code=True + ) # nosec B615 + except (ValueError, KeyError): + self.tokenizer = PreTrainedTokenizerFast.from_pretrained( + self.model_weights_dir, trust_remote_code=True + ) # nosec B615 jinja_file_path = os.path.join(self.model_weights_dir, "chat_template.jinja") with open(jinja_file_path, encoding="utf-8") as f: chat_template = f.read() @@ -89,6 +82,7 @@ def __init__( model_args=self.config, model_path=self.model_weights_dir, with_mtp=with_mtp, + temperature=temperature, top_p=top_p, top_k=top_k, use_topp=use_topp, @@ -111,12 +105,8 @@ def from_pretrained(self) -> None: self.decode_layer.from_pretrained(self.model_weights_dir) def extract_ffn_cache(self) -> tuple[dict[int, list], dict[int, set[str]]]: - """Extract MOE/MLP op objects and skip keys from current loaded weights. - - Returns: - Tuple of (cached_ffn_ops_per_device, skip_keys_per_device). - """ - from tilert.models.glm_5._dsa_v32.modules.end2end import ( + """Extract MOE/MLP op objects and skip keys from current loaded weights.""" + from tilert.models.glm_5.modules.end2end import ( _extract_ffn_ops, _get_moe_weight_keys, ) @@ -162,20 +152,7 @@ def generate( with_mtp: bool | None = None, prompt_tokens: list[int] | None = None, ) -> tuple[str, list[float], list[int], int]: - """Main function to load the model and perform single sequence generation. - - Args: - prompt: The input prompt string. - print_log: Whether to print generation logs. - with_mtp: Override MTP mode for this call. None uses self.with_mtp. - Requires MTP weights to have been loaded (self.with_mtp=True). - prompt_tokens: Pre-tokenized prompt tokens. If provided, skip tokenization - and use these tokens directly (useful for exact-length benchmarking). - - Returns: - Tuple of (result_text, time_list, accepted_counts, prompt_len). - accepted_counts is empty for non-MTP mode. - """ + """Main function to load the model and perform single sequence generation.""" active_mtp = with_mtp if with_mtp is not None else self.with_mtp if active_mtp and not self.with_mtp: raise ValueError("Cannot use MTP mode: MTP weights were not loaded") @@ -183,7 +160,7 @@ def generate( if active_mtp: return self._generate_with_mtp(prompt, print_log, prompt_tokens=prompt_tokens) result, time_list, prompt_len = self._generate_without_mtp( - prompt, print_log, with_mtp=active_mtp, prompt_tokens=prompt_tokens + prompt, print_log, prompt_tokens=prompt_tokens ) return result, time_list, [], prompt_len @@ -191,10 +168,9 @@ def _generate_without_mtp( self, prompt: str, print_log: bool = True, - with_mtp: bool = False, prompt_tokens: list[int] | None = None, ) -> tuple[str, list[float], int]: - """Standard generation without MTP.""" + """Standard generation without MTP (unified single-op decode).""" if prompt_tokens is None: messages = [{"role": "user", "content": prompt}] prompt_tokens = self.tokenizer.apply_chat_template( @@ -202,6 +178,7 @@ def _generate_without_mtp( tokenize=True, add_generation_prompt=True, enable_thinking=self.enable_thinking, + return_dict=False, ) max_seq_len = self.config.max_seq_len @@ -214,50 +191,65 @@ def _generate_without_mtp( tokens[0, :prompt_len] = torch.tensor( prompt_tokens, dtype=torch.long, device=self.default_device ) - prompt_mask = tokens != -1 - prev_pos = 0 - finished = torch.tensor( - [False] * self.batch_size, dtype=torch.bool, device=self.default_device - ) + ar_steps = max(1, min(1024, int(os.environ.get("GLM5_AR_N", "8")))) + return self._decode_without_mtp_ar(tokens, prompt_len, total_len, ar_steps, print_log) - time_list = [] - for cur_pos_val in range(1, total_len): - start_time = time.time() - multi_devices_results = self.decode_layer.forward( - tokens[0, prev_pos], with_mtp=with_mtp - ) - end_time = time.time() - time_list.append(end_time - start_time) + def _decode_without_mtp_ar( + self, + tokens: torch.Tensor, + prompt_len: int, + total_len: int, + ar_steps: int, + print_log: bool, + ) -> tuple[str, list[float], int]: + """w/o-MTP decode (unified single-op, unified-style).""" + time_list: list[float] = [] - intermediates, *_ = multi_devices_results[0] - next_token = intermediates[Idx.TOKEN_OUT][0][0] + self.decode_layer.set_prefill_valid_tokens(1, with_mtp=False) + for prev_pos in range(prompt_len - 1): + self.decode_layer.forward(tokens[0, prev_pos], with_mtp=False) + self.decode_layer.set_prefill_valid_tokens(0, with_mtp=False) - next_token = torch.where( - prompt_mask[0, cur_pos_val], tokens[0, cur_pos_val], next_token - ) - tokens[0, cur_pos_val] = next_token - is_stop_token = next_token.item() in self.stop_token_ids - finished |= torch.logical_and( - ~prompt_mask[0, cur_pos_val], - torch.tensor(is_stop_token, dtype=torch.bool, device=self.default_device), - ) - prev_pos = cur_pos_val - if cur_pos_val >= prompt_len: - decoded_tokens = self.tokenizer.decode( - [next_token.item()], skip_special_tokens=True - ) - if print_log: - print(decoded_tokens, end="", flush=True) + cur_pos = prompt_len - 1 + prev_token = tokens[0, prompt_len - 1].reshape(1).to(torch.int32) + finished = False + while cur_pos < total_len - 1 and not finished: + start_time = time.time() + self.decode_layer.show_hands_no_mtp(prev_token, ar_steps) + elapsed = time.time() - start_time - if finished.all(): - break + acc = self.decode_layer.ar_accepted_tokens_no_mtp(0).cpu() + n_tokens = int(acc[0].item()) + emitted = acc[1 : 1 + n_tokens].tolist() + per_step_time = elapsed / max(1, n_tokens) + + last_tok = int(prev_token[0].item()) + for tok in emitted: + if cur_pos + 1 >= total_len: + break + tokens[0, cur_pos + 1] = tok + cur_pos += 1 + last_tok = tok + if cur_pos >= prompt_len and print_log: + print( + self.tokenizer.decode([tok], skip_special_tokens=True), + end="", + flush=True, + ) + time_list.append(per_step_time) + if tok in self.stop_token_ids: + finished = True + break + prev_token = torch.tensor([last_tok], dtype=torch.int32, device=self.default_device) if print_log: print("\n") logger.info(f"--Number of tokens generated: {len(time_list)}") - - stats_time(time_list, "==== Performance ====") + if time_list: + total_t = sum(time_list) + tps = len(time_list) / total_t if total_t > 0 else 0 + logger.info(f"--Effective TPS (AR, ar_steps={ar_steps}): {tps:.2f} tokens/s") print("\n") self.decode_layer.reset_sequence() @@ -267,14 +259,12 @@ def _generate_without_mtp( toks = toks[prompt_len : prompt_len + self.max_new_tokens] stop_idx = len(toks) for i, tok in enumerate(toks): - if tok in self.stop_token_ids: + if tok == -1 or tok in self.stop_token_ids: stop_idx = i break toks = toks[:stop_idx] completion_tokens.append(toks) - decoded_tokens = self.tokenizer.batch_decode(completion_tokens, skip_special_tokens=True) - return f"{decoded_tokens[0]}\n" if decoded_tokens else "", time_list, prompt_len def _generate_with_mtp( @@ -287,8 +277,10 @@ def _generate_with_mtp( if prompt_tokens is None: prompt_tokens = self.tokenizer.apply_chat_template( [{"role": "user", "content": prompt}], + tokenize=True, add_generation_prompt=True, enable_thinking=self.enable_thinking, + return_dict=False, ) max_seq_len = self.config.max_seq_len @@ -303,8 +295,6 @@ def _generate_with_mtp( ) prefill_time_list = [] - decode_time_list = [] - decode_accepted_counts = [] cur_pos = 0 while cur_pos < prompt_len - 1: @@ -345,72 +335,82 @@ def _generate_with_mtp( self.decode_layer.set_prefill_valid_tokens(0) + ar_steps = max(1, min(1024, int(os.environ.get("GLM5_AR_N", "8")))) + return self._decode_ar(tokens, cur_pos, prompt_len, total_len, ar_steps, print_log) + + def _decode_ar( + self, + tokens: torch.Tensor, + cur_pos: int, + prompt_len: int, + total_len: int, + ar_steps: int, + print_log: bool, + ) -> tuple[str, list[float], list[int], int]: + """MTP decode (unified single-op, unified-style).""" + decode_time_list: list[float] = [] + decode_accepted_counts: list[int] = [] + + last_token = tokens[0, prompt_len - 1].item() + prev_draft = torch.full( + (1, self.mtp_seq_len), + last_token, + dtype=torch.int32, + device=self.default_device, + ) + finished = False while cur_pos < total_len - 1 and not finished: - if cur_pos == prompt_len - 1: - last_token = tokens[0, prompt_len - 1].item() - draft_tokens = torch.full( - (self.mtp_seq_len,), - last_token, - dtype=torch.long, - device=self.default_device, - ) - draft_tokens = draft_tokens.reshape(1, self.mtp_seq_len).to(torch.int32) - else: - draft_tokens = self.decode_layer.get_next_draft_tokens(0).reshape( - 1, self.mtp_seq_len - ) - start_time = time.time() - self.decode_layer.forward(draft_tokens, with_mtp=True) - end_time = time.time() - decode_time_list.append(end_time - start_time) - - num_accepted = self.decode_layer.get_num_accepted(0) - predicted_tokens = self.decode_layer.get_predicted_tokens(0).flatten() - decode_accepted_counts.append(num_accepted) + self.decode_layer.show_hands(prev_draft, ar_steps) + elapsed = time.time() - start_time + + acc = self.decode_layer.ar_accepted_tokens(0).cpu() + num = self.decode_layer.ar_num_accepted(0).cpu() + n_tokens = int(acc[0].item()) + n_steps = int(num[0].item()) + emitted = acc[1 : 1 + n_tokens].tolist() + per_step = num[1 : 1 + n_steps].tolist() + next_prev_draft = self.decode_layer.get_next_draft_tokens(0).reshape( + 1, self.mtp_seq_len + ) - num_output_tokens = num_accepted - for i in range(num_output_tokens): - if cur_pos + 1 + i >= total_len: + per_step_time = elapsed / max(1, len(per_step)) + offset = 0 + for na in per_step: + step_emit = emitted[offset : offset + na] + offset += na + for tok in step_emit: + if cur_pos + 1 >= total_len: + break + tokens[0, cur_pos + 1] = tok + cur_pos += 1 + if cur_pos >= prompt_len and print_log: + print( + self.tokenizer.decode([tok], skip_special_tokens=True), + end="", + flush=True, + ) + if tok in self.stop_token_ids: + finished = True + break + decode_time_list.append(per_step_time) + decode_accepted_counts.append(na) + if finished or cur_pos >= total_len - 1: break - new_token = int(predicted_tokens[i].item()) - tokens[0, cur_pos + 1 + i] = new_token - - if cur_pos + 1 + i >= prompt_len and print_log: - decoded_text = self.tokenizer.decode([new_token], skip_special_tokens=True) - print(decoded_text, end="", flush=True) - - if new_token in self.stop_token_ids: - finished = True - break - - cur_pos += num_accepted + prev_draft = next_prev_draft if print_log: print("\n") total_tokens = sum(decode_accepted_counts) logger.info(f"--Number of forward calls (decode): {len(decode_accepted_counts)}") logger.info(f"--Total tokens generated: {total_tokens}") - if len(decode_accepted_counts) > 0: - avg_accepted = sum(decode_accepted_counts) / len(decode_accepted_counts) - min_accepted = min(decode_accepted_counts) - max_accepted = max(decode_accepted_counts) - logger.info( - f"--Accepted tokens per call: mean={avg_accepted:.2f}, " - f"min={min_accepted}, max={max_accepted}" - ) - if decode_time_list: total_decode_time = sum(decode_time_list) effective_tps = total_tokens / total_decode_time if total_decode_time > 0 else 0 - avg_time_ms = total_decode_time / len(decode_time_list) * 1000 logger.info( - f"--Avg forward time: {avg_time_ms:.2f}ms, " - + f"({1000 / avg_time_ms:.2f} forwards/s)" + f"--Effective TPS (AR, ar_steps={ar_steps}): {effective_tps:.2f} tokens/s" ) - logger.info(f"--Effective TPS (with MTP): {effective_tps:.2f} tokens/s") - print("\n") self.decode_layer.reset_sequence() @@ -428,7 +428,6 @@ def _generate_with_mtp( completion_tokens.append(toks) decoded_tokens = self.tokenizer.batch_decode(completion_tokens, skip_special_tokens=True) - return ( f"{decoded_tokens[0]}\n" if decoded_tokens else "", decode_time_list, @@ -442,32 +441,7 @@ def inject_cache( start_pos: int = 0, end_pos: int | None = None, ) -> None: - """Inject external cache data into TileRT for P/D separation. - - This API allows injecting pre-computed KI/KV/PE cache data from an external - prefill system (e.g., SGLang), enabling prefill-decode disaggregation. - - Args: - layer_caches: List of (ki, kv, pe) tuples for each layer (0 to NUM_LAYERS-1). - Each tensor should be BF16 with shape [seqlen, dim] where: - - ki: [seqlen, 128] - compressed key (index_head_dim) - - kv: [seqlen, 512] - compressed key-value (kv_lora_rank) - - pe: [seqlen, 64] - position encoding cache (qk_rope_head_dim) - start_pos: Start position in cache to write (0-indexed). Defaults to 0. - end_pos: End position in cache (exclusive). If None, uses seqlen from tensors. - - Example: - >>> # Load cache from external prefill system - >>> layer_caches = [] # List of 78 (ki, kv, pe) tuples for GLM-5 - >>> for layer_id in range(78): - ... ki = load_ki_for_layer(layer_id) # [seqlen, 128] bf16 - ... kv = load_kv_for_layer(layer_id) # [seqlen, 512] bf16 - ... pe = load_pe_for_layer(layer_id) # [seqlen, 64] bf16 - ... layer_caches.append((ki, kv, pe)) - >>> generator.inject_cache(layer_caches, start_pos=0) - >>> generator.set_cur_pos(seqlen) # Set RoPE position - >>> # Continue generation from cache - """ + """Inject external cache data into TileRT for P/D separation.""" num_layers = len(layer_caches) if num_layers == 0: logger.warning("inject_cache called with empty layer_caches") @@ -497,53 +471,31 @@ def inject_cache( kv_src = kv[:cache_len].to(f"cuda:{device_id}") pe_src = pe[:cache_len].to(f"cuda:{device_id}") - caches[base_idx + 0][0, start_pos:end_pos, :].copy_(ki_src) - caches[base_idx + 1][0, start_pos:end_pos, :].copy_(kv_src) - caches[base_idx + 2][0, start_pos:end_pos, :].copy_(pe_src) + for _off, _src in ((0, ki_src), (1, kv_src), (2, pe_src)): + _dst = caches[base_idx + _off] + if _dst.size(1) < end_pos: + continue + _dst[0, start_pos:end_pos, :].copy_(_src) + + torch.cuda.synchronize(device_id) logger.info(f"Cache injection completed for {num_devices} devices") def set_cur_pos(self, cur_pos: int) -> None: - """Set the current position for RoPE. - - This should be called after inject_cache() to ensure the runtime position - matches the injected cache length, for correct RoPE position encoding - during continued generation. - - Args: - cur_pos: The current sequence position (typically the length of prefilled tokens). - - Example: - >>> generator.inject_cache(layer_caches, start_pos=0) - >>> generator.set_cur_pos(prefill_len) # Set position to prefill length - >>> # Now generate continues from the correct position - """ + """Set the current position for RoPE in C++ backend.""" if self.with_mtp: num_devices = self.decode_layer.num_devices for device_id in range(num_devices): intermediates, _, _, _ = self.decode_layer._get_device_result(device_id) cur_pos_tensor = intermediates[Idx.CUR_POS] cur_pos_tensor.fill_(cur_pos) + torch.cuda.synchronize(device_id) else: torch.ops.tilert.dsa_show_hands_set_cur_pos_glm5(cur_pos) logger.info(f"Set cur_pos to {cur_pos}") def inject_last_hidden_state(self, last_hidden_state: torch.Tensor) -> None: - """Inject the last hidden state for MTP mode. - - For MTP (Multi-Token Prediction), the MTP preprocess layer needs the - last hidden state from the main model's last token. - - Args: - last_hidden_state: [hidden_size] or [1, hidden_size] BF16 tensor. - The hidden state of the last token from prefill. - - Example: - >>> # After inject_cache, inject the last hidden state for MTP - >>> generator.inject_last_hidden_state(last_hidden_state) - >>> generator.set_cur_pos(prefill_len) - >>> # Then start generation - """ + """Inject the last hidden state for MTP mode.""" if not self.with_mtp: logger.warning("inject_last_hidden_state called but with_mtp is False, skipping") return @@ -557,5 +509,6 @@ def inject_last_hidden_state(self, last_hidden_state: torch.Tensor) -> None: lhs_tensor = intermediates[Idx.LAST_HIDDEN_STATES] lhs_src = last_hidden_state.to(f"cuda:{device_id}") lhs_tensor[0, 0, :].copy_(lhs_src.squeeze(0)) + torch.cuda.synchronize(device_id) logger.info(f"Injected last_hidden_state to {num_devices} devices") diff --git a/tilert/models/glm_5/model_args.py b/tilert/models/glm_5/model_args.py index 74e830c..b222729 100644 --- a/tilert/models/glm_5/model_args.py +++ b/tilert/models/glm_5/model_args.py @@ -12,43 +12,7 @@ @dataclass class ModelArgsGLM5(ModelArgs): - """ - Data class for defining model arguments and hyperparameters. - - Attributes: - arch_name (str): Architecture name. - max_batch_size (int): Maximum batch size. - max_seq_len (int): Maximum sequence length. - dtype (Literal["bf16", "fp8"]): Data type for computations. - scale_fmt (Optional[str]): Format for quantization scale. - vocab_size (int): Vocabulary size. - dim (int): Model dimension. - inter_dim (int): Intermediate dimension for MLP layers. - moe_inter_dim (int): Intermediate dimension for MoE layers. - n_layers (int): Number of transformer layers. - n_dense_layers (int): Number of dense layers in the model. - n_heads (int): Number of attention heads. - n_routed_experts (int): Number of routed experts for MoE layers. - n_shared_experts (int): Number of shared experts for MoE layers. - n_activated_experts (int): Number of activated experts in MoE layers. - n_expert_groups (int): Number of expert groups. - n_limited_groups (int): Number of limited groups for MoE routing. - score_func (Literal["softmax", "sigmoid"]): Scoring function for MoE routing. - route_scale (float): Scaling factor for routing scores. - q_lora_rank (int): LoRA rank for query projections. - kv_lora_rank (int): LoRA rank for key-value projections. - qk_nope_head_dim (int): Dimension for query-key projections without positional embeddings. - qk_rope_head_dim (int): Dimension for query-key projections with rotary embeddings. - v_head_dim (int): Dimension for value projections. - original_seq_len (Optional[int]): Original sequence length. - rope_theta (float): Base for rotary positional encoding. - rope_factor (Optional[float]): Scaling factor for extended sequence lengths. - beta_fast (Optional[int]): Fast beta correction factor. - beta_slow (Optional[int]): Slow beta correction factor. - mscale (float): Scaling factor for extended attention. - index_head_dim (int): Dimension for index head. - index_topk (int): Top-k for index head. - """ + """Data class for defining model arguments and hyperparameters.""" arch_name = "glm_5" @@ -68,7 +32,9 @@ class ModelArgsGLM5(ModelArgs): n_routed_experts: int = 256 n_shared_experts: int = 1 n_activated_experts: int = 8 - score_func: Literal["softmax", "sigmoid"] = "softmax" + n_expert_groups: int = 1 + n_limited_groups: int = 1 + score_func: Literal["softmax", "sigmoid"] = "sigmoid" route_scale: float = 2.5 q_lora_rank: int = 2048 diff --git a/tilert/models/glm_5/modules/__init__.py b/tilert/models/glm_5/modules/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tilert/models/glm_5/_dsa_v32/modules/dsa.py b/tilert/models/glm_5/modules/dsa.py similarity index 90% rename from tilert/models/glm_5/_dsa_v32/modules/dsa.py rename to tilert/models/glm_5/modules/dsa.py index 38a01c1..95c01ce 100644 --- a/tilert/models/glm_5/_dsa_v32/modules/dsa.py +++ b/tilert/models/glm_5/modules/dsa.py @@ -4,10 +4,10 @@ from tilert.models.base import SerializableTileRTModule from tilert.models.glm_5._dsa_v32.model_args import ModelArgs -from tilert.models.glm_5._dsa_v32.modules.mlp import MlpBlock -from tilert.models.glm_5._dsa_v32.modules.moe import MoeBlock from tilert.models.glm_5._dsa_v32.ops import RMSNormHeadProj -from tilert.models.glm_5._dsa_v32.temp_var_indices import TEMP_VARS_SIZE, Idx +from tilert.models.glm_5.modules.mlp import MlpBlock +from tilert.models.glm_5.modules.moe import MoeBlock +from tilert.models.glm_5.temp_var_indices import TEMP_VARS_SIZE, Idx class Dsa(SerializableTileRTModule): @@ -26,7 +26,7 @@ def __init__( num_devices=num_devices, remove_selected=True, ) - from tilert.models.glm_5._dsa_v32.modules.mla_v2 import ( + from tilert.models.glm_5.modules.mla_v2 import ( PureMlaV2, SparseSelectMlaV2, ) @@ -39,17 +39,17 @@ def __init__( if device_id == 0: self.v2_peer_bufs = torch.zeros(n_peers, dtype=torch.int64, device=dev) self.v2_partial_buf = torch.zeros( - model_args.max_batch_size, 4, model_args.dim, dtype=torch.bfloat16, device=dev + model_args.max_batch_size, 8, model_args.dim, dtype=torch.bfloat16, device=dev ) mla_kwargs = { "peer_bufs": self.v2_peer_bufs, "partial_buf": self.v2_partial_buf, } else: - max_seq_len = getattr(model_args, "num_mtp", 3) + 1 + max_seq_len = max(getattr(model_args, "num_mtp", 3) + 1, 8) topk = model_args.index_topk - self.v2_ll_buf = torch.zeros(max_seq_len * topk * 2, dtype=torch.int32, device=dev) - mla_kwargs = {"ll_buf": self.v2_ll_buf} + self.v2_recv_buf = torch.zeros(max_seq_len * topk * 2, dtype=torch.int32, device=dev) + mla_kwargs = {"recv_buf": self.v2_recv_buf} mla_num_devices: int | None = None if device_id != 0: @@ -139,7 +139,6 @@ def get_temp_vars( n_index_heads = self.model_args.index_n_heads max_seq_len = self.model_args.max_seq_len index_topk = self.model_args.index_topk - n_routed_experts = self.model_args.n_routed_experts n_activated_experts = self.model_args.n_activated_experts n_total_experts = self.model_args.n_activated_experts + self.model_args.n_shared_experts moe_inter_dim = self.model_args.moe_inter_dim // self.num_devices @@ -168,7 +167,10 @@ def get_temp_vars( temp_vars[Idx.O_LSE_ACC] = torch.empty(*batch_seq, n_local_heads, 32, **fp32_desc) temp_vars[Idx.PROJ_O] = torch.zeros(*batch_seq, n_local_heads, v_head_dim, **bf16_desc) temp_vars[Idx.UNPROJ_O] = torch.zeros(*batch_seq, dim, **bf16_desc) - temp_vars[Idx.SCORES] = torch.zeros(*batch_seq, n_routed_experts, **fp32_desc) + temp_vars[Idx.SCORES] = torch.full((4096,), float("nan"), **fp32_desc) + temp_vars[Idx.HIDDEN_MID] = torch.full( + (1, 8, n_activated_experts + 1, 256), float("nan"), **bf16_desc + ) temp_vars[Idx.X_MLP_IN] = torch.zeros(*batch_seq, dim, **bf16_desc) exp_up_gate = torch.zeros(*batch_seq, n_total_experts, moe_inter_dim, **bf16_desc) temp_vars[Idx.UP_GATE] = exp_up_gate @@ -222,6 +224,14 @@ def get_temp_vars( temp_vars[Idx.TOP_N_INDICES] = torch.zeros(*batch_seq, max_top_n, **int32_desc) temp_vars[Idx.LOGPROBS_FLAG] = torch.zeros(1, **int32_desc) + ar_max_steps = 1024 + temp_vars[Idx.AR_ACCEPTED_TOKENS] = torch.zeros(1 + ar_max_steps * seq_len, **int32_desc) + temp_vars[Idx.AR_NUM_ACCEPTED] = torch.zeros(1 + ar_max_steps, **int32_desc) + + temp_vars[Idx.GRAMMAR_BITMASK] = torch.full( + (*batch_seq, vocab_size // 32), -1, **int32_desc + ) + for i, t in enumerate(temp_vars): if t is None: raise RuntimeError(f"temp_vars[{i}] ({Idx(i).name}) was not initialized") diff --git a/tilert/models/glm_5/_dsa_v32/modules/end2end.py b/tilert/models/glm_5/modules/end2end.py similarity index 74% rename from tilert/models/glm_5/_dsa_v32/modules/end2end.py rename to tilert/models/glm_5/modules/end2end.py index 6b4e69c..f3d8375 100644 --- a/tilert/models/glm_5/_dsa_v32/modules/end2end.py +++ b/tilert/models/glm_5/modules/end2end.py @@ -9,14 +9,13 @@ import torch from safetensors import safe_open -from safetensors.torch import load_file from tilert import logger from tilert.models.base import TileRTModule from tilert.models.glm_5._dsa_v32.model_args import ModelArgs -from tilert.models.glm_5._dsa_v32.modules.dsa import Dsa -from tilert.models.glm_5._dsa_v32.modules.mtp import MTP -from tilert.models.glm_5._dsa_v32.temp_var_indices import Idx, validate_temp_vars_layout +from tilert.models.glm_5.modules.dsa import Dsa +from tilert.models.glm_5.modules.mtp import MTP +from tilert.models.glm_5.temp_var_indices import Idx, validate_temp_vars_layout from tilert.models.utils import precompute_freqs_cis from tilert.utils import get_profile_log_tensor @@ -26,6 +25,32 @@ DeviceResult = tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], torch.Tensor] +def _load_state_dicts_by_index( + model_path: str, + weight_file_map: dict[str, str], + weights_list: list[str], + device: str, + selective_only: bool = False, +) -> dict[str, torch.Tensor]: + """Load tensors treating the index (``weight_file_map``) as the per-key authority.""" + target_files = {weight_file_map[key] for key in weights_list} + weights_set = set(weights_list) + state_dicts: dict[str, torch.Tensor] = {} + for weight_file in sorted(target_files): + filepath = os.path.join(model_path, weight_file) + logger.info(f"Loading weights from {weight_file} to {device}") + with safe_open(filepath, framework="pt", device=device) as f: + for key in f.keys(): + if selective_only and key not in weights_set: + continue + if weight_file_map.get(key, weight_file) != weight_file: + continue + state_dicts[key] = f.get_tensor(key) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return state_dicts + + def _mark_weights_initialized(module: TileRTModule) -> None: """Recursively mark a module and all sub-ops as having initialized tilert weights.""" module.is_tilert_weights_init = True @@ -35,12 +60,9 @@ def _mark_weights_initialized(module: TileRTModule) -> None: def _extract_ffn_ops(dsa: "Dsa") -> list: - """Extract Moe/Mlp op objects from a Dsa's layer blocks. - - Returns a list of length n_layers where each element is a Moe or Mlp instance. - """ - from tilert.models.glm_5._dsa_v32.modules.mlp import MlpBlock - from tilert.models.glm_5._dsa_v32.modules.moe import MoeBlock + """Extract Moe/Mlp op objects from a Dsa's layer blocks.""" + from tilert.models.glm_5.modules.mlp import MlpBlock + from tilert.models.glm_5.modules.moe import MoeBlock ffn_ops = [] for block in dsa.exec_seq: @@ -61,8 +83,8 @@ def _extract_ffn_ops(dsa: "Dsa") -> list: def _get_moe_weight_keys(dsa: "Dsa") -> set[str]: """Get state_dict keys that belong exclusively to MOE/MLP ops in this Dsa.""" - from tilert.models.glm_5._dsa_v32.modules.mlp import MlpBlock - from tilert.models.glm_5._dsa_v32.modules.moe import MoeBlock + from tilert.models.glm_5.modules.mlp import MlpBlock + from tilert.models.glm_5.modules.moe import MoeBlock moe_keys: set[str] = set() mla_keys: set[str] = set() @@ -76,6 +98,11 @@ def _get_moe_weight_keys(dsa: "Dsa") -> set[str]: return moe_keys - mla_keys +def _glm5_suffix(is_glm5: "bool | str" = True) -> str: # noqa: U100 + """GLM5-native tree: the torch.ops name suffix is always ``_glm5``.""" + return "_glm5" + + def dsa_show_hands_prepare_money( params: list[torch.Tensor], temp_vars: list[torch.Tensor], @@ -83,11 +110,11 @@ def dsa_show_hands_prepare_money( profile_logs: torch.Tensor, forward_max_seq_len: int, with_mtp: bool = False, - is_glm5: bool = False, + is_glm5: "bool | str" = False, ) -> Any: """Prepare money for show hands""" mtp_flag = "_mtp_e2e" if with_mtp else "" - glm5_flag = "_glm5" if is_glm5 else "" + glm5_flag = _glm5_suffix(is_glm5) func_name = f"dsa{mtp_flag}_show_hands_prepare_money{glm5_flag}" if mtp_flag: return getattr(torch.ops.tilert, func_name)(params, temp_vars, cache_vars, profile_logs) @@ -96,69 +123,87 @@ def dsa_show_hands_prepare_money( ) -def dsa_show_hands(token_id: torch.Tensor, with_mtp: bool = False, is_glm5: bool = False) -> Any: - """Show hands with native MT""" +def dsa_show_hands( + token_id: torch.Tensor, + with_mtp: bool = False, + is_glm5: "bool | str" = False, + ar_steps: int = 1, +) -> Any: + """Show hands with native MT.""" mtp_flag = "_mtp_e2e" if with_mtp else "" - glm5_flag = "_glm5" if is_glm5 else "" - func_name = f"dsa{mtp_flag}_show_hands{glm5_flag}" - return getattr(torch.ops.tilert, func_name)(token_id) + glm5_flag = _glm5_suffix(is_glm5) + op = getattr(torch.ops.tilert, f"dsa{mtp_flag}_show_hands{glm5_flag}") + if is_glm5: + return op(token_id, int(ar_steps)) + return op(token_id) + + +def dsa_show_hands_accepted_tokens(dev: int, is_glm5: "bool | str" = True) -> torch.Tensor: + """Read w/o-MTP AR accepted-tokens flat buffer ([0]=count, [1:]=token stream).""" + glm5_flag = _glm5_suffix(is_glm5) + return getattr(torch.ops.tilert, f"dsa_show_hands_accepted_tokens{glm5_flag}")(dev) + + +def dsa_show_hands_num_accepted(dev: int, is_glm5: "bool | str" = True) -> torch.Tensor: + """Read w/o-MTP AR per-step num_accepted flat buffer ([0]=steps, [1:]=counts).""" + glm5_flag = _glm5_suffix(is_glm5) + return getattr(torch.ops.tilert, f"dsa_show_hands_num_accepted{glm5_flag}")(dev) + +def dsa_mtp_e2e_accepted_tokens(dev: int, is_glm5: "bool | str" = True) -> torch.Tensor: + """Read the AR accepted-tokens flat buffer ([0]=count, [1:]=token stream).""" + glm5_flag = _glm5_suffix(is_glm5) + return getattr(torch.ops.tilert, f"dsa_mtp_e2e_accepted_tokens{glm5_flag}")(dev) -def dsa_show_hands_reset(with_mtp: bool = False, is_glm5: bool = False) -> Any: + +def dsa_mtp_e2e_num_accepted(dev: int, is_glm5: "bool | str" = True) -> torch.Tensor: + """Read the AR per-step num_accepted flat buffer ([0]=steps, [1:]=counts).""" + glm5_flag = _glm5_suffix(is_glm5) + return getattr(torch.ops.tilert, f"dsa_mtp_e2e_num_accepted{glm5_flag}")(dev) + + +def dsa_show_hands_reset(with_mtp: bool = False, is_glm5: "bool | str" = False) -> Any: """Reset show one hand""" mtp_flag = "_mtp_e2e" if with_mtp else "" - glm5_flag = "_glm5" if is_glm5 else "" + glm5_flag = _glm5_suffix(is_glm5) func_name = f"dsa{mtp_flag}_show_hands_reset{glm5_flag}" return getattr(torch.ops.tilert, func_name)() -def dsa_show_hands_go_home(with_mtp: bool = False, is_glm5: bool = False) -> Any: +def dsa_show_hands_go_home(with_mtp: bool = False, is_glm5: "bool | str" = False) -> Any: """Go home""" mtp_flag = "_mtp_e2e" if with_mtp else "" - glm5_flag = "_glm5" if is_glm5 else "" + glm5_flag = _glm5_suffix(is_glm5) func_name = f"dsa{mtp_flag}_show_hands_go_home{glm5_flag}" return getattr(torch.ops.tilert, func_name)() def dsa_show_hands_set_sampling_seed( - seed: int, with_mtp: bool = False, is_glm5: bool = False + seed: int, with_mtp: bool = False, is_glm5: "bool | str" = False ) -> Any: - """Set the sampling seed (request-level, fixed for the entire request). - - Args: - seed: The sampling seed value. - """ + """Set the sampling seed (request-level, fixed for the entire request).""" mtp_flag = "_mtp_e2e" if with_mtp else "" - glm5_flag = "_glm5" if is_glm5 else "" + glm5_flag = _glm5_suffix(is_glm5) func_name = f"dsa{mtp_flag}_show_hands_set_sampling_seed{glm5_flag}" return getattr(torch.ops.tilert, func_name)(seed) def dsa_mtp_e2e_show_hands_set_prefill_valid_tokens( - num_valid_tokens: int, is_glm5: bool = False + num_valid_tokens: int, is_glm5: "bool | str" = False, with_mtp: bool = True ) -> Any: - """Set the number of valid (non-padding) tokens for prefill mode. - - This controls how many tokens are copied from draft_tokens to predicted_tokens - during prefill. Should be called before forward() when the chunk has padding. - - Args: - num_valid_tokens: Number of valid tokens in the chunk (1-4). - """ - mtp_flag = "_mtp_e2e" - glm5_flag = "_glm5" if is_glm5 else "" + """Select prefill (num_valid_tokens > 0) vs decode (0) mode.""" + mtp_flag = "_mtp_e2e" if with_mtp else "" + glm5_flag = _glm5_suffix(is_glm5) func_name = f"dsa{mtp_flag}_show_hands_set_prefill_valid_tokens{glm5_flag}" return getattr(torch.ops.tilert, func_name)(num_valid_tokens) -def dsa_mtp_e2e_show_hands_set_prefill_mtp_extra_token(token: int, is_glm5: bool = False) -> Any: - """Set the extra token for MTP[0] shifted input during prefill. - - Args: - token: The extra prompt token id (int32). - """ +def dsa_mtp_e2e_show_hands_set_prefill_mtp_extra_token( + token: int, is_glm5: "bool | str" = False +) -> Any: + """Set the extra token for MTP[0] shifted input during prefill.""" mtp_flag = "_mtp_e2e" - glm5_flag = "_glm5" if is_glm5 else "" + glm5_flag = _glm5_suffix(is_glm5) func_name = f"dsa{mtp_flag}_show_hands_set_prefill_mtp_extra_token{glm5_flag}" return getattr(torch.ops.tilert, func_name)(token) @@ -176,17 +221,23 @@ def __init__( top_p: float = 0.9, top_k: int = 256, use_topp: bool = False, + num_mtp: int = 3, ) -> None: validate_temp_vars_layout() print(f"Model args: {model_args.arch_name}") for k_arg, v_arg in model_args.__dict__.items(): print(f" - {k_arg}: {v_arg}") self.model_args = model_args - self.is_glm5 = self.model_args.arch_name == "glm_5" - assert self.model_args.arch_name in ["deepseek_v3_2", "glm_5"] + arch = self.model_args.arch_name + assert ( + arch == "glm_5" + ), f"glm_5-native ShowHandsDSALayer requires arch_name 'glm_5', got {arch!r}" + self.is_glm5: bool = True self.num_devices = 8 - self.forward_max_seq_len = 4 + assert num_mtp == 3, "num_mtp must be 3" + self.num_mtp = num_mtp + self.forward_max_seq_len = num_mtp + 1 self.model_path = model_path self.with_weight_conversion = with_weight_conversion @@ -222,30 +273,13 @@ def load_device_weights( if skip_keys: weights_list = [k for k in weights_list if k not in skip_keys] - target_files = set() - for weight_key in weights_list: - weight_file = weight_file_map[weight_key] - target_files.add(weight_file) - - state_dicts = {} - weights_set = set(weights_list) - for weight_file in target_files: - filepath = os.path.join(model_path, weight_file) - if skip_keys: - logger.info( - f"Selectively loading weights from {weight_file} for device {device_id}" - ) - with safe_open(filepath, framework="pt", device=f"cuda:{device_id}") as f: - for key in f.keys(): - if key in weights_set: - state_dicts[key] = f.get_tensor(key) - torch.cuda.empty_cache() - else: - logger.info(f"Loading weights from {weight_file} for device {device_id}") - state_dict = load_file(filepath, device=f"cuda:{device_id}") - state_dicts.update(state_dict) - del state_dict - torch.cuda.empty_cache() + state_dicts = _load_state_dicts_by_index( + model_path, + weight_file_map, + weights_list, + device=f"cuda:{device_id}", + selective_only=bool(skip_keys), + ) state_dicts["freqs_cis"] = self._gen_freqs_cis().to(device_id) return state_dicts @@ -339,16 +373,7 @@ def _init_weights( cached_ffn_ops_per_device: dict[int, list] | None = None, skip_keys_per_device: dict[int, set[str]] | None = None, ) -> None: - """Load the model weights from the given path or generate random weights. - - Args: - model_path: Path to the model weights directory. - cached_ffn_ops_per_device: Optional dict mapping device_id to cached FFN ops. - When provided, these ops are injected into the Dsa and their weights - are not re-loaded from disk. - skip_keys_per_device: Optional dict mapping device_id to safetensors keys - to skip during loading. Used together with cached_ffn_ops_per_device. - """ + """Load the model weights from the given path or generate random weights.""" self._v2_p2p: dict = {} def __load_weights(device_id: int, model_path: str | None) -> None: @@ -389,6 +414,7 @@ def __load_weights(device_id: int, model_path: str | None) -> None: dsa.init_tilert_weights(state_dicts) self._dsa_objects[device_id] = dsa params.extend(dsa.get_weights_list()) + torch.cuda.empty_cache() caches.extend(dsa.get_cache_vars()) if device_id == 0: @@ -397,7 +423,7 @@ def __load_weights(device_id: int, model_path: str | None) -> None: } else: self._v2_p2p[device_id] = { - "ll_buf": dsa.v2_ll_buf, + "recv_buf": dsa.v2_recv_buf, } intermediates.extend( self.generate_params_with_continuous_storage( @@ -428,12 +454,13 @@ def __load_weights(device_id: int, model_path: str | None) -> None: device=device_id, ) ) + intermediates[Idx.GRAMMAR_BITMASK].fill_(-1) base_params_count = len(params) base_caches_count = len(caches) if self.with_mtp: - from tilert.models.glm_5._dsa_v32.modules.mla_v2 import ( + from tilert.models.glm_5.modules.mla_v2 import ( PureMlaV2, SparseSelectMlaV2, ) @@ -446,7 +473,7 @@ def __load_weights(device_id: int, model_path: str | None) -> None: "peer_bufs": dsa.v2_peer_bufs, } else: - mtp_kwargs["mla_kwargs"] = {"ll_buf": dsa.v2_ll_buf} + mtp_kwargs["mla_kwargs"] = {"recv_buf": dsa.v2_recv_buf} mtp = MTP(self.model_args, device_id, self.num_devices, **mtp_kwargs) mtp.init_tilert_weights(state_dicts) params.extend(mtp.get_weights_list()) @@ -493,10 +520,10 @@ def _runner(dev_id: int) -> None: peer_bufs_cpu = torch.zeros(self.num_devices - 1, dtype=torch.int64) for i in range(self.num_devices - 1): dev_id = i + 1 - peer_bufs_cpu[i] = self._v2_p2p[dev_id]["ll_buf"].data_ptr() + peer_bufs_cpu[i] = self._v2_p2p[dev_id]["recv_buf"].data_ptr() gpu0["peer_bufs"].copy_(peer_bufs_cpu) logger.info( - "V2 P2P exchange complete: peer_bufs (ll_buf)=%s", + "V2 P2P exchange complete: peer_bufs (recv_buf)=%s", [hex(int(x)) for x in peer_bufs_cpu], ) @@ -554,18 +581,35 @@ def forward( with_mtp: bool | None = None, ) -> list[DeviceResult]: active_mtp = with_mtp if with_mtp is not None else self.with_mtp - dsa_show_hands(token_id.cpu(), active_mtp, self.is_glm5) + dsa_show_hands(token_id.cpu(), active_mtp, self.is_glm5, ar_steps=1) return [self._get_device_result(device_id) for device_id in range(self.num_devices)] - def set_sampling_seed(self, seed: int, with_mtp: bool | None = None) -> None: - """Set the sampling seed for top-p sampling. + def show_hands(self, prev_draft_tokens: torch.Tensor, ar_steps: int = 1) -> None: + """MTP decode (GLM5, unified-style single op).""" + dsa_show_hands(prev_draft_tokens.cpu(), True, self.is_glm5, ar_steps) + + def ar_accepted_tokens(self, dev: int = 0) -> torch.Tensor: + """AR accepted-tokens flat buffer ([0]=count, [1:]=token stream).""" + return dsa_mtp_e2e_accepted_tokens(dev, self.is_glm5) + + def ar_num_accepted(self, dev: int = 0) -> torch.Tensor: + """AR per-step num_accepted flat buffer ([0]=steps, [1:]=per-step counts).""" + return dsa_mtp_e2e_num_accepted(dev, self.is_glm5) + + def show_hands_no_mtp(self, prev_token: torch.Tensor, ar_steps: int = 1) -> None: + """w/o-MTP decode (GLM5, unified-style single op).""" + dsa_show_hands(prev_token.cpu(), False, self.is_glm5, ar_steps) - The seed is fixed for the entire request. Position provides per-step variation. + def ar_accepted_tokens_no_mtp(self, dev: int = 0) -> torch.Tensor: + """w/o-MTP AR accepted-tokens flat buffer ([0]=count, [1:]=token stream).""" + return dsa_show_hands_accepted_tokens(dev, self.is_glm5) - Args: - seed: The sampling seed value. - with_mtp: Override MTP mode for this call. Defaults to self.with_mtp. - """ + def ar_num_accepted_no_mtp(self, dev: int = 0) -> torch.Tensor: + """w/o-MTP AR per-step num_accepted flat buffer ([0]=steps, [1:]=counts).""" + return dsa_show_hands_num_accepted(dev, self.is_glm5) + + def set_sampling_seed(self, seed: int, with_mtp: bool | None = None) -> None: + """Set the sampling seed for top-p sampling.""" active_mtp = with_mtp if with_mtp is not None else self.with_mtp dsa_show_hands_set_sampling_seed(seed, active_mtp, self.is_glm5) @@ -595,84 +639,37 @@ def _get_device_result(self, device_id: int) -> DeviceResult: raise RuntimeError(f"Device {device_id} is not initialized") return device_result - def set_prefill_valid_tokens(self, num_valid_tokens: int) -> None: - """Set the number of valid tokens for prefill mode. - - This controls how many tokens are copied from draft_tokens to predicted_tokens - during prefill. Should be called before forward() when the chunk has padding. - - Args: - num_valid_tokens: Number of valid tokens in the chunk (1-4). - """ - dsa_mtp_e2e_show_hands_set_prefill_valid_tokens(num_valid_tokens, self.is_glm5) + def set_prefill_valid_tokens(self, num_valid_tokens: int, with_mtp: bool | None = None) -> None: + """Select prefill (num_valid_tokens > 0) vs decode (0) mode.""" + active_mtp = with_mtp if with_mtp is not None else self.with_mtp + dsa_mtp_e2e_show_hands_set_prefill_valid_tokens(num_valid_tokens, self.is_glm5, active_mtp) def set_prefill_mtp_extra_token(self, token: int) -> None: - """Set the extra token for MTP[0] shifted input during prefill. - - Args: - token: The prompt token at (cur_pos + mtp_seq_len). - """ + """Set the extra token for MTP[0] shifted input during prefill.""" dsa_mtp_e2e_show_hands_set_prefill_mtp_extra_token(token, self.is_glm5) def get_next_draft_tokens(self, device_id: int = 0) -> torch.Tensor: - """Get next_draft_tokens from the specified device. - - Args: - device_id: Device ID to get results from. - - Returns: - next_draft_tokens tensor of shape [1, MTP_SEQ_LEN]. - """ + """Get next_draft_tokens from the specified device.""" intermediates, _, _, _ = self._get_device_result(device_id) return intermediates[Idx.NEXT_DRAFT_TOKENS] def get_num_accepted(self, device_id: int = 0) -> int: - """Get number of accepted tokens from the specified device. - - Args: - device_id: Device ID to get results from. - - Returns: - Number of accepted tokens. - """ + """Get number of accepted tokens from the specified device.""" intermediates, _, _, _ = self._get_device_result(device_id) return int(intermediates[Idx.ACCEPTED_TOKENS][0].item()) def get_predicted_tokens(self, device_id: int = 0) -> torch.Tensor: - """Get predicted_tokens from the specified device. - - Args: - device_id: Device ID to get results from. - - Returns: - predicted_tokens tensor containing main model predictions. - """ + """Get predicted_tokens from the specified device.""" intermediates, _, _, _ = self._get_device_result(device_id) return intermediates[Idx.PREDICTED_TOKENS] def get_logits(self, device_id: int = 0) -> torch.Tensor: - """Get logits from the specified device. - - Args: - device_id: Device ID to get results from. - - Returns: - Logits tensor of shape [batch, seq_len, vocab_size] (FP32). - """ + """Get logits from the specified device.""" intermediates, _, _, _ = self._get_device_result(device_id) return intermediates[Idx.LOGITS_OUT] def get_top_n_logprobs(self, device_id: int = 0) -> tuple[torch.Tensor, torch.Tensor]: - """Get top-N log-probabilities and token IDs from the top_p kernel. - - Args: - device_id: Device ID to get results from. - - Returns: - Tuple of (log_probs, token_ids): - - log_probs: [batch, seq_len, 256] FP32 - - token_ids: [batch, seq_len, 256] INT32 - """ + """Get top-N log-probabilities and token IDs from the top_p kernel.""" intermediates, _, _, _ = self._get_device_result(device_id) return ( intermediates[Idx.TOP_N_LOG_PROBS], @@ -680,23 +677,12 @@ def get_top_n_logprobs(self, device_id: int = 0) -> tuple[torch.Tensor, torch.Te ) def get_token_logprob(self, device_id: int = 0) -> torch.Tensor: - """Get log-probability of the sampled token (from TOP_P_SCORES). - - Args: - device_id: Device ID to get results from. - - Returns: - Tensor of shape [batch, seq_len] (FP32). - """ + """Get log-probability of the sampled token (from TOP_P_SCORES).""" intermediates, _, _, _ = self._get_device_result(device_id) return intermediates[Idx.TOP_P_SCORES] def set_logprobs_enabled(self, enabled: bool) -> None: - """Enable or disable logprobs export in the top_p kernel. - - Args: - enabled: True to enable logprobs export, False to disable. - """ + """Enable or disable logprobs export in the top_p kernel.""" flag_val = 1 if enabled else 0 for device_id in range(self.num_devices): intermediates, _, _, _ = self._get_device_result(device_id) diff --git a/tilert/models/glm_5/_dsa_v32/modules/mla_v2.py b/tilert/models/glm_5/modules/mla_v2.py similarity index 74% rename from tilert/models/glm_5/_dsa_v32/modules/mla_v2.py rename to tilert/models/glm_5/modules/mla_v2.py index d9a9dd1..6a3c2fd 100644 --- a/tilert/models/glm_5/_dsa_v32/modules/mla_v2.py +++ b/tilert/models/glm_5/modules/mla_v2.py @@ -1,4 +1,4 @@ -"""MLA weight generator classes for device-group-specific pipelines.""" +"""V2 MLA weight generator classes for device-group-specific pipelines.""" import torch @@ -19,6 +19,7 @@ ) from tilert.models.glm_5._dsa_v32.ops.rmsnorm_projx_wqakis import ( RMSNormProjxWqakis, + RMSNormProjxWqakisAlgorithm, ) from tilert.models.glm_5._dsa_v32.ops.rmsnorm_projx_wqkva import ( RMSNormProjxWqkva, @@ -31,6 +32,8 @@ class SparseSelectMlaV2(SerializableTileRTModule): + """Device Group A (GPU 0): sparse selector MLA.""" + def __init__( self, model_args: ModelArgs, @@ -44,12 +47,13 @@ def __init__( self.rmsnorm_projx_wqakis = RMSNormProjxWqakis( model_args=model_args, device_id=device_id, num_devices=num_devices ) + self.rmsnorm_projx_wqakis.algorithm = RMSNormProjxWqakisAlgorithm.W8A16HMMA self.register_op(self.rmsnorm_projx_wqakis) self.rmsnorm_projq_wqi = RmsnormProjqWqi( model_args=model_args, device_id=device_id, num_devices=num_devices ) - self.rmsnorm_projq_wqi.algorithm = RmsnormProjqWqiAlgorithm.FP16MMA + self.rmsnorm_projq_wqi.algorithm = RmsnormProjqWqiAlgorithm.BF16MMA self.register_op(self.rmsnorm_projq_wqi) self.layernorm_rope_rotate = LayerNormRoPERotate( @@ -58,7 +62,10 @@ def __init__( self.register_op(self.layernorm_rope_rotate) self.projx_wis = ProjxWis( - model_args=model_args, device_id=device_id, num_devices=num_devices + model_args=model_args, + device_id=device_id, + num_devices=num_devices, + compute_kernel_type="bf16mma", ) self.register_op(self.projx_wis) @@ -70,7 +77,7 @@ def __init__( self.pe_cache: torch.Tensor | None = None def get_weights_list(self) -> list[torch.Tensor]: - """Return weight tensors.""" + """Return weight tensors in registration order.""" weights = super().get_weights_list() dev = f"cuda:{self.device_id}" @@ -79,7 +86,7 @@ def get_weights_list(self) -> list[torch.Tensor]: if self.partial_buf is None: self.partial_buf = torch.zeros( self.model_args.max_batch_size, - 4, + 8, self.model_args.dim, dtype=torch.bfloat16, device=dev, @@ -91,49 +98,51 @@ def get_weights_list(self) -> list[torch.Tensor]: return weights def get_cache_vars(self) -> list[torch.Tensor]: - """Return [ki_cache, kv_cache, pe_cache] matching DsaCacheVars layout.""" + """Return [k_cache, kv_cache, pe_cache] matching DsaCacheVars layout.""" cache_seq_len = self.model_args.max_seq_len + self.model_args.kv_cache_pad - bs_args = (self.model_args.max_batch_size, cache_seq_len) + bs = self.model_args.max_batch_size + dev = f"cuda:{self.device_id}" if self.ki_cache is None: ki_dim = self.model_args.index_head_dim - self.ki_cache = torch.zeros( - *bs_args, ki_dim, dtype=torch.bfloat16, device=f"cuda:{self.device_id}" - ) + self.ki_cache = torch.zeros(bs, cache_seq_len, ki_dim, dtype=torch.bfloat16, device=dev) if self.kv_cache is None: kv_dim = self.model_args.kv_lora_rank - self.kv_cache = torch.zeros( - *bs_args, kv_dim, dtype=torch.bfloat16, device=f"cuda:{self.device_id}" - ) + if getattr(self.model_args, "fp8_kv_cache", False): + self.kv_cache = torch.zeros( + bs, 1, kv_dim + (kv_dim // 128) * 4, dtype=torch.uint8, device=dev + ) + else: + self.kv_cache = torch.zeros(bs, 1, kv_dim, dtype=torch.bfloat16, device=dev) if self.pe_cache is None: - pe_dim = self.model_args.qk_rope_head_dim self.pe_cache = torch.zeros( - *bs_args, pe_dim, dtype=torch.bfloat16, device=f"cuda:{self.device_id}" + bs, 1, self.model_args.qk_rope_head_dim, dtype=torch.bfloat16, device=dev ) return [*super().get_cache_vars(), self.ki_cache, self.kv_cache, self.pe_cache] class PureMlaV2(SerializableTileRTModule): + """Device Group B (GPU 1-7): pure MLA.""" def __init__( self, model_args: ModelArgs, device_id: int, num_devices: int, - ll_buf: torch.Tensor | None = None, + recv_buf: torch.Tensor | None = None, ): super().__init__(model_args=model_args, device_id=device_id, num_devices=num_devices) self.rmsnorm_projx_wqkva = RMSNormProjxWqkva( model_args=model_args, device_id=device_id, num_devices=num_devices ) - self.rmsnorm_projx_wqkva.algorithm = RMSNormProjxWqkvaAlgorithm.DECOUPLED + self.rmsnorm_projx_wqkva.algorithm = RMSNormProjxWqkvaAlgorithm.W8A16HMMA self.register_op(self.rmsnorm_projx_wqkva) self.rmsnorm_projq_wqb = RmsnormProjqWqb( model_args=model_args, device_id=device_id, num_devices=num_devices ) - self.rmsnorm_projq_wqb.algorithm = RmsnormProjqWqbAlgorithm.FP16MMA + self.rmsnorm_projq_wqb.algorithm = RmsnormProjqWqbAlgorithm.BF16MMA self.register_op(self.rmsnorm_projq_wqb) self.rmsnorm_kv = KVRMSNorm( @@ -151,7 +160,7 @@ def __init__( ) self.register_op(self.projo_wkvb) - allreduce_algo = UnProjOAllReduceAlgorithm.FP16MMA + allreduce_algo = UnProjOAllReduceAlgorithm.BF16MMA self.unproj_o_allreduce = UnProjOAllReduce( model_args=model_args, device_id=device_id, @@ -160,14 +169,14 @@ def __init__( ) self.register_op(self.unproj_o_allreduce) - self.ll_buf = ll_buf + self.recv_buf = recv_buf self.ki_cache: torch.Tensor | None = None self.kv_cache: torch.Tensor | None = None self.pe_cache: torch.Tensor | None = None def init_random_weights(self) -> None: - """Initialize random weights for this module.""" + """Override to re-init ProjQWkvb/ProjOWkvb with HMMA-packed weights.""" super().init_random_weights() from tilert.models.common import init_func @@ -193,7 +202,7 @@ def init_random_weights(self) -> None: op.init_tilert_weights_hmma(per_dev) def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - """Load TileRT weights for this module from state_dict.""" + """Override to use HMMA-packed weights for ProjQWkvb and ProjOWkvb.""" self.projq_wqb.is_tilert_weights_init = True self.projo_wkvb.is_tilert_weights_init = True @@ -211,38 +220,47 @@ def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: op.init_tilert_weights_hmma(op_state_dict) def get_weights_list(self) -> list[torch.Tensor]: - """Return weight tensors.""" + """Return weight tensors in registration order.""" weights = super().get_weights_list() - if self.ll_buf is None: - max_seq_len = getattr(self.model_args, "num_mtp", 3) + 1 + if self.recv_buf is None: + max_seq_len = max(getattr(self.model_args, "num_mtp", 3) + 1, 8) topk = self.model_args.index_topk - self.ll_buf = torch.zeros( + self.recv_buf = torch.zeros( max_seq_len * topk * 2, dtype=torch.int32, device=f"cuda:{self.device_id}" ) - weights.append(self.ll_buf) + weights.append(self.recv_buf) return weights def get_cache_vars(self) -> list[torch.Tensor]: - """Return [ki_cache, kv_cache, pe_cache] matching DsaCacheVars layout.""" + """Return [k_cache, kv_cache, pe_cache] matching DsaCacheVars layout.""" cache_seq_len = self.model_args.max_seq_len + self.model_args.kv_cache_pad - bs_args = (self.model_args.max_batch_size, cache_seq_len) + bs = self.model_args.max_batch_size + dev = f"cuda:{self.device_id}" if self.ki_cache is None: - ki_dim = self.model_args.index_head_dim self.ki_cache = torch.zeros( - *bs_args, ki_dim, dtype=torch.bfloat16, device=f"cuda:{self.device_id}" + bs, 1, self.model_args.index_head_dim, dtype=torch.bfloat16, device=dev ) if self.kv_cache is None: kv_dim = self.model_args.kv_lora_rank - self.kv_cache = torch.zeros( - *bs_args, kv_dim, dtype=torch.bfloat16, device=f"cuda:{self.device_id}" - ) + if getattr(self.model_args, "fp8_kv_cache", False): + kv_merged = kv_dim + (kv_dim // 128) * 4 + self.kv_cache = torch.zeros( + bs, cache_seq_len, kv_merged, dtype=torch.uint8, device=dev + ) + else: + self.kv_cache = torch.zeros( + bs, cache_seq_len, kv_dim, dtype=torch.bfloat16, device=dev + ) if self.pe_cache is None: - pe_dim = self.model_args.qk_rope_head_dim self.pe_cache = torch.zeros( - *bs_args, pe_dim, dtype=torch.bfloat16, device=f"cuda:{self.device_id}" + bs, + cache_seq_len, + self.model_args.qk_rope_head_dim, + dtype=torch.bfloat16, + device=dev, ) return [*super().get_cache_vars(), self.ki_cache, self.kv_cache, self.pe_cache] diff --git a/tilert/models/glm_5/_dsa_v32/modules/mlp.py b/tilert/models/glm_5/modules/mlp.py similarity index 88% rename from tilert/models/glm_5/_dsa_v32/modules/mlp.py rename to tilert/models/glm_5/modules/mlp.py index 85fec25..79e8d69 100644 --- a/tilert/models/glm_5/_dsa_v32/modules/mlp.py +++ b/tilert/models/glm_5/modules/mlp.py @@ -1,13 +1,14 @@ from tilert.models.base import SerializableTileRTModule from tilert.models.glm_5._dsa_v32.model_args import ModelArgs -from tilert.models.glm_5._dsa_v32.modules.mla_v2 import PureMlaV2 as Mla from tilert.models.glm_5._dsa_v32.ops.down_allreduce import ( DownAllReduce, + DownAllReduceAlgorithm, ) from tilert.models.glm_5._dsa_v32.ops.rmsnorm_up_gate_silu import ( RMSNormUpGateSiLU, RMSNormUpGateSiLUAlgorithm, ) +from tilert.models.glm_5.modules.mla_v2 import PureMlaV2 as Mla class Mlp(SerializableTileRTModule): @@ -26,11 +27,14 @@ def __init__( device_id=device_id, num_devices=num_devices, ) - self.rmsnorm_mlp_up_gate_silu.algorithm = RMSNormUpGateSiLUAlgorithm.FP16MMA + self.rmsnorm_mlp_up_gate_silu.algorithm = RMSNormUpGateSiLUAlgorithm.BF16MMA_V2 self.register_op(self.rmsnorm_mlp_up_gate_silu) self.rmsnorm_mlp_down = DownAllReduce( - model_args=model_args, device_id=device_id, num_devices=num_devices + model_args=model_args, + device_id=device_id, + num_devices=num_devices, + algorithm=DownAllReduceAlgorithm.BF16MMA_V2, ) self.register_op(self.rmsnorm_mlp_down) diff --git a/tilert/models/glm_5/_dsa_v32/modules/moe.py b/tilert/models/glm_5/modules/moe.py similarity index 86% rename from tilert/models/glm_5/_dsa_v32/modules/moe.py rename to tilert/models/glm_5/modules/moe.py index 5410284..4e4b209 100644 --- a/tilert/models/glm_5/_dsa_v32/modules/moe.py +++ b/tilert/models/glm_5/modules/moe.py @@ -2,9 +2,9 @@ from tilert.models.base import SerializableTileRTModule from tilert.models.glm_5._dsa_v32.model_args import ModelArgs -from tilert.models.glm_5._dsa_v32.modules.mla_v2 import PureMlaV2 as Mla from tilert.models.glm_5._dsa_v32.ops.expert_down_allreduce import ( ExpertDownAllReduce, + ExpertDownAllReduceAlgorithm, ) from tilert.models.glm_5._dsa_v32.ops.expert_sel_up_gate_silu import ( ExpertSelectUpGateSiLU, @@ -13,6 +13,7 @@ from tilert.models.glm_5._dsa_v32.ops.rmsnorm_expert_proj import ( RMSNormExpertProj, ) +from tilert.models.glm_5.modules.mla_v2 import PureMlaV2 as Mla class Moe(SerializableTileRTModule): @@ -32,12 +33,16 @@ def __init__(self, model_args: ModelArgs, device_id: int, num_devices: int): model_args=model_args, device_id=device_id, num_devices=num_devices, - algorithm=ExpertSelectUpGateSiLUAlgorithm.FP16MMA, + algorithm=ExpertSelectUpGateSiLUAlgorithm.BF16MMA, ) self.register_op(self.exp_sel_up_gate_silu) + _expert_down_algo = ExpertDownAllReduceAlgorithm.BF16MMA self.expert_down_allreduce = ExpertDownAllReduce( - model_args=model_args, device_id=device_id, num_devices=num_devices + model_args=model_args, + device_id=device_id, + num_devices=num_devices, + algorithm=_expert_down_algo, ) self.register_op(self.expert_down_allreduce) diff --git a/tilert/models/glm_5/_dsa_v32/modules/mtp.py b/tilert/models/glm_5/modules/mtp.py similarity index 93% rename from tilert/models/glm_5/_dsa_v32/modules/mtp.py rename to tilert/models/glm_5/modules/mtp.py index ccfbdc8..1bfc1f1 100644 --- a/tilert/models/glm_5/_dsa_v32/modules/mtp.py +++ b/tilert/models/glm_5/modules/mtp.py @@ -2,9 +2,9 @@ from tilert.models.base import SerializableTileRTModule from tilert.models.glm_5._dsa_v32.model_args import ModelArgs -from tilert.models.glm_5._dsa_v32.modules.moe import MoeBlock -from tilert.models.glm_5._dsa_v32.modules.mtp_preprocess import MTPPreprocessLayer from tilert.models.glm_5._dsa_v32.ops import RMSNormHeadProj +from tilert.models.glm_5.modules.moe import MoeBlock +from tilert.models.glm_5.modules.mtp_preprocess import MTPPreprocessLayer class MTP(SerializableTileRTModule): diff --git a/tilert/models/glm_5/_dsa_v32/modules/mtp_preprocess.py b/tilert/models/glm_5/modules/mtp_preprocess.py similarity index 87% rename from tilert/models/glm_5/_dsa_v32/modules/mtp_preprocess.py rename to tilert/models/glm_5/modules/mtp_preprocess.py index debd75d..c7d2d83 100644 --- a/tilert/models/glm_5/_dsa_v32/modules/mtp_preprocess.py +++ b/tilert/models/glm_5/modules/mtp_preprocess.py @@ -70,18 +70,7 @@ class MTPPreprocessWeightsConverter(TilertWeightsConverter): """Converts ref-format weights to TileRT format for MTP preprocess.""" def convert_to_tilert(self, weights: list[torch.Tensor], device_id: int) -> list[torch.Tensor]: - """ - Convert ref weights to TileRT format for a specific device. - - Args: - weights: [embedding_rmsnorm_gamma, hidden_rmsnorm_gamma, eh_proj.weight] - Ref format: enorm.weight [7168], hnorm.weight [7168], - eh_proj.weight [7168, 14336]. - device_id: Target device ID for weight placement. - - Returns: - MTPPreprocessParams with converted weights for device_id. - """ + """Convert ref weights to TileRT format for a specific device.""" device = torch.device(f"cuda:{device_id}") embedding_rmsnorm_gamma, hidden_rmsnorm_gamma, eh_proj_weight = weights @@ -165,14 +154,7 @@ def init_reference_weights(self, state_dict: dict[str, torch.Tensor]) -> None: self.ref_eh_proj_weight = state_dict[self.ref_weights_alias.eh_proj] def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None: - """ - Load TileRT weights from state_dict. - - state_dict may use: - - Full keys: layer_{layer_id}_{alias}_dev_{device_id} - - Short keys: embedding_rmsnorm_gamma, hidden_rmsnorm_gamma, eh_proj_weights - - Ref keys: enorm.weight, hnorm.weight, eh_proj.weight (then convert) - """ + """Load TileRT weights from state_dict.""" converter = MTPPreprocessWeightsConverter(self.model_args, self.num_devices) params = converter.convert_to_tilert( [state_dict[k] for k in self.tilert_weights_alias()], self.device_id @@ -199,16 +181,7 @@ def golden_forward( x: torch.Tensor, last_hidden_states: torch.Tensor, ) -> torch.Tensor: - """ - Reference forward: enorm(x), hnorm(last_hidden), concat & eh_proj. - - Args: - x: [batch, seq_len, hidden_size] embedded tokens - last_hidden_states: [batch, seq_len, hidden_size] previous hidden - - Returns: - [batch, seq_len, hidden_size] projected hidden - """ + """Reference forward: enorm(x), hnorm(last_hidden), concat & eh_proj.""" assert self.ref_embedding_rmsnorm_gamma is not None assert self.ref_hidden_rmsnorm_gamma is not None assert self.ref_eh_proj_weight is not None diff --git a/tilert/models/glm_5/ops/__init__.py b/tilert/models/glm_5/ops/__init__.py new file mode 100644 index 0000000..888b650 --- /dev/null +++ b/tilert/models/glm_5/ops/__init__.py @@ -0,0 +1,7 @@ +"""Core operations for GLM5.""" + +from tilert.models.glm_5.ops.sparse_index_v3 import sparse_index_topk_v3 + +__all__ = [ + "sparse_index_topk_v3", +] diff --git a/tilert/models/glm_5/ops/sparse_index_v3.py b/tilert/models/glm_5/ops/sparse_index_v3.py new file mode 100644 index 0000000..8e58362 --- /dev/null +++ b/tilert/models/glm_5/ops/sparse_index_v3.py @@ -0,0 +1,51 @@ +"""GLM5 sparse index op Python wrapper.""" + +import torch + +__all__ = [ + "sparse_index_topk_v3", +] + + +def sparse_index_topk_v3( + q: torch.Tensor, # noqa: VNE001 + kv: torch.Tensor, + weights: torch.Tensor, + logits: torch.Tensor, + indices: torch.Tensor, + cur_pos: int, + profile_logs: torch.Tensor, +) -> None: + """GLM5 sparse index + top-k selection.""" + if q.dtype != torch.bfloat16: + raise ValueError("input must be a bfloat16 tensor.") + if kv.dtype != torch.bfloat16: + raise ValueError("kv must be a bfloat16 tensor.") + if weights.dtype != torch.bfloat16: + raise ValueError("weights must be a bfloat16 tensor.") + if logits.dtype != torch.float32: + raise ValueError("logits must be a float32 tensor.") + + seqlen = q.shape[-3] + head = q.shape[-2] + dim = q.shape[-1] + + if head != 32: + raise ValueError( + f"Unsupported head size: {head}. SparseIndexV3 fused op " + "supports head number of 32 (GLM5)." + ) + if dim != 128: + raise ValueError("dim must be 128, as we precompute scale inner kernel") + + device = q.device + if any(t.device != device for t in (kv, weights, logits, indices, profile_logs)): + raise ValueError( + "sparse_index inputs must be on the same device: " + f"q={device}, kv={kv.device}, weights={weights.device}, " + f"logits={logits.device}, profile_logs={profile_logs.device}" + ) + workspace = torch.zeros(seqlen, (200 * 1024 + 260), dtype=torch.int32, device=device) + torch.ops.tilert.sparse_index_topk_glm5_v3_op( + q, kv, weights, logits, cur_pos, indices, workspace, profile_logs + ) diff --git a/tilert/models/glm_5/params.py b/tilert/models/glm_5/params.py new file mode 100644 index 0000000..2721229 --- /dev/null +++ b/tilert/models/glm_5/params.py @@ -0,0 +1 @@ +"""GLM5 parameters and initializers.""" diff --git a/tilert/models/glm_5/_dsa_v32/temp_var_indices.py b/tilert/models/glm_5/temp_var_indices.py similarity index 73% rename from tilert/models/glm_5/_dsa_v32/temp_var_indices.py rename to tilert/models/glm_5/temp_var_indices.py index 3a7af62..de33438 100644 --- a/tilert/models/glm_5/_dsa_v32/temp_var_indices.py +++ b/tilert/models/glm_5/temp_var_indices.py @@ -1,13 +1,4 @@ -"""Named indices for DSA temporary variables. - -Lets Python code reference temp_vars by name instead of magic numbers. - -Usage:: - - from tilert.models.glm_5._dsa_v32.temp_var_indices import Idx - - token_out = intermediates[Idx.TOKEN_OUT] # equivalent to intermediates[25] -""" +"""Named indices for DSA temporary variables.""" from enum import IntEnum @@ -26,7 +17,7 @@ class DsaTempVarIdx(IntEnum): IDX_LOGITS = 8 IDX_SELECTS = 9 Q_NOPE = 10 - O = 11 # noqa: E741 + O = 11 # noqa: E741 β€” mirrors C++ DsaTempVars::O O_ACC = 12 O_LSE = 13 O_LSE_ACC = 14 @@ -71,24 +62,19 @@ class DsaTempVarIdx(IntEnum): TOP_N_LOG_PROBS = 53 TOP_N_INDICES = 54 LOGPROBS_FLAG = 55 + AR_ACCEPTED_TOKENS = 56 + AR_NUM_ACCEPTED = 57 + HIDDEN_MID = 58 + GRAMMAR_BITMASK = 59 -TEMP_VARS_SIZE = 56 +TEMP_VARS_SIZE = 60 Idx = DsaTempVarIdx def validate_temp_vars_layout() -> None: - """Validate the temporary-variable index enum. - - Checks: - 1. Enum member count equals TEMP_VARS_SIZE. - 2. Indices are contiguous 0..TEMP_VARS_SIZE-1 with no gaps or duplicates. - 3. (If the backend is loaded) the backend temp_vars_size matches TEMP_VARS_SIZE. - - Raises: - RuntimeError: If any validation check fails. - """ + """Validate the temporary-variable index enum.""" members = list(DsaTempVarIdx) if len(members) != TEMP_VARS_SIZE: @@ -112,7 +98,8 @@ def validate_temp_vars_layout() -> None: cpp_size = torch.ops.tilert.dsa_temp_vars_size() if cpp_size != TEMP_VARS_SIZE: raise RuntimeError( - f"TEMP_VARS_SIZE={TEMP_VARS_SIZE} != " f"backend temp_vars_size={cpp_size}" + f"Python TEMP_VARS_SIZE={TEMP_VARS_SIZE} != " + f"C++ DsaTempVars::temp_vars_size={cpp_size}" ) except (AttributeError, RuntimeError): pass diff --git a/tilert/pd_vllm/__init__.py b/tilert/pd_vllm/__init__.py new file mode 100644 index 0000000..2266c0c --- /dev/null +++ b/tilert/pd_vllm/__init__.py @@ -0,0 +1 @@ +"""vLLM-prefill + TileRT-decode PD disaggregation for GLM-5 and DeepSeek-V3.2.""" diff --git a/tilert/pd_vllm/decode_server.py b/tilert/pd_vllm/decode_server.py new file mode 100644 index 0000000..3372694 --- /dev/null +++ b/tilert/pd_vllm/decode_server.py @@ -0,0 +1,350 @@ +"""PD decode server (W6): HTTP orchestration around receive -> convert -> inject -> decode. + +Internal token-level API (the client-facing OpenAI layer lives in pd_router / +a later serving layer): + + POST /pd/decode {rid, first_token_id, max_tokens, sampling?, timeout_s?} + Waits for the wire transfer of `rid` to complete, converts, injects + into the engine, decodes, returns {"rid", "token_ids", "timing_ms"}. + GET /health {"status": "ok"} + GET /decode_status {"status": "idle"|"busy", "current_rid": ...} + +bs=1: a busy server answers 429 immediately (the router's gated dispatch +should make that unreachable). +""" + +import argparse +import contextlib +import json +import logging +import queue as queue_mod +import socket +import threading +import time +from typing import Any + +import uvicorn +from fastapi import FastAPI +from fastapi.responses import JSONResponse, StreamingResponse +from pydantic import BaseModel + +from tilert.pd_vllm.receive_server import ReceiveServer + +logger = logging.getLogger("pd_vllm.decode_server") + + +class DecodeBody(BaseModel): + rid: str + first_token_id: int + max_tokens: int = 256 + sampling: dict | None = None + timeout_s: float = 120.0 + stream: bool = False + + +def build_app(server: ReceiveServer, engine) -> FastAPI: + app = FastAPI() + lock = threading.Lock() + state: dict[str, Any] = {"current_rid": None} + + @app.get("/health") + def health(): + return {"status": "ok"} + + @app.get("/decode_status") + def decode_status(): + busy = lock.locked() + return {"status": "busy" if busy else "idle", "current_rid": state["current_rid"]} + + @app.post("/pd/cancel") + def pd_cancel(body: dict): + """Explicit kill switch: cancel the in-flight decode for `rid`. + + Deterministic cancel path β€” dead-connection detection at the + transport layer is unreliable (asyncio writes to a closed socket + do not raise), so the router calls this on client disconnect. + """ + rid = body.get("rid") + ev = state.get("cancel_event") + if rid and rid == state["current_rid"] and ev is not None: + ev.set() + logger.info("cancel requested for %s", rid) + return {"cancelled": rid} + return JSONResponse( + {"error": "no matching in-flight request", "current_rid": state["current_rid"]}, + status_code=404, + ) + + def _cleanup(): + try: + engine.reset() + except Exception: + logger.exception("engine reset failed") + server.release() + state["current_rid"] = None + state["cancel_event"] = None + lock.release() + + def _log_reqstat(body, req, n_tokens, timing): + logger.info( + "REQSTAT rid=%s seq=%d completion=%d %s", + body.rid, + req.seq_len, + n_tokens, + " ".join(f"{k}={v}" for k, v in timing.items()), + ) + + @app.post("/pd/decode") + def pd_decode(body: DecodeBody): + if not lock.acquire(blocking=False): + return JSONResponse( + {"error": "busy", "current_rid": state["current_rid"]}, status_code=429 + ) + state["current_rid"] = body.rid + t0 = time.time() + # phase 1: wire wait + convert + inject (common to both modes) + try: + # Drain until OUR rid arrives; drop stale completed entries + # (e.g. a transfer whose consumer never called /pd/decode). + req = None + deadline = time.time() + body.timeout_s + while time.time() < deadline: + try: + cand = server.completed.get(timeout=max(0.1, deadline - time.time())) + except queue_mod.Empty: + break + if cand.rid == body.rid: + req = cand + break + logger.warning( + "dropping unmatched request %s " "(waiting for %s)", cand.rid, body.rid + ) + server.release() + if req is None: + _cleanup() + return JSONResponse( + {"error": "kv_transfer_timeout", "rid": body.rid}, status_code=504 + ) + t_recv = time.time() + conv = server.profile.convert( + server.buffer, server.base_ptr, server.max_seq_len, req, server.profile.num_ranks + ) + t_conv = time.time() + engine.inject(conv) + t_inj = time.time() + except Exception as e: + logger.exception("prepare failed for %s", body.rid) + _cleanup() + return JSONResponse({"error": str(e), "rid": body.rid}, status_code=500) + + pre_timing = { + "wire_wait": round(1000 * (t_recv - t0), 1), + "convert": round(1000 * (t_conv - t_recv), 1), + "inject": round(1000 * (t_inj - t_conv), 1), + } + + # phase 2: decode + cancel = threading.Event() + state["cancel_event"] = cancel + + if not body.stream: + try: + tokens = engine.decode( + first_token_id=body.first_token_id, + max_tokens=body.max_tokens, + sampling=body.sampling, + cancel_event=cancel, + ) + timing = { + **pre_timing, + "decode": round(1000 * (time.time() - t_inj), 1), + **getattr(engine, "last_stats", {}), + } + _log_reqstat(body, req, len(tokens), timing) + return { + "rid": body.rid, + "token_ids": tokens, + "seq_len": req.seq_len, + "timing_ms": timing, + } + except Exception as e: + logger.exception("decode failed for %s", body.rid) + return JSONResponse({"error": str(e), "rid": body.rid}, status_code=500) + finally: + _cleanup() + + # streaming: ndjson lines {"t":[ids...]}* then {"done":true,...}; + # lock/engine ownership transfers to the generator. + q: queue_mod.Queue = queue_mod.Queue() + + def _run(): + try: + tokens = engine.decode( + first_token_id=body.first_token_id, + max_tokens=body.max_tokens, + sampling=body.sampling, + on_token=q.put, + cancel_event=cancel, + ) + q.put(("done", tokens)) + except Exception as e: # pragma: no cover + logger.exception("stream decode failed for %s", body.rid) + q.put(("error", str(e))) + + worker = threading.Thread(target=_run, name="pd-decode", daemon=True) + + async def _gen(): + # MUST be an async generator: on client disconnect starlette + # cancels the response task, and only async generators get the + # cancellation delivered into their frame so `finally` runs + # (a sync generator is silently abandoned -> the engine slot + # leaks forever; found by the streaming-cancel drill). + import asyncio + + import anyio + from starlette.concurrency import run_in_threadpool + + worker.start() + try: + batch: list[int] = [] + done_msg = None + last_activity = time.time() + while done_msg is None: + drained = False + while True: + try: + item = q.get_nowait() + except queue_mod.Empty: + break + drained = True + if isinstance(item, int): + batch.append(item) + else: + done_msg = item + break + if batch: + yield json.dumps({"t": batch}) + "\n" + batch = [] + if done_msg is None: + if drained: + last_activity = time.time() + elif time.time() - last_activity > 600: # noqa: R505 (exclusive branches) + yield json.dumps({"error": "decode stalled"}) + "\n" + return + else: + await asyncio.sleep(0.005) + kind, payload = done_msg + if kind == "done": + timing = { + **pre_timing, + "decode": round(1000 * (time.time() - t_inj), 1), + **getattr(engine, "last_stats", {}), + } + _log_reqstat(body, req, len(payload), timing) + yield json.dumps( + { + "done": True, + "n": len(payload), + "seq_len": req.seq_len, + "finish_reason": timing.get("finish_reason", "stop"), + "timing_ms": timing, + } + ) + "\n" + else: + yield json.dumps({"error": payload}) + "\n" + finally: + cancel.set() + # shield: cleanup must complete even inside a cancelled scope, + # and the worker must be joined before engine.reset() (the + # engine may be mid-decode_mtp on the GPU). + with anyio.CancelScope(shield=True): + await run_in_threadpool(worker.join, 120) + if worker.is_alive(): + logger.error("decode worker failed to stop for %s", body.rid) + _cleanup() + + return StreamingResponse(_gen(), media_type="application/x-ndjson") + + return app # noqa: R504 (assembled across the function) + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s") + ap = argparse.ArgumentParser() + ap.add_argument("--engine", choices=["stub", "tilert"], default="stub") + ap.add_argument("--model", default="glm5", help="model profile") + ap.add_argument("--max-seq-len", type=int, default=4096) + ap.add_argument("--ctrl-port", type=int, default=5556) + ap.add_argument("--http-port", type=int, default=5557) + ap.add_argument("--model-weights-dir", default="") + ap.add_argument("--with-mtp", action="store_true") + ap.add_argument( + "--transport", + choices=["mooncake", "nixl"], + default="mooncake", + help="RDMA data-plane backend " "(must match prefill's tilert_transport)", + ) + ap.add_argument( + "--kv-cache-dtype", + default="fp8_ds_mla", + help="MLA cache dtype (must match vLLM prefill); " "MLA-family profiles only", + ) + args = ap.parse_args() + + from tilert.pd_vllm.profiles import base as profiles + + profile = profiles.get_profile(args.model) + # MLA-family profiles (glm5/dsv32) need the cache dtype to size the receive + # buffer. + if hasattr(profile, "configure"): + profile.configure(args.kv_cache_dtype) + logger.info( + "profile %s MLA cache dtype = %s (layout v%d)", + profile.name, + args.kv_cache_dtype, + profile.layout_version, + ) + + if args.engine == "stub": + from tilert.pd_vllm.engine_iface import StubEngine + + engine: Any = StubEngine() + else: + logger.info( + "loading TileRT engine (profile=%s, weights=%s)...", + profile.name, + args.model_weights_dir, + ) + engine = profile.build_engine( + model_weights_dir=args.model_weights_dir, + max_seq_len=args.max_seq_len, + with_mtp=args.with_mtp, + ar_steps=8, + ) + logger.info("TileRT engine ready (cache window %d)", engine.max_seq_len) + + server = ReceiveServer( + profile, max_seq_len=args.max_seq_len, ctrl_port=args.ctrl_port, transport=args.transport + ) + app = build_app(server, engine) + logger.info( + "decode server on :%d (profile=%s, engine=%s, ctrl=:%d)", + args.http_port, + profile.name, + args.engine, + args.ctrl_port, + ) + # Bind dual-stack (IPv4 + IPv6) explicitly. uvicorn's host="::" is + # IPv6-only under some uvicorn/OS combinations, which leaves the decode + # HTTP endpoint unreachable from an IPv4 router. Mirror the control plane + # (receive_server) by clearing IPV6_V6ONLY on an AF_INET6 socket. + sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + with contextlib.suppress(OSError): + sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0) + sock.bind(("::", args.http_port)) + uvicorn.Server(uvicorn.Config(app, log_level="warning")).run(sockets=[sock]) + + +if __name__ == "__main__": + main() diff --git a/tilert/pd_vllm/engine_iface.py b/tilert/pd_vllm/engine_iface.py new file mode 100644 index 0000000..6ed9133 --- /dev/null +++ b/tilert/pd_vllm/engine_iface.py @@ -0,0 +1,55 @@ +"""Engine seam for the PD decode server (model-agnostic). + +``PDEngine`` is the interface the decode server drives; concrete adapters are +built by the active model profile (``profile.build_engine(...)``). +``StubEngine`` runs the whole serving path with no GPU / no tilert. +""" + +from collections.abc import Callable +from typing import Any, Protocol + + +class PDEngine(Protocol): + def inject(self, req: Any) -> None: + """Restore engine state to 'prefilled seq_len tokens' from req.""" + + def decode( + self, + first_token_id: int, + max_tokens: int, + sampling: dict | None, + on_token: Callable[[int], None] | None = None, + cancel_event=None, + ) -> list[int]: + """AR/MTP decode from first_token_id; returns completion ids. + + Includes first_token_id, excludes the stop token. on_token never fires + for stop tokens; cancel_event stops early; last_stats['finish_reason'] + is 'stop' | 'length' | 'cancelled'. + """ + + def reset(self) -> None: + """Release per-request state.""" + + +class StubEngine: + """Echo engine for plumbing tests: no GPU, no tilert.""" + + def __init__(self, fixed_tokens: tuple[int, ...] = (11, 22, 33)): + self._fixed = fixed_tokens + self.injected: Any = None + self.last_stats: dict = {} + + def inject(self, req: Any) -> None: + self.injected = req + + def decode(self, first_token_id, max_tokens, sampling, on_token=None, cancel_event=None): + out = ([int(first_token_id)] + list(self._fixed))[:max_tokens] + if on_token: + for t in out: + on_token(t) + self.last_stats = {"finish_reason": "stop"} + return out + + def reset(self) -> None: + self.injected = None diff --git a/tilert/pd_vllm/oai_parser.py b/tilert/pd_vllm/oai_parser.py new file mode 100644 index 0000000..931daa9 --- /dev/null +++ b/tilert/pd_vllm/oai_parser.py @@ -0,0 +1,230 @@ +"""OpenAI-semantics parser adapter over vLLM's parser engine (decision B1). + +Wraps ``vllm.parser`` (the NEW engine architecture in vllm >= 0.24; the old +``ReasoningParser``/``ToolParserManager`` API is superseded) into the +small surface the router needs: + + parser = make_parser("glm47", tokenizer, thinking=True) + parsed = parser.parse_complete(text) # non-streaming + sess = parser.stream() # per-request streaming + events = sess.feed(delta_text); sess.finish() # normalized event dicts + +Runs in the ROUTER environment only β€” that env must have vllm installed +(CPU-only import is fine; verified with CUDA_VISIBLE_DEVICES=""). The decode +node never imports vllm. +""" + +import logging +import uuid +from dataclasses import dataclass, field + +logger = logging.getLogger("pd_vllm.oai_parser") + + +@dataclass +class ToolCall: + call_id: str + name: str + arguments: str # JSON string (OpenAI convention) + + def to_openai(self, index: int) -> dict: + return { + "index": index, + "id": self.call_id, + "type": "function", + "function": {"name": self.name, "arguments": self.arguments}, + } + + +@dataclass +class Parsed: + reasoning_content: str | None + content: str | None + tool_calls: list[ToolCall] = field(default_factory=list) + + +def _new_call_id() -> str: + return f"call_{uuid.uuid4().hex[:24]}" + + +# family -> (config-builder import path, arg-converter import path). The +# glm47_moe parser engine uses the vllm.parser API shape (a `*_config(thinking)` +# builder + a `_*_arg_converter(raw, partial)`); the adapter picks the engine +# by family name. +_FAMILIES = { + "glm47": ("vllm.parser.glm47_moe", "glm47_moe_config", "_glm47_arg_converter"), +} + + +def make_parser(family: str, tokenizer, thinking: bool = True) -> "OaiParser": + if family not in _FAMILIES: + raise KeyError(f"unknown parser family {family!r}; " f"known: {sorted(_FAMILIES)}") + return OaiParser(family, tokenizer, thinking) + + +class OaiParser: + """Family-parameterized parser; one instance per model, ``stream()`` per request. + + Family is a vllm.parser engine (glm47). + """ + + def __init__(self, family: str, tokenizer, thinking: bool = True): + import importlib + + from vllm.parser.engine.events import EventType + from vllm.parser.engine.streaming_parser_engine import ( + StreamingParserEngine, + ) + + mod_name, cfg_name, conv_name = _FAMILIES[family] + mod = importlib.import_module(mod_name) + self._family = family + self._cfg_fn = getattr(mod, cfg_name) + self._Engine = StreamingParserEngine + self._ET = EventType + self._config = self._cfg_fn(thinking=thinking) + self._convert = getattr(mod, conv_name) + self._tok = tokenizer + + def with_thinking(self, thinking: bool) -> "OaiParser": + if thinking == (self._config.initial_state.name == "REASONING"): + return self + clone = object.__new__(OaiParser) + clone.__dict__.update(self.__dict__) + clone._config = self._cfg_fn(thinking=thinking) + return clone + + # ── non-streaming ──────────────────────────────────────────────────── + def parse_complete(self, text: str) -> Parsed: + engine = self._Engine(self._config, self._tok) + return self._reduce(engine.parse_complete(text)) + + def _reduce(self, events) -> Parsed: + ET = self._ET + reasoning, content = [], [] + slots: dict[int, dict] = {} + for e in events: + if e.type == ET.REASONING_CHUNK: + reasoning.append(e.value) + elif e.type == ET.TEXT_CHUNK: + content.append(e.value) + elif e.type in (ET.TOOL_NAME, ET.ARG_VALUE_CHUNK): + s = slots.setdefault(e.tool_index, {"name": [], "args": []}) + s["name" if e.type == ET.TOOL_NAME else "args"].append(e.value) + calls = [] + for i in sorted(slots): + name = "".join(slots[i]["name"]).strip() + if not name: + continue # unnamed fragment (heavy truncation) β€” drop + raw = "".join(slots[i]["args"]) + calls.append(ToolCall(_new_call_id(), name, self._convert(raw, True))) + r = "".join(reasoning) + c = "".join(content) + return Parsed(r if r else None, c if c else None, calls) + + # ── streaming ──────────────────────────────────────────────────────── + def stream(self) -> "OaiStream": + return OaiStream(self) + + +class OaiStream: + """Per-request streaming session. + + ``feed``/``finish`` return normalized event dicts: + {"kind": "reasoning", "text": ...} + {"kind": "content", "text": ...} + {"kind": "tool", "index": i, "id": ..., "name": ..., "arguments": ...} + + Reasoning/content stream through per delta. Tool calls are buffered and + emitted whole at TOOL_CALL_END (OpenAI clients accept arguments in any + fragmentation; whole-call emission sidesteps XMLβ†’JSON incremental + conversion). ``finish`` flushes a truncated trailing tool call with + partial-args conversion. + """ + + def __init__(self, parent: "OaiParser"): + self._p = parent + self._engine = parent._Engine(parent._config, parent._tok) + self._slots: dict[int, dict] = {} + self._emitted: set[int] = set() + + def feed(self, delta_text: str) -> list[dict]: + if not delta_text: + return [] + return self._consume(self._engine.feed(delta_text, [])) + + def finish(self) -> list[dict]: + out = self._consume(self._engine.finish()) + # flush truncated trailing tool call (never saw TOOL_CALL_END) + for i in sorted(self._slots): + if i in self._emitted: + continue + ev = self._flush_tool(i, partial=True) + if ev: + out.append(ev) + return out + + def _consume(self, events) -> list[dict]: + ET = self._p._ET + out: list[dict] = [] + for e in events: + if e.type == ET.REASONING_CHUNK: + out.append({"kind": "reasoning", "text": e.value}) + elif e.type == ET.TEXT_CHUNK: + out.append({"kind": "content", "text": e.value}) + elif e.type in (ET.TOOL_NAME, ET.ARG_VALUE_CHUNK): + s = self._slots.setdefault(e.tool_index, {"name": [], "args": []}) + s["name" if e.type == ET.TOOL_NAME else "args"].append(e.value) + elif e.type == ET.TOOL_CALL_END: + ev = self._flush_tool(e.tool_index, partial=False) + if ev: + out.append(ev) + return out + + def _flush_tool(self, index: int, partial: bool) -> dict | None: + s = self._slots.get(index) + if s is None or index in self._emitted: + return None + name = "".join(s["name"]).strip() + if not name: + return None + self._emitted.add(index) + args = self._p._convert("".join(s["args"]), partial) + return { + "kind": "tool", + "index": index, + "id": _new_call_id(), + "name": name, + "arguments": args, + } + + +class IncrementalDetok: + r"""Incremental tokenβ†’text for byte-level BPE tokenizers. + + Decodes a bounded trailing window; holds output while the window ends in + a partial multi-byte sequence (\\ufffd). Window folding is safe for + byte-level BPE: separate windows decode to concatenable byte streams. + Specials are KEPT (skip_special_tokens=False) β€” the parser consumes + etc.; the stop token never reaches the stream (engine adapter + suppresses it). + """ + + _FOLD = 256 + + def __init__(self, tokenizer): + self._tok = tokenizer + self._ids: list[int] = [] + self._emitted = 0 + + def push(self, ids: list[int]) -> str: + self._ids.extend(ids) + text = self._tok.decode(self._ids, skip_special_tokens=False) + if text.endswith("οΏ½"): + return "" + delta = text[self._emitted :] + self._emitted = len(text) + if len(self._ids) > self._FOLD: + self._ids = [] + self._emitted = 0 + return delta # noqa: R504 (self._emitted mutated after delta is computed) diff --git a/tilert/pd_vllm/pd_router.py b/tilert/pd_vllm/pd_router.py new file mode 100644 index 0000000..87a61e5 --- /dev/null +++ b/tilert/pd_vllm/pd_router.py @@ -0,0 +1,456 @@ +"""PD router (W6): client-facing entry over vLLM prefill + TileRT decode. + +Does OpenAI-semantics output parsing (reasoning + tool calls), streaming and +non-streaming. + +Flow per request (phase-1 hybrid, see design doc): + 1. pick a free decode node (in-memory busy tracking; all busy -> 429) + 2. forward to vLLM with max_tokens=1 + logprobs and inject + kv_transfer_params {tilert_host, tilert_ctrl_port} β€” the connector + claims the request and RDMA-sends state to the decode node + 3. extract rid + first_token_id from the vLLM response + (requires vLLM serve launched with --return-tokens-as-token-ids) + 4. call the decode node (/pd/decode; stream or not) and assemble the + OpenAI response: reasoning_content / content / tool_calls via the + vLLM parser engine (decision B1 β€” this process's env has vllm + installed, CPU-only; the decode node does not). + +Environment: run in a vllm-equipped env with CUDA_VISIBLE_DEVICES="" (the +router must never touch GPUs). --parser none falls back to raw passthrough. + +Run: + CUDA_VISIBLE_DEVICES= python -m tilert.pd_vllm.pd_router \ + --vllm-url http://prefill-node:8000 \ + --decode decode-node:5556:5557 --port 23333 \ + --model-path /path/to/GLM-5.1 --parser glm47 +""" + +import argparse +import json +import logging +import threading +import time + +import requests +import uvicorn +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse, StreamingResponse + +from tilert.pd_vllm.wire import derive_rid + +logger = logging.getLogger("pd_vllm.router") + + +class DecodeNode: + def __init__(self, host: str, ctrl_port: int, http_port: int): + self.host = host + self.ctrl_port = ctrl_port + self.http_port = http_port + self.busy = False + + @property + def http_base(self) -> str: + return f"http://{self.host}:{self.http_port}" + + +class Pool: + def __init__(self, nodes: list[DecodeNode]): + self.nodes = nodes + self._lock = threading.Lock() + + def acquire(self) -> DecodeNode | None: + with self._lock: + for n in self.nodes: + if not n.busy: + n.busy = True + return n + return None + + def release(self, node: DecodeNode) -> None: + with self._lock: + node.busy = False + + +def first_token_from_logprobs(resp: dict, is_chat: bool) -> int: + """Parse 'token_id:N' (vLLM --return-tokens-as-token-ids) from logprobs.""" + choice = resp["choices"][0] + lp = choice.get("logprobs") or {} + tok: str | None = None + if is_chat: + content = lp.get("content") or [] + if content: + tok = content[0].get("token") + else: + toks = lp.get("tokens") or [] + if toks: + tok = toks[0] + if tok and tok.startswith("token_id:"): + return int(tok.split(":", 1)[1]) + raise ValueError( + f"cannot extract first token id from logprobs ({tok!r}); launch vLLM " + f"with --return-tokens-as-token-ids and request logprobs" + ) + + +def _thinking_enabled(body: dict) -> bool: + ctk = body.get("chat_template_kwargs") or {} + return bool(ctk.get("enable_thinking", True)) + + +class RouterCtx: + """Immutable per-process context (tokenizer, parser factory, config).""" + + def __init__(self, vllm_url: str, pool: Pool, tokenizer, parser_name: str): + self.vllm_url = vllm_url + self.pool = pool + self.tokenizer = tokenizer + self.parser_name = parser_name + self._parsers = {} + if parser_name != "none": + if tokenizer is None: + raise SystemExit("--parser requires --model-path (tokenizer)") + from tilert.pd_vllm.oai_parser import make_parser + + self._parsers[True] = make_parser(parser_name, tokenizer, thinking=True) + self._parsers[False] = self._parsers[True].with_thinking(False) + logger.info("parser '%s' ready (thinking variants cached)", parser_name) + + def parser(self, thinking: bool): + return self._parsers.get(thinking) + + +def build_app(ctx: RouterCtx) -> FastAPI: + app = FastAPI() + pool = ctx.pool + + @app.get("/health") + def health(): + return {"status": "ok", "decode_free": sum(1 for n in pool.nodes if not n.busy)} + + @app.get("/pool_status") + def pool_status(): + return {"nodes": [{"host": n.host, "busy": n.busy} for n in pool.nodes]} + + # ── shared prefill step ────────────────────────────────────────────── + def _prefill(path, body, node): + prefill_body = dict(body) + prefill_body["max_tokens"] = 1 + prefill_body["stream"] = False + if path.endswith("chat/completions"): + prefill_body["logprobs"] = True + prefill_body["top_logprobs"] = 1 + else: + prefill_body["logprobs"] = 1 + prefill_body["kv_transfer_params"] = { + "tilert_host": node.host, + "tilert_ctrl_port": node.ctrl_port, + } + r = requests.post(f"{ctx.vllm_url}{path}", json=prefill_body, timeout=600) + r.raise_for_status() + return r.json() + + def _sampling_of(body): + return {k: body[k] for k in ("temperature", "top_p", "top_k") if k in body} + + def _max_tokens_of(body): + return int(body.get("max_tokens") or body.get("max_completion_tokens") or 256) + + # ── non-streaming ──────────────────────────────────────────────────── + def _handle(path: str, body: dict): + is_chat = path.endswith("chat/completions") + node = pool.acquire() + if node is None: + return JSONResponse({"error": "all decode nodes busy"}, status_code=429) + t0 = time.time() + try: + prefill = _prefill(path, body, node) + t_prefill = time.time() + rid = derive_rid(prefill["id"]) + first_token_id = first_token_from_logprobs(prefill, is_chat) + + dr = requests.post( + f"{node.http_base}/pd/decode", + json={ + "rid": rid, + "first_token_id": first_token_id, + "max_tokens": _max_tokens_of(body), + "sampling": _sampling_of(body), + }, + timeout=600, + ) + dr.raise_for_status() + decode = dr.json() + token_ids = decode["token_ids"] + timing = decode.get("timing_ms", {}) + finish = timing.get("finish_reason", "stop") + if finish == "cancelled": + finish = "stop" + + choice: dict = {"index": 0, "finish_reason": finish} + parser = ctx.parser(_thinking_enabled(body)) if is_chat else None + if parser is not None: + text = ctx.tokenizer.decode(token_ids, skip_special_tokens=False) + parsed = parser.parse_complete(text) + msg = {"role": "assistant", "content": parsed.content or ""} + if parsed.reasoning_content: + msg["reasoning_content"] = parsed.reasoning_content + if parsed.tool_calls: + msg["tool_calls"] = [c.to_openai(i) for i, c in enumerate(parsed.tool_calls)] + choice["finish_reason"] = "tool_calls" + choice["message"] = msg + else: + text = ( + ctx.tokenizer.decode(token_ids, skip_special_tokens=True) + if ctx.tokenizer + else None + ) + if is_chat: + choice["message"] = {"role": "assistant", "content": text} + else: + choice["text"] = text + choice["token_ids"] = token_ids + + return JSONResponse( + { + "id": prefill["id"], + "object": "chat.completion" if is_chat else "text_completion", + "created": int(time.time()), + "model": prefill.get("model"), + "choices": [choice], + "usage": { + "prompt_tokens": (prefill.get("usage") or {}).get("prompt_tokens"), + "completion_tokens": len(token_ids), + }, + "pd_timing_ms": { + "prefill": round(1000 * (t_prefill - t0), 1), + **timing, + }, + } + ) + except Exception as e: + logger.exception("pd request failed") + return JSONResponse({"error": str(e)}, status_code=502) + finally: + pool.release(node) + + # ── streaming (chat only) ──────────────────────────────────────────── + async def _handle_stream(path: str, body: dict, request: Request): + from starlette.concurrency import run_in_threadpool + + node = pool.acquire() + if node is None: + return JSONResponse({"error": "all decode nodes busy"}, status_code=429) + + try: + prefill = await run_in_threadpool(_prefill, path, body, node) + rid = derive_rid(prefill["id"]) + first_token_id = first_token_from_logprobs(prefill, True) + except Exception as e: + pool.release(node) + logger.exception("pd stream request failed before streaming") + return JSONResponse({"error": str(e)}, status_code=502) + + chunk_id = prefill["id"] + model = prefill.get("model") + prompt_tokens = (prefill.get("usage") or {}).get("prompt_tokens") + parser = ctx.parser(_thinking_enabled(body)) + + def _chunk(delta: dict, finish=None, usage=None) -> str: + payload = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], + } + if usage is not None: + payload["usage"] = usage + return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n" + + def _event_delta(ev: dict) -> dict: + if ev["kind"] == "reasoning": + return {"reasoning_content": ev["text"]} + if ev["kind"] == "content": + return {"content": ev["text"]} + return { + "tool_calls": [ + { + "index": ev["index"], + "id": ev["id"], + "type": "function", + "function": {"name": ev["name"], "arguments": ev["arguments"]}, + } + ] + } + + def _fire_cancel(): + try: + requests.post(f"{node.http_base}/pd/cancel", json={"rid": rid}, timeout=5) + except Exception: + logger.warning("cancel POST failed for %s", rid) + + async def _gen(): + import anyio + import httpx + + from tilert.pd_vllm.oai_parser import IncrementalDetok + + n_tokens = 0 + saw_tool = False + finish_reason = "stop" + client_gone = False + completed_ok = False + detok = IncrementalDetok(ctx.tokenizer) + sess = parser.stream() if parser else None + client = httpx.AsyncClient(timeout=httpx.Timeout(600, read=600)) + try: + yield _chunk({"role": "assistant"}) + async with client.stream( + "POST", + f"{node.http_base}/pd/decode", + json={ + "rid": rid, + "first_token_id": first_token_id, + "max_tokens": _max_tokens_of(body), + "sampling": _sampling_of(body), + "stream": True, + }, + ) as resp: + resp.raise_for_status() + async for line in resp.aiter_lines(): + # Deterministic client-liveness check: writes to a + # dead socket do NOT raise (verified by drill), so + # poll the ASGI disconnect state every line. + if await request.is_disconnected(): + client_gone = True + logger.info("client disconnected, cancelling %s", rid) + break + if not line: + continue + msg = json.loads(line) + if "t" in msg: + n_tokens += len(msg["t"]) + text = detok.push(msg["t"]) + if not text: + continue + if sess is None: + yield _chunk({"content": text}) + continue + for ev in sess.feed(text): + if ev["kind"] == "tool": + saw_tool = True + yield _chunk(_event_delta(ev)) + elif "done" in msg: + finish_reason = msg.get("finish_reason", "stop") + if finish_reason == "cancelled": + finish_reason = "stop" + elif "error" in msg: + yield _chunk({"content": f"\n[decode error: {msg['error']}]"}) + finish_reason = "stop" + if client_gone: + logger.info("client gone mid-stream for %s", rid) + return # finally fires the cancel + if sess is not None: + for ev in sess.finish(): + if ev["kind"] == "tool": + saw_tool = True + yield _chunk(_event_delta(ev)) + if saw_tool: + finish_reason = "tool_calls" + yield _chunk( + {}, + finish=finish_reason, + usage={ + "prompt_tokens": prompt_tokens, + "completion_tokens": n_tokens, + }, + ) + yield "data: [DONE]\n\n" + completed_ok = True + except Exception: + logger.exception("stream failed mid-flight for %s", rid) + finally: + # Runs under cancellation too (client disconnect cancels this + # task). Order matters: release first (sync, can't be + # cancelled), then best-effort cancel via a plain thread + # (an await here could be cancelled before firing), then a + # shielded aclose. + pool.release(node) + if not completed_ok: + threading.Thread(target=_fire_cancel, daemon=True).start() + with anyio.CancelScope(shield=True): + await client.aclose() + + return StreamingResponse(_gen(), media_type="text/event-stream") + + @app.post("/v1/chat/completions") + async def chat(request: Request): + from starlette.concurrency import run_in_threadpool + + body = await request.json() + if body.get("stream"): + return await _handle_stream("/v1/chat/completions", body, request) + # blocking work off the event loop (decode can take minutes) + return await run_in_threadpool(_handle, "/v1/chat/completions", body) + + @app.post("/v1/completions") + async def completions(request: Request): + from starlette.concurrency import run_in_threadpool + + body = await request.json() + if body.get("stream"): + return JSONResponse( + {"error": "streaming is supported on /v1/chat/completions"}, status_code=400 + ) + return await run_in_threadpool(_handle, "/v1/completions", body) + + return app # noqa: R504 (assembled across the function) + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s") + ap = argparse.ArgumentParser() + ap.add_argument("--vllm-url", required=True) + ap.add_argument( + "--decode", nargs="+", required=True, help="decode nodes as host:ctrl_port:http_port" + ) + ap.add_argument("--host", default="0.0.0.0") # nosec B104 (bind-all by design) + ap.add_argument("--port", type=int, default=23333) + ap.add_argument( + "--model-path", default="", help="tokenizer path (required unless --parser none)" + ) + ap.add_argument( + "--parser", + choices=["glm47", "none"], + default="glm47", + help="output parser (reasoning + tool calls)", + ) + args = ap.parse_args() + + nodes = [] + for spec in args.decode: + host, cport, hport = spec.rsplit(":", 2) + nodes.append(DecodeNode(host, int(cport), int(hport))) + + tokenizer = None + if args.model_path: + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained( + args.model_path, trust_remote_code=True + ) # nosec B615 + + ctx = RouterCtx(args.vllm_url, Pool(nodes), tokenizer, args.parser) + app = build_app(ctx) + logger.info( + "router on :%d -> vllm=%s, %d decode node(s), parser=%s", + args.port, + args.vllm_url, + len(nodes), + args.parser, + ) + uvicorn.run(app, host=args.host, port=args.port, log_level="warning") + + +if __name__ == "__main__": + main() diff --git a/tilert/pd_vllm/prefill_connector.py b/tilert/pd_vllm/prefill_connector.py new file mode 100644 index 0000000..abd4819 --- /dev/null +++ b/tilert/pd_vllm/prefill_connector.py @@ -0,0 +1,341 @@ +"""TileRT PD producer connector for vLLM prefill (model-agnostic framework). + +Loaded into vLLM via the official plugin surface: + + --kv-transfer-config '{ + "kv_connector": "TileRTConnector", + "kv_connector_module_path": "tilert.pd_vllm.prefill_connector", + "kv_role": "kv_producer", + "kv_connector_extra_config": {"tilert_host": "", + "tilert_ctrl_port": 5556, + "tilert_model": "glm5"} + }' + +Claim discipline (MultiConnector-safe): only requests whose +``kv_transfer_params`` carry ``tilert_host`` are claimed; everything else is a +strict no-op so a native connector can coexist. + +The connector owns the model-agnostic plumbing (claim, chunked-prefill +tracking, worker init, staging, background send, TCP handshake); all per-model +extraction / layout / RDMA planning is delegated to the selected model profile +(``tilert_model``, default ``glm5``). +""" + +import logging +import queue +import threading +from dataclasses import dataclass, field + +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorBase_V1, + KVConnectorMetadata, + SupportsHMA, +) + +from tilert.pd_vllm import wire +from tilert.pd_vllm.profiles import base as profiles +from tilert.pd_vllm.wire import derive_rid + +logger = logging.getLogger("pd_vllm.connector") + + +@dataclass +class _ReqMeta: + req_id: str + rid: str + num_tokens: int + last_prompt_token: int + block_ids_per_group: list + tilert_host: str + tilert_ctrl_port: int + sampling: dict | None = None + + +@dataclass +class TileRTMetadata(KVConnectorMetadata): + requests: list = field(default_factory=list) + + +@dataclass +class _Pending: + """Scheduler-side chunked-prefill accumulation.""" + + req_id: str + prompt_token_ids: list + total_tokens: int + block_ids_per_group: list + params: dict + + +class TileRTConnector(KVConnectorBase_V1, SupportsHMA): + # ══════════════════════════ init ══════════════════════════ + + def __init__(self, vllm_config, role, kv_cache_config=None): + super().__init__(vllm_config, role, kv_cache_config) + extra = vllm_config.kv_transfer_config.kv_connector_extra_config or {} + self._default_host = extra.get("tilert_host") + self._default_port = int(extra.get("tilert_ctrl_port", 5556)) + self._sync_send = bool(extra.get("tilert_sync_send", False)) + self._max_seq = int(extra.get("tilert_max_seq_len", vllm_config.model_config.max_model_len)) + self._profile = profiles.get_profile(extra.get("tilert_model", "glm5")) + self._transport_name = extra.get("tilert_transport", "mooncake") + + # scheduler-side + self._pending: dict[str, _Pending] = {} + + # worker-side (lazy) + self._kv_caches: dict = {} + self._reg = None # profile registration (layer map) + self._tp_rank: int | None = None + self._transport = None + self._staging = None + self._send_q: queue.Queue = queue.Queue() + self._sender_thread: threading.Thread | None = None + + logger.info( + "TileRTConnector: role=%s profile=%s target=%s:%s sync=%s", + role, + self._profile.name, + self._default_host, + self._default_port, + self._sync_send, + ) + + # ══════════════════════ scheduler side ═════════════════════ + + @staticmethod + def _claim(params) -> dict | None: + """Return kv_transfer_params if this request is ours, else None.""" + if params and isinstance(params, dict) and params.get("tilert_host"): + return params + return None + + def _params_of(self, new_req) -> dict | None: + sp = getattr(new_req, "sampling_params", None) + extra = getattr(sp, "extra_args", None) if sp is not None else None + if extra: + return self._claim(extra.get("kv_transfer_params")) + return None + + def get_num_new_matched_tokens(self, request, num_computed_tokens): + return 0, False + + def update_state_after_alloc(self, request, blocks, num_external_tokens): + pass + + def build_connector_meta(self, scheduler_output) -> KVConnectorMetadata: + meta = TileRTMetadata() + num_sched = scheduler_output.num_scheduled_tokens or {} + + for req_id in scheduler_output.finished_req_ids: + self._pending.pop(req_id, None) + for req_id in getattr(scheduler_output, "preempted_req_ids", None) or []: + self._pending.pop(req_id, None) + + for new_req in scheduler_output.scheduled_new_reqs: + params = self._params_of(new_req) + if params is None: + continue # not ours β€” strict no-op (MultiConnector safety) + token_ids = list(new_req.prompt_token_ids or []) + if not token_ids: + continue + groups = [list(g) for g in new_req.block_ids] + n = num_sched.get(new_req.req_id, 0) + if new_req.num_computed_tokens + n >= len(token_ids): + meta.requests.append(self._emit(new_req.req_id, token_ids, groups, params)) + else: + self._pending[new_req.req_id] = _Pending( + req_id=new_req.req_id, + prompt_token_ids=token_ids, + total_tokens=len(token_ids), + block_ids_per_group=groups, + params=params, + ) + + cached = scheduler_output.scheduled_cached_reqs + for i, req_id in enumerate(getattr(cached, "req_ids", []) or []): + p = self._pending.get(req_id) + if p is None: + continue + new_blocks = cached.new_block_ids[i] + if new_blocks is not None: + for gi, g in enumerate(new_blocks): + if gi < len(p.block_ids_per_group) and g: + p.block_ids_per_group[gi].extend(g) + n = num_sched.get(req_id, 0) + if cached.num_computed_tokens[i] + n >= p.total_tokens: + meta.requests.append( + self._emit(req_id, p.prompt_token_ids, p.block_ids_per_group, p.params) + ) + del self._pending[req_id] + return meta + + def _emit(self, req_id, token_ids, groups, params) -> _ReqMeta: + host = params.get("tilert_host") or self._default_host + assert host is not None, "claimed a request with no tilert_host" + m = _ReqMeta( + req_id=req_id, + rid=derive_rid(req_id), + num_tokens=len(token_ids), + last_prompt_token=int(token_ids[-1]), + block_ids_per_group=groups, + tilert_host=host, + tilert_ctrl_port=int(params.get("tilert_ctrl_port", self._default_port)), + sampling=params.get("sampling"), + ) + logger.info( + "claimed %s (rid=%s, %d tokens) -> %s:%d", + req_id, + m.rid, + m.num_tokens, + m.tilert_host, + m.tilert_ctrl_port, + ) + return m + + def request_finished(self, request, block_ids): + self._pending.pop(getattr(request, "request_id", ""), None) + return False, None + + def request_finished_all_groups(self, request, block_ids): + return self.request_finished(request, block_ids) + + # ══════════════════════ worker side ════════════════════════ + + def register_kv_caches(self, kv_caches): + self._kv_caches = kv_caches + cfg = getattr(self, "_kv_cache_config", None) + self._reg = self._profile.classify_layers(kv_caches, cfg) + + def _ensure_worker_ready(self) -> None: + if self._transport is not None: + return + import torch + from vllm.distributed import get_tensor_model_parallel_rank + + self._tp_rank = int(get_tensor_model_parallel_rank()) + + from tilert.pd_vllm.transport import make_transport + + hostname = wire.local_ip() + total = self._profile.staging_bytes(self._reg, self._tp_rank, self._max_seq) + dev = torch.cuda.current_device() + self._staging = torch.zeros(total, dtype=torch.uint8, device=f"cuda:{dev}") + + self._transport = make_transport(self._transport_name) + self._transport.init(hostname) + self._transport.register(self._staging.data_ptr(), total, dev) + + if not self._sync_send: + self._sender_thread = threading.Thread( + target=self._sender_loop, name="tilert-pd-sender", daemon=True + ) + self._sender_thread.start() + logger.info( + "worker ready: rank=%d transport=%s staging=%.1f MB profile=%s", + self._tp_rank, + self._transport.name, + total / 1e6, + self._profile.name, + ) + + def start_load_kv(self, forward_context, **kwargs): + pass + + def wait_for_layer_load(self, layer_name): + pass + + def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs): + pass + + def wait_for_save(self): + metadata = self._get_connector_metadata() + if not isinstance(metadata, TileRTMetadata) or not metadata.requests: + return + self._ensure_worker_ready() + assert self._tp_rank is not None # set by _ensure_worker_ready + if self._tp_rank not in self._profile.sender_ranks: + return # this rank does not participate (e.g. replicated MLA) + for m in metadata.requests: + try: + sections = self._profile.extract( + self._reg, m, self._tp_rank, self._staging, self._max_seq + ) + except Exception: + logger.exception("extraction failed for %s", m.rid) + continue + job = {"meta": m, "sections": sections, "seq": sections["seq"]} + if self._sync_send: + self._send(job) + else: + self._send_q.put(job) + + def get_finished(self, finished_req_ids): + return None, None + + # ── background send ── + + def _sender_loop(self) -> None: + while True: + job = self._send_q.get() + try: + self._send(job) + except Exception: + logger.exception("send failed for %s", job["meta"].rid) + + def _send(self, job: dict) -> None: + import socket as _socket + import time as _time + + # _send only runs after wait_for_save() -> _ensure_worker_ready() + assert self._transport is not None and self._staging is not None + assert self._tp_rank is not None + m: _ReqMeta = job["meta"] + seq = job["seq"] + t0 = _time.time() + conn = _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM) + try: + conn.setsockopt(_socket.IPPROTO_TCP, _socket.TCP_NODELAY, 1) + conn.settimeout(60) + conn.connect((m.tilert_host, m.tilert_ctrl_port)) + hello = wire.recv_msg(conn) + assert hello.get("magic") == wire.MAGIC, f"bad hello: {hello}" + assert hello.get("layout_version") == self._profile.layout_version, ( + f"layout version mismatch: {hello.get('layout_version')} " + f"vs {self._profile.layout_version}" + ) + assert hello.get("transport") == self._transport.name, ( + f"transport mismatch: decode={hello.get('transport')} " + f"vs prefill={self._transport.name}" + ) + remote_max_seq = int(hello["max_seq_len"]) + assert seq <= remote_max_seq, f"seq {seq} exceeds decode max_seq_len {remote_max_seq}" + + wire.send_msg( + conn, + { + "rid": m.rid, + "rank": self._tp_rank, + "seq_len": seq, + "last_prompt_token": m.last_prompt_token, + "sampling": m.sampling, + }, + ) + + base = self._staging.data_ptr() + srcs, dsts, lens = self._profile.rdma_plan( + hello, job["sections"], self._tp_rank, seq, base + ) + self._transport.write(hello, srcs, dsts, lens) + + wire.send_msg(conn, {"done": True, "rid": m.rid, "rank": self._tp_rank}) + logger.info( + "sent %s: rank=%d seq=%d %.1f MB in %.1f ms", + m.rid, + self._tp_rank, + seq, + sum(lens) / 1e6, + 1000 * (_time.time() - t0), + ) + finally: + conn.close() diff --git a/tilert/pd_vllm/profiles/__init__.py b/tilert/pd_vllm/profiles/__init__.py new file mode 100644 index 0000000..fe663bf --- /dev/null +++ b/tilert/pd_vllm/profiles/__init__.py @@ -0,0 +1 @@ +"""Model profiles for the PD data plane.""" diff --git a/tilert/pd_vllm/profiles/base.py b/tilert/pd_vllm/profiles/base.py new file mode 100644 index 0000000..f1231d7 --- /dev/null +++ b/tilert/pd_vllm/profiles/base.py @@ -0,0 +1,97 @@ +"""ModelProfile seam: everything model-specific in the PD data plane. + +The framework (prefill connector plumbing, receive server + control plane, +decode server orchestration, router) is model-agnostic and calls into the +active profile for the parts that differ between models: + + GLM-5 : replicated MLA latent KV + NSA KI index + MTP draft + DeepSeek-V3.2 : replicated MLA latent KV + NSA KI index + MTP draft +""" + +from __future__ import annotations + +from typing import Any, Protocol + + +class ModelProfile(Protocol): + name: str + num_ranks: int + sender_ranks: frozenset + + @property + def layout_version(self) -> int: ... + + # ── receive side (decode node) ─────────────────────────────────────── + def buffer_bytes(self, max_seq_len: int) -> int: + """Total receive-buffer size for one request slot.""" + + def hello_layout(self, base_ptr: int, max_seq_len: int) -> dict[str, int]: + """Region base addresses, merged into the hello message. + + Tells the sender where to RDMA-write each section. + """ + + def convert( + self, buffer: Any, base_ptr: int, max_seq_len: int, received: Any, num_devices: int + ) -> Any: + """Received buffer -> native per-device tensors (ConvertedRequest).""" + + # ── prefill side (vLLM connector worker) ───────────────────────────── + def classify_layers(self, kv_caches: dict, kv_cache_config: Any) -> Any: + """Inspect registered kv_caches and return an opaque registration. + + The framework passes it back to ``staging_bytes``/``extract``. Raise on + an unexpected layer set (e.g. missing speculative layer). + """ + + def staging_bytes(self, reg: Any, tp_rank: int, max_seq_len: int) -> int: + """Per-rank staging-buffer size.""" + + def extract(self, reg: Any, req_meta: Any, tp_rank: int, staging, max_seq_len: int) -> Any: + """Copy this rank's KV out of the paged caches into ``staging``. + + Runs inside the forward window; returns opaque ``sections``. + """ + + def rdma_plan( + self, hello: dict, sections: Any, tp_rank: int, seq_len: int, staging_base: int + ) -> tuple[list, list, list]: + """(src_ptrs, dst_ptrs, lengths) for one mooncake batch write.""" + + # ── engine (decode node) ───────────────────────────────────────────── + def build_engine( + self, model_weights_dir: str, max_seq_len: int, with_mtp: bool, ar_steps: int + ) -> Any: + """Construct the decode engine adapter (inject/decode/reset).""" + + +_REGISTRY: dict[str, ModelProfile] = {} +_ALIASES = { + "glm5": "glm5", + "glm_5": "glm5", + "glm-5": "glm5", + "dsv32": "dsv32", + "deepseek_v3_2": "dsv32", + "deepseek-v3.2": "dsv32", + "dsv3.2": "dsv32", + "v32": "dsv32", +} + + +def register(profile: ModelProfile) -> None: + _REGISTRY[profile.name] = profile + + +def get_profile(name: str) -> ModelProfile: + canon = _ALIASES.get(name, name) + if canon not in _REGISTRY: + # lazy import so a profile's heavy deps load only when selected + if canon == "glm5": + from tilert.pd_vllm.profiles import glm5 # noqa: F401 + elif canon == "dsv32": + from tilert.pd_vllm.profiles import dsv32 # noqa: F401 + if canon not in _REGISTRY: + raise KeyError( + f"unknown model profile {name!r}; " f"accepted keys (incl. aliases): {sorted(_ALIASES)}" + ) + return _REGISTRY[canon] diff --git a/tilert/pd_vllm/profiles/dsv32.py b/tilert/pd_vllm/profiles/dsv32.py new file mode 100644 index 0000000..5691772 --- /dev/null +++ b/tilert/pd_vllm/profiles/dsv32.py @@ -0,0 +1,41 @@ +"""DeepSeek-V3.2 profile β€” thin config over the shared MLA+NSA data plane.""" + +from __future__ import annotations + +from tilert.pd_vllm.profiles import base +from tilert.pd_vllm.profiles.mla_nsa import ( + MlaNsaEngineAdapter, + MlaNsaProfile, +) + +NUM_LAYERS = 62 # 61 main + 1 MTP draft (HF: 61 hidden + 1 nextn) +LAYOUT_VERSION = 11 # dsv32 wire family (distinct from glm5's 10) + + +def _build_engine(model_weights_dir, max_seq_len, with_mtp, ar_steps): + import tilert + + if hasattr(tilert, "load_backend"): + tilert.load_backend("deepseek_v3_2") # multi-backend builds only + from tilert.models.deepseek_v3_2.generator import DSAv32Generator + from tilert.models.deepseek_v3_2.model_args import ModelArgs + + gen = DSAv32Generator( + model_args=ModelArgs(), + max_new_tokens=max(max_seq_len - 256, 4096 - 256), + model_weights_dir=model_weights_dir, + with_mtp=with_mtp, + use_topp=True, + ) + gen.from_pretrained() + return MlaNsaEngineAdapter(gen, with_mtp) + + +base.register( + MlaNsaProfile( + name="dsv32", + num_layers=NUM_LAYERS, + layout_version=LAYOUT_VERSION, + engine_factory=_build_engine, + ) +) diff --git a/tilert/pd_vllm/profiles/glm5.py b/tilert/pd_vllm/profiles/glm5.py new file mode 100644 index 0000000..c98d7e9 --- /dev/null +++ b/tilert/pd_vllm/profiles/glm5.py @@ -0,0 +1,49 @@ +"""GLM-5 profile β€” thin config over the shared MLA+NSA data plane. + +GLM-5 = 79 cache layers (78 main + 1 MTP draft), MLA latent KV + NSA KI index. +All layout / convert / extract / RDMA logic lives in ``mla_nsa``; this file +only pins the layer count, wire version, and the GLM5Generator engine build. +""" + +from __future__ import annotations + +from tilert.pd_vllm.profiles import base +from tilert.pd_vllm.profiles.mla_nsa import ( + MlaNsaEngineAdapter, + MlaNsaProfile, +) + +NUM_LAYERS = 79 # 78 main + 1 MTP draft +LAYOUT_VERSION = 10 # glm5 wire family + + +def _build_engine(model_weights_dir, max_seq_len, with_mtp, ar_steps): + import tilert + + # multi-backend builds (tilert>=0.1.x) load the per-model .so on demand; + # single-backend builds auto-register on import and lack load_backend. + if hasattr(tilert, "load_backend"): + tilert.load_backend("glm5") + from tilert.models.glm_5.generator import GLM5Generator + from tilert.models.glm_5.model_args import ModelArgsGLM5 + + gen = GLM5Generator( + model_args=ModelArgsGLM5(), + max_new_tokens=max(max_seq_len - 256, 4096 - 256), + model_weights_dir=model_weights_dir, + with_mtp=with_mtp, + use_topp=True, + enable_thinking=False, + ) + gen.from_pretrained() + return MlaNsaEngineAdapter(gen, with_mtp) + + +base.register( + MlaNsaProfile( + name="glm5", + num_layers=NUM_LAYERS, + layout_version=LAYOUT_VERSION, + engine_factory=_build_engine, + ) +) diff --git a/tilert/pd_vllm/profiles/mla_nsa.py b/tilert/pd_vllm/profiles/mla_nsa.py new file mode 100644 index 0000000..a270f7c --- /dev/null +++ b/tilert/pd_vllm/profiles/mla_nsa.py @@ -0,0 +1,486 @@ +"""Shared MLA + NSA-KI data plane for the DeepSeek-family models.""" + +from __future__ import annotations + +import logging +import re +from dataclasses import dataclass + +import torch + +from tilert.pd_vllm import wire + +logger = logging.getLogger("pd_vllm.profile.mla_nsa") + +KV_LORA_RANK = 512 +QK_ROPE_HEAD_DIM = 64 +INDEX_HEAD_DIM = 128 +KI_QUANT_BLOCK = 128 +KV_QUANT_BLOCK = 128 # per-128 fp8 scale on the kv latent +PAGE_SIZE = 64 + +# The MLA KV cache dtype is a launch choice (vLLM ``--kv-cache-dtype``), NOT +# tied to the (fp8) model weights β€” both are supported and selected at runtime: +# +# fp8_ds_mla : cache tensor [nblk, page, 656] u8; per token 512 fp8 kv_c + +# 16 B (4 fp32) scale + 128 B bf16 k_pe. Split into a 528-B +# kv_merged plane + 128-B pe plane; kv dequantized fp8->bf16 on +# the decode side. (recommended, aligns with SGLang fp8) +# bf16 : cache tensor [nblk, page, 576] bf16; per token 512 bf16 kv_c + +# 64 bf16 k_pe = 1024-B kv plane + 128-B pe plane, no dequant. +KV_FP8_BYTES = KV_LORA_RANK # 512 (fp8, 1 B each) +KV_SCALE_BYTES = KV_LORA_RANK // KV_QUANT_BLOCK * 4 # 16 (4 fp32 scales) +KV_BYTES_FP8 = KV_FP8_BYTES + KV_SCALE_BYTES # 528 B/token +KV_BYTES_BF16 = KV_LORA_RANK * 2 # 1024 B/token +PE_BPT = QK_ROPE_HEAD_DIM * 2 # 128 B/token bf16 (both) +MLA_BPT_FP8 = KV_BYTES_FP8 + PE_BPT # 656 (fp8 cache stride) +MLA_BPT_BF16 = (KV_LORA_RANK + QK_ROPE_HEAD_DIM) * 2 # 1152 (bf16 cache stride) +_VERSION_BF16_OFFSET = 40 # bf16 layout_version = base + 40 +KI_PAGE_BYTES = ( + PAGE_SIZE * INDEX_HEAD_DIM + PAGE_SIZE * INDEX_HEAD_DIM // KI_QUANT_BLOCK * 4 +) # 8448 + + +def _max_pages(max_seq_len: int) -> int: + return (max_seq_len + PAGE_SIZE - 1) // PAGE_SIZE + + +def _hadamard(x: torch.Tensor) -> torch.Tensor: + """Hadamard rotation of the last dim (scale d^-0.5), matching TileRT's indexer. + + Uses fast_hadamard_transform if present, else a scipy matmul. + """ + d = x.shape[-1] + try: + from fast_hadamard_transform import hadamard_transform + + return hadamard_transform(x, scale=d**-0.5) + except Exception: + from scipy.linalg import hadamard as _h + + H = torch.from_numpy(_h(d).astype("float32")).to(x.device) * (d**-0.5) + return (x.float() @ H).to(x.dtype) + + +@dataclass +class ConvertedRequest: + rid: str + seq_len: int + last_prompt_token: int + first_token_id: int | None + sampling: dict | None + layers: list # [(ki[seq,128], kv[seq,512], pe[seq,64]) bf16] x num_layers + + +@dataclass +class _Reg: + mla_layers: list # [(lid, name, kv_t, gi)] sorted + ki_layers: list # [(lid, name, ki_t, gi)] sorted + + +class MlaNsaProfile: + """Config-driven MLA+NSA profile. + + ``engine_factory(weights, max_seq, with_mtp, ar_steps) -> adapter`` builds + the model-specific engine. + """ + + num_ranks = wire.NUM_RANKS + sender_ranks = frozenset({0}) # MLA latent replicated across TP + + def __init__( + self, name: str, num_layers: int, layout_version: int, engine_factory, mla_fp8: bool = True + ): + self.name = name + self.num_layers = num_layers + self._base_version = layout_version + self._engine_factory = engine_factory + self.mla_fp8 = mla_fp8 # fp8_ds_mla (True) vs bf16 (False) MLA cache + + def configure(self, kv_cache_dtype: str) -> MlaNsaProfile: + """Select the MLA cache dtype (decode side; prefill auto-detects).""" + d = (kv_cache_dtype or "").lower() + if d in ("fp8_ds_mla", "fp8", "fp8_e4m3"): + self.mla_fp8 = True + elif d in ("bf16", "bfloat16", "auto"): + self.mla_fp8 = False + else: + raise ValueError( + f"unknown kv_cache_dtype {kv_cache_dtype!r}; " f"want fp8_ds_mla or bf16" + ) + return self + + @property + def layout_version(self) -> int: + # distinct wire version per cache dtype so a mismatched pairing + # (prefill fp8 vs decode bf16) is rejected at hello, not corrupted + return self._base_version + (0 if self.mla_fp8 else _VERSION_BF16_OFFSET) + + @property + def _kv_bpt(self) -> int: + return KV_BYTES_FP8 if self.mla_fp8 else KV_BYTES_BF16 + + @property + def _mla_bpt(self) -> int: + return MLA_BPT_FP8 if self.mla_fp8 else MLA_BPT_BF16 + + # ── plane sizing (depends on num_layers + cache dtype) ── + def _kv_plane(self, max_seq_len: int) -> int: + return self.num_layers * max_seq_len * self._kv_bpt + + def _pe_plane(self, max_seq_len: int) -> int: + return self.num_layers * max_seq_len * PE_BPT + + def _ki_plane(self, max_seq_len: int) -> int: + return self.num_layers * _max_pages(max_seq_len) * KI_PAGE_BYTES + + # ── receive side ── + def buffer_bytes(self, max_seq_len: int) -> int: + return ( + self._kv_plane(max_seq_len) + self._pe_plane(max_seq_len) + self._ki_plane(max_seq_len) + ) + + def hello_layout(self, base_ptr: int, max_seq_len: int) -> dict[str, int]: + kv = base_ptr + pe = kv + self._kv_plane(max_seq_len) + ki = pe + self._pe_plane(max_seq_len) + return {"kv_base": kv, "pe_base": pe, "ki_base": ki} + + @torch.inference_mode() + def convert(self, buffer, base_ptr, max_seq_len, received, num_devices=1): + seq = received.seq_len + npages = _max_pages(seq) + pe_base = self._kv_plane(max_seq_len) + ki_base = pe_base + self._pe_plane(max_seq_len) + kv_bpt = self._kv_bpt + layers = [] + for lid in range(self.num_layers): + ko = lid * max_seq_len * kv_bpt + kv_raw = buffer[ko : ko + seq * kv_bpt].view(seq, kv_bpt) + if self.mla_fp8: + kv = self._dequant_kv(kv_raw, seq) # fp8+scale -> bf16 512 + else: + kv = ( + kv_raw.view(torch.bfloat16).view(seq, KV_LORA_RANK).contiguous() + ) # already bf16 + po = pe_base + lid * max_seq_len * PE_BPT + pe = ( + buffer[po : po + seq * PE_BPT] + .view(torch.bfloat16) + .view(seq, QK_ROPE_HEAD_DIM) + .contiguous() + ) + io = ki_base + lid * _max_pages(max_seq_len) * KI_PAGE_BYTES + ki_raw = buffer[io : io + npages * KI_PAGE_BYTES].view(npages, KI_PAGE_BYTES) + layers.append((self._dequant_ki(ki_raw, seq), kv, pe)) + torch.cuda.synchronize() + return ConvertedRequest( + rid=received.rid, + seq_len=seq, + last_prompt_token=received.last_prompt_token, + first_token_id=received.first_token_id, + sampling=received.sampling, + layers=layers, + ) + + @staticmethod + def _dequant_kv(kv_raw: torch.Tensor, seq_len: int) -> torch.Tensor: + """Dequantize kv_merged [seq,528] u8 (512 fp8 + 4 fp32 scale) -> bf16 [seq,512]. + + Per-128-block scale: kv[:, b*128:(b+1)*128] *= scale[:, b]. + """ + nblk = KV_LORA_RANK // KV_QUANT_BLOCK + fp8 = ( + kv_raw[:, :KV_FP8_BYTES] + .reshape(-1) + .view(torch.float8_e4m3fn) + .reshape(seq_len, KV_LORA_RANK) + ) + scale = ( + kv_raw[:, KV_FP8_BYTES:] + .reshape(-1) + .contiguous() + .view(torch.float32) + .reshape(seq_len, nblk) + ) + fp32 = fp8.float().view(seq_len, nblk, KV_QUANT_BLOCK) + deq = (fp32 * scale.unsqueeze(-1)).view(seq_len, KV_LORA_RANK) + return deq.to(torch.bfloat16) + + @staticmethod + def _dequant_ki(ki_raw: torch.Tensor, seq_len: int) -> torch.Tensor: + npages = ki_raw.shape[0] + fp8_bytes = PAGE_SIZE * INDEX_HEAD_DIM + ki_fp8 = ( + ki_raw[:, :fp8_bytes] + .reshape(-1) + .view(torch.float8_e4m3fn) + .reshape(npages * PAGE_SIZE, INDEX_HEAD_DIM) + ) + scale = ( + ki_raw[:, fp8_bytes:] + .reshape(-1) + .contiguous() + .view(torch.float32) + .reshape(npages * PAGE_SIZE, INDEX_HEAD_DIM // KI_QUANT_BLOCK) + ) + deq = (ki_fp8[:seq_len].float() * scale[:seq_len]).to(torch.bfloat16) + return _hadamard(deq) + + # ── prefill side ── + def classify_layers(self, kv_caches: dict, kv_cache_config) -> _Reg: + group_of = {} + for gi, g in enumerate(getattr(kv_cache_config, "kv_cache_groups", []) or []): + for ln in getattr(g, "layer_names", []): + group_of[ln] = gi + + def lid_of(name): + m = re.search(r"\.(\d+)\.", name) + base_i = int(m.group(1)) if m else -1 + return self.num_layers - 1 if name.startswith("mtp.") else base_i + + mla, ki = [], [] + for name, cache in kv_caches.items(): + t = cache[0] if isinstance(cache, (tuple, list)) else cache + gi = group_of.get(name, -1) + if "indexer" in name.lower() or "index_k" in name.lower(): + ki.append((lid_of(name), name, t, gi)) + else: + mla.append((lid_of(name), name, t, gi)) + mla.sort(key=lambda x: x[0]) + ki.sort(key=lambda x: x[0]) + if len(mla) != self.num_layers or len(ki) != self.num_layers: + raise RuntimeError( + f"{self.name} classify: {len(mla)} MLA + {len(ki)} KI layers " + f"(expected {self.num_layers} each); check --speculative-config" + f" and the vLLM layer naming" + ) + # auto-detect MLA cache dtype from the actual cache stride (the prefill + # cache is ground truth; the decode side is told via --kv-cache-dtype) + t0 = mla[0][2] + bpt = t0.shape[-1] * t0.element_size() + if bpt == MLA_BPT_FP8: + self.mla_fp8 = True + elif bpt == MLA_BPT_BF16: + self.mla_fp8 = False + else: + raise RuntimeError( + f"{self.name}: unexpected MLA cache stride {bpt} B/token; " + f"expected {MLA_BPT_FP8} (fp8_ds_mla) or {MLA_BPT_BF16} (bf16)" + ) + logger.info( + "%s registered %d MLA + %d KI layers, MLA cache=%s", + self.name, + len(mla), + len(ki), + "fp8_ds_mla" if self.mla_fp8 else "bf16", + ) + return _Reg(mla_layers=mla, ki_layers=ki) + + def staging_bytes(self, reg, tp_rank, max_seq_len): + if tp_rank not in self.sender_ranks: + return 4 + return self.buffer_bytes(max_seq_len) + + @torch.inference_mode() + def extract(self, reg: _Reg, m, tp_rank, staging, max_seq_len): + torch.cuda.synchronize() + seq = m.num_tokens + npages = _max_pages(seq) + mla_ids = m.block_ids_per_group[reg.mla_layers[0][3]] + bt = torch.tensor(mla_ids, dtype=torch.long) + offs = torch.arange(PAGE_SIZE) + slots = (offs.reshape(1, -1) + bt.reshape(-1, 1) * PAGE_SIZE).flatten()[:seq] + ki_ids = m.block_ids_per_group[reg.ki_layers[0][3]] + ki_bt = torch.tensor(ki_ids[:npages], dtype=torch.long) + + pe_base = self._kv_plane(max_seq_len) + ki_base = pe_base + self._pe_plane(max_seq_len) + kv_bpt, mla_bpt = self._kv_bpt, self._mla_bpt + for lid in range(self.num_layers): + # raw-byte split of the MLA cache row: works for both dtypes + # (fp8 656 -> 528+128, bf16 1152 -> 1024+128), no conversion here + kv_t = reg.mla_layers[lid][2] + raw = kv_t if kv_t.dtype == torch.uint8 else kv_t.view(torch.uint8) + flat = raw.reshape(-1, mla_bpt) # [ntok, mla_bpt] u8 + rows = flat[slots.to(flat.device)] # [seq, mla_bpt] u8 + kv_merged = rows[:, :kv_bpt].contiguous() # kv_c (+scale) + pe = rows[:, kv_bpt:].contiguous() # 64 bf16 (128 B) + ko = lid * max_seq_len * kv_bpt + po = pe_base + lid * max_seq_len * PE_BPT + staging[ko : ko + seq * kv_bpt].copy_(kv_merged.flatten()) + staging[po : po + seq * PE_BPT].copy_(pe.flatten()) + + ki_t = reg.ki_layers[lid][2] + ki_pages = ki_t[ki_bt.to(ki_t.device)].reshape(npages, -1) + io = ki_base + lid * _max_pages(max_seq_len) * KI_PAGE_BYTES + staging[io : io + npages * KI_PAGE_BYTES].copy_( + ki_pages.contiguous().view(torch.uint8).flatten() + ) + torch.cuda.synchronize() + return {"seq": seq, "npages": npages, "stage_max": max_seq_len} + + def rdma_plan(self, hello, sections, tp_rank, seq_len, base): + remote_max = int(hello["max_seq_len"]) + stage_max = sections["stage_max"] + npages = sections["npages"] + srcs, dsts, lens = [], [], [] + s_pe = self._kv_plane(stage_max) + s_ki = s_pe + self._pe_plane(stage_max) + kv_bpt = self._kv_bpt + r_kv, r_pe, r_ki = (int(hello["kv_base"]), int(hello["pe_base"]), int(hello["ki_base"])) + for lid in range(self.num_layers): + srcs.append(base + lid * stage_max * kv_bpt) + dsts.append(r_kv + lid * remote_max * kv_bpt) + lens.append(seq_len * kv_bpt) + srcs.append(base + s_pe + lid * stage_max * PE_BPT) + dsts.append(r_pe + lid * remote_max * PE_BPT) + lens.append(seq_len * PE_BPT) + srcs.append(base + s_ki + lid * _max_pages(stage_max) * KI_PAGE_BYTES) + dsts.append(r_ki + lid * _max_pages(remote_max) * KI_PAGE_BYTES) + lens.append(npages * KI_PAGE_BYTES) + return srcs, dsts, lens + + def build_engine(self, model_weights_dir, max_seq_len, with_mtp, ar_steps): + return self._engine_factory(model_weights_dir, max_seq_len, with_mtp, ar_steps) + + +class MlaNsaEngineAdapter: + """Shared decode adapter for GLM-5 / DSV3.2 (same inject + 3-phase MTP). + + ``generator`` is a ready ``from_pretrained``'d GLM5Generator / + DSAv32Generator; both expose inject_cache / set_cur_pos / decode_layer + with forward / get_next_draft_tokens / get_num_accepted / + get_predicted_tokens / reset_sequence, and share DSV3.2's TOKEN_OUT index. + """ + + def __init__(self, generator, with_mtp: bool): + import torch as _torch + + self._torch = _torch + self.gen = generator + self.with_mtp = with_mtp + self.mtp_seq_len = getattr(generator, "mtp_seq_len", 4) + self.max_seq_len = getattr(generator.decode_layer, "max_seq_len", 200000) + self.last_stats: dict = {} + self.stop_ids = self._resolve_stop_ids(generator) + + @staticmethod + def _resolve_stop_ids(generator) -> set: + # GLM-5 exposes a stop_token_ids set; DSV3.2 exposes only eos_id. + sids = getattr(generator, "stop_token_ids", None) + if sids: + return set(sids) + eos = getattr(generator, "eos_id", None) + return {int(eos)} if eos is not None else set() + + def inject(self, req) -> None: + self.gen.inject_cache(req.layers, start_pos=0) + self.gen.set_cur_pos(req.seq_len - 1) + self._last_prompt_token = req.last_prompt_token + self._seq_len = req.seq_len + + def decode(self, first_token_id, max_tokens, sampling, on_token=None, cancel_event=None): + sampling = sampling or {} + temp = float(sampling.get("temperature", 1.0)) + if temp < 1e-5: + self.gen.update_sampling_params(temperature=1.0, top_p=1.0, top_k=1, use_topp=False) + else: + self.gen.update_sampling_params( + temperature=temp, + top_p=float(sampling.get("top_p", 0.95)), + top_k=int(sampling.get("top_k", 256)), + use_topp=True, + ) + budget = min(int(max_tokens), self.max_seq_len - self._seq_len - 1) + if budget <= 0: + self.last_stats = {"finish_reason": "length"} + return [int(first_token_id)] + if self.with_mtp: + return self._decode_mtp(first_token_id, budget, on_token, cancel_event) + return self._decode_standard(first_token_id, budget, on_token, cancel_event) + + def _decode_mtp(self, first_token_id, budget, on_token, cancel_event): + dl = self.gen.decode_layer + T = self.mtp_seq_len + stop_ids = self.stop_ids + torch = self._torch + tokens = [int(first_token_id)] + if on_token: + on_token(int(first_token_id)) + if int(first_token_id) in stop_ids: + self.last_stats = {"finish_reason": "stop"} + return [] + dl.set_prefill_valid_tokens(0) + draft = torch.full((1, T), int(self._last_prompt_token), dtype=torch.int32, device="cuda:0") + accepted, finish, fwd, finished = [], "length", 0, False + while not finished and len(tokens) < budget: + if cancel_event is not None and cancel_event.is_set(): + finish = "cancelled" + break + if fwd == 1: + draft = torch.full((1, T), int(first_token_id), dtype=torch.int32, device="cuda:0") + elif fwd > 1: + draft = dl.get_next_draft_tokens(0).reshape(1, T) + dl.forward(draft) + n_acc = dl.get_num_accepted(0) + pred = dl.get_predicted_tokens(0).flatten() + if fwd == 0: + fwd += 1 + continue + accepted.append(n_acc) + fwd += 1 + for i in range(n_acc): + if len(tokens) >= budget: + break + tok = int(pred[i].item()) + if tok in stop_ids: + finished = True + finish = "stop" + break + tokens.append(tok) + if on_token: + on_token(tok) + dl.reset_sequence() + self.last_stats = { + "finish_reason": finish, + "mtp_accept_mean": round(sum(accepted) / max(1, len(accepted)), 3), + "mtp_verify_calls": len(accepted), + } + return tokens + + def _decode_standard(self, first_token_id, budget, on_token, cancel_event): + from tilert.models.deepseek_v3_2.temp_var_indices import Idx + + dl = self.gen.decode_layer + stop_ids = self.stop_ids + torch = self._torch + tokens = [int(first_token_id)] + if on_token: + on_token(int(first_token_id)) + if int(first_token_id) in stop_ids: + self.last_stats = {"finish_reason": "stop"} + return [] + finish = "length" + cur = torch.tensor(int(first_token_id), dtype=torch.long, device="cuda:0") + while len(tokens) < budget: + if cancel_event is not None and cancel_event.is_set(): + finish = "cancelled" + break + res = dl.forward(cur) + intermediates, *_ = res[0] + nxt = intermediates[Idx.TOKEN_OUT][0][0] + tok = int(nxt.item()) + if tok in stop_ids: + finish = "stop" + break + tokens.append(tok) + if on_token: + on_token(tok) + cur = nxt + dl.reset_sequence() + self.last_stats = {"finish_reason": finish} + return tokens + + def reset(self) -> None: + pass diff --git a/tilert/pd_vllm/receive_server.py b/tilert/pd_vllm/receive_server.py new file mode 100644 index 0000000..3ad06f2 --- /dev/null +++ b/tilert/pd_vllm/receive_server.py @@ -0,0 +1,190 @@ +"""Decode-side receive server (W4): Mooncake buffer + TCP control plane.""" + +import contextlib +import logging +import queue +import socket +import threading +import time +from dataclasses import dataclass, field + +import torch + +from tilert.pd_vllm import wire + +logger = logging.getLogger("pd_vllm.receive") + + +@dataclass +class ReceivedRequest: + rid: str + seq_len: int + last_prompt_token: int + first_token_id: int | None + sampling: dict | None + done_ranks: set = field(default_factory=set) + t_first_conn: float = 0.0 + t_complete: float = 0.0 + + +class ReceiveServer: + def __init__( + self, + profile, + max_seq_len: int, + ctrl_port: int = 5556, + hostname: str | None = None, + device: str = "cuda:0", + request_timeout: float = 120.0, + transport: str = "mooncake", + ): + self.profile = profile + self.max_seq_len = max_seq_len + self.ctrl_port = ctrl_port + self.device = device + self.request_timeout = request_timeout + + total = profile.buffer_bytes(max_seq_len) + logger.info( + "allocating receive buffer: %.2f GB on %s (profile=%s)", + total / 1024**3, + device, + profile.name, + ) + self.buffer = torch.zeros(total, dtype=torch.uint8, device=device) + self.base_ptr = self.buffer.data_ptr() + self._hello_layout = profile.hello_layout(self.base_ptr, max_seq_len) + + # RDMA transport (mooncake default / nixl), single cuda:0 registration + from tilert.pd_vllm.transport import make_transport + + if hostname is None: + hostname = wire.local_ip() + dev_id = torch.device(device).index or 0 + self._transport = make_transport(transport) + self._transport.init(hostname) + self._transport.register(self.base_ptr, total, dev_id) + self._transport_meta = self._transport.local_meta() + logger.info( + "transport=%s ready, buffer registered (%.2f GB)", self._transport.name, total / 1024**3 + ) + + self._lock = threading.Lock() + self._current: ReceivedRequest | None = None + self.completed: queue.Queue[ReceivedRequest] = queue.Queue() + + # dual-stack: accept IPv4 (v4-mapped) and IPv6, incl. link-local peers + # (e.g. an IPv6-only decode node reached over fe80::.../bond0) + self._srv = socket.socket(socket.AF_INET6, socket.SOCK_STREAM) + self._srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + with contextlib.suppress(OSError): + self._srv.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0) + self._srv.bind(("::", ctrl_port)) + self._srv.listen(32) + self._stop = threading.Event() + self._thread = threading.Thread( + target=self._accept_loop, name="pd-recv-accept", daemon=True + ) + self._thread.start() + logger.info("control plane listening on :%d", ctrl_port) + + # ── public ─────────────────────────────────────────────────────────── + + def release(self) -> None: + """Mark the single receive slot free (call after inject/decode).""" + with self._lock: + self._current = None + + def close(self) -> None: + self._stop.set() + with contextlib.suppress(OSError): + self._srv.close() + + # ── accept / per-connection handling ───────────────────────────────── + + def _accept_loop(self) -> None: + while not self._stop.is_set(): + try: + conn, addr = self._srv.accept() + except OSError: + break + t = threading.Thread(target=self._handle, args=(conn, addr), daemon=True) + t.start() + + def _handle(self, conn: socket.socket, addr) -> None: + try: + conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + conn.settimeout(self.request_timeout) + # `busy` in hello is advisory (a same-rid rank must still proceed); + # the authoritative accept/reject happens once the rid is known. + with self._lock: + advisory_busy = self._current is not None and self._current.t_complete == 0.0 + wire.send_msg( + conn, + wire.hello_msg( + self._transport.name, + self._transport_meta, + self.max_seq_len, + self.profile.layout_version, + self._hello_layout, + busy=advisory_busy, + ), + ) + + req = wire.recv_msg(conn) + rid, rank = req["rid"], int(req["rank"]) + if req.get("seq_len", 0) > self.max_seq_len: + wire.send_msg(conn, {"error": "seq_len exceeds max_seq_len"}) + return + + with self._lock: + cur = self._current + if cur is None or cur.rid != rid: + if ( + cur is not None + and cur.t_complete == 0.0 + and time.time() - cur.t_first_conn < self.request_timeout + ): + # busy with a different in-flight rid + wire.send_msg(conn, {"error": "busy", "busy_rid": cur.rid}) + logger.warning("rejecting %s (busy with %s)", rid, cur.rid) + return + self._current = cur = ReceivedRequest( + rid=rid, + seq_len=int(req["seq_len"]), + last_prompt_token=int(req.get("last_prompt_token", 0)), + first_token_id=req.get("first_token_id"), + sampling=req.get("sampling"), + t_first_conn=time.time(), + ) + logger.info("request %s: seq_len=%d", rid, cur.seq_len) + + # wait for this rank's done (RDMA happens meanwhile) + done = wire.recv_msg(conn) + if not done.get("done"): + logger.warning("rank %d sent non-done message: %s", rank, done) + return + with self._lock: + cur = self._current + if cur is None or cur.rid != rid: + return + cur.done_ranks.add(rank) + logger.info( + "request %s: rank %d done (%d/%d)", + rid, + rank, + len(cur.done_ranks), + len(self.profile.sender_ranks), + ) + if cur.done_ranks >= set(self.profile.sender_ranks): + cur.t_complete = time.time() + self.completed.put(cur) + logger.info( + "request %s: all ranks done in %.1f ms", + rid, + 1000 * (cur.t_complete - cur.t_first_conn), + ) + except Exception: + logger.exception("connection from %s failed", addr) + finally: + conn.close() diff --git a/tilert/pd_vllm/transport.py b/tilert/pd_vllm/transport.py new file mode 100644 index 0000000..7d5e779 --- /dev/null +++ b/tilert/pd_vllm/transport.py @@ -0,0 +1,109 @@ +"""Pluggable RDMA transport for the PD data plane: Mooncake (default) or NIXL.""" + +from __future__ import annotations + +import base64 +import os + + +class Transport: + name = "?" + + def init(self, host: str) -> None: ... + def register(self, ptr: int, nbytes: int, dev_id: int) -> None: ... + def local_meta(self) -> dict: ... # type: ignore[empty-body] + def write(self, remote_meta: dict, srcs, dsts, lens) -> None: ... + + +class MooncakeTransport(Transport): + """serve_sglang precedent: one TransferEngine, P2P handshake, sync write.""" + + name = "mooncake" + + def init(self, host: str) -> None: + from mooncake.engine import TransferEngine + + self.engine = TransferEngine() + ret = self.engine.initialize(host, "P2PHANDSHAKE", "rdma", "") + if ret != 0: + raise RuntimeError(f"Mooncake engine init failed: {ret}") + self.session_id = f"{host}:{self.engine.get_rpc_port()}" + + def register(self, ptr: int, nbytes: int, dev_id: int) -> None: + ret = self.engine.batch_register_memory([ptr], [nbytes]) + if ret != 0: + raise RuntimeError(f"Mooncake register failed: {ret}") + + def local_meta(self) -> dict: + return {"session_id": self.session_id} + + def write(self, remote_meta: dict, srcs, dsts, lens) -> None: + ret = self.engine.batch_transfer_sync_write(remote_meta["session_id"], srcs, dsts, lens) + if ret != 0: + raise RuntimeError(f"mooncake write failed: {ret}") + + +class NixlTransport(Transport): + """NIXL agent over the UCX backend (GPUDirect RDMA). + + Registers VRAM regions with 4-tuple descriptors, exchanges agent metadata + via the hello, and issues WRITE transfers built from (src,dst,len) triples. + """ + + name = "nixl" + _MAX_POLL = 2_000_000 + + def init(self, host: str) -> None: + from nixl._api import nixl_agent, nixl_agent_config + + # agent name must be globally unique across the two peers + self._agent = nixl_agent(f"{host}:{os.getpid()}", nixl_agent_config(backends=["UCX"])) + self._remotes: dict[bytes, str] = {} # remote meta -> remote name + self._dev = 0 + + def register(self, ptr: int, nbytes: int, dev_id: int) -> None: + self._dev = dev_id + self._agent.register_memory([(ptr, nbytes, dev_id, "")], "VRAM") + + def local_meta(self) -> dict: + return { + "nixl_meta": base64.b64encode(self._agent.get_agent_metadata()).decode(), + "nixl_dev": self._dev, + } + + def write(self, remote_meta: dict, srcs, dsts, lens) -> None: + meta_b = base64.b64decode(remote_meta["nixl_meta"]) + rname = self._remotes.get(meta_b) + if rname is None: + rname = self._agent.add_remote_agent(meta_b) + self._remotes[meta_b] = rname + rdev = int(remote_meta.get("nixl_dev", 0)) + ld = self._agent.get_xfer_descs( + [(int(s), int(n), self._dev) for s, n in zip(srcs, lens)], "VRAM" + ) + rd = self._agent.get_xfer_descs( + [(int(d), int(n), rdev) for d, n in zip(dsts, lens)], "VRAM" + ) + h = self._agent.initialize_xfer("WRITE", ld, rd, rname) + try: + st = self._agent.transfer(h) + polls = 0 + while st not in ("DONE", "ERR"): + st = self._agent.check_xfer_state(h) + polls += 1 + if polls > self._MAX_POLL: + raise RuntimeError("nixl xfer timed out") + if st == "ERR": + raise RuntimeError("nixl xfer failed") + finally: + self._agent.release_xfer_handle(h) + + +_BACKENDS = {"mooncake": MooncakeTransport, "nixl": NixlTransport} + + +def make_transport(name: str | None) -> Transport: + key = (name or "mooncake").lower() + if key not in _BACKENDS: + raise ValueError(f"unknown transport {name!r}; " f"choices: {sorted(_BACKENDS)}") + return _BACKENDS[key]() diff --git a/tilert/pd_vllm/wire.py b/tilert/pd_vllm/wire.py new file mode 100644 index 0000000..284850c --- /dev/null +++ b/tilert/pd_vllm/wire.py @@ -0,0 +1,92 @@ +"""Shared control-plane protocol for vLLM-prefill -> TileRT-decode PD.""" + +import json +import socket +import struct + +MAGIC = "tilert-pd" + +NUM_RANKS = 8 +EXPECTED_RANKS = tuple(range(NUM_RANKS)) + + +def local_ip(probe_addr: str | None = None) -> str: + """Best-effort local IP for the mooncake session identity.""" + import os + + probe = probe_addr or os.environ.get("TILERT_PD_PROBE_ADDR", "8.8.8.8") + s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + try: + s.connect((probe, 1)) + return s.getsockname()[0] + finally: + s.close() + + +def derive_rid(request_id: str) -> str: + """Map a vLLM request/response id to the client-visible rid. + + Shared by the prefill connector (internal id) and the router (response id) + so both agree. + """ + rid = request_id + for prefix in ("chatcmpl-", "cmpl-"): + if rid.startswith(prefix): + rid = rid[len(prefix) :] + break + parts = rid.rsplit("-", 1) + if len(parts) == 2 and len(parts[1]) <= 8 and all(c in "0123456789abcdef" for c in parts[1]): + rid = parts[0] + parts = rid.rsplit("-", 1) + if len(parts) == 2 and parts[1].isdigit() and len(parts[1]) <= 3: + rid = parts[0] + return rid + + +def send_msg(sock: socket.socket, obj: dict) -> None: + data = json.dumps(obj).encode() + sock.sendall(struct.pack("!I", len(data)) + data) + + +def recv_msg(sock: socket.socket) -> dict: + hdr = _recv_exact(sock, 4) + (n,) = struct.unpack("!I", hdr) + if n > 16 << 20: + raise ValueError(f"control message too large: {n}") + return json.loads(_recv_exact(sock, n).decode()) + + +def _recv_exact(sock: socket.socket, n: int) -> bytes: + buf = b"" + while len(buf) < n: + chunk = sock.recv(n - len(buf)) + if not chunk: + raise ConnectionError("connection closed mid-message") + buf += chunk + return buf + + +def hello_msg( + transport: str, + transport_meta: dict, + max_seq_len: int, + layout_version: int, + layout: dict, + busy: bool, +) -> dict: + """Build the common hello envelope. + + ``transport`` names the RDMA backend and ``transport_meta`` carries its + connection info (mooncake: session_id; nixl: nixl_meta/nixl_dev). + ``layout`` carries profile-specific region base addresses (e.g. gdn_base / + gqa_k_base / kv_base). + """ + return { + "magic": MAGIC, + "layout_version": layout_version, + "transport": transport, + "max_seq_len": max_seq_len, + "busy": busy, + **transport_meta, + **layout, + }