diff --git a/.flake8 b/.flake8
index a247806..3b308d2 100644
--- a/.flake8
+++ b/.flake8
@@ -38,4 +38,17 @@ per-file-ignores =
# 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
+ #
+ # Since v0.1.6 the Python sources are recovered from the published wheel,
+ # which ships them without comments: per-line noqa's are not available and
+ # long string literals would need code edits to re-wrap. The remaining codes
+ # on the pd_vllm line and the entries below relax those checks per exported
+ # path instead. flake8 applies only the longest matching pattern, so each
+ # entry repeats the codes it still needs (e.g. F401 for __init__.py).
+ tilert/pd_vllm/*:U100,E704,B007,D205,D209,E501,E741,F401,R501,R504,R505,SIM105,SIM115,U101,VNE002,VNE003
+ tilert/models/glm_5_2_rocm/*:E501,F401,R504,SIM113,U100,VNE002
+ tilert/models/misc_rocm/*:E501,F401,U100
+ tilert/models/base_rocm.py:E501,U100
+ tilert/__init__.py:E501,F401
+ tilert/generate.py:E501
+ tilert/utils.py:E501
diff --git a/README.md b/README.md
index bd88e06..ada7261 100644
--- a/README.md
+++ b/README.md
@@ -20,7 +20,9 @@ ______________________________________________________________________
## π° News
-- π **2026-07-14 Β· [v0.1.5](https://github.com/tile-ai/TileRT/releases/tag/v0.1.5.post3) Released**. Introduce [**PD (prefillβdecode) disaggregation**](https://www.tilert.ai/blog/tilert-vllm-disaggregation.html) β vLLM prefill + TileRT decode, behind an OpenAI-compatible endpoint. Supported on GLM-5/5.1 and DeepSeek-V3.2.
+- π **2026-09-24 Β· [TileRT Takes the Top Spot on AgentX](https://www.tilert.ai/blog/tilert-amd-agentx.html)**. On SemiAnalysis's InferenceX **AgentX** leaderboard, TileRT powers **GLM-5.3** to **469 tok/s** of single-user generation on **million-token agent sessions** with **8Γ AMD Instinct MI355X** β **#1 in single-user performance**. Running the original **FP8** weights, it leads the second-place NVIDIA GB300 NVL72 FP4 submission by more than 100 tok/s. [v0.1.6](https://github.com/tile-ai/TileRT/releases/tag/v0.1.6.post2) brings ROCm support to the public release: GLM-5.2/5.3 on AMD Instinct MI350X/MI355X.
+
+- π **2026-07-14 Β· [v0.1.5](https://github.com/tile-ai/TileRT/releases/tag/v0.1.5.post2) Released**. Introduce [**PD (prefillβdecode) disaggregation**](https://www.tilert.ai/blog/tilert-vllm-disaggregation.html) β 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.
@@ -56,9 +58,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Γ 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).
+ GLM-5.2/5.3-FP8 token generation speed on 8Γ AMD Instinct MI350X with TileRT v0.1.6 (greedy decoding, bf16 KV and indexer caches). Output length 1K, input length 1Kβ1M. Bars compare TileRT without MTP, with MTP at depth 3 and average acceptance length 3.2, and the peak under best-case MTP acceptance (4.0).
______________________________________________________________________
@@ -66,28 +68,31 @@ ______________________________________________________________________
## Installation
> \[!IMPORTANT\]
-> 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.
+> TileRT v0.1.6 is distributed as a **pre-built binary wheel** that carries three backends: two CUDA backends (DeepSeek-V3.2, GLM-5/5.1) and one ROCm backend (GLM-5.2/5.3). Each backend is linked against the exact PyTorch ABI listed below, and `tilert.load_backend()` refuses a PyTorch of the wrong version or flavor (CUDA vs ROCm). Other combinations of Python, CUDA/ROCm, or PyTorch versions are **untested and not guaranteed to work** β please reproduce this environment for a supported setup.
+
+### Build environment of the v0.1.6 wheel
-### Build environment of the v0.1.5 wheel
+The official `tilert==0.1.6.post2` wheel on PyPI was compiled against the following stacks. Treat these as **hard requirements**, not lower bounds (`transformers` / `tokenizers` are lower bounds).
-The official `tilert==0.1.5.post3` wheel on PyPI was compiled against the following stack. Treat these as **hard requirements**, not lower bounds (`transformers` / `tokenizers` are lower bounds since v0.1.5.post2).
+| Component | NVIDIA β DeepSeek-V3.2, GLM-5/5.1 | AMD β GLM-5.2/5.3 |
+| ---------------- | --------------------------------------------------- | --------------------------------------------------- |
+| GPU | 8Γ NVIDIA **B200** | 8Γ AMD Instinct **MI350X** / **MI355X** |
+| Driver / runtime | NVIDIA driver supporting the **CUDA 13.2** runtime | **ROCm 7.14.0** |
+| Operating System | Linux **x86_64**, glibc **β₯ 2.28** (manylinux_2_28) | Linux **x86_64**, glibc **β₯ 2.28** (manylinux_2_28) |
+| Python | **3.12** | **3.12** |
+| PyTorch | **`torch==2.11.0+cu130`** | **`torch==2.12.0+rocm7.14.0`** |
+| `transformers` | **`>= 4.46.3`** | **`>= 4.46.3`** |
+| `tokenizers` | **`>= 0.20.3`** | **`>= 0.20.3`** |
-| Component | Pinned version |
-| ---------------- | --------------------------------------------------- |
-| GPU | 8Γ NVIDIA **B200** |
-| NVIDIA driver | Supports **CUDA 13.2** runtime |
-| Operating System | Linux **x86_64**, glibc **β₯ 2.28** (manylinux_2_28) |
-| Python | **3.12** |
-| PyTorch | **`torch==2.11.0+cu130`** |
-| `transformers` | **`>= 4.46.3`** |
-| `tokenizers` | **`>= 0.20.3`** |
+### Recommended: pre-built Docker images
-### Recommended: pre-built Docker image
+The pinned build environments above are preinstalled in our official
+images β this is the **recommended way to run v0.1.6** and avoids any
+version drift on the host.
-The pinned build environment above is preinstalled in our official image
-β 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:
+**NVIDIA (B200).** The image is mirrored to two registries; pull from
+whichever is reachable, then launch a container with all 8 B200 GPUs
+attached:
```bash
# GitHub Container Registry
@@ -95,29 +100,45 @@ docker pull ghcr.io/tile-ai/tilert:cu132-latest
# Docker Hub
docker pull tileai/tilert:cu132-latest
+
+docker run --rm -it --gpus all --ipc=host \
+ -v "$PWD":/workspace -w /workspace \
+ ghcr.io/tile-ai/tilert:cu132-latest
```
-Launch a container with all 8 B200 GPUs attached, then install the
-wheel inside:
+**AMD (MI350X / MI355X).** Pull the ROCm image and launch a container with the
+GPUs attached:
```bash
-docker run --rm -it --gpus all --ipc=host \
+docker pull ghcr.io/tile-ai/tilert-rocm-decode:0.1.6
+
+docker run --rm -it --ipc=host \
+ --device=/dev/kfd --device=/dev/dri --group-add video \
-v "$PWD":/workspace -w /workspace \
- ghcr.io/tile-ai/tilert:cu132-latest
+ ghcr.io/tile-ai/tilert-rocm-decode:0.1.6
+```
-# Inside the container β install from PyPI:
-pip install tilert==0.1.5.post3
+The ROCm image sets `HF_HUB_OFFLINE=1`; run `unset HF_HUB_OFFLINE` before
+downloading a checkpoint from inside the container.
+
+Inside either container, install the wheel. It accepts both PyTorch builds
+(`torch>=2.11,<2.13`), so pip keeps the one preinstalled in the image:
+
+```bash
+# From PyPI:
+pip install tilert==0.1.6.post2
# 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.5.post3/tilert-0.1.5.post3-cp312-cp312-manylinux_2_28_x86_64.whl
+pip install https://github.com/tile-ai/TileRT/releases/download/v0.1.6.post2/tilert-0.1.6.post2-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.5.post3 / torch 2.11.0+cu130 / cuda 13.0
+python -c "import tilert, torch; print('tilert', tilert.__version__, '/ torch', torch.__version__)"
+# NVIDIA β expected: tilert 0.1.6.post2 / torch 2.11.0+cu130
+# AMD β expected: tilert 0.1.6.post2 / torch 2.12.0+rocm7.14.0
```
Proceed to [Getting Started](#getting-started) to download and convert model weights.
@@ -130,11 +151,12 @@ Starting from release v0.1.3, TileRT no longer requires downloading pre-converte
### Step 2: Shard Weights with `weight_converter`
-The converter ships inside the `tilert` wheel. It rewrites the official HF
-checkpoint into TileRT's per-device layout β 8 shards, one per B200, with
-keys suffixed `*_dev_{0..7}` and a fresh `model.safetensors.index.json`.
-The runtime loads these shards directly; the original checkpoint is no
-longer needed after conversion.
+The converters ship inside the `tilert` wheel. For the NVIDIA models,
+`weight_converter` rewrites the official HF checkpoint into TileRT's
+per-device layout β 8 shards, one per B200, with keys suffixed
+`*_dev_{0..7}` and a fresh `model.safetensors.index.json`. The runtime
+loads these shards directly; the original checkpoint is no longer needed
+after conversion.
For **DeepSeek-V3.2**:
@@ -154,35 +176,51 @@ python -m tilert.models.preprocess.weight_converter \
--save_dir "/path/to/GLM-5-FP8-TileRT"
```
-`--model_dir` is the directory of the downloaded HF checkpoint;
-`--save_dir` is where the sharded TileRT-format weights will land.
+For **GLM-5.2/5.3** (AMD MI350X / MI355X), the ROCm port has its own converter:
-### Step 3: Register the Sharded Weights Path
+```bash
+python -m tilert.models.glm_5_2_rocm.weight_converter \
+ --model_dir "/path/to/GLM-5.2-FP8" \
+ --save_dir "/path/to/GLM-5.2-FP8-TileRT" \
+ --num_mtp 1
+```
-Either pass `--model-weights-dir ` on every `tilert.generate`
-invocation, or register the path once in `~/.tilert/config.toml` so the
-CLI picks it up automatically:
+`--num_mtp 1` also converts the MTP module, which MTP decoding needs;
+leave it out to convert the base model only.
-```toml
-[weights]
-deepseek_v3_2 = "/path/to/DeepSeek-V3.2-TileRT"
-glm5 = "/path/to/GLM-5-FP8-TileRT"
-```
+`--model_dir` is the directory of the downloaded HF checkpoint;
+`--save_dir` is where the sharded TileRT-format weights will land.
### Running the Generation Example
-The simplest entry point is the bundled CLI. Pick `--model deepseek_v3_2`
-or `--model glm5`; weights resolve from `~/.tilert/config.toml` or from
-an explicit `--model-weights-dir`:
+The simplest entry point is the bundled CLI. Pick `--model deepseek_v3_2`,
+`--model glm5`, or `--model glm5_2_rocm`, and point `--model-weights-dir`
+(required) at the converted weights:
```bash
-python -m tilert.generate --model deepseek_v3_2 --max-new-tokens 1000
+python -m tilert.generate --model glm5_2_rocm \
+ --model-weights-dir /path/to/GLM-5.2-FP8-TileRT \
+ --prompt "Tell me three jokes." --max-new-tokens 1000
```
+On AMD, `--max-seq-len` (default 8192) sets the KV-cache length; the
+NVIDIA models take it from their model args.
+
+> \[!WARNING\]
+> In the published `tilert==0.1.6.post2` wheel, the CLI fails for the
+> NVIDIA models (`--model deepseek_v3_2` / `--model glm5`) once the weights
+> are loaded, with
+> `TypeError: generate() got an unexpected keyword argument 'max_new_tokens'`.
+> The fix is in this repository's [`tilert/generate.py`](tilert/generate.py)
+> but not yet in a published wheel; until it is, drive the NVIDIA models
+> through the Python API below.
+
> \[!NOTE\]
-> 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
+> v0.1.6 ships **three independent backend libraries** β
+> `libtilert_dsv32.so` and `libtilert_glm5.so` (CUDA), and
+> `libtilert_glm52_rocm.so` (ROCm) β and loads exactly one per Python
+> process via `tilert.load_backend(model_type)`, which also checks the
+> running PyTorch's version and flavor (CUDA vs ROCm). Run different models in
> separate processes β they cannot coexist in a single interpreter.
To drive generation programmatically, load the backend first, then build
@@ -219,7 +257,10 @@ completion = generator.generate(prompt)
(For **GLM-5**, swap in `tilert.load_backend("glm5")` and
`from tilert.models.glm_5.generator import GLM5Generator` with
-`ModelArgsGLM5`.)
+`ModelArgsGLM5`. For **GLM-5.2/5.3** on AMD, use
+`tilert.load_backend("glm5_2_rocm")` and
+`from tilert.models.glm_5_2_rocm.generator import Glm52Generator` with
+`ModelArgsGlm52`; it takes `num_mtp=` (0, 1 or 3) instead of `with_mtp=`.)
For example, TileRT may generate:
@@ -238,12 +279,18 @@ This example demonstrates basic single-step autoregressive generation using the
### Running the Generation Example with Multi-Token Prediction (MTP)
-TileRT also supports Multi-Token Prediction (MTP), which allows the model to generate multiple tokens per forward pass and reduces sequential decoding depth. Enable it from the CLI with `--with-mtp`:
+TileRT also supports Multi-Token Prediction (MTP), which allows the model to generate multiple tokens per forward pass and reduces sequential decoding depth. Enable it from the CLI with `--num-mtp`, the draft depth (`0`, the default, disables it):
```bash
-python -m tilert.generate --model deepseek_v3_2 --with-mtp --max-new-tokens 1000
+python -m tilert.generate --model glm5_2_rocm \
+ --model-weights-dir /path/to/GLM-5.2-FP8-TileRT \
+ --num-mtp 3 --max-new-tokens 1000
```
+GLM-5.2/5.3 run MTP at depth 1 or 3 and need weights converted with
+`--num_mtp 1`; the NVIDIA models run any `--num-mtp` above 0 at their
+fixed depth of 3.
+
Or programmatically, pass `with_mtp=True` to the generator:
```python
@@ -306,7 +353,7 @@ Of course! Here are 10 short jokes for you.
This example highlights how MTP enables TileRT to efficiently generate longer outputs by accepting multiple tokens per decoding step, while preserving the same Python API interface.
-For the full list of CLI flags (sampling, batching, benchmark modes, β¦), run `python -m tilert.generate --help`.
+For the full list of CLI flags (prompt files, sampling, MTP depth, β¦), run `python -m tilert.generate --help`.
## Disaggregated Serving: vLLM Prefill + TileRT Decode
@@ -316,10 +363,11 @@ TileRT can run as the **decode engine behind a vLLM prefill**, integrated throug
- 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.
+- On the **decode** and **router** hosts, install `tilert` with the `pd` extra: `pip install "tilert[pd]==0.1.6.post2"` pulls in FastAPI, Uvicorn, HTTPX, Requests and Pydantic.
- **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`.
+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`. For GLM-5.2/5.3 on AMD MI350X / MI355X, the decode server's `--model glm5_2` profile (`glm5_3` is accepted as an alias) runs the ROCm engine; start that decode node from the ROCm image.
### Topology A: vLLM prefill β TileRT decode
@@ -365,6 +413,10 @@ Send OpenAI requests to `http://:23333/v1/chat/completions`. The router
A decode engine serves one sequence at a time, so the router reserves a node per request and answers `429` while they are all busy. Add `--queue-timeout ` to make a request wait for a free node instead of failing: useful when a single client fans out into concurrent sub-conversations β an agentic session spawning sub-agents, say β and the burst is wider than the pool but short-lived. Waits longer than 0.1 s are logged. The default, `0`, keeps the fail-fast behaviour.
+The router resolves sampling defaults the way vLLM does: `--generation-config auto` (the default) reads `generation_config.json` under `--model-path`, `--generation-config vllm` uses neutral defaults, and `--default-temperature` / `--default-top-p` / `--default-top-k` override single values. The resolved values are sent explicitly to both the prefill and the decode leg, so the two cannot disagree. Decode nodes advertise what they support (logprobs, penalties, `ignore_eos`), and a request that asks for something its node cannot honour is rejected with an error.
+
+On the decode node, `--pd-buffer-device cpu` keeps the PD receive buffer in pinned host memory registered with the transport: it frees VRAM at the cost of one extra host-to-device copy per request. The default is `cuda`; `TILERT_PD_BUFFER_DEVICE` sets it too.
+
### 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.
diff --git a/assets/glm5_3fp8_benchmark.png b/assets/glm5_3fp8_benchmark.png
new file mode 100644
index 0000000..b24e2b4
Binary files /dev/null and b/assets/glm5_3fp8_benchmark.png differ
diff --git a/pyproject.toml b/pyproject.toml
index 1d1600e..feac475 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -15,19 +15,30 @@ classifiers = [
]
dependencies = [
- # 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.
- "torch==2.11.0",
+ # The v0.1.6 wheel carries backends for two torch ABIs, hence the range:
+ # the CUDA backends (DeepSeek-V3.2, GLM-5) need ``torch==2.11.0+cu130``
+ # from PyTorch's cu130 index (``--index-url
+ # https://download.pytorch.org/whl/cu130``), the ROCm backend (GLM-5.2/5.3)
+ # needs ``torch==2.12.0+rocm7.14.0``. ``tilert.load_backend`` refuses a
+ # torch of the wrong version or flavor, so install the matching one first.
+ "torch>=2.11,<2.13",
"transformers>=4.46.3",
"tokenizers>=0.20.3",
+ "safetensors",
"numpy",
"scipy",
"einops",
]
[project.optional-dependencies]
+# Runtime dependencies of the PD router and decode server (``tilert.pd_vllm``).
+pd = [
+ "fastapi",
+ "uvicorn",
+ "httpx",
+ "requests",
+ "pydantic",
+]
dev = [
"pytest>=7.0.0",
"black==25.1.0",
@@ -65,7 +76,10 @@ 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.5 wheel.
+# Note: this repository ships the public Python sources of the v0.1.6.post2
+# wheel. They are recovered from the wheel itself, which ships them without
+# docstrings or comments, and reformatted to this repo's lint config;
+# ``tilert/generate.py`` also carries a CLI fix that no published wheel has yet.
# 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.
@@ -113,13 +127,41 @@ 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.
+# Relax annotation-completeness there. Since v0.1.6 the sources are recovered
+# from the published wheel, which ships them without comments, so per-line
+# ``# type: ignore``s are not available either: the error codes that occur in
+# the exported code are disabled, and the remaining checks still apply.
[[tool.mypy.overrides]]
module = "tilert.pd_vllm.*"
disallow_untyped_defs = false
disallow_incomplete_defs = false
warn_return_any = false
+disable_error_code = [
+ "arg-type",
+ "assignment",
+ "attr-defined",
+ "func-returns-value",
+ "index",
+ "misc",
+ "return-value",
+ "type-var",
+ "union-attr",
+ "unreachable",
+ "var-annotated",
+]
+
+# The ROCm port (GLM-5.2/5.3 on MI350X/MI355X) comes from the same wheel export.
+[[tool.mypy.overrides]]
+module = [
+ "tilert.models.glm_5_2_rocm.*",
+ "tilert.models.misc_rocm.*",
+ "tilert.models.base_rocm",
+ "tilert.models.common_mxfp4",
+]
+disallow_untyped_defs = false
+disallow_incomplete_defs = false
+warn_return_any = false
+disable_error_code = ["assignment", "attr-defined", "dict-item", "index", "union-attr"]
[tool.bandit]
exclude_dirs = ["tests", "3rd-party"]
@@ -202,7 +244,7 @@ args = ["--keep-percent-format", "--py311-plus"]
[tool.codespell]
skip = "3rd-party,*.git,*.pyc,build,dist,*.egg-info,.mypy_cache,.pytest_cache,"
check-hidden = true
-ignore-words-list = "ThrID"
+ignore-words-list = "ThrID,pres"
[tool.mdformat]
number = true
diff --git a/requirements.txt b/requirements.txt
index ade8509..9ea62fe 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,15 +1,22 @@
-# 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:
+# Runtime dependencies of the v0.1.6 wheel. The wheel carries backends for two
+# torch ABIs (hence the range below), so install the torch build that matches
+# your GPUs first:
#
+# # NVIDIA (DeepSeek-V3.2, GLM-5/5.1): PyTorch's cu130 build. PyPI's default
+# # ``torch`` is a different CUDA build and will not load the cu130-linked
+# # tilert binaries.
# pip install --index-url https://download.pytorch.org/whl/cu130 torch==2.11.0
+#
+# # AMD (GLM-5.2/5.3): torch==2.12.0+rocm7.14.0, preinstalled in the ROCm image
+# # (ghcr.io/tile-ai/tilert-rocm-decode:0.1.6).
+#
# pip install -r requirements.txt
#
-# The recommended path remains the prebuilt Docker image (see README).
-torch==2.11.0
+# The recommended path remains the prebuilt Docker images (see README).
+torch>=2.11,<2.13
transformers>=4.46.3
tokenizers>=0.20.3
+safetensors
numpy
scipy
einops
diff --git a/tilert/__init__.py b/tilert/__init__.py
index d34ce51..6c9666b 100644
--- a/tilert/__init__.py
+++ b/tilert/__init__.py
@@ -1,15 +1,4 @@
-"""TileRT Python package.
-
-Two backend libraries ship with TileRT β one per model family:
-
- - ``libtilert_dsv32.so`` (DeepSeek-V3.2)
- - ``libtilert_glm5.so`` (GLM-5)
-
-They are NOT loaded at import time. The caller selects a backend via
-``load_backend(model_type)`` (done automatically by ``tilert.generate``).
-Only one backend may be loaded per process β both register the ``tilert``
-torch-op namespace. Run DSv3.2 and GLM-5 in separate processes.
-"""
+"""TileRT: tile-based runtime for ultra-low-latency LLM inference."""
import ctypes
import logging
@@ -22,7 +11,6 @@
if not hasattr(torch, "ops"):
raise RuntimeError("PyTorch is required but torch.ops is not available")
-
try:
__version__ = pkg_version("tilert")
except PackageNotFoundError:
@@ -30,31 +18,47 @@
def init_logging() -> logging.Logger:
- """Initialize logging configuration."""
+ level_name = os.environ.get("TILERT_LOG_LEVEL", "INFO").upper()
logging.basicConfig(
- level=logging.DEBUG,
+ level=getattr(logging, level_name, logging.INFO),
format="%(filename)s:%(lineno)d [%(levelname)s]: %(message)s",
)
return logging.getLogger(__name__)
logger = init_logging()
-
_BACKENDS = {
"deepseek_v3_2": "libtilert_dsv32.so",
"glm5": "libtilert_glm5.so",
+ "glm5_2_rocm": "libtilert_glm52_rocm.so",
+}
+_TORCH_FOR_LIB = {
+ "libtilert_dsv32.so": ("2.11", False),
+ "libtilert_glm5.so": ("2.11", False),
+ "libtilert_glm52_rocm.so": ("2.12", True),
}
+
+def _check_torch(so_name: str) -> None:
+ wanted = _TORCH_FOR_LIB.get(so_name)
+ if wanted is None:
+ return
+ version, needs_rocm = wanted
+ is_rocm = getattr(torch.version, "hip", None) is not None
+ if is_rocm != needs_rocm:
+ raise RuntimeError(
+ f"{so_name} needs a {('ROCm' if needs_rocm else 'CUDA')} build of torch; this interpreter has torch {torch.__version__}"
+ )
+ if not torch.__version__.startswith(f"{version}."):
+ raise RuntimeError(
+ f"{so_name} was built against torch {version}; this interpreter has torch {torch.__version__}"
+ )
+
+
_loaded_backend: str | None = None
def load_backend(model_type: str) -> None:
- """Load the backend for ``model_type`` (lazy, once per process).
-
- DeepSeek-V3.2 and GLM-5 ship as separate libraries; the matching one is
- loaded on first use. Loading a second, different backend in the same
- process raises (both libraries define the ``tilert`` op namespace).
- """
global _loaded_backend
so_name = _BACKENDS.get(model_type)
if so_name is None:
@@ -62,10 +66,10 @@ def load_backend(model_type: str) -> None:
if _loaded_backend is not None:
if _loaded_backend != so_name:
raise RuntimeError(
- f"TileRT backend '{_loaded_backend}' already loaded; cannot load "
- f"'{so_name}' in the same process. Run {model_type} in a fresh process."
+ f"TileRT backend '{_loaded_backend}' already loaded; cannot load '{so_name}' in the same process. Run {model_type} in a fresh process."
)
return
+ _check_torch(so_name)
pkg_dir = Path(__file__).parent
lib_path = pkg_dir / so_name
if not lib_path.exists():
@@ -76,16 +80,9 @@ def load_backend(model_type: str) -> None:
ctypes.CDLL(str(lib_path), mode=ctypes.RTLD_GLOBAL | os.RTLD_LAZY)
torch.ops.load_library(str(lib_path))
_loaded_backend = so_name
- logger.info(
- "Loaded TileRT backend %s (%s) for model_type=%s", so_name, lib_path.name, model_type
- )
+ logger.info("Loaded TileRT backend %s for model_type=%s", lib_path.name, model_type)
from .tilert_init import tilert_init # noqa: E402
-__all__ = [
- "logger",
- "load_backend",
- "tilert_init",
- "__version__",
-]
+__all__ = ["logger", "load_backend", "tilert_init", "__version__"]
diff --git a/tilert/generate.py b/tilert/generate.py
index bfcd97f..3beb965 100644
--- a/tilert/generate.py
+++ b/tilert/generate.py
@@ -1,299 +1,172 @@
-"""Text generation script for TileRT."""
+"""TileRT offline generation CLI."""
+import argparse
+import importlib.util
+import sys
import time
-from argparse import ArgumentParser
-from typing import TYPE_CHECKING
-
-import tilert
-
-if TYPE_CHECKING:
- from tilert.models.deepseek_v3_2.generator import DSAv32Generator
- from tilert.models.glm_5.generator import GLM5Generator
-from tilert.benchmark import BenchMode
-from tilert.benchmark import coding_prompt as coding_bench
-from tilert.benchmark import long_prompt as long_bench
-from tilert.benchmark import merge_stats, print_summary_table
-from tilert.benchmark import short_prompt as short_bench
-from tilert.benchmark.config import get_weights_dir
+from pathlib import Path
+from typing import Any
+import torch
-def get_generator(
- model_type: str,
- max_new_tokens: int,
- temperature: float,
- model_weights_dir: str,
- with_mtp: bool,
- top_p: float = 0.9,
- top_k: int = 256,
- enable_thinking: bool = False,
- sampling_seed: int = 42,
-) -> "DSAv32Generator | GLM5Generator":
- """Load the matching backend .so and build the generator for ``model_type``.
-
- DeepSeek-V3.2 and GLM-5 ship as separate libraries; only one backend loads
- per process. Generators are imported lazily after the backend is loaded.
- """
- tilert.load_backend(model_type)
-
- if model_type == "deepseek_v3_2":
- from tilert.models.deepseek_v3_2.generator import DSAv32Generator
- from tilert.models.deepseek_v3_2.model_args import ModelArgs as DSAv32ModelArgs
+import tilert
- return DSAv32Generator(
- model_args=DSAv32ModelArgs(),
- max_new_tokens=max_new_tokens,
- temperature=temperature,
- model_weights_dir=model_weights_dir,
- with_mtp=with_mtp,
- top_p=top_p,
- top_k=top_k,
- use_topp=top_p < 1.0,
- sampling_seed=sampling_seed,
- enable_thinking=enable_thinking,
+_MODEL_PACKAGES: dict[str, str] = {
+ "deepseek_v3_2": "tilert.models.deepseek_v3_2",
+ "glm5": "tilert.models.glm_5",
+ "glm5_2_rocm": "tilert.models.glm_5_2_rocm",
+}
+
+
+def _release_products() -> set[str] | None:
+ try:
+ op = torch.ops.tilert.supported_models.default
+ keyset = torch._C.DispatchKeySet(torch._C.DispatchKey.CPU)
+ return set(op.redispatch(keyset))
+ except Exception:
+ return None
+
+
+def available_models() -> list[str]:
+ present = [
+ model
+ for model, package in _MODEL_PACKAGES.items()
+ if importlib.util.find_spec(package) is not None
+ ]
+ products = _release_products()
+ if products is None:
+ return present
+ filtered = [m for m in present if _MODEL_PACKAGES[m].rsplit(".", 1)[-1] in products]
+ return filtered or present
+
+
+def get_generator(model: str, weights_dir: str, args: argparse.Namespace) -> Any:
+ tilert.load_backend(model)
+ use_topp = args.top_p < 1.0
+ if model == "glm5_2_rocm":
+ from tilert.models.glm_5_2_rocm.generator import Glm52Generator
+ from tilert.models.glm_5_2_rocm.model_args import ModelArgsGlm52
+
+ return Glm52Generator(
+ model_args=ModelArgsGlm52(),
+ model_weights_dir=weights_dir,
+ max_new_tokens=args.max_new_tokens,
+ temperature=args.temperature,
+ top_p=args.top_p,
+ use_topp=use_topp,
+ sampling_seed=args.sampling_seed,
+ num_mtp=args.num_mtp,
+ max_seq_len=args.max_seq_len,
)
-
- if model_type == "glm5":
+ if model == "glm5":
from tilert.models.glm_5.generator import GLM5Generator
from tilert.models.glm_5.model_args import ModelArgsGLM5
return GLM5Generator(
model_args=ModelArgsGLM5(),
- max_new_tokens=max_new_tokens,
- temperature=temperature,
- model_weights_dir=model_weights_dir,
- with_mtp=with_mtp,
- top_p=top_p,
- top_k=top_k,
- use_topp=top_p < 1.0,
- enable_thinking=enable_thinking,
- sampling_seed=sampling_seed,
+ model_weights_dir=weights_dir,
+ max_new_tokens=args.max_new_tokens,
+ temperature=args.temperature,
+ top_p=args.top_p,
+ top_k=args.top_k,
+ use_topp=use_topp,
+ sampling_seed=args.sampling_seed,
+ with_mtp=args.num_mtp > 0,
)
+ if model == "deepseek_v3_2":
+ from tilert.models.deepseek_v3_2.generator import DSAv32Generator
+ from tilert.models.deepseek_v3_2.model_args import ModelArgs
- raise ValueError(f"unsupported model_type: {model_type!r}")
+ return DSAv32Generator(
+ model_args=ModelArgs(),
+ model_weights_dir=weights_dir,
+ max_new_tokens=args.max_new_tokens,
+ temperature=args.temperature,
+ top_p=args.top_p,
+ top_k=args.top_k,
+ use_topp=use_topp,
+ sampling_seed=args.sampling_seed,
+ enable_thinking=args.enable_thinking,
+ with_mtp=args.num_mtp > 0,
+ )
+ raise SystemExit(
+ f"[generate] model {model!r} is not offered by this build; available: {available_models()}"
+ )
-def parse_args(): # type: ignore
- parser = ArgumentParser(description="Command-line interface for text generation.")
- parser.add_argument(
- "--model-weights-dir",
- type=str,
- default=None,
- help="Path to model weights directory (resolved from ~/.tilert/config.toml if omitted)",
- )
- parser.add_argument(
- "--model",
- type=str,
- default="deepseek_v3_2",
- choices=["deepseek_v3_2", "glm5"],
- help="Model type to use (default: deepseek_v3_2).",
- )
- parser.add_argument("--max-new-tokens", type=int, default=4000, help="Max tokens to generate")
- parser.add_argument("--temperature", type=float, default=1.0, help="Sampling temperature")
- parser.add_argument(
- "--top-p",
- type=float,
- default=1.0,
- help="Top-p (nucleus) sampling threshold. Use < 1.0 to enable top-p sampling (e.g. 0.9)",
- )
- parser.add_argument("--top-k", type=int, default=256, help="Top-k sampling threshold")
- parser.add_argument("--interactive", action="store_true")
- parser.add_argument(
- "--with-mtp",
- action="store_true",
- help="Enable MTP (Multi-Token Prediction) for speculative decoding",
- )
- parser.add_argument(
- "--use-random-weights",
- action="store_true",
- help="Use random weights instead of pretrained (for testing MTP without real weights)",
+def _prompts(args: argparse.Namespace) -> list[str]:
+ if args.prompt_file:
+ text = Path(args.prompt_file).read_text(encoding="utf-8")
+ out = [p.strip() for p in text.split("\n\n") if p.strip()]
+ if not out:
+ raise SystemExit(f"[generate] no prompts in {args.prompt_file}")
+ return out
+ return [args.prompt]
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+ models = available_models()
+ if not models:
+ raise SystemExit(
+ "[generate] this install ships no model package; the engine library may have failed to load (see the warning from `import tilert`)"
+ )
+ p = argparse.ArgumentParser(
+ prog="python -m tilert.generate", description=__doc__.splitlines()[0]
)
- parser.add_argument(
- "--enable-thinking",
- action="store_true",
- help="Enable thinking mode in chat template",
+ p.add_argument("--model", choices=models, default=models[0])
+ p.add_argument("--model-weights-dir", required=True, help="converted weights dir")
+ src = p.add_mutually_exclusive_group()
+ src.add_argument("--prompt", default="Hello! Tell me about yourself.")
+ src.add_argument(
+ "--prompt-file", help="file of prompts, separated by blank lines; each is generated in turn"
)
- parser.add_argument(
- "--sampling-seed",
+ p.add_argument("--max-new-tokens", type=int, default=200)
+ p.add_argument(
+ "--max-seq-len",
type=int,
- default=42,
- help="Sampling seed for top-p sampling (fixed per request, default: 42)",
- )
- parser.add_argument(
- "--model-name",
- type=str,
- default=None,
- help="Override display name for benchmark tables",
- )
- parser.add_argument(
- "--tag",
- type=str,
- default=None,
- help="Tag for regression_plots/ directory (default: auto-detect from git state)",
- )
- parser.add_argument(
- "--modes",
- type=str,
- default=None,
- help="Comma-separated mode filters: top-k1,top-p0.95 (default: all)",
+ default=8192,
+ help="KV cache length; honoured by the ROCm models (the CUDA ones take it from their model args)",
)
- parser.add_argument(
- "--workloads",
- type=str,
- default=None,
- help="Comma-separated workload filters: short,coding,long (default: all)",
+ p.add_argument("--temperature", type=float, default=0.7)
+ p.add_argument("--top-p", type=float, default=0.95, help="1.0 selects greedy (argmax) decoding")
+ p.add_argument("--top-k", type=int, default=20)
+ p.add_argument("--sampling-seed", type=int, default=42)
+ p.add_argument(
+ "--num-mtp", type=int, default=0, help="MTP draft depth; 0 disables speculative decoding"
)
- parser.add_argument(
- "--enable-logprobs",
- action="store_true",
- help="Enable kernel-level top-256 logprobs export (for benchmarking overhead)",
- )
- return parser.parse_args()
-
-
-if __name__ == "__main__":
- """
- Usage (run as a module; --model-weights-dir may be omitted if the path is
- registered under ~/.tilert/config.toml). Run DeepSeek-V3.2 and GLM-5 in
- separate processes β the two backends cannot coexist in one interpreter.
-
- # DeepSeek-V3.2 β standard generation with pretrained weights:
- python -m tilert.generate --model deepseek_v3_2 \
- --model-weights-dir /path/to/DeepSeek-V3.2-TileRT \
- --max-new-tokens 1000 2>&1 | tee test.log
-
- # DeepSeek-V3.2 β MTP generation with random weights (for testing):
- python -m tilert.generate --model deepseek_v3_2 --with-mtp --use-random-weights \
- --model-weights-dir /path/to/DeepSeek-V3.2-TileRT \
- --max-new-tokens 1000 2>&1 | tee test.log
-
- # DeepSeek-V3.2 β MTP generation with pretrained weights:
- python -m tilert.generate --model deepseek_v3_2 --with-mtp \
- --model-weights-dir /path/to/DeepSeek-V3.2-TileRT \
- --max-new-tokens 1000 2>&1 | tee test.log
-
- # GLM-5 β standard generation:
- python -m tilert.generate --model glm5 \
- --model-weights-dir /path/to/GLM-5-FP8-TileRT \
- --max-new-tokens 1000 2>&1 | tee test.log
-
- # GLM-5 β MTP generation:
- python -m tilert.generate --model glm5 --with-mtp \
- --model-weights-dir /path/to/GLM-5-FP8-TileRT \
- --max-new-tokens 1000 2>&1 | tee test.log
- """
- args = parse_args()
-
- config_key = args.model
- model_name = args.model.upper()
- if args.model_name:
- model_name = args.model_name
- model_weights_dir = get_weights_dir(config_key, cli_override=args.model_weights_dir)
-
- if args.interactive:
- with_mtp = args.with_mtp
- else:
- with_mtp = True
+ p.add_argument("--enable-thinking", action="store_true")
+ p.add_argument("--quiet", action="store_true", help="only print completions")
+ return p.parse_args(argv)
- generator = get_generator(
- model_type=args.model,
- max_new_tokens=args.max_new_tokens,
- temperature=args.temperature,
- model_weights_dir=model_weights_dir,
- with_mtp=with_mtp,
- top_p=args.top_p,
- top_k=args.top_k,
- enable_thinking=args.enable_thinking,
- sampling_seed=args.sampling_seed,
- )
+def main(argv: list[str] | None = None) -> int:
+ args = parse_args(argv)
+ prompts = _prompts(args)
t0 = time.monotonic()
- if args.use_random_weights:
- print("Initializing random weights...")
- if hasattr(generator, "init"):
- generator.init() # type: ignore[union-attr]
- generator.init_random_weights()
- else:
- print("Loading pretrained weights...")
- generator.from_pretrained()
- load_time = time.monotonic() - t0
-
- if args.enable_logprobs:
- if hasattr(generator.decode_layer, "set_logprobs_enabled"):
- generator.decode_layer.set_logprobs_enabled(True) # type: ignore[union-attr]
- print("Logprobs export enabled (top-256)")
- else:
- print(f"Warning: logprobs not supported for {type(generator).__name__}")
-
- if args.interactive:
- print("Welcome to the TileRT interactive mode! Type '/exit' to exit.")
- while True:
- prompt = input(">>> ")
- if prompt == "/exit":
- break
- _ = generator.generate(prompt) # type: ignore[has-type]
- else:
+ generator = get_generator(args.model, args.model_weights_dir, args)
+ generator.init()
+ generator.from_pretrained()
+ if not args.quiet:
+ print(f"[generate] {args.model} loaded in {time.monotonic() - t0:.1f}s")
+ try:
+ for i, prompt in enumerate(prompts):
+ if not args.quiet:
+ print(f"\n=== prompt {i + 1}/{len(prompts)} ===\n{prompt}\n--- output ---")
+ t = time.monotonic()
+ text, times, accepts, prompt_len = generator.generate(prompt, print_log=False)
+ wall = time.monotonic() - t
+ print(text)
+ if not args.quiet:
+ n_tok = sum(accepts) if accepts else len(times)
+ rate = n_tok / wall if wall > 0 else 0.0
+ line = f"[generate] {prompt_len} prompt tok -> {n_tok} tok in {wall:.1f}s ({rate:.1f} tok/s)"
+ if accepts:
+ line += f", mean accept {sum(accepts) / len(accepts):.2f}"
+ print(line)
+ finally:
+ generator.cleanup()
+ return 0
- bench_top_p = args.top_p if args.top_p < 1.0 else 0.95
- modes = [
- BenchMode(with_mtp=False, label="top-k1 w/o MTP"),
- BenchMode(with_mtp=True, label="top-k1 w/ MTP"),
- BenchMode(
- with_mtp=True,
- label=f"top-p{bench_top_p} w/ MTP",
- use_topp=True,
- top_p=bench_top_p,
- top_k=args.top_k,
- temperature=args.temperature,
- ),
- ]
-
- if args.modes:
- allowed = {m.strip() for m in args.modes.split(",")}
- modes = [m for m in modes if any(a in m.label for a in allowed)]
- if not modes:
- raise SystemExit(
- f"Error: --modes '{args.modes}' matched no benchmark modes. "
- f"Valid tokens: top-k1, top-p0.95"
- )
-
- t0 = time.monotonic()
- workload_runners = []
- allowed_workloads = (
- {w.strip() for w in args.workloads.split(",")}
- if args.workloads
- else {"short", "coding", "long"}
- )
- if "short" in allowed_workloads:
- workload_runners.append(short_bench.run)
- if "coding" in allowed_workloads:
- workload_runners.append(coding_bench.run)
- if "long" in allowed_workloads:
- workload_runners.append(long_bench.run)
- if not workload_runners:
- raise SystemExit(
- f"Error: --workloads '{args.workloads}' matched no workloads. "
- f"Valid values: short, coding, long"
- )
- all_bench_results = [
- runner(generator, modes) for runner in workload_runners # type: ignore[arg-type]
- ]
- bench_time = time.monotonic() - t0
- all_bench_stats = [stats for stats, _ in all_bench_results]
-
- print_summary_table(
- merge_stats(all_bench_stats),
- model_name=model_name,
- )
-
- total = load_time + bench_time
- print(f"\n## {model_name} Timing")
- print()
- print("| Phase | Time |")
- print("|-------|------|")
- print(f"| Loading | {load_time:.1f}s |")
- print(f"| Benchmark | {bench_time:.1f}s |")
- print(f"| **Total** | **{total:.1f}s** |")
-
- print("Cleaning up...")
- generator.cleanup()
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/tilert/models/base_rocm.py b/tilert/models/base_rocm.py
new file mode 100644
index 0000000..bede0d2
--- /dev/null
+++ b/tilert/models/base_rocm.py
@@ -0,0 +1,156 @@
+"""Base classes for TileRT model modules."""
+
+from abc import ABC
+from enum import Enum
+from typing import Any, ClassVar
+
+import torch
+import torch.nn as nn
+
+from tilert import logger
+
+__all__ = ["SerializableTileRTModule", "TileRTModule", "TilertWeightsConverter"]
+
+
+class TilertWeightsConverter:
+ """Tilert weights converter: dispatches to ``convert_to_``."""
+
+ def __init__(self, model_args: Any, num_devices: int):
+ self.model_args = model_args
+ self.num_devices = num_devices
+
+ def dispatch(self, algorithm: Enum, weights: list[torch.Tensor]) -> Any:
+ dispatch_method = getattr(self, f"convert_to_{algorithm.value}")
+ return dispatch_method(weights)
+
+
+class TileRTModule(nn.Module, ABC):
+ """Base class for all TileRT modules."""
+
+ _SUPPORTED_ALGORITHMS: ClassVar[dict[str, list[Enum]]] = {}
+
+ @classmethod
+ def get_supported_algorithms(cls, arch_name: str) -> list[Enum]:
+ if arch_name not in cls._SUPPORTED_ALGORITHMS:
+ raise ValueError(
+ f"{cls.__name__} does not support arch '{arch_name}'. Supported: {list(cls._SUPPORTED_ALGORITHMS.keys())}"
+ )
+ return cls._SUPPORTED_ALGORITHMS[arch_name]
+
+ def __init__(
+ self,
+ op_name: str = "",
+ layer_idx: int = 0,
+ model_args: Any | None = None,
+ num_devices: int = 1,
+ device_id: int = 0,
+ *args: Any,
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(*args, **kwargs)
+ self.model_args = model_args
+ self.num_devices = num_devices
+ self.device_id = device_id
+ self.algorithm: Enum | None = None
+ self.layer_idx = layer_idx
+ self.is_tilert_weights_init = False
+ self.op_name = type(self).__name__ if op_name == "" else op_name
+
+ def get_cache_vars(self) -> list[torch.Tensor]:
+ return []
+
+ def get_tilert_weights_alias(self) -> list[str]:
+ return list(self.tilert_weights_alias())
+
+ def set_algorithm(self, algorithm: Enum) -> None:
+ if self._SUPPORTED_ALGORITHMS:
+ arch = self.model_args.arch_name
+ supported = self.get_supported_algorithms(arch)
+ if algorithm not in supported:
+ raise ValueError(
+ f"{type(self).__name__}: algorithm {algorithm} not supported for arch '{arch}'. Supported: {supported}"
+ )
+ self.algorithm = algorithm
+
+ def golden_forward(self, *args: Any, **kwargs: Any) -> Any:
+ raise RuntimeError("golden_forward is not available in release builds")
+
+ def tilert_forward(self, *args: Any, **kwargs: Any) -> Any:
+ del args, kwargs
+ raise NotImplementedError("Tilert forward not implemented")
+
+ def init_tilert_vars(self, batch_size: int, seq_len: int) -> None:
+ del batch_size, seq_len
+
+
+class SerializableTileRTModule(TileRTModule):
+ """Composite module: ordered sequence of sub-ops with prefixed weight keys."""
+
+ def __init__(
+ self,
+ model_args: Any,
+ device_id: int = 0,
+ num_devices: int = 1,
+ remove_selected: bool = False,
+ ):
+ super().__init__(
+ type(self).__name__, model_args=model_args, device_id=device_id, num_devices=num_devices
+ )
+ self.remove_selected = remove_selected
+ self.exec_seq: list[TileRTModule] = []
+ self.prefix_seq: list[str] = []
+ self.suffix_seq: list[str] = []
+ self.retain_weights_seq: list[bool] = []
+
+ def get_cache_vars(self) -> list[torch.Tensor]:
+ cache_vars = []
+ for op in self.exec_seq:
+ cache_vars.extend(op.get_cache_vars())
+ return cache_vars
+
+ def register_op(
+ self, op: TileRTModule, prefix: str = "", suffix: str = "", retain_weights: bool = False
+ ) -> None:
+ self.exec_seq.append(op)
+ self.prefix_seq.append(prefix)
+ self.suffix_seq.append(suffix)
+ self.retain_weights_seq.append(retain_weights)
+
+ def get_tilert_weights_alias(self) -> list[str]:
+ weights_alias: list[str] = []
+ for op in self.exec_seq:
+ weights_alias.extend(op.get_tilert_weights_alias())
+ return weights_alias
+
+ def get_weights_list(self) -> list[torch.Tensor]:
+ weights = []
+ for op in self.exec_seq:
+ weights.extend(op.get_weights_list())
+ return weights
+
+ def init_tilert_weights(self, state_dict: dict[str, torch.Tensor]) -> None:
+ for op, prefix, suffix, retain_weights in zip(
+ self.exec_seq, self.prefix_seq, self.suffix_seq, self.retain_weights_seq
+ ):
+ if op.is_tilert_weights_init:
+ logger.debug(f"Skipping init_tilert_weights for {op.op_name} (already initialized)")
+ continue
+ keys_to_remove = set()
+ op_state_dict = {}
+ for op_key in op.get_tilert_weights_alias():
+ original_key = f"{prefix}{op_key}{suffix}"
+ if original_key in state_dict:
+ op_state_dict[op_key] = state_dict[original_key]
+ if self.remove_selected:
+ keys_to_remove.add(original_key)
+ op.init_tilert_weights(op_state_dict)
+ if self.remove_selected and (not retain_weights):
+ for k in keys_to_remove:
+ del state_dict[k]
+
+ def init_random_weights(self) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def init_tilert_vars(self, batch_size: int, seq_len: int) -> None:
+ for op in self.exec_seq:
+ op.init_tilert_vars(batch_size, seq_len)
diff --git a/tilert/models/common_mxfp4.py b/tilert/models/common_mxfp4.py
new file mode 100644
index 0000000..66c527e
--- /dev/null
+++ b/tilert/models/common_mxfp4.py
@@ -0,0 +1,419 @@
+from __future__ import annotations
+
+import os
+
+import torch
+
+__all__ = [
+ "E2M1_TABLE",
+ "dequant_mxfp4_to_bf16",
+ "marlin_pack_fp4_tiles",
+ "make_bf16_prefolded_scale",
+ "make_random_mxfp4_weight",
+ "standard_bf16_scale_bytes",
+ "build_ug_weights_mma_natural",
+ "build_down_weights_mma_natural",
+]
+E2M1_TABLE = torch.tensor(
+ [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0],
+ dtype=torch.float32,
+)
+
+
+def dequant_mxfp4_to_bf16(
+ fp4_nibbles: torch.Tensor, e8m0: torch.Tensor, block_size: int = 32
+) -> torch.Tensor:
+ assert fp4_nibbles.dtype == torch.uint8, fp4_nibbles.dtype
+ assert e8m0.dtype == torch.uint8, e8m0.dtype
+ assert fp4_nibbles.shape[-1] % block_size == 0
+ assert fp4_nibbles.shape[:-1] == e8m0.shape[:-1]
+ assert e8m0.shape[-1] == fp4_nibbles.shape[-1] // block_size
+ table = E2M1_TABLE.to(fp4_nibbles.device)
+ vals_f32 = table[fp4_nibbles.long()]
+ scale_f32 = torch.pow(2.0, e8m0.to(torch.float32) - 127.0)
+ scale_f32 = scale_f32.repeat_interleave(block_size, dim=-1)
+ return (vals_f32 * scale_f32).to(torch.bfloat16)
+
+
+def make_bf16_prefolded_scale(e8m0: torch.Tensor) -> torch.Tensor:
+ assert e8m0.dtype == torch.uint8
+ bits = e8m0.to(torch.int32) + 126 << 7
+ return bits.to(torch.int32).bitwise_and(65535).to(torch.int16).view(torch.uint16)
+
+
+def standard_bf16_scale_bytes(e8m0: torch.Tensor) -> torch.Tensor:
+ bits = (e8m0.to(torch.int32) << 7).contiguous()
+ return bits.view(torch.uint8).reshape(*bits.shape, 4)[..., :2].contiguous()
+
+
+def _use_sm90_lop3_packing(device: torch.device) -> bool:
+ if os.environ.get("TILERT_FORCE_SM90_PACKING", "") not in ("", "0"):
+ return True
+ if device.type != "cuda" or not torch.cuda.is_available():
+ return False
+ idx = device.index if device.index is not None else torch.cuda.current_device()
+ major = int(torch.cuda.get_device_capability(idx)[0])
+ return major < 10
+
+
+def _make_mma_natural_u32(nibble_tile: torch.Tensor) -> torch.Tensor:
+ assert nibble_tile.shape[-2:] == (16, 16)
+ pre = nibble_tile.shape[:-2]
+ device = nibble_tile.device
+ nib = nibble_tile.to(torch.int64) & 15
+ flat = nib.reshape(*pre, 256)
+ bidx, lid = torch.meshgrid(
+ torch.arange(4, device=device), torch.arange(32, device=device), indexing="ij"
+ )
+ m_off = (bidx & 1) * 8
+ k_off = (bidx >> 1) * 8
+ m = lid // 4 + m_off
+ k_lo = lid % 4 * 2 + k_off
+ k_hi = k_lo + 1
+ idx_lo = (m * 16 + k_lo).reshape(-1)
+ idx_hi = (m * 16 + k_hi).reshape(-1)
+ lo_vals = flat.index_select(-1, idx_lo).reshape(*pre, 4, 32)
+ hi_vals = flat.index_select(-1, idx_hi).reshape(*pre, 4, 32)
+ if _use_sm90_lop3_packing(device):
+ lo_sh = torch.tensor([12, 8, 4, 0], dtype=torch.int64, device=device)
+ hi_sh = torch.tensor([28, 24, 20, 16], dtype=torch.int64, device=device)
+ view = [1] * len(pre) + [4, 1]
+ out = (lo_vals << lo_sh.view(*view)).sum(dim=-2) + (hi_vals << hi_sh.view(*view)).sum(
+ dim=-2
+ )
+ return out.to(torch.int32).view(torch.uint32)
+ bytes_ = lo_vals | hi_vals << 4
+ shifts = torch.tensor([0, 8, 16, 24], dtype=torch.int64, device=device)
+ out = (bytes_ << shifts.view(*[1] * len(pre), 4, 1)).sum(dim=-2)
+ return out.to(torch.int32).view(torch.uint32)
+
+
+def build_ug_weights_mma_natural(
+ gate_fp4: torch.Tensor,
+ gate_e8m0: torch.Tensor,
+ up_fp4: torch.Tensor,
+ up_e8m0: torch.Tensor,
+ dim: int,
+ moe_inter_dim: int,
+ *,
+ sms_per_expert: int = 16,
+ k_per_page: int | None = None,
+) -> torch.Tensor:
+ e = gate_fp4.shape[0]
+ outer_iters = 2 * moe_inter_dim // 16 // 32
+ n_slices = sms_per_expert * outer_iters
+ rh = 16
+ kp = k_per_page if k_per_page is not None else _ug_k_per_page(dim)
+ assert dim % kp == 0 and kp % 256 == 0, f"bad k_per_page {kp} for dim {dim}"
+ np_ = dim // kp
+ sc = kp // 32
+ mat_bytes_per_page = 2 * (kp // 16) * 32 * 4
+ scale_bytes_per_page = 32 * sc * 2
+ g_fp4 = gate_fp4.reshape(e, n_slices, rh, np_, kp)
+ u_fp4 = up_fp4.reshape(e, n_slices, rh, np_, kp)
+ g_e8m0 = gate_e8m0.reshape(e, n_slices, rh, np_, sc)
+ u_e8m0 = up_e8m0.reshape(e, n_slices, rh, np_, sc)
+ page_fp4 = torch.cat([g_fp4, u_fp4], dim=2).permute(0, 1, 3, 2, 4).contiguous()
+ page_e8m0 = torch.cat([g_e8m0, u_e8m0], dim=2).permute(0, 1, 3, 2, 4).contiguous()
+ nib = page_fp4.reshape(e, n_slices, np_, 2, 16, kp // 16, 16)
+ nib = nib.permute(0, 1, 2, 3, 5, 4, 6).contiguous()
+ u32 = _make_mma_natural_u32(nib)
+ mat_bytes = u32.view(torch.uint8).reshape(e, n_slices, np_, mat_bytes_per_page)
+ scale_bytes = standard_bf16_scale_bytes(page_e8m0).reshape(
+ e, n_slices, np_, scale_bytes_per_page
+ )
+ return torch.cat([mat_bytes, scale_bytes], dim=-1).contiguous()
+
+
+def _ug_k_per_page(dim: int) -> int:
+ if dim % 2048 == 0:
+ return 2048
+ if dim % 1792 == 0:
+ return 1792
+ return 1024
+
+
+def build_down_weights_mma_natural(
+ mat_fp4: torch.Tensor,
+ mat_e8m0: torch.Tensor,
+ dim: int,
+ moe_inter_dim: int,
+ *,
+ down_num_sms: int = 128,
+) -> torch.Tensor:
+ e = mat_fp4.shape[0]
+ dev = mat_fp4.device
+ odpb = dim // down_num_sms
+ pad = (odpb + 15) // 16 * 16
+ n_mt = pad // 16
+ n_kt = moe_inter_dim // 16
+ n_sc = moe_inter_dim // 32
+ mat_bytes = n_mt * n_kt * 32 * 4
+ scale_bytes = pad * n_sc * 2
+ mat = mat_fp4.reshape(e, down_num_sms, odpb, moe_inter_dim)
+ e8 = mat_e8m0.reshape(e, down_num_sms, odpb, n_sc)
+ pad_rows = pad - odpb
+ if pad_rows > 0:
+ mat = torch.cat(
+ [
+ mat,
+ torch.zeros(
+ e, down_num_sms, pad_rows, moe_inter_dim, dtype=torch.uint8, device=dev
+ ),
+ ],
+ dim=2,
+ )
+ e8 = torch.cat(
+ [e8, torch.full((e, down_num_sms, pad_rows, n_sc), 127, dtype=torch.uint8, device=dev)],
+ dim=2,
+ )
+ nib = mat.reshape(e, down_num_sms, n_mt, 16, n_kt, 16)
+ nib = nib.permute(0, 1, 2, 4, 3, 5).contiguous()
+ u32 = _make_mma_natural_u32(nib)
+ mat_b = u32.view(torch.uint8).reshape(e, down_num_sms, mat_bytes)
+ scale_b = standard_bf16_scale_bytes(e8).reshape(e, down_num_sms, scale_bytes)
+ return torch.cat([mat_b, scale_b], dim=-1).contiguous()
+
+
+SF_VEC_NVFP4 = 16
+
+
+def build_ug_weights_mma_natural_nvfp4(
+ gate_fp4: torch.Tensor,
+ gate_sf: torch.Tensor,
+ up_fp4: torch.Tensor,
+ up_sf: torch.Tensor,
+ dim: int,
+ moe_inter_dim: int,
+ *,
+ sms_per_expert: int = 16,
+) -> torch.Tensor:
+ e = gate_fp4.shape[0]
+ outer_iters = 2 * moe_inter_dim // 16 // 32
+ n_slices = sms_per_expert * outer_iters
+ rh = 16
+ kp = _ug_k_per_page(dim)
+ np_ = dim // kp
+ sc = kp // SF_VEC_NVFP4
+ mat_bytes_per_page = 2 * (kp // 16) * 32 * 4
+ scale_bytes_per_page = 32 * sc
+ g_fp4 = gate_fp4.reshape(e, n_slices, rh, np_, kp)
+ u_fp4 = up_fp4.reshape(e, n_slices, rh, np_, kp)
+ g_sf = gate_sf.reshape(e, n_slices, rh, np_, sc)
+ u_sf = up_sf.reshape(e, n_slices, rh, np_, sc)
+ page_fp4 = torch.cat([g_fp4, u_fp4], dim=2).permute(0, 1, 3, 2, 4).contiguous()
+ page_sf = torch.cat([g_sf, u_sf], dim=2).permute(0, 1, 3, 2, 4).contiguous()
+ nib = page_fp4.reshape(e, n_slices, np_, 2, 16, kp // 16, 16)
+ nib = nib.permute(0, 1, 2, 3, 5, 4, 6).contiguous()
+ u32 = _make_mma_natural_u32(nib)
+ mat_bytes = u32.view(torch.uint8).reshape(e, n_slices, np_, mat_bytes_per_page)
+ scale_bytes = page_sf.reshape(e, n_slices, np_, scale_bytes_per_page)
+ return torch.cat([mat_bytes, scale_bytes], dim=-1).contiguous()
+
+
+def build_down_weights_mma_natural_nvfp4(
+ mat_fp4: torch.Tensor,
+ mat_sf: torch.Tensor,
+ dim: int,
+ moe_inter_dim: int,
+ *,
+ down_num_sms: int = 128,
+) -> torch.Tensor:
+ e = mat_fp4.shape[0]
+ dev = mat_fp4.device
+ odpb = dim // down_num_sms
+ pad = (odpb + 15) // 16 * 16
+ n_mt = pad // 16
+ n_kt = moe_inter_dim // 16
+ n_sc = moe_inter_dim // SF_VEC_NVFP4
+ mat_bytes = n_mt * n_kt * 32 * 4
+ scale_bytes = pad * n_sc
+ mat = mat_fp4.reshape(e, down_num_sms, odpb, moe_inter_dim)
+ sf = mat_sf.reshape(e, down_num_sms, odpb, n_sc)
+ pad_rows = pad - odpb
+ if pad_rows > 0:
+ mat = torch.cat(
+ [
+ mat,
+ torch.zeros(
+ e, down_num_sms, pad_rows, moe_inter_dim, dtype=torch.uint8, device=dev
+ ),
+ ],
+ dim=2,
+ )
+ sf = torch.cat(
+ [sf, torch.zeros((e, down_num_sms, pad_rows, n_sc), dtype=torch.uint8, device=dev)],
+ dim=2,
+ )
+ nib = mat.reshape(e, down_num_sms, n_mt, 16, n_kt, 16)
+ nib = nib.permute(0, 1, 2, 4, 3, 5).contiguous()
+ u32 = _make_mma_natural_u32(nib)
+ mat_b = u32.view(torch.uint8).reshape(e, down_num_sms, mat_bytes)
+ scale_b = sf.reshape(e, down_num_sms, scale_bytes)
+ return torch.cat([mat_b, scale_b], dim=-1).contiguous()
+
+
+_UG_KPAGE_NVFP4 = 256
+_UG_KTILE_NVFP4 = 64
+_UG_SFA_ATOM_NVFP4 = 512
+
+
+def _sfa_perm_nvfp4(device: torch.device) -> torch.Tensor:
+ return torch.tensor([m % 32 * 4 + m // 32 for m in range(128)], dtype=torch.long, device=device)
+
+
+def _interleave_mat_nvfp4(nib: torch.Tensor) -> torch.Tensor:
+ e, m, k = nib.shape
+ assert m % 8 == 0 and k % 32 == 0
+ packed = (nib[:, :, 0::2] | nib[:, :, 1::2] << 4).to(torch.uint8)
+ ku128 = k // 32
+ inter = packed.reshape(e, m // 8, 8, ku128, 16).permute(0, 1, 3, 2, 4).contiguous()
+ return inter.reshape(e, -1)
+
+
+def _sfa_atoms_nvfp4(sb: torch.Tensor) -> torch.Tensor:
+ e, m, n_blk = sb.shape
+ assert m == 128 and n_blk % 4 == 0
+ n_atom = n_blk // 4
+ u32 = sb.to(torch.int32).reshape(e, 128, n_atom, 4)
+ packed = u32[..., 0] | u32[..., 1] << 8 | u32[..., 2] << 16 | u32[..., 3] << 24
+ atom = torch.zeros(e, n_atom, 128, dtype=torch.int32, device=sb.device)
+ atom[:, :, _sfa_perm_nvfp4(sb.device)] = packed.permute(0, 2, 1).contiguous()
+ return atom.reshape(e, -1).view(torch.uint8)
+
+
+def build_ug_weights_utcmma_nvfp4(
+ gate_nib: torch.Tensor,
+ gate_sf: torch.Tensor,
+ up_nib: torch.Tensor,
+ up_sf: torch.Tensor,
+ dim: int,
+ moe_inter_dim: int,
+) -> torch.Tensor:
+ e = gate_nib.shape[0]
+ ug_mtiles = moe_inter_dim // 128
+ n_pages = dim // _UG_KPAGE_NVFP4
+ ug_ktiles = _UG_KPAGE_NVFP4 // _UG_KTILE_NVFP4
+ mat_b = 128 * _UG_KPAGE_NVFP4 // 2
+ sfa_b = _UG_SFA_ATOM_NVFP4 * ug_ktiles
+ page_b = 2 * (mat_b + sfa_b)
+ sc_per_page = _UG_KPAGE_NVFP4 // SF_VEC_NVFP4
+ out = torch.zeros(e, ug_mtiles, n_pages, page_b, dtype=torch.uint8, device=gate_nib.device)
+ for h in range(ug_mtiles):
+ r = slice(h * 128, h * 128 + 128)
+ for p in range(n_pages):
+ ks = slice(p * _UG_KPAGE_NVFP4, p * _UG_KPAGE_NVFP4 + _UG_KPAGE_NVFP4)
+ ss = slice(p * sc_per_page, p * sc_per_page + sc_per_page)
+ out[:, h, p, 0:mat_b] = _interleave_mat_nvfp4(gate_nib[:, r, ks])
+ out[:, h, p, mat_b : 2 * mat_b] = _interleave_mat_nvfp4(up_nib[:, r, ks])
+ out[:, h, p, 2 * mat_b : 2 * mat_b + sfa_b] = _sfa_atoms_nvfp4(gate_sf[:, r, ss])
+ out[:, h, p, 2 * mat_b + sfa_b : page_b] = _sfa_atoms_nvfp4(up_sf[:, r, ss])
+ return out.reshape(e, -1).contiguous()
+
+
+def build_down_weights_utcmma_nvfp4(
+ mat_nib: torch.Tensor,
+ mat_sf: torch.Tensor,
+ dim: int,
+ moe_inter_dim: int,
+ *,
+ down_num_sms: int = 128,
+) -> torch.Tensor:
+ e = mat_nib.shape[0]
+ dev = mat_nib.device
+ odpb = dim // down_num_sms
+ n_sc = moe_inter_dim // SF_VEC_NVFP4
+ mat_b = odpb * moe_inter_dim // 2
+ sfa_b = _UG_SFA_ATOM_NVFP4 * (moe_inter_dim // _UG_KTILE_NVFP4)
+ per = mat_b + sfa_b
+ out = torch.zeros(e, down_num_sms, per, dtype=torch.uint8, device=dev)
+ for c in range(down_num_sms):
+ r = slice(c * odpb, c * odpb + odpb)
+ sb_pad = torch.zeros(e, 128, n_sc, dtype=torch.uint8, device=dev)
+ sb_pad[:, :odpb] = mat_sf[:, r, :]
+ out[:, c, 0:mat_b] = _interleave_mat_nvfp4(mat_nib[:, r, :])
+ out[:, c, mat_b:per] = _sfa_atoms_nvfp4(sb_pad)
+ return out.contiguous()
+
+
+def marlin_pack_fp4_tiles(fp4_nibbles: torch.Tensor) -> torch.Tensor:
+ assert fp4_nibbles.dtype == torch.uint8
+ assert fp4_nibbles.shape[-2] % 16 == 0
+ assert fp4_nibbles.shape[-1] % 16 == 0
+ M, K = fp4_nibbles.shape[-2:]
+ pre = fp4_nibbles.shape[:-2]
+ tiles = fp4_nibbles.reshape(*pre, M // 16, 16, K // 16, 16)
+ tiles = tiles.permute(*range(len(pre)), -4, -2, -3, -1).contiguous()
+ device = fp4_nibbles.device
+ lid = torch.arange(32, device=device)
+ row_t = (lid // 4).to(torch.long)
+ col_t = (lid % 4 * 2).to(torch.long)
+
+ def _gather(row_offset: int, col_offset: int) -> torch.Tensor:
+ rows = (row_t + row_offset).view(32)
+ cols = (col_t + col_offset).view(32)
+ return tiles[..., rows, cols]
+
+ nib0 = _gather(0, 0)
+ nib1 = _gather(0, 8)
+ nib2 = _gather(8, 0)
+ nib3 = _gather(8, 8)
+ nib4 = _gather(0, 1)
+ nib5 = _gather(0, 9)
+ nib6 = _gather(8, 1)
+ nib7 = _gather(8, 9)
+
+ def _u32(x: torch.Tensor) -> torch.Tensor:
+ return x.to(torch.int32).bitwise_and(15)
+
+ packed = (
+ _u32(nib0) << 0
+ | _u32(nib1) << 4
+ | _u32(nib2) << 8
+ | _u32(nib3) << 12
+ | _u32(nib4) << 16
+ | _u32(nib5) << 20
+ | _u32(nib6) << 24
+ | _u32(nib7) << 28
+ )
+ return packed.to(torch.int32).view(torch.uint32)
+
+
+def make_random_mxfp4_weight(
+ shape: tuple[int, ...],
+ *,
+ e8m0_min: int = 110,
+ e8m0_max: int = 120,
+ block_size: int = 32,
+ device: str | torch.device = "cuda",
+ generator: torch.Generator | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ assert shape[-1] % block_size == 0
+ fp4 = torch.randint(0, 16, shape, dtype=torch.uint8, device=device, generator=generator)
+ scale_shape = (*shape[:-1], shape[-1] // block_size)
+ e8m0 = torch.randint(
+ e8m0_min, e8m0_max + 1, scale_shape, dtype=torch.uint8, device=device, generator=generator
+ )
+ dequant = dequant_mxfp4_to_bf16(fp4, e8m0, block_size)
+ return (fp4, e8m0, dequant)
+
+
+def _unpack_fp4_nibbles_last(packed: torch.Tensor) -> torch.Tensor:
+ if packed.dtype != torch.uint8:
+ packed = packed.contiguous().view(torch.uint8)
+ lo = packed.bitwise_and(15)
+ hi = packed.bitwise_right_shift(4).bitwise_and(15)
+ target_shape = list(packed.shape)
+ target_shape[-1] = packed.shape[-1] * 2
+ return torch.stack([lo, hi], dim=-1).reshape(target_shape).contiguous()
+
+
+def _pack_fp4_nibbles_last(unpacked: torch.Tensor) -> torch.Tensor:
+ assert (
+ unpacked.shape[-1] % 2 == 0
+ ), f"_pack_fp4_nibbles_last requires even last dim, got shape {tuple(unpacked.shape)}"
+ if unpacked.dtype != torch.uint8:
+ unpacked = unpacked.to(torch.uint8)
+ lo = unpacked[..., 0::2]
+ hi = unpacked[..., 1::2]
+ return (lo | hi << 4).contiguous()
diff --git a/tilert/models/glm_5_2_rocm/__init__.py b/tilert/models/glm_5_2_rocm/__init__.py
new file mode 100644
index 0000000..8a0c64a
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/__init__.py
@@ -0,0 +1 @@
+"""GLM-5 / GLM-5.2 model ops (ROCm port, MI350X/CDNA4 target)."""
diff --git a/tilert/models/glm_5_2_rocm/checkpoint_config.py b/tilert/models/glm_5_2_rocm/checkpoint_config.py
new file mode 100644
index 0000000..978e8f1
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/checkpoint_config.py
@@ -0,0 +1,207 @@
+"""HF ``config.json`` <-> ``ModelArgsGlm52`` cross-check."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from typing import Any
+
+from tilert.models.glm_5_2_rocm.model_args import (
+ KIND_DENSE,
+ KIND_MOE_SHARED,
+ ModelArgsGlm52,
+ layer_kind,
+)
+
+__all__ = [
+ "CONVERTER_VERSION",
+ "ConfigMismatch",
+ "describe_hf_config",
+ "load_hf_config",
+ "rope_theta_of",
+ "sha256_file",
+ "validate_hf_config",
+]
+CONVERTER_VERSION = 3
+
+
+class ConfigMismatch(ValueError):
+ """The checkpoint's config.json contradicts the model args / kernels."""
+
+
+def load_hf_config(path_or_dir: str) -> dict[str, Any]:
+ path = path_or_dir
+ if os.path.isdir(path):
+ path = os.path.join(path, "config.json")
+ with open(path) as f:
+ return json.load(f)
+
+
+def sha256_file(path: str) -> str:
+ h = hashlib.sha256()
+ with open(path, "rb") as f:
+ for chunk in iter(lambda: f.read(1 << 20), b""):
+ h.update(chunk)
+ return h.hexdigest()
+
+
+def rope_theta_of(cfg: dict[str, Any]) -> float | None:
+ rp = cfg.get("rope_parameters")
+ if isinstance(rp, dict) and "rope_theta" in rp:
+ return float(rp["rope_theta"])
+ if "rope_theta" in cfg:
+ return float(cfg["rope_theta"])
+ return None
+
+
+def _rope_scaling_of(cfg: dict[str, Any]) -> Any:
+ rp = cfg.get("rope_parameters")
+ if isinstance(rp, dict):
+ t = rp.get("rope_type", "default")
+ return None if t in (None, "default") else t
+ return cfg.get("rope_scaling")
+
+
+def validate_hf_config(
+ cfg: dict[str, Any],
+ args: ModelArgsGlm52 | None = None,
+ *,
+ n_layers: int | None = None,
+ num_mtp: int | None = None,
+ max_seq_len: int | None = None,
+) -> None:
+ args = args or ModelArgsGlm52()
+ n_layers = args.n_layers if n_layers is None else n_layers
+ num_mtp = args.num_mtp if num_mtp is None else num_mtp
+ max_seq_len = args.max_seq_len if max_seq_len is None else max_seq_len
+ bad: list[str] = []
+
+ def want(key: str, expect: Any) -> None:
+ got = cfg.get(key, "")
+ if got != expect:
+ bad.append(f"{key}: checkpoint {got!r} != expected {expect!r}")
+
+ want("hidden_size", args.dim)
+ want("vocab_size", args.vocab_size)
+ want("intermediate_size", args.inter_dim)
+ want("moe_intermediate_size", args.moe_inter_dim)
+ want("num_attention_heads", args.n_heads)
+ want("q_lora_rank", args.q_lora_rank)
+ want("kv_lora_rank", args.kv_lora_rank)
+ want("qk_nope_head_dim", args.qk_nope_head_dim)
+ want("qk_rope_head_dim", args.qk_rope_head_dim)
+ want("v_head_dim", args.v_head_dim)
+ want("index_topk", args.index_topk)
+ want("index_head_dim", args.index_head_dim)
+ want("index_n_heads", args.index_n_heads)
+ want("n_routed_experts", args.n_routed_experts)
+ want("n_shared_experts", args.n_shared_experts)
+ want("num_experts_per_tok", args.n_activated_experts)
+ want("routed_scaling_factor", args.route_scale)
+ want("first_k_dense_replace", args.n_dense_layers)
+ want("rms_norm_eps", args.eps)
+ want("scoring_func", "sigmoid")
+ want("topk_method", "noaux_tc")
+ want("norm_topk_prob", True)
+ want("n_group", 1)
+ want("topk_group", 1)
+ if "indexer_rope_interleave" in cfg:
+ want("indexer_rope_interleave", True)
+ theta = rope_theta_of(cfg)
+ if theta is None:
+ bad.append("rope_theta: missing (neither rope_parameters.rope_theta nor rope_theta)")
+ elif theta != args.rope_theta:
+ bad.append(f"rope_theta: checkpoint {theta!r} != expected {args.rope_theta!r}")
+ scaling = _rope_scaling_of(cfg)
+ if scaling is not None:
+ bad.append(f"rope scaling {scaling!r}: the port implements plain RoPE only")
+ depth = cfg.get("num_hidden_layers", "")
+ if not isinstance(depth, int) or depth < n_layers:
+ bad.append(f"num_hidden_layers: checkpoint {depth!r} < wanted {n_layers}")
+ if n_layers == args.n_layers and depth != args.n_layers:
+ bad.append(f"num_hidden_layers: checkpoint {depth!r} != {args.n_layers}")
+ nextn = cfg.get("num_nextn_predict_layers", 0)
+ if num_mtp > 0 and nextn != 1:
+ bad.append(
+ f"num_nextn_predict_layers: checkpoint {nextn!r}, the port runs exactly one MTP module"
+ )
+ share = cfg.get("index_share_for_mtp_iteration")
+ if num_mtp > 1 and share is not None and (share is not True):
+ bad.append(
+ f"index_share_for_mtp_iteration: checkpoint {share!r}, chained MTP drafts (num_mtp {num_mtp}) share MTP[0]'s selection"
+ )
+ max_pos = cfg.get("max_position_embeddings", "")
+ if not isinstance(max_pos, int):
+ bad.append(f"max_position_embeddings: {max_pos!r}")
+ elif max_seq_len > max_pos:
+ bad.append(f"max_seq_len {max_seq_len} > max_position_embeddings {max_pos}")
+ itypes = cfg.get("indexer_types")
+ freq = cfg.get("index_topk_freq")
+ skip_off = cfg.get("index_skip_topk_offset")
+ if isinstance(itypes, list):
+ n = min(len(itypes), n_layers)
+ exp = ["shared" if layer_kind(i) == KIND_MOE_SHARED else "full" for i in range(n)]
+ if itypes[:n] != exp:
+ bad.append("indexer_types: does not follow the (i-2) % 4 full/shared rule")
+ elif freq is not None or skip_off is not None:
+ if not (isinstance(freq, int) and isinstance(skip_off, int) and (freq > 0)):
+ bad.append(f"index_topk_freq / index_skip_topk_offset: {freq!r} / {skip_off!r}")
+ else:
+ for i in range(n_layers):
+ full = max(i - skip_off + 1, 0) % freq == 0
+ if full != (layer_kind(i) != KIND_MOE_SHARED):
+ bad.append(
+ f"index_topk_freq {freq} / index_skip_topk_offset {skip_off}: layer {i} disagrees with the kernels' (i-2) % 4 full/shared rule"
+ )
+ break
+ else:
+ bad.append(
+ "indexer layer pattern: neither indexer_types nor index_topk_freq/index_skip_topk_offset present"
+ )
+ mtypes = cfg.get("mlp_layer_types")
+ moe_freq = cfg.get("moe_layer_freq")
+ if isinstance(mtypes, list):
+ n = min(len(mtypes), n_layers)
+ exp = ["dense" if layer_kind(i) == KIND_DENSE else "sparse" for i in range(n)]
+ if mtypes[:n] != exp:
+ bad.append("mlp_layer_types: does not follow the 3-dense-then-MoE rule")
+ elif moe_freq is not None:
+ if moe_freq != 1:
+ bad.append(f"moe_layer_freq: checkpoint {moe_freq!r} != 1 (every layer)")
+ else:
+ bad.append("mlp layer pattern: neither mlp_layer_types nor moe_layer_freq present")
+ q = cfg.get("quantization_config")
+ if not isinstance(q, dict):
+ bad.append("quantization_config: missing (the port expects an fp8 checkpoint)")
+ else:
+ if q.get("quant_method") != "fp8" or q.get("fmt") != "e4m3":
+ qm, fmt = (q.get("quant_method"), q.get("fmt"))
+ bad.append(f"quantization_config: {qm!r}/{fmt!r} != fp8/e4m3")
+ if list(q.get("weight_block_size", [])) != [128, 128]:
+ bad.append(f"weight_block_size: {q.get('weight_block_size')!r} != [128, 128]")
+ if bad:
+ raise ConfigMismatch(
+ "checkpoint config.json does not match ModelArgsGlm52 / the kernels:\n "
+ + "\n ".join(bad)
+ )
+
+
+def describe_hf_config(cfg: dict[str, Any]) -> dict[str, Any]:
+ q = cfg.get("quantization_config") or {}
+ return {
+ "hf_model_type": cfg.get("model_type"),
+ "hf_architectures": cfg.get("architectures"),
+ "num_hidden_layers": cfg.get("num_hidden_layers"),
+ "num_nextn_predict_layers": cfg.get("num_nextn_predict_layers"),
+ "max_position_embeddings": cfg.get("max_position_embeddings"),
+ "rope_theta": rope_theta_of(cfg),
+ "index_topk": cfg.get("index_topk"),
+ "index_share_for_mtp_iteration": cfg.get("index_share_for_mtp_iteration"),
+ "quant": {
+ "quant_method": q.get("quant_method"),
+ "fmt": q.get("fmt"),
+ "weight_block_size": q.get("weight_block_size"),
+ },
+ "eos_token_id": cfg.get("eos_token_id"),
+ }
diff --git a/tilert/models/glm_5_2_rocm/end2end.py b/tilert/models/glm_5_2_rocm/end2end.py
new file mode 100644
index 0000000..aaf6cef
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/end2end.py
@@ -0,0 +1,624 @@
+"""GLM-5.2 Show Hands: the TP8 e2e Python wrapper (8x MI350X, one process)."""
+
+from __future__ import annotations
+
+import contextlib
+import dataclasses
+import json
+import math
+import os
+import threading
+import weakref
+from concurrent.futures import ThreadPoolExecutor
+
+import torch
+from safetensors import safe_open
+
+from tilert import logger
+from tilert.models.glm_5_2_rocm.checkpoint_config import (
+ CONVERTER_VERSION,
+ load_hf_config,
+ sha256_file,
+ validate_hf_config,
+)
+from tilert.models.glm_5_2_rocm.model_args import KIND_MOE_SHARED, ModelArgsGlm52, layer_kind
+from tilert.models.glm_5_2_rocm.ops import (
+ index_collective,
+ llm_preprocess,
+ top1_allreduce,
+ top_p,
+ unprojo_allreduce,
+)
+from tilert.models.glm_5_2_rocm.ops.rmsnorm_head_proj import TOP1_WS_WORDS
+from tilert.models.glm_5_2_rocm.weight_converter import (
+ fp8_ki_enabled,
+ fp8_kv_enabled,
+ load_rank_params,
+ random_rank_params,
+)
+
+__all__ = ["Glm52ShowHands", "Idx", "TEMP_VARS_SIZE"]
+HIDDEN = 6144
+ROPE_DIM = 64
+Q_LORA = 2048
+KV_LORA = 512
+KV_FP8_ROW = KV_LORA + KV_LORA // 128 * 4
+NOPE = 192
+V_HEAD = 256
+INDEX_HEADS = 32
+INDEX_DIM = 128
+TOPK = 2048
+HIST_ROW = 520 + 2 * 64 * 256 + 256 * 256 + 257
+HIST_MAX_SAMPLES = 8
+TIE_PAIRS = 4096 * 64
+RADIX = 256
+MOE_SLOTS = 9
+TOP_K = 8
+VOCAB_SHARD = 19360
+NUM_PES = 8
+_lease_lock = threading.Lock()
+_lease_owner: weakref.ReferenceType | None = None
+
+
+def _lease_holder() -> Glm52ShowHands | None:
+ return _lease_owner() if _lease_owner is not None else None
+
+
+class Idx:
+ """Temp-var indices -- must mirror the engine's own temp table."""
+
+ X = 0
+ ROPE_FREQS = 1
+ Q_PE = 2
+ Q_NOPE = 3
+ FLASH_ACC = 4
+ FLASH_MAX = 5
+ FLASH_SUM = 6
+ UNPROJ_O = 7
+ EXP_OUT = 8
+ NORM_HIDDEN = 9
+ SCORES = 10
+ MOE_PROBS = 11
+ MOE_IDX = 12
+ HIDDEN_MID = 13
+ DENSE_MID = 14
+ LOGITS = 15
+ NORM_OUT = 16
+ TOKEN_OUT = 17
+ PROB_OUT = 18
+ TOPP_SEND = 19
+ TOKEN_ID = 20
+ CUR_POS = 21
+ SAMPLING_SEED = 22
+ SAMPLING_POS = 23
+ AR_ACC = 24
+ AR_NUM = 25
+ IDX_SCORES = 26
+ IDX_IQ_RT = 27
+ IDX_LOGITS = 28
+ IDX_HIST = 29
+ IDX_TIE = 30
+ IDX_SELECTS = 31
+ SYM_ATTN = 32
+ SYM_ATTN_TAB = 33
+ SYM_FFN = 34
+ SYM_FFN_TAB = 35
+ SYM_TOP1 = 36
+ SYM_TOP1_TAB = 37
+ SYM_TOPP = 38
+ SYM_TOPP_TAB = 39
+ XFER_BUF = 40
+ XFER_TAB = 41
+ MTP_EH = 42
+ MTP_HIDDEN = 43
+ MTP_TOKENS = 44
+ DRAFT_TOKENS = 45
+ NEXT_DRAFT = 46
+ NUM_ACCEPTED = 47
+ LAST_TOKEN = 48
+ LAST_HIDDEN = 49
+ LAYER_TRACE = 50
+ MTP_CUR_POS = 51
+ SCORE_LINES = 52
+ MOE_FLAGS = 53
+ MID_PAIRS = 54
+ Q_PAIRS = 55
+ KV_PAIRS = 56
+ PE_PAIRS = 57
+ M1_PAIRS = 58
+ M56_PAIRS = 59
+ M56_FLAGS = 60
+ PROJ_PAIRS = 61
+ M5_TRIPLE_SEN = 62
+ MLA_AR_QLINES = 63
+ MLA_AR_KVNEW = 64
+ MLA_AR_PENEW = 65
+ MLA_MOE_HLINES = 66
+ SS_Q_PAIRS = 67
+ SS_KI_PAIRS = 68
+ SS_IQ_PAIRS = 69
+ IDX_IQ_RT8 = 70
+ IDX_IQ_SCALE = 71
+ TOP1_WS = 72
+
+
+TEMP_VARS_SIZE = 73
+
+
+def validate_temp_vars_layout() -> None:
+ got = int(torch.ops.tilert.glm52_temp_vars_size())
+ if got != TEMP_VARS_SIZE:
+ raise RuntimeError(
+ f"temp-var layout drift: the loaded library expects {got} slots, this package builds {TEMP_VARS_SIZE}; the library and the python package are from different builds."
+ )
+
+
+class Glm52ShowHands:
+ """TP8 show-hands wrapper: weights, temps, caches, and the engine calls."""
+
+ def __init__(
+ self,
+ model_args: ModelArgsGlm52 | None = None,
+ temperature: float = 1.0,
+ top_p_val: float = 0.95,
+ sampling_seed: int = 42,
+ use_topp: bool = False,
+ max_seq_len: int | None = None,
+ n_layers: int | None = None,
+ num_mtp: int | None = None,
+ ) -> None:
+ self.args = dataclasses.replace(model_args) if model_args else ModelArgsGlm52()
+ if max_seq_len is not None:
+ self.args.max_seq_len = max_seq_len
+ self.n_layers = self.args.n_layers if n_layers is None else n_layers
+ self.num_mtp = self.args.num_mtp if num_mtp is None else num_mtp
+ if self.num_mtp not in (0, 1, 3):
+ raise ValueError(
+ f"num_mtp must be 0 (plain), 1 (seq-2 verify) or 3 (seq-4 verify + chained drafts); got {self.num_mtp}"
+ )
+ if (
+ isinstance(self.n_layers, bool)
+ or not isinstance(self.n_layers, int)
+ or (not 1 <= self.n_layers <= self.args.n_layers)
+ ):
+ raise ValueError(
+ f"n_layers must be an int in [1, {self.args.n_layers}] (got {self.n_layers!r})"
+ )
+ if not 0.0 < top_p_val <= 1.0:
+ raise ValueError(f"top_p must be in (0, 1] (got {top_p_val})")
+ if not (math.isfinite(temperature) and temperature > 0.0):
+ raise ValueError(f"temperature must be finite and > 0 (got {temperature})")
+ self.temperature = temperature
+ self.top_p = top_p_val
+ self.sampling_seed = sampling_seed
+ self.use_topp = use_topp
+ self.npes = self.args.num_devices
+ self.max_samples = self.num_mtp + 1
+ self._temps: list[list[torch.Tensor]] = []
+ self._caches: list[list[torch.Tensor]] = []
+ self._prepared = False
+
+ def _build_temp_vars(self, rank: int) -> list[torch.Tensor]:
+ a = self.args
+ dev = f"cuda:{rank}"
+ S = self.max_samples
+ L = a.max_seq_len
+ heads = a.local_heads
+ bf16 = {"dtype": torch.bfloat16, "device": dev}
+ f32 = {"dtype": torch.float32, "device": dev}
+ i32 = {"dtype": torch.int32, "device": dev}
+ i64 = {"dtype": torch.int64, "device": dev}
+ u8 = {"dtype": torch.uint8, "device": dev}
+ splits = self.num_mtp + 1
+ from tilert.models.glm_5_2_rocm.ops.flash_sparse_mla import split_tile_n
+
+ n_splits = TOPK // split_tile_n()
+ t: list[torch.Tensor] = [torch.empty(0)] * TEMP_VARS_SIZE
+ t[Idx.X] = torch.zeros(S, HIDDEN, **bf16)
+ t[Idx.ROPE_FREQS] = torch.zeros(S, ROPE_DIM, **f32)
+ t[Idx.Q_PE] = torch.zeros(S, heads * ROPE_DIM, **bf16)
+ t[Idx.Q_NOPE] = torch.zeros(S, heads * KV_LORA, **bf16)
+ t[Idx.FLASH_ACC] = torch.zeros(S, heads, n_splits, KV_LORA, **f32)
+ t[Idx.FLASH_MAX] = torch.zeros(S, heads, n_splits, **f32)
+ t[Idx.FLASH_SUM] = torch.zeros(S, heads, n_splits, **f32)
+ t[Idx.UNPROJ_O] = torch.zeros(S, HIDDEN, **bf16)
+ t[Idx.EXP_OUT] = torch.zeros(S, HIDDEN, **bf16)
+ t[Idx.NORM_HIDDEN] = torch.zeros(S, HIDDEN, **bf16)
+ t[Idx.SCORES] = torch.zeros(S, a.n_routed_experts, **f32)
+ t[Idx.LAYER_TRACE] = torch.zeros(self.n_layers + 1, S, HIDDEN, **bf16)
+ t[Idx.MTP_CUR_POS] = torch.zeros(a.max_batch_size, **i32)
+ t[Idx.SCORE_LINES] = torch.zeros(S, 32, 32, **i32)
+ t[Idx.MOE_FLAGS] = torch.zeros(512, **i32)
+ t[Idx.MID_PAIRS] = torch.zeros(S, MOE_SLOTS, 256, **i32)
+ t[Idx.Q_PAIRS] = torch.zeros(S, Q_LORA // 2, 2, **i32)
+ t[Idx.KV_PAIRS] = torch.zeros(S, KV_LORA // 2, 2, **i32)
+ t[Idx.PE_PAIRS] = torch.zeros(S, ROPE_DIM // 2, 2, **i32)
+ t[Idx.M1_PAIRS] = torch.zeros(S, heads * 256 // 2, 2, **i32)
+ t[Idx.M56_PAIRS] = torch.zeros(8, heads, KV_LORA, **i32)
+ t[Idx.M56_FLAGS] = torch.zeros(16 * heads, **i32)
+ t[Idx.PROJ_PAIRS] = torch.zeros(heads * 16 * 4 * 16 * (1 + 8), **i32)
+ t[Idx.M5_TRIPLE_SEN] = torch.zeros(2 * 8 * (TOPK // 64), **i32)
+ t[Idx.MLA_AR_QLINES] = torch.zeros(heads * (KV_LORA + ROPE_DIM) // 16 * 4 * 16, **i32)
+ t[Idx.MLA_AR_KVNEW] = torch.zeros(4 * KV_LORA * 2, **i32)
+ t[Idx.MLA_AR_PENEW] = torch.zeros(4 * ROPE_DIM * 2, **i32)
+ t[Idx.MLA_MOE_HLINES] = torch.zeros(4 * 256 * 16, **i32)
+ t[Idx.SS_Q_PAIRS] = torch.zeros(S, Q_LORA // 2, 2, **i32)
+ t[Idx.SS_KI_PAIRS] = torch.zeros(S, INDEX_DIM // 2, 2, **i32)
+ t[Idx.SS_IQ_PAIRS] = torch.zeros(S, INDEX_HEADS * INDEX_DIM // 2, 2, **i32)
+ t[Idx.MOE_PROBS] = torch.zeros(S, TOP_K, **f32)
+ t[Idx.MOE_IDX] = torch.zeros(S, TOP_K, **i32)
+ t[Idx.HIDDEN_MID] = torch.zeros(S, MOE_SLOTS, a.moe_inter_shard, **bf16)
+ t[Idx.DENSE_MID] = torch.zeros(S, a.dense_inter_shard, **bf16)
+ t[Idx.LOGITS] = torch.zeros(S, VOCAB_SHARD, **f32)
+ t[Idx.NORM_OUT] = torch.zeros(S, HIDDEN, **bf16)
+ t[Idx.TOKEN_OUT] = torch.zeros(S, **i32)
+ t[Idx.PROB_OUT] = torch.zeros(S, **f32)
+ t[Idx.TOPP_SEND] = torch.zeros(S, top_p.SEND_BYTES, **u8)
+ t[Idx.TOKEN_ID] = torch.zeros(a.max_batch_size, **i32)
+ t[Idx.CUR_POS] = torch.zeros(a.max_batch_size, **i32)
+ t[Idx.SAMPLING_SEED] = torch.full((S,), self.sampling_seed, **i64)
+ t[Idx.SAMPLING_POS] = torch.zeros(S, **i64)
+ t[Idx.AR_ACC] = torch.zeros(a.max_batch_size, L + 1, **i32)
+ t[Idx.AR_NUM] = torch.zeros(a.max_batch_size, L + 1, **i32)
+ t[Idx.IDX_SCORES] = torch.zeros(S, INDEX_HEADS, **bf16)
+ t[Idx.IDX_IQ_RT] = torch.zeros(S, INDEX_HEADS * INDEX_DIM, **bf16)
+ t[Idx.IDX_IQ_RT8] = torch.zeros(S, INDEX_HEADS * INDEX_DIM, **u8)
+ t[Idx.IDX_IQ_SCALE] = torch.zeros(S, INDEX_HEADS, **f32)
+ t[Idx.IDX_LOGITS] = torch.zeros(S, L, **f32)
+ t[Idx.IDX_HIST] = torch.zeros(HIST_MAX_SAMPLES * HIST_ROW, **i32)
+ t[Idx.IDX_TIE] = torch.zeros(S * 2 * TIE_PAIRS, **i32)
+ t[Idx.IDX_SELECTS] = torch.zeros(S, TOPK, **i32)
+ t[Idx.SYM_ATTN] = unprojo_allreduce.sym_buffer(S, dev)
+ t[Idx.SYM_ATTN_TAB] = torch.zeros(NUM_PES, **i64)
+ t[Idx.SYM_FFN] = unprojo_allreduce.sym_buffer(S, dev)
+ t[Idx.SYM_FFN_TAB] = torch.zeros(NUM_PES, **i64)
+ t[Idx.SYM_TOP1] = top1_allreduce.sym_buffer(S, dev)
+ t[Idx.SYM_TOP1_TAB] = torch.zeros(NUM_PES, **i64)
+ t[Idx.SYM_TOPP] = top_p.sym_buffer(S, dev)
+ t[Idx.SYM_TOPP_TAB] = torch.zeros(NUM_PES, **i64)
+ t[Idx.XFER_BUF] = torch.zeros(index_collective.xfer_buf_bytes(S, TOPK) // 4, **i32)
+ t[Idx.XFER_TAB] = torch.zeros(NUM_PES, **i64)
+ t[Idx.MTP_EH] = torch.zeros(S, HIDDEN, **bf16)
+ t[Idx.MTP_HIDDEN] = torch.zeros(S, HIDDEN, **bf16)
+ t[Idx.MTP_TOKENS] = torch.zeros(S, **i32)
+ t[Idx.DRAFT_TOKENS] = torch.zeros(a.max_batch_size * S, **i32)
+ t[Idx.NEXT_DRAFT] = torch.zeros(a.max_batch_size * S, **i32)
+ t[Idx.NUM_ACCEPTED] = torch.zeros(a.max_batch_size, **i32)
+ t[Idx.LAST_TOKEN] = torch.zeros(a.max_batch_size, **i32)
+ ws_words = TOP1_WS_WORDS
+ t[Idx.TOP1_WS] = torch.zeros(ws_words, dtype=torch.int64, device=dev)
+ t[Idx.LAST_HIDDEN] = torch.zeros(a.max_batch_size, HIDDEN, **bf16)
+ assert all((x.numel() > 0 or i == 0 for i, x in enumerate(t))), "temp gap"
+ del splits
+ return t
+
+ def _build_caches(self, rank: int) -> list[torch.Tensor]:
+ a = self.args
+ dev = f"cuda:{rank}"
+ L = a.max_seq_len
+ B = a.max_batch_size
+ bf16 = {"dtype": torch.bfloat16, "device": dev}
+ caches: list[torch.Tensor] = []
+ n_extra = 1 if self.num_mtp > 0 else 0
+
+ def kv_cache() -> torch.Tensor:
+ if fp8_kv_enabled():
+ return torch.zeros(B, L, KV_FP8_ROW, dtype=torch.uint8, device=dev)
+ return torch.zeros(B, L, KV_LORA, **bf16)
+
+ if rank == 0:
+ for _ in range(self.n_layers):
+ caches.append(kv_cache())
+ caches.append(torch.zeros(B, L, ROPE_DIM, **bf16))
+ n_full = sum(1 for i in range(self.n_layers) if layer_kind(i) != KIND_MOE_SHARED)
+ for _ in range(n_full + n_extra):
+ if fp8_ki_enabled():
+ caches.append(
+ torch.zeros(B * L * (INDEX_DIM + 4), dtype=torch.uint8, device=dev)
+ )
+ else:
+ caches.append(torch.zeros(B, L, INDEX_DIM, **bf16))
+ else:
+ for _ in range(self.n_layers + n_extra):
+ caches.append(kv_cache())
+ caches.append(torch.zeros(B, L, ROPE_DIM, **bf16))
+ return caches
+
+ _SYM_PAIRS = (
+ (Idx.SYM_ATTN, Idx.SYM_ATTN_TAB),
+ (Idx.SYM_FFN, Idx.SYM_FFN_TAB),
+ (Idx.SYM_TOP1, Idx.SYM_TOP1_TAB),
+ (Idx.SYM_TOPP, Idx.SYM_TOPP_TAB),
+ (Idx.XFER_BUF, Idx.XFER_TAB),
+ )
+
+ def _link_sym_tables(self) -> None:
+ for buf_i, tab_i in self._SYM_PAIRS:
+ ptrs = [self._temps[r][buf_i].data_ptr() for r in range(self.npes)]
+ for r in range(self.npes):
+ self._temps[r][tab_i].copy_(
+ torch.tensor(ptrs, dtype=torch.int64), non_blocking=False
+ )
+
+ def _acquire_lease(self) -> None:
+ global _lease_owner
+ with _lease_lock:
+ holder = _lease_holder()
+ if holder is self:
+ raise RuntimeError(
+ "this Glm52ShowHands instance is already prepared; call cleanup() before loading again"
+ )
+ if holder is not None:
+ raise RuntimeError(
+ "another Glm52ShowHands instance is live in this process (the show-hands state is a singleton); call cleanup() on it first -- this instance was not prepared and the live one is untouched"
+ )
+ _lease_owner = weakref.ref(self)
+
+ def owns_state(self) -> bool:
+ with _lease_lock:
+ return _lease_holder() is self
+
+ def _prepare(self, params_by_rank: list[list[torch.Tensor]]) -> None:
+ assert self.owns_state(), "_prepare without the lease"
+ self._params = params_by_rank
+ for r in range(self.npes):
+ torch.ops.tilert.glm52_prepare_rank(
+ r, params_by_rank[r], self._temps[r], self._caches[r]
+ )
+ torch.ops.tilert.glm52_show_hands_prepare_money(
+ self.npes,
+ self.n_layers,
+ self.args.max_seq_len,
+ self.num_mtp,
+ self.use_topp,
+ self.temperature,
+ self.top_p,
+ )
+ self._prepared = True
+
+ def _alloc_state(self) -> None:
+ validate_temp_vars_layout()
+ if self.args.max_seq_len < TOPK:
+ raise ValueError(
+ f"max_seq_len ({self.args.max_seq_len}) must be >= the sparse selection budget ({TOPK}); the flash kernel rejects a cache shorter than topk"
+ )
+ n = int(unprojo_allreduce.enable_peer_access(self.npes))
+ if n < self.npes:
+ raise RuntimeError(
+ f"peer access enabled on only {n} of {self.npes} devices; the collectives need all 8 GPUs peer-mapped"
+ )
+ self._temps = []
+ self._caches = []
+ for r in range(self.npes):
+ with torch.cuda.device(r):
+ self._temps.append(self._build_temp_vars(r))
+ self._caches.append(self._build_caches(r))
+ self._link_sym_tables()
+
+ def _make_freqs_cis(self) -> torch.Tensor:
+ return llm_preprocess.make_freqs_cis(
+ self.args.max_seq_len, theta=self.args.rope_theta, device="cpu"
+ )
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def from_pretrained(self, weights_dir: str) -> None:
+ self._acquire_lease()
+ try:
+ self._from_pretrained_impl(weights_dir)
+ except BaseException:
+ self.cleanup()
+ raise
+
+ def _init_random_weights_impl(self, seed: int) -> None:
+ self._alloc_state()
+ freqs_cis = self._make_freqs_cis()
+ params = []
+ for r in range(self.npes):
+ with torch.cuda.device(r):
+ params.append(
+ random_rank_params(
+ self.args,
+ r,
+ f"cuda:{r}",
+ self.n_layers,
+ seed=seed,
+ num_mtp=self.num_mtp,
+ freqs_cis=freqs_cis,
+ )
+ )
+ logger.info("rank %d: random weights packed", r)
+ self._prepare(params)
+
+ def _from_pretrained_impl(self, weights_dir: str) -> None:
+ meta_path = os.path.join(weights_dir, "tilert_meta.json")
+ cfg_path = os.path.join(weights_dir, "shared", "config.json")
+ if not os.path.isfile(meta_path):
+ raise ValueError(f"{weights_dir}: no tilert_meta.json (not a conversion?)")
+ if not os.path.isfile(cfg_path):
+ raise ValueError(f"{weights_dir}: no shared/config.json to validate against")
+ hf_cfg = load_hf_config(cfg_path)
+ validate_hf_config(
+ hf_cfg,
+ self.args,
+ n_layers=self.n_layers,
+ num_mtp=self.num_mtp,
+ max_seq_len=self.args.max_seq_len,
+ )
+ if True:
+ with open(meta_path) as f:
+ meta = json.load(f)
+ stamp_hint = f"; stamp it with `python -m tilert.models.glm_5_2_rocm.weight_converter --stamp-provenance --model_dir --save_dir {weights_dir}`"
+ if meta.get("model") != "glm_5_2":
+ raise ValueError(
+ f"{weights_dir} holds a {meta.get('model')!r} conversion, not glm_5_2"
+ + ("" if "model" in meta else stamp_hint)
+ )
+ version = meta.get("converter_version")
+ if version != CONVERTER_VERSION:
+ raise ValueError(
+ f"{weights_dir}: converter_version {version!r} != this loader's {CONVERTER_VERSION}"
+ + stamp_hint
+ )
+ want_sha = meta.get("source_config_sha256")
+ if not want_sha:
+ raise ValueError(
+ f"{weights_dir}: tilert_meta.json carries no source_config_sha256 (no provenance)"
+ + stamp_hint
+ )
+ if want_sha != sha256_file(cfg_path):
+ raise ValueError(
+ f"{weights_dir}: shared/config.json differs from the checkpoint config the conversion was cut from (sha256 mismatch)"
+ )
+ if meta.get("n_layers", self.n_layers) < self.n_layers:
+ raise ValueError(
+ f"{weights_dir} holds only {meta['n_layers']} layers, this run wants {self.n_layers}"
+ )
+ if self.num_mtp > 0 and meta.get("num_mtp", 0) < 1:
+ raise ValueError(
+ f"{weights_dir} was converted without the MTP module (rerun the converter with --num_mtp 1) or construct with num_mtp=0"
+ )
+ if not meta.get("attn_fp8_lossless", False):
+ raise ValueError(
+ f"{weights_dir} holds RE-QUANTIZED Wq_b/Wkv_b tensors (128-row scales); the attention kernels read 64-row-stripe scales over the checkpoint's own fp8 bytes. Repair in place with: python -m tilert.models.glm_5_2_rocm.weight_converter --augment-attn-lossless --model_dir --save_dir {weights_dir}"
+ )
+ if not meta.get("attn_tp8", False):
+ raise ValueError(
+ f"{weights_dir} was converted without the attn_tp8 set, which the pure-MLA TP8 runner reads on every shared layer. Either re-run the converter or augment the existing directory with weight_converter.augment_attn_tp8"
+ )
+ self._alloc_state()
+ with safe_open(
+ os.path.join(weights_dir, "shared", "embed.safetensors"), framework="pt"
+ ) as f:
+ embed = f.get_tensor("embed")
+ freqs_cis = self._make_freqs_cis()
+
+ def load_one(r: int) -> list[torch.Tensor]:
+ torch.cuda.set_device(r)
+ out = load_rank_params(
+ weights_dir,
+ self.args,
+ r,
+ f"cuda:{r}",
+ n_layers=self.n_layers,
+ num_mtp=self.num_mtp,
+ embed=embed,
+ freqs_cis=freqs_cis,
+ )
+ logger.info("rank %d: weights loaded and packed", r)
+ return out
+
+ with ThreadPoolExecutor(max_workers=self.npes) as pool:
+ params = list(pool.map(load_one, range(self.npes)))
+ self._prepare(params)
+
+ def _check_token(self, token_id: torch.Tensor) -> None:
+ if not isinstance(token_id, torch.Tensor):
+ raise TypeError(f"token_id must be an int32 [1] tensor, got {token_id!r}")
+ if token_id.dtype != torch.int32 or token_id.numel() != 1:
+ raise ValueError(
+ f"token_id must be int32 [1], got {token_id.dtype} {tuple(token_id.shape)}"
+ )
+ tok = int(token_id.item())
+ if not 0 <= tok < self.args.vocab_size:
+ raise ValueError(f"token_id {tok} outside [0, {self.args.vocab_size})")
+
+ def _check_steps(self, n: object) -> int:
+ if isinstance(n, bool) or not isinstance(n, int):
+ raise TypeError(f"step count must be an int, got {n!r}")
+ if not 1 <= n <= self.args.max_seq_len:
+ raise ValueError(f"step count {n} outside [1, max_seq_len = {self.args.max_seq_len}]")
+ return n
+
+ def forward(self, token_id: torch.Tensor) -> None:
+ assert self._prepared, "call init_random_weights() or from_pretrained()"
+ self._check_token(token_id)
+ torch.ops.tilert.glm52_show_hands(token_id)
+
+ def step(self, token_id: int) -> int:
+ tok = torch.tensor([token_id], dtype=torch.int32)
+ self.forward(tok)
+ return int(self._temps[0][Idx.TOKEN_OUT][0].item())
+
+ def prefill(self, token_id: torch.Tensor, next_token: int) -> None:
+ assert self._prepared, "call init_random_weights() or from_pretrained()"
+ self._check_token(token_id)
+ if isinstance(next_token, bool) or not isinstance(next_token, int):
+ raise TypeError(f"next_token must be an int, got {next_token!r}")
+ if not -1 <= next_token < self.args.vocab_size:
+ raise ValueError(f"next_token {next_token} outside [-1, {self.args.vocab_size})")
+ torch.ops.tilert.glm52_show_hands_prefill(
+ token_id, torch.tensor([next_token], dtype=torch.int32)
+ )
+
+ def decode_n(self, n: int) -> None:
+ torch.ops.tilert.glm52_show_hands_decode_n(self._check_steps(n))
+
+ def mtp_n(self, n: int) -> int:
+ return int(torch.ops.tilert.glm52_show_hands_mtp_n(self._check_steps(n)))
+
+ def seed_draft(self, token: int, draft: int) -> None:
+ torch.ops.tilert.glm52_seed_draft(token, draft)
+
+ @property
+ def accepted_count(self) -> int:
+ return int(self._temps[0][Idx.AR_ACC][0, 0].item())
+
+ def accepted_tokens(self, start: int = 0, end: int | None = None) -> list[int]:
+ end = self.accepted_count if end is None else end
+ if end <= start:
+ return []
+ row = self._temps[0][Idx.AR_ACC][0, 1 + start : 1 + end]
+ return row.cpu().tolist()
+
+ @property
+ def step_count(self) -> int:
+ return int(self._temps[0][Idx.AR_NUM][0, 0].item())
+
+ def accepted_step_counts(self, start: int = 0, end: int | None = None) -> list[int]:
+ end = self.step_count if end is None else end
+ if end <= start:
+ return []
+ row = self._temps[0][Idx.AR_NUM][0, 1 + start : 1 + end]
+ return row.cpu().tolist()
+
+ @property
+ def token_out(self) -> int:
+ return int(self._temps[0][Idx.TOKEN_OUT][0].item())
+
+ def update_sampling(self, use_topp: bool, temperature: float, top_p_val: float) -> None:
+ if not 0.0 < top_p_val <= 1.0:
+ raise ValueError(f"top_p must be in (0, 1] (got {top_p_val})")
+ if not (math.isfinite(temperature) and temperature > 0.0):
+ raise ValueError(f"temperature must be finite and > 0 (got {temperature})")
+ torch.ops.tilert.glm52_update_sampling(use_topp, temperature, top_p_val)
+ self.use_topp = use_topp
+ self.temperature = temperature
+ self.top_p = top_p_val
+
+ def reset_sequence(self) -> None:
+ torch.ops.tilert.glm52_show_hands_reset()
+
+ def set_cur_pos(self, cur_pos: int) -> None:
+ torch.ops.tilert.glm52_show_hands_set_cur_pos(cur_pos)
+
+ def cleanup(self) -> None:
+ global _lease_owner
+ self._prepared = False
+ try:
+ with _lease_lock:
+ owns = _lease_holder() is self
+ if owns:
+ torch.ops.tilert.glm52_show_hands_go_home()
+ finally:
+ with _lease_lock:
+ if _lease_holder() is self:
+ _lease_owner = None
+ self._params = []
+ self._temps = []
+ self._caches = []
+
+ def __del__(self) -> None:
+ with contextlib.suppress(Exception):
+ self.cleanup()
diff --git a/tilert/models/glm_5_2_rocm/generator.py b/tilert/models/glm_5_2_rocm/generator.py
new file mode 100644
index 0000000..4e801c8
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/generator.py
@@ -0,0 +1,311 @@
+"""Glm52Generator: prompt in, text out, over the TP8 show-hands chain."""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+
+import torch
+from transformers import AutoTokenizer
+
+from tilert import logger
+from tilert.models.glm_5_2_rocm.end2end import Glm52ShowHands
+from tilert.models.glm_5_2_rocm.model_args import ModelArgsGlm52
+from tilert.tilert_init import tilert_init
+
+__all__ = ["Glm52Generator"]
+STOP_TOKENS: dict[str, int] = {
+ "<|endoftext|>": 154820,
+ "<|user|>": 154827,
+ "<|assistant|>": 154828,
+ "<|observation|>": 154829,
+}
+TOP_K_FIXED = 256
+
+
+def _check_max_new_tokens(n: object) -> int:
+ if isinstance(n, bool) or not isinstance(n, int) or n < 0:
+ raise ValueError(f"max_new_tokens must be a non-negative int (got {n!r})")
+ return n
+
+
+class Glm52Generator:
+ """Single-process, 8-GPU generator for GLM-5.2."""
+
+ def __init__(
+ self,
+ model_weights_dir: str = "",
+ model_args: ModelArgsGlm52 | None = None,
+ max_new_tokens: int = 64,
+ temperature: float = 1.0,
+ top_p: float = 0.95,
+ sampling_seed: int = 42,
+ use_topp: bool = False,
+ n_layers: int | None = None,
+ num_mtp: int | None = None,
+ max_seq_len: int | None = None,
+ ) -> None:
+ self.model_weights_dir = model_weights_dir
+ self.max_new_tokens = _check_max_new_tokens(max_new_tokens)
+ self.decode_layer = Glm52ShowHands(
+ model_args=model_args,
+ temperature=temperature,
+ top_p_val=top_p,
+ sampling_seed=sampling_seed,
+ use_topp=use_topp,
+ max_seq_len=max_seq_len,
+ n_layers=n_layers,
+ num_mtp=num_mtp,
+ )
+ self.config = self.decode_layer.args
+ self.use_mtp = self.decode_layer.num_mtp > 0
+ self._tokenizer: AutoTokenizer | None = None
+ self._stop_token_ids: set[int] = set()
+ self.last_completion_tokens: list[int] = []
+
+ def _tok_dir(self) -> str:
+ shared = os.path.join(self.model_weights_dir, "shared")
+ if os.path.isfile(os.path.join(shared, "tokenizer_config.json")):
+ return shared
+ return self.model_weights_dir
+
+ def _ensure_tokenizer(self) -> None:
+ if self._tokenizer is not None:
+ return
+ tok_dir = self._tok_dir()
+ if not os.path.isfile(os.path.join(tok_dir, "tokenizer_config.json")):
+ raise FileNotFoundError(
+ f"No tokenizer_config.json under {self.model_weights_dir}. Re-run the converter so the output dir is self-contained."
+ )
+ self._tokenizer = AutoTokenizer.from_pretrained(tok_dir)
+ eos = self._tokenizer.eos_token_id
+ if eos is not None:
+ self._stop_token_ids.add(eos)
+ for name, known_id in STOP_TOKENS.items():
+ ids = self._tokenizer.encode(name, add_special_tokens=False)
+ if len(ids) == 1:
+ self._stop_token_ids.add(int(ids[0]))
+ continue
+ added = getattr(self._tokenizer, "added_tokens_encoder", {})
+ if name in added:
+ self._stop_token_ids.add(int(added[name]))
+ continue
+ logger.warning(
+ "tokenizer does not resolve stop token %r; using GLM-5.2's id %d", name, known_id
+ )
+ self._stop_token_ids.add(known_id)
+ gc_path = os.path.join(tok_dir, "generation_config.json")
+ if os.path.isfile(gc_path):
+ with open(gc_path) as f:
+ gc = json.load(f)
+ eos_ids = gc.get("eos_token_id", [])
+ if isinstance(eos_ids, int):
+ eos_ids = [eos_ids]
+ self._stop_token_ids.update(int(t) for t in eos_ids)
+ logger.info(f"Stop token IDs: {sorted(self._stop_token_ids)}")
+
+ @property
+ def tokenizer(self) -> AutoTokenizer:
+ self._ensure_tokenizer()
+ assert self._tokenizer is not None
+ return self._tokenizer
+
+ @property
+ def stop_token_ids(self) -> set[int]:
+ self._ensure_tokenizer()
+ return self._stop_token_ids
+
+ def init(self) -> None:
+ tilert_init()
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def from_pretrained(self) -> None:
+ self.decode_layer.from_pretrained(self.model_weights_dir)
+
+ def cleanup(self) -> None:
+ self.decode_layer.cleanup()
+
+ def update_sampling_params(
+ self,
+ temperature: float = 1.0,
+ top_p: float = 0.95,
+ top_k: int = TOP_K_FIXED,
+ use_topp: bool = True,
+ ) -> None:
+ if isinstance(top_k, bool) or top_k != TOP_K_FIXED:
+ raise ValueError(
+ f"top_k is fixed at {TOP_K_FIXED} (per-rank top-256 candidate set, 2048 globally) by the top-p kernels; got {top_k!r}"
+ )
+ self.decode_layer.update_sampling(use_topp, temperature, top_p)
+
+ @torch.inference_mode()
+ def generate(
+ self,
+ prompt: str,
+ print_log: bool = True,
+ prompt_tokens: list[int] | None = None,
+ max_new_tokens: int | None = None,
+ raw: bool = False,
+ with_mtp: bool | None = None,
+ ) -> tuple[str, list[float], list[int], int]:
+ use_mtp = self.use_mtp if with_mtp is None else with_mtp
+ if use_mtp:
+ assert self.decode_layer.num_mtp > 0, "constructed with num_mtp=0"
+ self._ensure_tokenizer()
+ assert self._tokenizer is not None
+ n_new = (
+ self.max_new_tokens if max_new_tokens is None else _check_max_new_tokens(max_new_tokens)
+ )
+ if prompt_tokens is None:
+ if raw:
+ prompt_tokens = self._tokenizer.encode(prompt)
+ else:
+ messages = [{"role": "user", "content": prompt}]
+ res = self._tokenizer.apply_chat_template(
+ messages, tokenize=True, add_generation_prompt=True
+ )
+ prompt_tokens = res["input_ids"] if hasattr(res, "keys") else res
+ prompt_len = len(prompt_tokens)
+ pos_limit = self.decode_layer.args.max_seq_len
+ if not prompt_tokens:
+ raise ValueError("empty prompt: nothing to condition the first token on")
+ vocab = self.config.vocab_size
+ for i, tid in enumerate(prompt_tokens):
+ if isinstance(tid, bool) or not isinstance(tid, int) or (not 0 <= tid < vocab):
+ raise ValueError(f"prompt token {i} = {tid!r} is not an int in [0, {vocab})")
+ if prompt_len > pos_limit:
+ raise ValueError(
+ f"prompt of {prompt_len} tokens does not fit the {pos_limit}-token cache (max_seq_len)"
+ )
+ self.last_completion_tokens = []
+ if n_new == 0:
+ return ("", [], [], prompt_len)
+ self.decode_layer.reset_sequence()
+ times: list[float] = []
+ t0 = time.time()
+ mtp_prefill = use_mtp
+ for i, tid in enumerate(prompt_tokens):
+ tok = torch.tensor([tid], dtype=torch.int32)
+ if mtp_prefill:
+ nxt = prompt_tokens[i + 1] if i + 1 < prompt_len else -1
+ self.decode_layer.prefill(tok, nxt)
+ else:
+ self.decode_layer.forward(tok)
+ prefill_s = time.time() - t0
+ base = self.decode_layer.accepted_count
+ first = self.decode_layer.token_out
+ completion: list[int] = [first]
+ accept_lens: list[int] = []
+ if print_log:
+ print(self._tokenizer.decode([first], skip_special_tokens=True), end="", flush=True)
+ stopped = first in self.stop_token_ids
+ if use_mtp:
+ self.decode_layer.seed_draft(first, first)
+ ar_steps = max(1, min(1024, int(os.environ.get("GLM5_AR_N", "8"))))
+ mtp_seq = self.decode_layer.num_mtp + 1
+ chain_slack = max(0, self.decode_layer.num_mtp - 1)
+ step_base = self.decode_layer.step_count
+ steps_done = 0
+ produced = 0
+ while not stopped and produced < n_new - 1:
+ room = pos_limit - (prompt_len + produced)
+ k = min(ar_steps, (room - chain_slack) // mtp_seq)
+ if k < 1:
+ logger.warning(
+ "stopping decode: position %d + %d would reach the %d-token cache limit",
+ prompt_len + produced,
+ mtp_seq + chain_slack,
+ pos_limit,
+ )
+ break
+ t0 = time.time()
+ got = self.decode_layer.mtp_n(k)
+ dt = time.time() - t0
+ new = self.decode_layer.accepted_tokens(base + produced)
+ per_step = self.decode_layer.accepted_step_counts(step_base + steps_done)
+ steps_done += len(per_step)
+ step_dt = dt / max(1, len(per_step))
+ offset = 0
+ for acc in per_step:
+ step_toks = new[offset : offset + acc]
+ offset += acc
+ taken = 0
+ for tok in step_toks:
+ if len(completion) >= n_new:
+ stopped = True
+ break
+ completion.append(tok)
+ taken += 1
+ if print_log:
+ print(
+ self._tokenizer.decode([tok], skip_special_tokens=True),
+ end="",
+ flush=True,
+ )
+ if tok in self.stop_token_ids:
+ stopped = True
+ break
+ if taken > 0:
+ times.append(step_dt)
+ accept_lens.append(taken)
+ if stopped:
+ break
+ produced += got
+ else:
+ chunk = 8
+ produced = 0
+ while not stopped and produced < n_new - 1:
+ room = pos_limit - (prompt_len + produced)
+ if room < 1:
+ logger.warning(
+ "stopping decode: position %d reached the %d-token cache limit",
+ prompt_len + produced,
+ pos_limit,
+ )
+ break
+ n = min(chunk, n_new - 1 - produced, room)
+ t0 = time.time()
+ self.decode_layer.decode_n(n)
+ dt = time.time() - t0
+ new = self.decode_layer.accepted_tokens(base + produced)
+ for tok in new:
+ completion.append(tok)
+ times.append(dt / n)
+ if print_log:
+ print(
+ self._tokenizer.decode([tok], skip_special_tokens=True),
+ end="",
+ flush=True,
+ )
+ if tok in self.stop_token_ids:
+ stopped = True
+ break
+ produced += n
+ if print_log:
+ print("\n")
+ n_tok = sum(accept_lens) if accept_lens else len(times)
+ logger.info(
+ f"--Tokens generated: {len(completion)} (prefill {prompt_len} tok in {prefill_s:.2f}s)"
+ )
+ if times and n_tok:
+ per_tok = sum(times) / n_tok
+ logger.info("==== Performance ====")
+ logger.info(f"--Average time per token: {per_tok * 1000:.4f} ms")
+ logger.info(f"--Effective tokens per second: {1 / per_tok:.2f}")
+ if accept_lens:
+ logger.info(
+ "--MTP mean accept length: %.3f over %d steps",
+ sum(accept_lens) / len(accept_lens),
+ len(accept_lens),
+ )
+ print("\n")
+ for i, tok in enumerate(completion):
+ if tok in self.stop_token_ids:
+ completion = completion[:i]
+ break
+ self.last_completion_tokens = completion
+ text = self._tokenizer.decode(completion, skip_special_tokens=True)
+ return (text, times, accept_lens, prompt_len)
diff --git a/tilert/models/glm_5_2_rocm/model_args.py b/tilert/models/glm_5_2_rocm/model_args.py
new file mode 100644
index 0000000..465bfc0
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/model_args.py
@@ -0,0 +1,63 @@
+"""Model arguments for GLM-5.2 (TP8 decode on 8x MI350X)."""
+
+from dataclasses import dataclass
+
+__all__ = ["ModelArgsGlm52", "layer_kind", "full_layer_ordinals"]
+KIND_DENSE = 0
+KIND_MOE_FULL = 1
+KIND_MOE_SHARED = 2
+
+
+def layer_kind(i: int) -> int:
+ if i < 3:
+ return KIND_DENSE
+ return KIND_MOE_FULL if (i - 2) % 4 == 0 else KIND_MOE_SHARED
+
+
+def full_layer_ordinals(n_layers: int) -> list[int]:
+ return [i for i in range(n_layers) if layer_kind(i) != KIND_MOE_SHARED]
+
+
+@dataclass
+class ModelArgsGlm52:
+ """GLM-5.2 model arguments (decode bring-up scope: batch=1, seq=1)."""
+
+ arch_name = "glm_5_2"
+ max_batch_size: int = 1
+ max_seq_len: int = 4096
+ vocab_size: int = 154880
+ dim: int = 6144
+ inter_dim: int = 12288
+ moe_inter_dim: int = 2048
+ n_layers: int = 78
+ n_dense_layers: int = 3
+ n_heads: int = 64
+ q_lora_rank: int = 2048
+ kv_lora_rank: int = 512
+ qk_nope_head_dim: int = 192
+ qk_rope_head_dim: int = 64
+ v_head_dim: int = 256
+ index_n_heads: int = 32
+ index_head_dim: int = 128
+ index_topk: int = 2048
+ n_routed_experts: int = 256
+ n_shared_experts: int = 1
+ n_activated_experts: int = 8
+ route_scale: float = 2.5
+ rope_theta: float = 8000000.0
+ eps: float = 1e-05
+ num_devices: int = 8
+ local_heads: int = 10
+ num_mtp: int = 0
+
+ @property
+ def vocab_shard(self) -> int:
+ return self.vocab_size // self.num_devices
+
+ @property
+ def dense_inter_shard(self) -> int:
+ return self.inter_dim // self.num_devices
+
+ @property
+ def moe_inter_shard(self) -> int:
+ return self.moe_inter_dim // self.num_devices
diff --git a/tilert/models/glm_5_2_rocm/ops/__init__.py b/tilert/models/glm_5_2_rocm/ops/__init__.py
new file mode 100644
index 0000000..2ae2839
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/__init__.py
@@ -0,0 +1 @@
+pass
diff --git a/tilert/models/glm_5_2_rocm/ops/down_allreduce.py b/tilert/models/glm_5_2_rocm/ops/down_allreduce.py
new file mode 100644
index 0000000..4db3327
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/down_allreduce.py
@@ -0,0 +1,47 @@
+"""GLM-5.2 dense MLP DownAllReduce op wrapper (layers 0-2)."""
+
+import os
+
+import torch
+
+from tilert.models.glm_5_2_rocm.ops import unprojo_allreduce as m7
+
+DENSE_INTER = 1536
+HIDDEN = m7.HIDDEN
+PROTO_0, PROTO_1 = (m7.PROTO_0, m7.PROTO_1)
+sym_bytes = m7.sym_bytes
+sym_buffer = m7.sym_buffer
+sym_table = m7.sym_table
+enable_peer_access = m7.enable_peer_access
+
+
+class DownAllReduceGlm5(m7.UnprojOAllReduceGlm5):
+ """One rank's dense-MLP Wdown shard."""
+
+ def __init__(self, device: str = "cuda:0") -> None:
+ self.num_heads = 0
+ self.k = DENSE_INTER
+ self.device = device
+ self.w_fp8 = None
+ self.scales = None
+ self.packed = None
+
+ def tilert_forward(
+ self,
+ mid: torch.Tensor,
+ residual: torch.Tensor | None = None,
+ proto: int = PROTO_0,
+ sym: torch.Tensor | None = None,
+ mype: int = 0,
+ npes: int = 1,
+ flag: int = 1,
+ out: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ assert self.packed is not None and self.scales is not None
+ os.environ["TILERT_GLM5_AR_PROTO"] = str(proto)
+ if out is None:
+ out = torch.empty(mid.shape[0], HIDDEN, dtype=torch.bfloat16, device=mid.device)
+ torch.ops.tilert.glm5_down_allreduce_op(
+ mid, self.packed, self.scales, residual, sym, mype, npes, flag, out
+ )
+ return out
diff --git a/tilert/models/glm_5_2_rocm/ops/eh_proj_allreduce.py b/tilert/models/glm_5_2_rocm/ops/eh_proj_allreduce.py
new file mode 100644
index 0000000..61ba679
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/eh_proj_allreduce.py
@@ -0,0 +1,137 @@
+"""GLM-5.2 MtpPreProcess op wrapper: packers, golden, forward."""
+
+import os
+
+import torch
+
+HIDDEN = 6144
+K_IN = 1536
+NUM_PES = 8
+EPS = 1e-05
+PROTO_0, PROTO_1 = (0, 1)
+MAX_SEQ = 8
+
+
+def swizzle_128_bf16(w: torch.Tensor) -> torch.Tensor:
+ rows, k = w.shape
+ assert rows % 16 == 0 and k % 32 == 0
+ w16 = w.to(torch.bfloat16).view(torch.uint16)
+ rg = torch.arange(rows // 16)
+ kc = torch.arange(k // 32)
+ lane = torch.arange(64)
+ i = torch.arange(8)
+ RG, KC, L, II = torch.meshgrid(rg, kc, lane, i, indexing="ij")
+ packed = w16[RG * 16 + L % 16, KC * 32 + L // 16 * 8 + II]
+ return packed.reshape(-1).contiguous().view(torch.uint8)
+
+
+def swizzle_256_bf16(w: torch.Tensor) -> torch.Tensor:
+ rows, k = w.shape
+ assert rows % 8 == 0 and k % 64 == 0
+ w16 = w.to(torch.bfloat16).view(torch.uint16)
+ rg = torch.arange(rows // 8)
+ kc = torch.arange(k // 64)
+ lane = torch.arange(64)
+ s = torch.arange(2)
+ i = torch.arange(4)
+ RG, KC, L, S, II = torch.meshgrid(rg, kc, lane, s, i, indexing="ij")
+ rows_ix = RG * 8 + (L >> 5) * 4 + (L & 3)
+ ks = KC * 64 + S * 32 + (L >> 2 & 7) * 4 + II
+ packed = w16[rows_ix, ks]
+ return packed.reshape(-1).contiguous().view(torch.uint8)
+
+
+def rms_inv_ref(x: torch.Tensor) -> torch.Tensor:
+ xf = x.float()
+ return torch.rsqrt(xf.square().mean(dim=-1) + EPS)
+
+
+class EhProjAllReduceGlm5:
+ """The full W_eh [6144, 12288] plus per-rank packed slices."""
+
+ def __init__(self, device: str = "cuda:0") -> None:
+ self.device = device
+ self.w: torch.Tensor | None = None
+ self.e_gamma: torch.Tensor | None = None
+ self.h_gamma: torch.Tensor | None = None
+ self.packed: dict[int, torch.Tensor] = {}
+ self._dev_gammas: tuple[torch.Tensor, torch.Tensor] | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def init_reference_weights(
+ self, w: torch.Tensor, e_gamma: torch.Tensor, h_gamma: torch.Tensor
+ ) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def _rank_slice_bf16(self, mype: int) -> torch.Tensor:
+ assert self.w is not None
+ ws = self.w[:, mype * K_IN : (mype + 1) * K_IN]
+ return ws.to(torch.bfloat16)
+
+ def packed_weights(self, mype: int) -> torch.Tensor:
+ if mype not in self.packed:
+ ws = self._rank_slice_bf16(mype)
+ self.packed[mype] = swizzle_128_bf16(ws).to(self.device)
+ return self.packed[mype]
+
+ def gammas(self) -> tuple[torch.Tensor, torch.Tensor]:
+ if self._dev_gammas is None:
+ self._dev_gammas = (self.e_gamma.to(self.device), self.h_gamma.to(self.device))
+ return self._dev_gammas
+
+ def partial_golden(
+ self, embed_x: torch.Tensor, last_hidden: torch.Tensor, mype: int
+ ) -> torch.Tensor:
+ assert self.w is not None
+ is_h = mype >= NUM_PES // 2
+ src = (last_hidden if is_h else embed_x).float().cpu()
+ gamma = self.h_gamma if is_h else self.e_gamma
+ g0 = mype % (NUM_PES // 2) * K_IN
+ x_slice = src[:, g0 : g0 + K_IN]
+ act = (x_slice * gamma[g0 : g0 + K_IN][None, :]).to(torch.bfloat16).float()
+ wf = self._rank_slice_bf16(mype).float().cpu()
+ rinv = rms_inv_ref(src.to(torch.bfloat16))
+ return (act @ wf.T * rinv[:, None]).to(torch.bfloat16)
+
+ @staticmethod
+ def reduce_golden(partials: list[torch.Tensor]) -> torch.Tensor:
+ acc = torch.zeros(partials[0].shape, dtype=torch.float32)
+ for p in partials:
+ acc += p.float().cpu()
+ return acc.to(torch.bfloat16)
+
+ def golden_forward(
+ self, embed_x: torch.Tensor, last_hidden: torch.Tensor, mype: int
+ ) -> torch.Tensor:
+ raise RuntimeError("golden_forward is not available in release builds")
+
+ def tilert_forward(
+ self,
+ embed_x: torch.Tensor,
+ last_hidden: torch.Tensor,
+ proto: int = PROTO_0,
+ sym: torch.Tensor | None = None,
+ mype: int = 0,
+ npes: int = 1,
+ flag: int = 1,
+ out: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ os.environ["TILERT_GLM5_AR_PROTO"] = str(proto)
+ e_gamma, h_gamma = self.gammas()
+ if out is None:
+ out = torch.empty(embed_x.shape[0], HIDDEN, dtype=torch.bfloat16, device=embed_x.device)
+ torch.ops.tilert.glm5_eh_proj_allreduce_op(
+ embed_x,
+ last_hidden,
+ e_gamma,
+ h_gamma,
+ self.packed_weights(mype),
+ sym,
+ mype,
+ npes,
+ flag,
+ out,
+ )
+ return out
diff --git a/tilert/models/glm_5_2_rocm/ops/flash_sparse_mla.py b/tilert/models/glm_5_2_rocm/ops/flash_sparse_mla.py
new file mode 100644
index 0000000..83f9da8
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/flash_sparse_mla.py
@@ -0,0 +1,79 @@
+"""GLM-5.2 FlashSparseMLA: shapes, the index wire helpers and the torch golden shared by the fused-tail wrappers and their tests."""
+
+import math
+
+import torch
+
+KV_LORA_RANK = 512
+PE_DIM = 64
+QK_DIM = KV_LORA_RANK + PE_DIM
+TILE_N = 64
+TOPK_DEFAULT = 2048
+SUPPORTED_HEADS = (8, 10, 16)
+SUPPORTED_SAMPLES = (1, 2, 4, 8)
+XCDS = 8
+GLM5_SOFTMAX_SCALE = (192 + 64) ** (-0.5)
+
+
+def split_tile_n() -> int:
+ return 32
+
+
+def softmax_scale(
+ qk_nope_head_dim: int = 192, qk_rope_head_dim: int = 64, rope_factor: float | None = None
+) -> float:
+ scale = (qk_nope_head_dim + qk_rope_head_dim) ** (-0.5)
+ mscale = 1.0 if rope_factor is None else 0.1 * math.log(rope_factor) + 1.0
+ return scale * mscale * mscale
+
+
+def pack_xfer(indices: torch.Tensor, flag: int) -> torch.Tensor:
+ flat = indices.reshape(-1).to(torch.int32)
+ assert flat.numel() % 2 == 0
+ out = torch.empty(flat.numel() * 2, dtype=torch.int32, device=flat.device)
+ out[0::4] = flat[0::2]
+ out[2::4] = flat[1::2]
+ fw = torch.tensor(flag, dtype=torch.int64).to(torch.int32)
+ out[1::4] = fw
+ out[3::4] = fw
+ return out
+
+
+def xfer_send(indices: torch.Tensor, xfer_buf: torch.Tensor, flag: int) -> None:
+ torch.ops.tilert.glm5_xfer_send_indices_op(
+ indices.reshape(-1).contiguous().to(torch.int32), xfer_buf, flag
+ )
+
+
+def xfer_buffer(samples: int, topk: int, device: str | torch.device) -> torch.Tensor:
+ return torch.zeros(samples * topk * 2, dtype=torch.int32, device=device)
+
+
+def golden_attention(
+ q_nope: torch.Tensor,
+ q_pe: torch.Tensor,
+ kv_cache: torch.Tensor,
+ pe_cache: torch.Tensor,
+ indices: torch.Tensor | None,
+ cur_pos: int,
+ topk: int = TOPK_DEFAULT,
+ scale: float = GLM5_SOFTMAX_SCALE,
+) -> torch.Tensor:
+ seq, heads, _ = q_nope.shape
+ seqlen_kv = kv_cache.shape[0]
+ kvf = kv_cache.float()
+ pef = pe_cache.float()
+ dev = q_nope.device
+ out = torch.zeros(seq, heads, KV_LORA_RANK, dtype=torch.float32, device=dev)
+ for s in range(seq):
+ kv_len = cur_pos + 1 + s
+ scores = (q_nope[s].float() @ kvf.T + q_pe[s].float() @ pef.T) * scale
+ mask = torch.full((seqlen_kv,), float("-inf"), device=dev)
+ if kv_len > topk:
+ assert indices is not None, "sparse step needs a selection"
+ mask[indices[s].long()] = 0.0
+ else:
+ mask[:kv_len] = 0.0
+ probs = (scores + mask).softmax(dim=-1, dtype=torch.float32)
+ out[s] = probs.to(torch.bfloat16).float() @ kvf
+ return out.reshape(seq, heads * KV_LORA_RANK).to(torch.bfloat16)
diff --git a/tilert/models/glm_5_2_rocm/ops/fp8_kv_cache.py b/tilert/models/glm_5_2_rocm/ops/fp8_kv_cache.py
new file mode 100644
index 0000000..0231a6a
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/fp8_kv_cache.py
@@ -0,0 +1,81 @@
+"""GLM-5.2 fp8 kv latent cache: the 528-byte row format."""
+
+from __future__ import annotations
+
+import torch
+
+KV_LORA_RANK = 512
+SCALE_BLOCK_SIZE = 128
+NUM_SCALE_BLOCKS = KV_LORA_RANK // SCALE_BLOCK_SIZE
+FP8_MAX = 448.0
+AMAX_FLOOR = 0.0001
+SCALE_BYTES = NUM_SCALE_BLOCKS * 4
+KV_ROW_BYTES = KV_LORA_RANK + SCALE_BYTES
+__all__ = [
+ "KV_LORA_RANK",
+ "SCALE_BLOCK_SIZE",
+ "NUM_SCALE_BLOCKS",
+ "KV_ROW_BYTES",
+ "quant_kv_to_fp8_blocked",
+ "dequant_fp8_blocked",
+ "quant_dequant",
+ "pack_kv_528",
+ "unpack_kv_528",
+ "quant_pack_kv",
+ "dequant_kv_528",
+]
+
+
+def quant_kv_to_fp8_blocked(
+ x: torch.Tensor, block_size: int = SCALE_BLOCK_SIZE
+) -> tuple[torch.Tensor, torch.Tensor]:
+ dim = x.shape[-1]
+ assert dim % block_size == 0, f"dim {dim} not divisible by {block_size}"
+ nb = dim // block_size
+ xf = x.float()
+ lead = xf.shape[:-1]
+ xb = xf.reshape(*lead, nb, block_size)
+ amax = xb.abs().amax(dim=-1).clamp(min=AMAX_FLOOR)
+ scale = (amax * (1.0 / FP8_MAX)).contiguous()
+ inv = 1.0 / scale
+ q = (xb * inv.unsqueeze(-1)).clamp(min=-FP8_MAX, max=FP8_MAX)
+ fp8 = q.reshape(*lead, dim).to(torch.float8_e4m3fn)
+ return (fp8.contiguous(), scale)
+
+
+def dequant_fp8_blocked(
+ fp8: torch.Tensor, scale: torch.Tensor, block_size: int = SCALE_BLOCK_SIZE
+) -> torch.Tensor:
+ dim = fp8.shape[-1]
+ nb = dim // block_size
+ lead = fp8.shape[:-1]
+ xb = fp8.float().reshape(*lead, nb, block_size)
+ return (xb * scale.unsqueeze(-1)).reshape(*lead, dim).to(torch.bfloat16)
+
+
+def quant_dequant(x: torch.Tensor, block_size: int = SCALE_BLOCK_SIZE) -> torch.Tensor:
+ fp8, scale = quant_kv_to_fp8_blocked(x, block_size)
+ return dequant_fp8_blocked(fp8, scale, block_size)
+
+
+def pack_kv_528(fp8: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+ lead = fp8.shape[:-1]
+ fp8_u8 = fp8.reshape(-1, KV_LORA_RANK).view(torch.uint8)
+ scale_u8 = scale.reshape(-1, NUM_SCALE_BLOCKS).float().contiguous().view(torch.uint8)
+ return torch.cat([fp8_u8, scale_u8], dim=-1).reshape(*lead, KV_ROW_BYTES).contiguous()
+
+
+def unpack_kv_528(buf: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ lead = buf.shape[:-1]
+ flat = buf.reshape(-1, KV_ROW_BYTES)
+ fp8 = flat[:, :KV_LORA_RANK].contiguous().view(torch.float8_e4m3fn)
+ scale = flat[:, KV_LORA_RANK:].contiguous().view(torch.float32)
+ return (fp8.reshape(*lead, KV_LORA_RANK), scale.reshape(*lead, NUM_SCALE_BLOCKS))
+
+
+def quant_pack_kv(kv: torch.Tensor) -> torch.Tensor:
+ return pack_kv_528(*quant_kv_to_fp8_blocked(kv))
+
+
+def dequant_kv_528(buf: torch.Tensor) -> torch.Tensor:
+ return dequant_fp8_blocked(*unpack_kv_528(buf))
diff --git a/tilert/models/glm_5_2_rocm/ops/fused_moe_allreduce.py b/tilert/models/glm_5_2_rocm/ops/fused_moe_allreduce.py
new file mode 100644
index 0000000..3ce6a00
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/fused_moe_allreduce.py
@@ -0,0 +1,143 @@
+"""GLM-5.2 MoE monokernel (v4) op wrapper: banks, packers, golden, forward."""
+
+import os
+
+import torch
+
+from tilert.models.glm_5_2_rocm.ops.moe_router import HIDDEN, NUM_EXPERTS, MoeRouterGlm5
+from tilert.models.glm_5_2_rocm.ops.upgate_silu import (
+ MOE_SLOTS,
+ NUM_MOE_WEIGHTS,
+ TOP_K,
+ MoeUpGateSiluGlm5,
+)
+
+GRID_BLOCKS = 256
+SCORE_LINE_WORDS = 32
+FUSED_SAMPLES = (1, 2, 4)
+
+
+class MoeFrontBanksGlm5:
+ """Router bank + ONE MoE up/gate bank (expert 0 = shared) + one gamma: the monokernel's front-half weights and their golden."""
+
+ def __init__(self, device: str = "cuda:0", num_weights: int = NUM_MOE_WEIGHTS):
+ self.device = device
+ self.router = MoeRouterGlm5(device=device)
+ self.moe = MoeUpGateSiluGlm5(device=device, num_weights=num_weights)
+ self.packed_m4: torch.Tensor | None = None
+
+ @property
+ def gamma(self) -> torch.Tensor | None:
+ return self.router.gamma
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def router_packed(self) -> torch.Tensor:
+ if self.packed_m4 is None:
+ from tilert.models.glm_5_2_rocm.ops.eh_proj_allreduce import swizzle_256_bf16
+
+ assert self.router.w is not None
+ self.packed_m4 = swizzle_256_bf16(self.router.w).to(self.device)
+ return self.packed_m4
+
+ def golden_forward(
+ self, hidden: torch.Tensor, indices: torch.Tensor
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ raise RuntimeError("golden_forward is not available in release builds")
+
+
+class FusedMoeAllreduceGlm5:
+ """The WHOLE MoE chain in one launch: MoeFrontBanksGlm5's banks plus a down bank (K-split packing) and the all-reduce plumbing."""
+
+ def __init__(self, device: str = "cuda:0", num_weights: int = NUM_MOE_WEIGHTS):
+ from tilert.models.glm_5_2_rocm.ops.moe_down_allreduce import MoeDownBankGlm5
+
+ self.device = device
+ self.front = MoeFrontBanksGlm5(device=device, num_weights=num_weights)
+ self.down = MoeDownBankGlm5(device=device, num_weights=num_weights)
+ self.score_lines = torch.zeros(
+ max(FUSED_SAMPLES), 32, SCORE_LINE_WORDS, dtype=torch.int32, device=device
+ )
+ self.flags = torch.zeros(2 * GRID_BLOCKS, dtype=torch.int32, device=device)
+ self.mid_pairs = torch.zeros(
+ max(FUSED_SAMPLES), MOE_SLOTS, 256, dtype=torch.int32, device=device
+ )
+ self._sen_tag = 0
+ self._down_v4: torch.Tensor | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def golden_forward(
+ self,
+ hidden: torch.Tensor,
+ indices: torch.Tensor,
+ probs: torch.Tensor,
+ residual: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ raise RuntimeError("golden_forward is not available in release builds")
+
+ def pack_down_v4(self) -> torch.Tensor:
+ if self._down_v4 is None:
+ assert self.down.w_fp8 is not None
+ w8 = self.down.w_fp8.view(torch.uint8)
+ e = w8.shape[0]
+ v = w8.view(e, 256, 6, 4, 8, 2, 16).permute(0, 1, 4, 5, 2, 3, 6)
+ self._down_v4 = v.reshape(e, -1).contiguous().view(-1).to(self.device)
+ return self._down_v4
+
+ def tilert_forward(
+ self,
+ hidden: torch.Tensor,
+ bias: torch.Tensor,
+ residual: torch.Tensor | None = None,
+ proto: int = 0,
+ sym: torch.Tensor | None = None,
+ mype: int = 0,
+ npes: int = 1,
+ flag: int = 1,
+ out: torch.Tensor | None = None,
+ sen_tag: int | None = None,
+ ):
+ assert self.front.moe.packed is not None
+ os.environ["TILERT_GLM5_AR_PROTO"] = str(proto)
+ s_n = hidden.shape[0]
+ assert s_n in FUSED_SAMPLES
+ if sen_tag is None:
+ self._sen_tag += 1
+ sen_tag = self._sen_tag
+ dev = hidden.device
+ norm = torch.empty(s_n, HIDDEN, dtype=torch.bfloat16, device=dev)
+ scores = torch.empty(s_n, NUM_EXPERTS, dtype=torch.float32, device=dev)
+ probs = torch.zeros(s_n, TOP_K, dtype=torch.float32, device=dev)
+ indices = torch.zeros(s_n, TOP_K, dtype=torch.int32, device=dev)
+ mid = torch.zeros(s_n, MOE_SLOTS, 256, dtype=torch.bfloat16, device=dev)
+ if out is None:
+ out = torch.empty(s_n, HIDDEN, dtype=torch.bfloat16, device=dev)
+ torch.ops.tilert.glm5_fused_moe_allreduce_v4_op(
+ hidden,
+ self.front.router.gamma,
+ self.front.router_packed(),
+ self.front.moe.packed,
+ self.front.moe.scales,
+ bias,
+ self.pack_down_v4(),
+ self.down.scales,
+ residual,
+ sym,
+ mype,
+ npes,
+ flag,
+ norm,
+ scores,
+ self.score_lines[:s_n],
+ self.flags,
+ probs,
+ indices,
+ mid,
+ out,
+ self.mid_pairs[:s_n],
+ sen_tag,
+ )
+ return (norm, scores, mid, probs, indices, out)
diff --git a/tilert/models/glm_5_2_rocm/ops/fused_moe_w8a8.py b/tilert/models/glm_5_2_rocm/ops/fused_moe_w8a8.py
new file mode 100644
index 0000000..f2b2f1f
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/fused_moe_w8a8.py
@@ -0,0 +1,210 @@
+"""GLM-5.2 W8A8 MoE monokernel (v4 skeleton): the v4 op's banks + the fp8 packings + the quantized-math goldens."""
+
+import torch
+
+from tilert.models.glm_5_2_rocm.ops.fused_moe_allreduce import (
+ FUSED_SAMPLES,
+ GRID_BLOCKS,
+ FusedMoeAllreduceGlm5,
+)
+from tilert.models.glm_5_2_rocm.ops.moe_router import HIDDEN, NUM_EXPERTS
+from tilert.models.glm_5_2_rocm.ops.upgate_silu import (
+ MOE_SLOTS,
+ NUM_MOE_WEIGHTS,
+ SCALE_BLK,
+ TOP_K,
+ _silu,
+)
+from tilert.models.glm_5_2_rocm.ops.upgate_silu_w8a8 import (
+ quant_act_row,
+ quant_mid_rows,
+ swizzle_pair_interleaved_k128,
+)
+
+
+def swizzle_down_k128(w_fp8: torch.Tensor) -> torch.Tensor:
+ rows, k = w_fp8.shape
+ assert rows % 24 == 0 and k == 256
+ w8 = w_fp8.view(torch.uint8)
+ out = torch.zeros(rows // 24 * 6144, dtype=torch.uint8)
+ lane = torch.arange(64)
+ h = torch.arange(2)
+ i = torch.arange(16)
+ L, H, II = torch.meshgrid(lane, h, i, indexing="ij")
+ for blk in range(rows // 24):
+ r0 = blk * 24
+ base = blk * 6144
+ for c in range(2):
+ ks = c * 128 + L // 16 * 32 + H * 16 + II
+ main = w8[r0 + L % 16, ks]
+ idx = base + c * 2048 + H * 1024 + L * 16 + II
+ out[idx.reshape(-1)] = main.reshape(-1)
+ m8 = torch.arange(8)
+ g = torch.arange(4)
+ G, M, H2, I2 = torch.meshgrid(g, m8, h, i, indexing="ij")
+ for c in range(2):
+ ks = c * 128 + G * 32 + H2 * 16 + I2
+ tail = w8[r0 + 16 + M, ks]
+ idx = base + 4096 + c * 1024 + H2 * 512 + (G * 8 + M) * 16 + I2
+ out[idx.reshape(-1)] = tail.reshape(-1)
+ return out.contiguous()
+
+
+class FusedMoeAllreduceW8A8Glm5(FusedMoeAllreduceGlm5):
+ """The v4 monokernel's banks + the fp8 packings + the W8A8 op."""
+
+ def __init__(self, device: str = "cuda:0", num_weights: int = NUM_MOE_WEIGHTS):
+ super().__init__(device=device, num_weights=num_weights)
+ self.packed_k128: torch.Tensor | None = None
+ self.down_k128: torch.Tensor | None = None
+
+ def pack_k128(self) -> torch.Tensor:
+ if self.packed_k128 is None:
+ moe = self.front.moe
+ assert moe.w_fp8 is not None
+ ps = [
+ swizzle_pair_interleaved_k128(moe.w_fp8[e], moe.inter)
+ for e in range(moe.num_weights)
+ ]
+ self.packed_k128 = torch.cat(ps).contiguous().to(self.device)
+ return self.packed_k128
+
+ def pack_down_k128(self) -> torch.Tensor:
+ if self.down_k128 is None:
+ assert self.down.w_fp8 is not None
+ ps = [swizzle_down_k128(self.down.w_fp8[e]) for e in range(self.down.w_fp8.shape[0])]
+ self.down_k128 = torch.cat(ps).contiguous().to(self.device)
+ return self.down_k128
+
+ def golden_down_w8a8(
+ self,
+ mid_bf16: torch.Tensor,
+ probs: torch.Tensor,
+ indices: torch.Tensor,
+ residual: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ assert self.down.w_fp8 is not None and self.down.scales is not None
+ q, ds = quant_mid_rows(mid_bf16.cpu())
+ qf = q.float()
+ wq = self.down.w_fp8.float().cpu()
+ scales = self.down.scales.cpu()
+ idx = indices.cpu()
+ pr = probs.cpu()
+ s_n = mid_bf16.shape[0]
+ out = torch.zeros(s_n, HIDDEN, dtype=torch.float32)
+ row_blk = torch.arange(HIDDEN) // 128
+ for s in range(s_n):
+ for slot in range(MOE_SLOTS):
+ e = 0 if slot == 0 else 1 + int(idx[s, slot - 1])
+ w = 1.0 if slot == 0 else float(pr[s, slot - 1])
+ acc = torch.zeros(HIDDEN, dtype=torch.float32)
+ for kb in range(2):
+ ks = slice(kb * 128, (kb + 1) * 128)
+ part = wq[e][:, ks] @ qf[s, slot, ks]
+ a_sc = float(ds[s, slot, kb])
+ acc += part * (scales[e][row_blk, kb] * w * a_sc)
+ out[s] += acc
+ out = out.to(torch.bfloat16).float()
+ if residual is not None:
+ out += residual.float().cpu()
+ return out.to(torch.bfloat16).to(mid_bf16.device)
+
+ def golden_forward_w8a8(
+ self,
+ hidden: torch.Tensor,
+ indices: torch.Tensor,
+ probs: torch.Tensor,
+ residual: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ moe = self.front.moe
+ assert moe.w_fp8 is not None and moe.scales is not None
+ norm_ref, scores_ref = self.front.router.golden_forward(hidden)
+ act8, a_scales = quant_act_row(norm_ref.cpu())
+ a = act8.float()
+ wq = moe.w_fp8.float().cpu()
+ scales = moe.scales.cpu()
+ idx = indices.cpu()
+ s_n = a.shape[0]
+ inter = moe.inter
+ mid = torch.zeros(s_n, MOE_SLOTS, inter, dtype=torch.float32)
+ n_kb = HIDDEN // SCALE_BLK
+ for s in range(s_n):
+ for slot in range(MOE_SLOTS):
+ e = 0 if slot == 0 else 1 + int(idx[s, slot - 1])
+ acc = torch.zeros(2 * inter, dtype=torch.float32)
+ for kb in range(n_kb):
+ ks = slice(kb * SCALE_BLK, (kb + 1) * SCALE_BLK)
+ part = wq[e][:, ks] @ a[s, ks]
+ row_blk = torch.arange(2 * inter) // SCALE_BLK
+ acc += part * (scales[e][row_blk, kb] * float(a_scales[s, kb]))
+ gate, up = (acc[:inter], acc[inter:])
+ mid[s, slot] = _silu(gate) * up
+ mid_bf16 = mid.to(torch.bfloat16).to(hidden.device)
+ out_ref = self.golden_down_w8a8(mid_bf16, probs, indices, residual)
+ return (norm_ref, scores_ref, mid_bf16, out_ref)
+
+ def tilert_forward_w8a8_v4(
+ self,
+ hidden: torch.Tensor,
+ bias: torch.Tensor,
+ residual: torch.Tensor | None = None,
+ proto: int = 0,
+ sym: torch.Tensor | None = None,
+ mype: int = 0,
+ npes: int = 1,
+ flag: int = 1,
+ out: torch.Tensor | None = None,
+ sen_tag: int | None = None,
+ ):
+ import os
+
+ os.environ["TILERT_GLM5_AR_PROTO"] = str(proto)
+ s_n = hidden.shape[0]
+ assert s_n in FUSED_SAMPLES
+ if sen_tag is None:
+ self._sen_tag += 1
+ sen_tag = self._sen_tag
+ dev = hidden.device
+ norm = torch.empty(s_n, HIDDEN, dtype=torch.bfloat16, device=dev)
+ scores = torch.empty(s_n, NUM_EXPERTS, dtype=torch.float32, device=dev)
+ probs = torch.zeros(s_n, TOP_K, dtype=torch.float32, device=dev)
+ indices = torch.zeros(s_n, TOP_K, dtype=torch.int32, device=dev)
+ mid = torch.zeros(s_n, MOE_SLOTS, 256, dtype=torch.bfloat16, device=dev)
+ if out is None:
+ out = torch.empty(s_n, HIDDEN, dtype=torch.bfloat16, device=dev)
+ torch.ops.tilert.glm5_fused_moe_allreduce_w8a8_v4_op(
+ hidden,
+ self.front.router.gamma,
+ self.front.router_packed(),
+ self.pack_k128(),
+ self.front.moe.scales,
+ bias,
+ self.pack_down_k128(),
+ self.down.scales,
+ residual,
+ sym,
+ mype,
+ npes,
+ flag,
+ norm,
+ scores,
+ self.score_lines[:s_n],
+ self.flags,
+ probs,
+ indices,
+ mid,
+ out,
+ self.mid_pairs[:s_n],
+ sen_tag,
+ )
+ return (norm, scores, mid, probs, indices, out)
+
+
+__all__ = [
+ "FusedMoeAllreduceW8A8Glm5",
+ "GRID_BLOCKS",
+ "quant_act_row",
+ "quant_mid_rows",
+ "swizzle_down_k128",
+ "swizzle_pair_interleaved_k128",
+]
diff --git a/tilert/models/glm_5_2_rocm/ops/index_collective.py b/tilert/models/glm_5_2_rocm/ops/index_collective.py
new file mode 100644
index 0000000..82c4f9c
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/index_collective.py
@@ -0,0 +1,61 @@
+"""GLM-5.2 indexer collective op wrappers: the GPU0 index broadcast and its all-reduce leg."""
+
+import os
+
+import torch
+
+HIDDEN = 6144
+NUM_PES = 8
+PROTO_0, PROTO_1 = (0, 1)
+BLOCKS = 256
+ROWS_PER_BLOCK = 24
+MAX_SEQ = 4
+
+
+def xfer_buf_bytes(samples: int, topk: int) -> int:
+ return int(torch.ops.tilert.glm5_broadcast_xfer_buf_bytes(samples, topk))
+
+
+def xfer_buffer(samples: int, topk: int, device) -> torch.Tensor:
+ return torch.zeros(xfer_buf_bytes(samples, topk), dtype=torch.uint8, device=device)
+
+
+def peer_table(buffers: list[torch.Tensor], device) -> torch.Tensor:
+ return torch.tensor([b.data_ptr() for b in buffers], dtype=torch.int64, device=device)
+
+
+def broadcast(indices: torch.Tensor, sym: torch.Tensor, mype: int, npes: int, flag: int) -> None:
+ torch.ops.tilert.glm5_broadcast_selected_token_ids_op(indices, sym, mype, npes, flag)
+
+
+def broadcast_golden(indices: torch.Tensor, flag: int) -> torch.Tensor:
+ pairs = indices.reshape(-1).to(torch.int32).cpu().reshape(-1, 2)
+ out = torch.empty(pairs.shape[0], 4, dtype=torch.int32)
+ out[:, 0] = pairs[:, 0]
+ out[:, 1] = flag
+ out[:, 2] = pairs[:, 1]
+ out[:, 3] = flag
+ return out.reshape(-1)
+
+
+def padded_allreduce_add(
+ out: torch.Tensor,
+ residual: torch.Tensor | None = None,
+ proto: int = PROTO_0,
+ sym: torch.Tensor | None = None,
+ mype: int = 0,
+ npes: int = 1,
+ flag: int = 1,
+) -> torch.Tensor:
+ os.environ["TILERT_GLM5_AR_PROTO"] = str(proto)
+ torch.ops.tilert.glm5_padded_allreduce_add_op(residual, sym, mype, npes, flag, out)
+ return out
+
+
+def reduce_golden(partials: list[torch.Tensor], residual: torch.Tensor | None) -> torch.Tensor:
+ acc = torch.zeros(partials[0].shape, dtype=torch.float32)
+ for p in partials:
+ acc += p.float().cpu()
+ if residual is not None:
+ acc += residual.float().cpu()
+ return acc.to(torch.bfloat16)
diff --git a/tilert/models/glm_5_2_rocm/ops/llm_preprocess.py b/tilert/models/glm_5_2_rocm/ops/llm_preprocess.py
new file mode 100644
index 0000000..2e56430
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/llm_preprocess.py
@@ -0,0 +1,45 @@
+"""GLM-5.2 LlmPreprocess op wrapper: golden + forward."""
+
+import torch
+
+HIDDEN = 6144
+ROPE_DIM = 64
+SUPPORTED_SAMPLES = (1, 2, 4, 8)
+
+
+def make_freqs_cis(max_pos: int, theta: float = 8000000.0, device: str = "cuda:0") -> torch.Tensor:
+ inv = 1.0 / theta ** (torch.arange(0, ROPE_DIM, 2, dtype=torch.float32) / ROPE_DIM)
+ ang = torch.outer(torch.arange(max_pos), inv)
+ cis = torch.polar(torch.ones_like(ang), ang)
+ return torch.view_as_real(cis).reshape(max_pos, ROPE_DIM).contiguous().to(device)
+
+
+class LlmPreprocessGlm5:
+ """Embedding + rope-freq row gather at cur_pos."""
+
+ OP_NAME = "glm5_llm_preprocess_op"
+
+ def __init__(self, device: str = "cuda:0"):
+ self.device = device
+ self.embed: torch.Tensor | None = None
+ self.freqs_cis: torch.Tensor | None = None
+
+ def init_random_weights(self, vocab: int, max_pos: int, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def golden_forward(
+ self, token_id: torch.Tensor, cur_pos: torch.Tensor, seq_len: int
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ raise RuntimeError("golden_forward is not available in release builds")
+
+ def tilert_forward(
+ self, token_id: torch.Tensor, cur_pos: torch.Tensor, seq_len: int
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ assert self.embed is not None and self.freqs_cis is not None
+ samples = token_id.numel()
+ x = torch.empty(samples, HIDDEN, dtype=torch.bfloat16, device=self.device)
+ freqs = torch.empty(samples, ROPE_DIM, dtype=torch.float32, device=self.device)
+ torch.ops.tilert.glm5_llm_preprocess_op(
+ token_id, self.embed, self.freqs_cis, cur_pos, x, freqs, seq_len
+ )
+ return (x, freqs)
diff --git a/tilert/models/glm_5_2_rocm/ops/moe_down_allreduce.py b/tilert/models/glm_5_2_rocm/ops/moe_down_allreduce.py
new file mode 100644
index 0000000..2207bed
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/moe_down_allreduce.py
@@ -0,0 +1,108 @@
+"""GLM-5.2 MoE down-projection bank, golden and symmetric-buffer helpers."""
+
+import torch
+
+HIDDEN = 6144
+EXPERT_DIM = 256
+TOP_K = 8
+SLOTS = TOP_K + 1
+SCALE_BLK = 128
+NUM_PES = 8
+FP8_MAX = 448.0
+PROTO_0, PROTO_1 = (0, 1)
+
+
+def quantize_fp8_block(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ rows, k = w.shape
+ assert rows % SCALE_BLK == 0 and k % SCALE_BLK == 0
+ sr, sc = (rows // SCALE_BLK, k // SCALE_BLK)
+ blocks = w.float().view(sr, SCALE_BLK, sc, SCALE_BLK)
+ amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12)
+ scales = amax / FP8_MAX
+ q = (blocks / scales).to(torch.float8_e4m3fn)
+ return (q.reshape(rows, k).contiguous(), scales.view(sr, sc).contiguous())
+
+
+def swizzle_m4(w_fp8: torch.Tensor) -> torch.Tensor:
+ rows, k = w_fp8.shape
+ assert rows % 8 == 0 and k % 128 == 0
+ w8 = w_fp8.view(torch.uint8)
+ rg = torch.arange(rows // 8)
+ kc = torch.arange(k // 128)
+ lane = torch.arange(64)
+ s = torch.arange(4)
+ i = torch.arange(4)
+ RG, KC, L, S, II = torch.meshgrid(rg, kc, lane, s, i, indexing="ij")
+ rows_ix = RG * 8 + (L >> 5) * 4 + (L & 3)
+ ks = KC * 128 + S * 32 + (L >> 2 & 7) * 4 + II
+ return w8[rows_ix, ks].reshape(-1).contiguous()
+
+
+def sym_bytes(samples: int) -> int:
+ return int(torch.ops.tilert.glm5_moe_sym_bytes(samples))
+
+
+def sym_buffer(samples: int, device) -> torch.Tensor:
+ return torch.zeros(sym_bytes(samples), dtype=torch.uint8, device=device)
+
+
+def sym_table(buffers: list[torch.Tensor], device) -> torch.Tensor:
+ assert len(buffers) == NUM_PES
+ return torch.tensor([b.data_ptr() for b in buffers], dtype=torch.int64, device=device)
+
+
+class MoeDownBankGlm5:
+ """One rank's expert down-projection bank (fp8 + block scales) and the down golden. Packing is the monokernel wrapper's job (pack_down_v4)."""
+
+ def __init__(self, device: str = "cuda:0", num_weights: int = 257) -> None:
+ self.device = device
+ self.num_weights = num_weights
+ self.w_fp8: torch.Tensor | None = None
+ self.scales: torch.Tensor | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def init_reference_weights(self, w: torch.Tensor) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def partial_golden(
+ self, hidden_mid: torch.Tensor, probs: torch.Tensor, indices: torch.Tensor
+ ) -> torch.Tensor:
+ assert self.w_fp8 is not None and self.scales is not None
+ mid = hidden_mid.float().cpu()
+ wq = self.w_fp8.float().cpu()
+ scales = self.scales.cpu()
+ pr = probs.float().cpu()
+ idx = indices.cpu()
+ s_n = mid.shape[0]
+ acc = torch.zeros(s_n, HIDDEN, dtype=torch.float32)
+ row_scale = scales.repeat_interleave(SCALE_BLK, dim=1)
+ for s in range(s_n):
+ for slot in range(SLOTS):
+ e = 0 if slot == 0 else 1 + int(idx[s, slot - 1])
+ w = 1.0 if slot == 0 else float(pr[s, slot - 1])
+ a = mid[s, slot]
+ for kb in range(EXPERT_DIM // SCALE_BLK):
+ k0, k1 = (kb * SCALE_BLK, (kb + 1) * SCALE_BLK)
+ part = wq[e, :, k0:k1] @ a[k0:k1]
+ acc[s] += part * (row_scale[e, :, kb] * w)
+ return acc.to(torch.bfloat16)
+
+ @staticmethod
+ def reduce_golden(partials: list[torch.Tensor], residual: torch.Tensor | None) -> torch.Tensor:
+ acc = torch.zeros(partials[0].shape, dtype=torch.float32)
+ for p in partials:
+ acc += p.float().cpu()
+ if residual is not None:
+ acc += residual.float().cpu()
+ return acc.to(torch.bfloat16)
+
+ def golden_forward(
+ self,
+ hidden_mid: torch.Tensor,
+ probs: torch.Tensor,
+ indices: torch.Tensor,
+ residual: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ raise RuntimeError("golden_forward is not available in release builds")
diff --git a/tilert/models/glm_5_2_rocm/ops/moe_router.py b/tilert/models/glm_5_2_rocm/ops/moe_router.py
new file mode 100644
index 0000000..29824f1
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/moe_router.py
@@ -0,0 +1,68 @@
+"""GLM-5.2 router weights container + routing golden."""
+
+import torch
+
+HIDDEN = 6144
+NUM_EXPERTS = 256
+TOP_K = 8
+ROUTE_SCALE = 2.5
+EPS = 1e-05
+_CHUNK_K = 32
+
+
+def swizzle_router_bf16(w: torch.Tensor) -> torch.Tensor:
+ rows, k = w.shape
+ assert rows % 16 == 0 and k % _CHUNK_K == 0
+ w16 = w.to(torch.bfloat16).view(torch.uint16)
+ tile = torch.arange(rows // 16)
+ kc = torch.arange(k // _CHUNK_K)
+ lane = torch.arange(64)
+ i = torch.arange(8)
+ T, KC, L, II = torch.meshgrid(tile, kc, lane, i, indexing="ij")
+ packed = w16[T * 16 + L % 16, KC * _CHUNK_K + L // 16 * 8 + II]
+ return packed.reshape(-1).contiguous().view(torch.uint8)
+
+
+def select_topk_golden(
+ partials: torch.Tensor, bias: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+ logits = partials.float().cpu().sum(dim=0)
+ b = bias.float().cpu()
+ scores = torch.sigmoid(logits)
+ ranked = scores + b[None, :]
+ s_n = scores.shape[0]
+ probs = torch.zeros(s_n, TOP_K, dtype=torch.float32)
+ idx = torch.zeros(s_n, TOP_K, dtype=torch.int32)
+ for s in range(s_n):
+ r = ranked[s].clone()
+ vals = []
+ for k in range(TOP_K):
+ e = int(torch.argmax(r))
+ idx[s, k] = e
+ vals.append(float(scores[s, e]))
+ r[e] = -float("inf")
+ total = 0.0
+ for v in vals:
+ total += v
+ for k, v in enumerate(vals):
+ probs[s, k] = v * (ROUTE_SCALE / total)
+ return (probs, idx)
+
+
+class MoeRouterGlm5:
+ """One rank's logical router tensor + gamma (weights container/golden)."""
+
+ def __init__(self, device: str = "cuda:0") -> None:
+ self.device = device
+ self.w: torch.Tensor | None = None
+ self.packed: torch.Tensor | None = None
+ self.gamma: torch.Tensor | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def init_reference_weights(self, w: torch.Tensor, gamma: torch.Tensor) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def golden_forward(self, hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ raise RuntimeError("golden_forward is not available in release builds")
diff --git a/tilert/models/glm_5_2_rocm/ops/mtp.py b/tilert/models/glm_5_2_rocm/ops/mtp.py
new file mode 100644
index 0000000..df4f406
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/mtp.py
@@ -0,0 +1,129 @@
+"""GLM-5.2 MTP step-tail helpers: the CPU goldens for the fused tail."""
+
+import torch
+
+HIDDEN = 6144
+SIMULATE_TOKEN_ID = 100
+
+
+def xfer_buf_bytes(producer_samples: int, topk: int) -> int:
+ return int(torch.ops.tilert.glm5_broadcast_xfer_buf_bytes(producer_samples, topk))
+
+
+def verify_golden(
+ draft_tokens: torch.Tensor,
+ predicted_tokens: torch.Tensor,
+ mtp0_tokens: torch.Tensor,
+ mtp0_hidden: torch.Tensor,
+ cur_pos: torch.Tensor,
+ idx_source: torch.Tensor | None = None,
+ sim_accept: torch.Tensor | None = None,
+) -> dict[str, torch.Tensor]:
+ d = draft_tokens.cpu()
+ p = predicted_tokens.cpu().clone()
+ m = mtp0_tokens.cpu().clone()
+ h = mtp0_hidden.cpu()
+ batch, mtp_seq_len = d.shape
+ num_verified = mtp_seq_len - 1
+ out = {
+ "num_accepted": torch.zeros(batch, dtype=torch.int32),
+ "cur_pos": cur_pos.cpu().clone(),
+ "last_token": torch.zeros(batch, dtype=torch.int32),
+ "last_hidden": torch.zeros(batch, HIDDEN, dtype=h.dtype),
+ "next_draft_tokens": torch.zeros(batch, mtp_seq_len, dtype=torch.int32),
+ }
+ if idx_source is not None:
+ topk = idx_source.shape[1]
+ out["idx_selects"] = torch.zeros(batch, topk, dtype=torch.int32)
+ for b in range(batch):
+ if sim_accept is not None:
+ acc = min(max(int(sim_accept.cpu()[b]), 0), mtp_seq_len - 1)
+ p[b, :] = SIMULATE_TOKEN_ID
+ m[b, :] = SIMULATE_TOKEN_ID
+ else:
+ acc = 0
+ for i in range(num_verified):
+ if d[b, i + 1] == p[b, i]:
+ acc += 1
+ else:
+ break
+ out["num_accepted"][b] = acc + 1
+ out["cur_pos"][b] += acc + 1
+ out["last_token"][b] = m[b, acc]
+ out["next_draft_tokens"][b, 0] = p[b, acc]
+ out["next_draft_tokens"][b, 1] = m[b, acc]
+ out["last_hidden"][b] = h[b * mtp_seq_len + acc]
+ if idx_source is not None:
+ out["idx_selects"][b] = idx_source.cpu()[b * mtp_seq_len + acc]
+ out["predicted_tokens"] = p
+ out["mtp0_tokens"] = m
+ return out
+
+
+def assemble_accepted_golden(
+ num_accepted: torch.Tensor, predicted: torch.Tensor, ar_acc: torch.Tensor, ar_num: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+ acc = ar_acc.cpu().clone()
+ num = ar_num.cpu().clone()
+ na = num_accepted.cpu()
+ pred = predicted.cpu()
+ acc_stride = acc.shape[1]
+ num_stride = num.shape[1]
+ for b in range(pred.shape[0]):
+ n = int(na[b])
+ base = int(acc[b, 0]) + 1
+ for i in range(n):
+ if base + i < acc_stride:
+ acc[b, base + i] = pred[b, i]
+ acc[b, 0] = min(int(acc[b, 0]) + n, acc_stride - 1)
+ k = int(num[b, 0])
+ if 1 + k < num_stride:
+ num[b, 1 + k] = n
+ num[b, 0] = min(k + 1, num_stride - 1)
+ return (acc, num)
+
+
+def step_tail_golden(
+ draft_tokens: torch.Tensor,
+ predicted_tokens: torch.Tensor,
+ mtp0_tokens: torch.Tensor,
+ mtp0_hidden: torch.Tensor,
+ cur_pos: torch.Tensor,
+ next_draft_tokens: torch.Tensor,
+ idx_selects: torch.Tensor,
+ ar_acc: torch.Tensor,
+ ar_num: torch.Tensor,
+ idx_source: torch.Tensor | None = None,
+ local_gather: bool = False,
+ sim_accept: torch.Tensor | None = None,
+) -> dict[str, torch.Tensor]:
+ batch, mtp_seq_len = draft_tokens.shape
+ src = None
+ if idx_source is not None:
+ src = idx_source
+ elif local_gather:
+ src = idx_selects[: batch * mtp_seq_len]
+ v = verify_golden(
+ draft_tokens, predicted_tokens, mtp0_tokens, mtp0_hidden, cur_pos, src, sim_accept
+ )
+ acc, num = assemble_accepted_golden(v["num_accepted"], v["predicted_tokens"], ar_acc, ar_num)
+ positions = v["cur_pos"].long()[:, None] + torch.arange(mtp_seq_len)[None, :]
+ rotated = next_draft_tokens.cpu().clone()
+ rotated[:, :2] = v["next_draft_tokens"][:, :2]
+ idx = idx_selects.cpu().clone()
+ if src is not None:
+ idx[:batch] = v["idx_selects"]
+ return {
+ "draft_tokens": rotated.clone(),
+ "predicted_tokens": v["predicted_tokens"],
+ "mtp0_tokens": v["mtp0_tokens"],
+ "num_accepted": v["num_accepted"],
+ "cur_pos": v["cur_pos"],
+ "last_token": v["last_token"],
+ "last_hidden": v["last_hidden"],
+ "next_draft_tokens": rotated,
+ "idx_selects": idx,
+ "ar_acc": acc,
+ "ar_num": num,
+ "positions": positions,
+ }
diff --git a/tilert/models/glm_5_2_rocm/ops/proj_wkvb.py b/tilert/models/glm_5_2_rocm/ops/proj_wkvb.py
new file mode 100644
index 0000000..367d7eb
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/proj_wkvb.py
@@ -0,0 +1,85 @@
+"""GLM-5.2 projection op wrappers: quantizer, weight packer, golden, forward."""
+
+import torch
+
+KV_LORA_RANK = 512
+QK_NOPE_DIM = 192
+V_HEAD_DIM = 256
+M4_SCALE_BLK_K = 64
+M6_SCALE_BLK_K = 128
+SCALE_BLK_M = 64
+FP8_MAX = 448.0
+SUPPORTED_SAMPLES = (1, 2, 4, 8)
+SUPPORTED_HEADS = (8, 10)
+NUM_HEADS = 10
+_CHUNK = 64
+
+
+def quantize_fp8_block(w: torch.Tensor, scale_blk_k: int) -> tuple[torch.Tensor, torch.Tensor]:
+ rows, k = w.shape
+ assert rows % SCALE_BLK_M == 0 and k % scale_blk_k == 0
+ scale_rows, scale_cols = (rows // SCALE_BLK_M, k // scale_blk_k)
+ blocks = w.float().view(scale_rows, SCALE_BLK_M, scale_cols, scale_blk_k)
+ amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12)
+ scales = amax / FP8_MAX
+ q = (blocks / scales).to(torch.float8_e4m3fn)
+ return (q.reshape(rows, k).contiguous(), scales.view(scale_rows, scale_cols).contiguous())
+
+
+def swizzle_weights_contig(w_fp8: torch.Tensor) -> torch.Tensor:
+ rows, k = w_fp8.shape
+ assert rows % 16 == 0 and k % _CHUNK == 0
+ w8 = w_fp8.view(torch.uint8)
+ rg = torch.arange(rows // 16)
+ kc = torch.arange(k // _CHUNK)
+ lane = torch.arange(64)
+ sp = torch.arange(2)
+ i = torch.arange(8)
+ RG, KC, L, SP, II = torch.meshgrid(rg, kc, lane, sp, i, indexing="ij")
+ rows_ix = RG * 16 + L % 16
+ ks = KC * _CHUNK + SP * 32 + L // 16 * 8 + II
+ return w8[rows_ix, ks].reshape(-1).contiguous()
+
+
+class _ProjWkvbGlm5:
+ """Shared op class: golden + tilert forwards share one set of weights."""
+
+ K: int
+ OUT_DIM: int
+ SCALE_BLK_K: int
+ OP_NAME: str
+
+ def __init__(self, device: str = "cuda:0", num_heads: int = NUM_HEADS):
+ assert num_heads in SUPPORTED_HEADS
+ self.device = device
+ self.num_heads = num_heads
+ self.act_dim = num_heads * self.K
+ self.rows = num_heads * self.OUT_DIM
+ self.w_fp8: torch.Tensor | None = None
+ self.scales: torch.Tensor | None = None
+ self.packed: torch.Tensor | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def init_reference_weights(self, w: torch.Tensor) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def golden_forward(self, act: torch.Tensor) -> torch.Tensor:
+ raise RuntimeError("golden_forward is not available in release builds")
+
+
+class ProjQWkvbGlm5(_ProjWkvbGlm5):
+ """Project q_nope_down [S, H*192] to q_nope [S, H*512]."""
+
+ K = QK_NOPE_DIM
+ OUT_DIM = KV_LORA_RANK
+ SCALE_BLK_K = M4_SCALE_BLK_K
+
+
+class ProjOWkvbGlm5(_ProjWkvbGlm5):
+ """Project o [S, H*512] to proj_o [S, H*256] (the flat [S, 2560] the tail reads)."""
+
+ K = KV_LORA_RANK
+ OUT_DIM = V_HEAD_DIM
+ SCALE_BLK_K = M6_SCALE_BLK_K
diff --git a/tilert/models/glm_5_2_rocm/ops/pure_mla_allreduce.py b/tilert/models/glm_5_2_rocm/ops/pure_mla_allreduce.py
new file mode 100644
index 0000000..cc9f72a
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/pure_mla_allreduce.py
@@ -0,0 +1,245 @@
+"""GLM-5.2 attention block in ONE launch (pure_mla_allreduce)."""
+
+import torch
+
+from tilert.models.glm_5_2_rocm.ops.flash_sparse_mla import GLM5_SOFTMAX_SCALE, KV_LORA_RANK, PE_DIM
+from tilert.models.glm_5_2_rocm.ops.proj_wkvb import ProjQWkvbGlm5
+from tilert.models.glm_5_2_rocm.ops.pure_mla_m567 import PureMlaM567Glm5
+from tilert.models.glm_5_2_rocm.ops.qkv_rope import QkvRopeGlm5
+from tilert.models.glm_5_2_rocm.ops.rmsnorm_kv import RmsnormKvGlm5
+from tilert.models.glm_5_2_rocm.ops.rmsnorm_projq_wqb import RmsnormProjQWqbGlm5
+from tilert.models.glm_5_2_rocm.ops.rmsnorm_projx_wqkva import KV_DIM
+from tilert.models.glm_5_2_rocm.ops.rmsnorm_projx_wqkva import PE_DIM as M0_PE_DIM
+from tilert.models.glm_5_2_rocm.ops.rmsnorm_projx_wqkva import Q_DIM, RmsnormProjXWqkvaGlm5
+from tilert.models.glm_5_2_rocm.ops.unprojo_allreduce import HIDDEN
+
+Q_LINE_WORDS = 16
+Q_LINE_MAX_SEQ = 4
+MAX_SEQ = 4
+
+
+def q_tiles(num_heads: int) -> int:
+ return num_heads * (KV_LORA_RANK + PE_DIM) // 16
+
+
+def q_words(num_heads: int) -> int:
+ return q_tiles(num_heads) * Q_LINE_MAX_SEQ * Q_LINE_WORDS
+
+
+class Exchange:
+ """The attention block's scratch -- allocate ONCE and share across layers."""
+
+ def __init__(self, samples: int, heads: int = 10, device: torch.device | str = "cuda"):
+ self.samples = samples
+ self.heads = heads
+ self.q_pairs = torch.zeros(samples, Q_DIM // 2, 2, dtype=torch.int32, device=device)
+ self.kv_pairs = torch.zeros(samples, KV_DIM // 2, 2, dtype=torch.int32, device=device)
+ self.pe_pairs = torch.zeros(samples, M0_PE_DIM // 2, 2, dtype=torch.int32, device=device)
+ self.m1_pairs = torch.zeros(samples, heads * 256 // 2, 2, dtype=torch.int32, device=device)
+
+
+class Legs:
+ """The three intra-launch legs -- one allocation shared by every layer."""
+
+ def __init__(self, num_heads: int, device: str = "cuda:0") -> None:
+ i32 = {"dtype": torch.int32, "device": device}
+ self.qlines = torch.zeros(q_words(num_heads), **i32)
+ self.kvnew_pairs = torch.zeros(MAX_SEQ * KV_LORA_RANK * 2, **i32)
+ self.penew_pairs = torch.zeros(MAX_SEQ * PE_DIM * 2, **i32)
+
+
+class PureMlaAllReduceGlm5:
+ """One rank's whole attention block: the projection weights, the output shards, the exchange, the tail's scratch and the three legs."""
+
+ def __init__(
+ self,
+ device: str = "cuda:0",
+ num_heads: int = 10,
+ topk: int = 2048,
+ scale: float = GLM5_SOFTMAX_SCALE,
+ ) -> None:
+ assert num_heads in (8, 10), "num_heads must be 8 or 10"
+ self.device = device
+ self.num_heads = num_heads
+ self.topk = topk
+ self.scale = scale
+ self.m0 = RmsnormProjXWqkvaGlm5(device=device)
+ self.m1 = RmsnormProjQWqbGlm5(device=device, num_heads=num_heads)
+ self.m3 = RmsnormKvGlm5(device=device)
+ self.m4 = ProjQWkvbGlm5(device=device, num_heads=num_heads)
+ self.tail = PureMlaM567Glm5(device=device, num_heads=num_heads, topk=topk, scale=scale)
+ self._ex: dict[int, Exchange] = {}
+ self.legs = Legs(num_heads, device)
+ self._tag = 0
+
+ @property
+ def m6(self):
+ return self.tail.m6
+
+ @property
+ def m7(self):
+ return self.tail.m7
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def exchange(self, samples: int) -> Exchange:
+ if samples not in self._ex:
+ self._ex[samples] = Exchange(samples, heads=self.num_heads, device=self.device)
+ return self._ex[samples]
+
+ def next_tag(self) -> int:
+ self._tag += 1
+ return self._tag
+
+ def alloc_partials(self, samples: int) -> tuple:
+ return self.tail.alloc_partials(samples)
+
+ def forward(
+ self,
+ hidden_in: torch.Tensor,
+ cur_pos: torch.Tensor,
+ rope_freqs: torch.Tensor,
+ pe_cache: torch.Tensor,
+ kv_cache: torch.Tensor,
+ indices: torch.Tensor | None,
+ partials: tuple,
+ residual: torch.Tensor | None = None,
+ sym: torch.Tensor | None = None,
+ mype: int = 0,
+ npes: int = 1,
+ tag: int | None = None,
+ out: torch.Tensor | None = None,
+ xfer_buf: torch.Tensor | None = None,
+ flag: int = 0,
+ timeline: torch.Tensor | None = None,
+ q_nope: torch.Tensor | None = None,
+ q_pe: torch.Tensor | None = None,
+ reuse_selection: int = 0,
+ ) -> torch.Tensor:
+ acc, pmax, psum = partials
+ samples = hidden_in.shape[0]
+ dev = hidden_in.device
+ if out is None:
+ out = torch.empty(samples, HIDDEN, dtype=torch.bfloat16, device=dev)
+ if q_nope is None:
+ q_nope = torch.empty(samples, self.m4.rows, dtype=torch.bfloat16, device=dev)
+ if q_pe is None:
+ q_pe = torch.empty(samples, self.m1.pe_dim, dtype=torch.bfloat16, device=dev)
+ ex = self.exchange(samples)
+ idx = None if indices is None else indices.contiguous().to(torch.int32)
+ torch.ops.tilert.glm5_pure_mla_allreduce_op(
+ hidden_in,
+ self.m0.gamma_arg,
+ self.m0.packed,
+ self.m0.scales,
+ cur_pos,
+ pe_cache,
+ ex.q_pairs,
+ self.m1.gamma_arg,
+ self.m1.packed,
+ self.m1.scales,
+ q_pe,
+ ex.kv_pairs,
+ ex.pe_pairs,
+ ex.m1_pairs,
+ self.m3.gamma,
+ kv_cache,
+ self.m4.packed,
+ self.m4.scales,
+ q_nope,
+ rope_freqs,
+ samples,
+ self.legs.qlines,
+ self.legs.kvnew_pairs,
+ self.legs.penew_pairs,
+ idx,
+ xfer_buf,
+ flag,
+ acc,
+ pmax,
+ psum,
+ self.tail.sen_a,
+ self.tail.sen_b,
+ self.tail.sen_proj,
+ self.next_tag() if tag is None else tag,
+ self.m6.packed,
+ self.m6.scales,
+ self.m7.packed,
+ self.m7.scales,
+ residual,
+ sym,
+ mype,
+ npes,
+ out,
+ self.topk,
+ self.scale,
+ timeline,
+ reuse_selection,
+ )
+ return out
+
+ def golden_m01234(
+ self,
+ hidden_in: torch.Tensor,
+ cur_pos: torch.Tensor,
+ rope_freqs: torch.Tensor,
+ pe_cache: torch.Tensor,
+ kv_cache: torch.Tensor,
+ seq_len: int,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ pe_out = pe_cache.clone()
+ q_down, kv = self.m0.golden_forward(hidden_in, cur_pos, seq_len, pe_out)
+ q_nope_down, q_pe_raw = self.m1.golden_forward(q_down)
+ q_pe, pe_out = QkvRopeGlm5(num_heads=self.num_heads).golden_forward(
+ q_pe_raw, pe_out, rope_freqs, cur_pos, seq_len
+ )
+ kv_rows = self.m3.golden_forward(kv)
+ kv_out = kv_cache.clone()
+ batch = cur_pos.numel()
+ for b in range(batch):
+ sp = int(cur_pos[b])
+ kv_out[b, sp : sp + seq_len] = kv_rows[b * seq_len : (b + 1) * seq_len]
+ q_nope = self.m4.golden_forward(q_nope_down)
+ return (q_nope, q_pe, pe_out, kv_out)
+
+ def golden_tail(
+ self,
+ q_nope: torch.Tensor,
+ q_pe: torch.Tensor,
+ kv_cache: torch.Tensor,
+ pe_cache: torch.Tensor,
+ indices: torch.Tensor | None,
+ cur_pos: torch.Tensor,
+ residual: torch.Tensor | None,
+ partials: list[torch.Tensor] | None = None,
+ ) -> torch.Tensor:
+ from tilert.models.glm_5_2_rocm.ops.flash_sparse_mla import golden_attention
+
+ samples = q_nope.shape[0]
+ attn_o = golden_attention(
+ q_nope.view(samples, self.num_heads, KV_LORA_RANK),
+ q_pe.view(samples, self.num_heads, PE_DIM),
+ kv_cache,
+ pe_cache,
+ indices,
+ int(cur_pos.reshape(-1)[0].item()),
+ topk=self.topk,
+ scale=self.scale,
+ )
+ part = self.m7.partial_golden(self.m6.golden_forward(attn_o))
+ parts = [part] if partials is None else partials
+ return self.m7.reduce_golden(parts, residual).to(q_nope.device)
+
+ def golden_forward(
+ self,
+ hidden_in: torch.Tensor,
+ cur_pos: torch.Tensor,
+ rope_freqs: torch.Tensor,
+ pe_cache: torch.Tensor,
+ kv_cache: torch.Tensor,
+ indices: torch.Tensor | None,
+ residual: torch.Tensor | None,
+ seq_len: int | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ raise RuntimeError("golden_forward is not available in release builds")
diff --git a/tilert/models/glm_5_2_rocm/ops/pure_mla_m567.py b/tilert/models/glm_5_2_rocm/ops/pure_mla_m567.py
new file mode 100644
index 0000000..7c500de
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/pure_mla_m567.py
@@ -0,0 +1,82 @@
+"""GLM-5.2 attention tail (split + combine, ProjOWkvb, UnprojOAllReduce): the tail half's weight shards and scratch, as ``pure_mla_allreduce`` composes them."""
+
+import torch
+
+from tilert.models.glm_5_2_rocm.ops.flash_sparse_mla import GLM5_SOFTMAX_SCALE, KV_LORA_RANK, TILE_N
+from tilert.models.glm_5_2_rocm.ops.proj_wkvb import ProjOWkvbGlm5
+from tilert.models.glm_5_2_rocm.ops.unprojo_allreduce import UnprojOAllReduceGlm5
+
+XCDS = 8
+M56_MAX_SEQ = 8
+TILES_PER_HEAD = 16
+MAX_SEQ = 4
+BLOCKS, STAMPS = (256, 19)
+STAMP_NAMES = (
+ "start",
+ "split",
+ "merge",
+ "m6poll",
+ "m6mfma",
+ "m6pub",
+ "publish",
+ "leadspin",
+ "leadflag",
+ "xflag",
+ "fill",
+ "gemv",
+ "exchange",
+ "end",
+ "gather",
+ "fix",
+ "q",
+ "score",
+ "idx",
+)
+LINE_WORDS = 16
+XCDS_LEADERS = 8
+
+
+def proj_words(num_heads: int) -> int:
+ nlines = num_heads * TILES_PER_HEAD * MAX_SEQ
+ return nlines * LINE_WORDS + XCDS * nlines * 16
+
+
+class PureMlaM567Glm5:
+ """One rank's tail banks: the projection shard, the W_o shard and the tail's scratch (zero-init once, never reset)."""
+
+ def __init__(
+ self,
+ device: str = "cuda:0",
+ num_heads: int = 10,
+ topk: int = 2048,
+ scale: float = GLM5_SOFTMAX_SCALE,
+ ) -> None:
+ assert num_heads in (8, 10), "weight layout is H = 8 or 10"
+ assert topk > 0 and topk % TILE_N == 0 and (topk <= 2048)
+ self.device = device
+ self.num_heads = num_heads
+ self.topk = topk
+ self.scale = scale
+ self.num_splits_max = topk // TILE_N
+ self.m6 = ProjOWkvbGlm5(device=device, num_heads=num_heads)
+ self.m7 = UnprojOAllReduceGlm5(num_heads=num_heads, device=device)
+ i32 = {"dtype": torch.int32, "device": device}
+ self.sen_a = torch.zeros(2 * XCDS * self.num_splits_max, **i32)
+ self.sen_b = torch.zeros(M56_MAX_SEQ * num_heads * KV_LORA_RANK, **i32)
+ self.sen_proj = torch.zeros(proj_words(num_heads), **i32)
+ self._tag = 0
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def next_tag(self) -> int:
+ self._tag += 1
+ return self._tag
+
+ def alloc_partials(self, samples: int) -> tuple:
+ shape = (samples, self.num_heads, self.num_splits_max)
+ return (
+ torch.empty(shape + (KV_LORA_RANK,), dtype=torch.float32, device=self.device),
+ torch.empty(shape + (1,), dtype=torch.float32, device=self.device),
+ torch.empty(shape + (1,), dtype=torch.float32, device=self.device),
+ )
diff --git a/tilert/models/glm_5_2_rocm/ops/pure_mla_moe_layer.py b/tilert/models/glm_5_2_rocm/ops/pure_mla_moe_layer.py
new file mode 100644
index 0000000..0d94b06
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/pure_mla_moe_layer.py
@@ -0,0 +1,296 @@
+"""GLM-5.2 ONE WHOLE MoE LAYER IN ONE LAUNCH (pure_mla_moe_layer)."""
+
+import torch
+
+from tilert.models.glm_5_2_rocm.ops.flash_sparse_mla import GLM5_SOFTMAX_SCALE
+from tilert.models.glm_5_2_rocm.ops.fused_moe_allreduce import FUSED_SAMPLES, FusedMoeAllreduceGlm5
+from tilert.models.glm_5_2_rocm.ops.fused_moe_w8a8 import FusedMoeAllreduceW8A8Glm5
+from tilert.models.glm_5_2_rocm.ops.moe_router import NUM_EXPERTS
+from tilert.models.glm_5_2_rocm.ops.pure_mla_allreduce import PureMlaAllReduceGlm5
+from tilert.models.glm_5_2_rocm.ops.unprojo_allreduce import HIDDEN
+from tilert.models.glm_5_2_rocm.ops.upgate_silu import MOE_SLOTS, NUM_MOE_WEIGHTS, TOP_K
+
+BLOCKS = 256
+MAX_SEQ = 4
+HLINE_WORDS = 16
+MOE_V4 = 0
+MOE_W8A8 = 1
+
+
+def hline_words() -> int:
+ return MAX_SEQ * BLOCKS * HLINE_WORDS
+
+
+class Legs:
+ """The attention -> MoE handoff legs -- one allocation for every layer."""
+
+ def __init__(self, device: str = "cuda:0") -> None:
+ i32 = {"dtype": torch.int32, "device": device}
+ self.hlines = torch.zeros(hline_words(), **i32)
+
+
+class PureMlaMoeLayerGlm5:
+ """One rank's whole MoE layer: the attention block's weights and scratch (``PureMlaAllReduceGlm5``), the MoE banks and workspaces (``FusedMoeAllreduceGlm5`` / the w8a8 subclass) and the handoff legs."""
+
+ def __init__(
+ self,
+ device: str = "cuda:0",
+ num_heads: int = 10,
+ topk: int = 2048,
+ scale: float = GLM5_SOFTMAX_SCALE,
+ moe_w8a8: bool = False,
+ num_weights: int = NUM_MOE_WEIGHTS,
+ ) -> None:
+ self.device = device
+ self.num_heads = num_heads
+ self.moe_w8a8 = moe_w8a8
+ self.mla = PureMlaAllReduceGlm5(device=device, num_heads=num_heads, topk=topk, scale=scale)
+ self.moe = (
+ FusedMoeAllreduceW8A8Glm5(device=device, num_weights=num_weights)
+ if moe_w8a8
+ else FusedMoeAllreduceGlm5(device=device, num_weights=num_weights)
+ )
+ self.legs = Legs(device)
+ self._moe_scratch: dict[int, dict[str, torch.Tensor]] = {}
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def next_tag(self) -> int:
+ return self.mla.next_tag()
+
+ def next_sen_tag(self) -> int:
+ self.moe._sen_tag += 1
+ return self.moe._sen_tag
+
+ def alloc_partials(self, samples: int) -> tuple:
+ return self.mla.alloc_partials(samples)
+
+ def moe_banks(self) -> tuple[torch.Tensor, torch.Tensor]:
+ self.moe.front.router_packed()
+ if self.moe_w8a8:
+ return (self.moe.pack_k128(), self.moe.pack_down_k128())
+ return (self.moe.front.moe.packed, self.moe.pack_down_v4())
+
+ def moe_scratch(self, samples: int) -> dict[str, torch.Tensor]:
+ if samples not in self._moe_scratch:
+ dev = self.device
+ self._moe_scratch[samples] = {
+ "q_nope": torch.empty(samples, self.mla.m4.rows, dtype=torch.bfloat16, device=dev),
+ "q_pe": torch.empty(samples, self.mla.m1.pe_dim, dtype=torch.bfloat16, device=dev),
+ "norm": torch.empty(samples, HIDDEN, dtype=torch.bfloat16, device=dev),
+ "scores": torch.empty(samples, NUM_EXPERTS, dtype=torch.float32, device=dev),
+ "probs": torch.zeros(samples, TOP_K, dtype=torch.float32, device=dev),
+ "indices": torch.zeros(samples, TOP_K, dtype=torch.int32, device=dev),
+ "mid": torch.zeros(samples, MOE_SLOTS, 256, dtype=torch.bfloat16, device=dev),
+ }
+ return self._moe_scratch[samples]
+
+ def forward(
+ self,
+ hidden_in: torch.Tensor,
+ cur_pos: torch.Tensor,
+ rope_freqs: torch.Tensor,
+ pe_cache: torch.Tensor,
+ kv_cache: torch.Tensor,
+ indices: torch.Tensor | None,
+ partials: tuple,
+ bias: torch.Tensor,
+ residual: torch.Tensor | None = None,
+ sym_attn: torch.Tensor | None = None,
+ sym_ffn: torch.Tensor | None = None,
+ mype: int = 0,
+ npes: int = 1,
+ tag: int | None = None,
+ sen_tag: int | None = None,
+ ffn_flag: int | None = None,
+ unproj_o: torch.Tensor | None = None,
+ x_out: torch.Tensor | None = None,
+ xfer_buf: torch.Tensor | None = None,
+ flag: int = 0,
+ timeline: torch.Tensor | None = None,
+ reuse_selection: int = 0,
+ ) -> torch.Tensor:
+ samples = hidden_in.shape[0]
+ assert samples in FUSED_SAMPLES
+ dev = hidden_in.device
+ mla = self.mla
+ acc, pmax, psum = partials
+ if unproj_o is None:
+ unproj_o = torch.empty(samples, HIDDEN, dtype=torch.bfloat16, device=dev)
+ if x_out is None:
+ x_out = torch.empty(samples, HIDDEN, dtype=torch.bfloat16, device=dev)
+ ex = mla.exchange(samples)
+ idx = None if indices is None else indices.contiguous().to(torch.int32)
+ t = self.next_tag() if tag is None else tag
+ st = self.next_sen_tag() if sen_tag is None else sen_tag
+ ff = st if ffn_flag is None else ffn_flag
+ sc = self.moe_scratch(samples)
+ q_nope, q_pe = (sc["q_nope"], sc["q_pe"])
+ ug_w, down_w = self.moe_banks()
+ torch.ops.tilert.glm5_pure_mla_moe_layer_op(
+ hidden_in,
+ mla.m0.gamma_arg,
+ mla.m0.packed,
+ mla.m0.scales,
+ cur_pos,
+ pe_cache,
+ ex.q_pairs,
+ mla.m1.gamma_arg,
+ mla.m1.packed,
+ mla.m1.scales,
+ q_pe,
+ ex.kv_pairs,
+ ex.pe_pairs,
+ ex.m1_pairs,
+ mla.m3.gamma,
+ kv_cache,
+ mla.m4.packed,
+ mla.m4.scales,
+ q_nope,
+ rope_freqs,
+ samples,
+ mla.legs.qlines,
+ mla.legs.kvnew_pairs,
+ mla.legs.penew_pairs,
+ idx,
+ xfer_buf,
+ flag,
+ acc,
+ pmax,
+ psum,
+ mla.tail.sen_a,
+ mla.tail.sen_b,
+ mla.tail.sen_proj,
+ t,
+ mla.m6.packed,
+ mla.m6.scales,
+ mla.m7.packed,
+ mla.m7.scales,
+ residual,
+ sym_attn,
+ mype,
+ npes,
+ unproj_o,
+ mla.topk,
+ mla.scale,
+ [
+ self.moe.front.router.gamma,
+ self.moe.front.router_packed(),
+ ug_w,
+ self.moe.front.moe.scales,
+ bias,
+ down_w,
+ self.moe.down.scales,
+ ],
+ sym_ffn,
+ ff,
+ [
+ sc["norm"],
+ sc["scores"],
+ self.moe.score_lines[:samples],
+ self.moe.flags,
+ sc["probs"],
+ sc["indices"],
+ sc["mid"],
+ x_out,
+ self.moe.mid_pairs[:samples],
+ self.legs.hlines,
+ ],
+ st,
+ MOE_W8A8 if self.moe_w8a8 else MOE_V4,
+ timeline,
+ reuse_selection,
+ )
+ return x_out
+
+ def chain_forward(
+ self,
+ hidden_in: torch.Tensor,
+ cur_pos: torch.Tensor,
+ rope_freqs: torch.Tensor,
+ pe_cache: torch.Tensor,
+ kv_cache: torch.Tensor,
+ indices: torch.Tensor | None,
+ partials: tuple,
+ bias: torch.Tensor,
+ residual: torch.Tensor | None = None,
+ sym_attn: torch.Tensor | None = None,
+ sym_ffn: torch.Tensor | None = None,
+ mype: int = 0,
+ npes: int = 1,
+ tag: int | None = None,
+ sen_tag: int | None = None,
+ ffn_flag: int | None = None,
+ unproj_o: torch.Tensor | None = None,
+ x_out: torch.Tensor | None = None,
+ xfer_buf: torch.Tensor | None = None,
+ flag: int = 0,
+ timeline: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ samples = hidden_in.shape[0]
+ dev = hidden_in.device
+ if unproj_o is None:
+ unproj_o = torch.empty(samples, HIDDEN, dtype=torch.bfloat16, device=dev)
+ if x_out is None:
+ x_out = torch.empty(samples, HIDDEN, dtype=torch.bfloat16, device=dev)
+ t = self.next_tag() if tag is None else tag
+ st = self.next_sen_tag() if sen_tag is None else sen_tag
+ ff = st if ffn_flag is None else ffn_flag
+ sc = self.moe_scratch(samples)
+ ug_w, down_w = self.moe_banks()
+ router_w = self.moe.front.router_packed()
+ self.mla.forward(
+ hidden_in,
+ cur_pos,
+ rope_freqs,
+ pe_cache,
+ kv_cache,
+ indices,
+ partials,
+ residual,
+ sym=sym_attn,
+ mype=mype,
+ npes=npes,
+ tag=t,
+ out=unproj_o,
+ xfer_buf=xfer_buf,
+ flag=flag,
+ q_nope=sc["q_nope"],
+ q_pe=sc["q_pe"],
+ timeline=timeline,
+ )
+ moe_op = (
+ torch.ops.tilert.glm5_fused_moe_allreduce_w8a8_v4_op
+ if self.moe_w8a8
+ else torch.ops.tilert.glm5_fused_moe_allreduce_v4_op
+ )
+ moe_op(
+ unproj_o,
+ self.moe.front.router.gamma,
+ router_w,
+ ug_w,
+ self.moe.front.moe.scales,
+ bias,
+ down_w,
+ self.moe.down.scales,
+ unproj_o,
+ sym_ffn,
+ mype,
+ npes,
+ ff,
+ sc["norm"],
+ sc["scores"],
+ self.moe.score_lines[:samples],
+ self.moe.flags,
+ sc["probs"],
+ sc["indices"],
+ sc["mid"],
+ x_out,
+ self.moe.mid_pairs[:samples],
+ st,
+ )
+ return (x_out, unproj_o)
+
+
+__all__ = ["MOE_V4", "MOE_W8A8", "Legs", "PureMlaMoeLayerGlm5", "hline_words"]
diff --git a/tilert/models/glm_5_2_rocm/ops/qkv_rope.py b/tilert/models/glm_5_2_rocm/ops/qkv_rope.py
new file mode 100644
index 0000000..c852d58
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/qkv_rope.py
@@ -0,0 +1,55 @@
+"""GLM-5.2 QkvRope op wrapper: golden + forward for the MLA pe-path RoPE."""
+
+import torch
+
+ROPE_DIM = 64
+NUM_HEADS = 10
+SUPPORTED_HEADS = (8, 10)
+SUPPORTED_SAMPLES = (1, 2, 4, 8)
+
+
+def make_rope_freqs(
+ samples: int, device: str = "cuda:0", seed: int = 0, positions: list[int] | None = None
+) -> torch.Tensor:
+ gen = torch.Generator(device="cpu").manual_seed(seed)
+ if positions is None:
+ ang = torch.rand(samples, ROPE_DIM // 2, generator=gen, dtype=torch.float32)
+ ang = ang * (2 * torch.pi)
+ else:
+ from tilert.models.glm_5_2_rocm.ops.llm_preprocess import make_freqs_cis
+
+ assert len(positions) == samples
+ table = make_freqs_cis(max(positions) + 1, device="cpu")
+ return table[torch.tensor(positions)].contiguous().to(device)
+ freqs = torch.stack([torch.cos(ang), torch.sin(ang)], dim=-1)
+ return freqs.reshape(samples, ROPE_DIM).contiguous().to(device)
+
+
+def rotate_golden(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+ samples = x.shape[0]
+ xf = x.float().reshape(samples, -1, ROPE_DIM // 2, 2)
+ fr = freqs.float().reshape(samples, 1, ROPE_DIM // 2, 2)
+ a, b = (xf[..., 0], xf[..., 1])
+ c, d = (fr[..., 0], fr[..., 1])
+ out = torch.stack([a * c - b * d, a * d + b * c], dim=-1)
+ return out.reshape(samples, -1).to(torch.bfloat16)
+
+
+class QkvRopeGlm5:
+ """Rotate q_pe in place and the pe_cache row at cur_pos in place."""
+
+ def __init__(self, device: str = "cuda:0", num_heads: int = NUM_HEADS):
+ assert num_heads in SUPPORTED_HEADS
+ self.device = device
+ self.num_heads = num_heads
+ self.pe_dim = num_heads * ROPE_DIM
+
+ def golden_forward(
+ self,
+ q_pe: torch.Tensor,
+ pe_cache: torch.Tensor,
+ freqs: torch.Tensor,
+ cur_pos: torch.Tensor,
+ seq_len: int,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ raise RuntimeError("golden_forward is not available in release builds")
diff --git a/tilert/models/glm_5_2_rocm/ops/rmsnorm_head_proj.py b/tilert/models/glm_5_2_rocm/ops/rmsnorm_head_proj.py
new file mode 100644
index 0000000..e6046cb
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/rmsnorm_head_proj.py
@@ -0,0 +1,62 @@
+"""GLM-5.2 RmsnormHeadProj op wrapper: golden + forward."""
+
+import torch
+
+HIDDEN = 6144
+TOP1_WS_WORDS = 16
+VOCAB_SHARD = 19360
+TILES = VOCAB_SHARD // 16
+EPS = 1e-05
+SUPPORTED_SAMPLES = (1, 2, 4, 8)
+
+
+def swizzle_head_bf16(w: torch.Tensor) -> torch.Tensor:
+ assert w.shape == (VOCAB_SHARD, HIDDEN)
+ w16 = w.contiguous().to(torch.bfloat16).view(torch.uint16)
+ w16 = w16.reshape(TILES, 16, HIDDEN // 32, 4, 8)
+ packed = w16.permute(0, 2, 3, 1, 4).contiguous()
+ return packed.reshape(-1).contiguous().view(torch.uint8)
+
+
+class RmsnormHeadProjGlm5:
+ """Final rmsnorm + bf16 head projection to f32 logits."""
+
+ OP_NAME = "glm5_rmsnorm_head_proj_op"
+
+ def __init__(self, device: str = "cuda:0"):
+ self.device = device
+ self.gamma: torch.Tensor | None = None
+ self.head: torch.Tensor | None = None
+ self.packed: torch.Tensor | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def init_reference_weights(self, gamma: torch.Tensor, head: torch.Tensor) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def golden_norm(self, hidden: torch.Tensor) -> torch.Tensor:
+ assert self.gamma is not None
+ x = hidden.float()
+ ssq = (x * x).sum(dim=-1, keepdim=True)
+ rms_inv = torch.rsqrt(ssq / HIDDEN + EPS)
+ return (x * self.gamma[None, :] * rms_inv).to(torch.bfloat16)
+
+ def golden_forward(self, hidden: torch.Tensor) -> torch.Tensor:
+ raise RuntimeError("golden_forward is not available in release builds")
+
+ def tilert_forward(
+ self, hidden: torch.Tensor, with_norm_out: bool = True
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ assert self.gamma is not None and self.packed is not None
+ samples = hidden.shape[0]
+ logits = torch.empty(samples, VOCAB_SHARD, dtype=torch.float32, device=self.device)
+ norm_out = (
+ torch.empty(samples, HIDDEN, dtype=torch.bfloat16, device=self.device)
+ if with_norm_out
+ else None
+ )
+ torch.ops.tilert.glm5_rmsnorm_head_proj_op(
+ hidden, self.gamma, self.packed, logits, norm_out
+ )
+ return (logits, norm_out)
diff --git a/tilert/models/glm_5_2_rocm/ops/rmsnorm_kv.py b/tilert/models/glm_5_2_rocm/ops/rmsnorm_kv.py
new file mode 100644
index 0000000..df73544
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/rmsnorm_kv.py
@@ -0,0 +1,28 @@
+"""GLM-5.2 RmsnormKv op wrapper: golden + forward for the kv-cache RMSNorm."""
+
+import torch
+
+KV_DIM = 512
+EPS = 1e-05
+SUPPORTED_SAMPLES = (1, 2, 4, 8)
+
+
+class RmsnormKvGlm5:
+ """Normalize kv rows into the bf16 kv cache at cur_pos."""
+
+ def __init__(self, device: str = "cuda:0"):
+ self.device = device
+ self.gamma: torch.Tensor | None = None
+
+ def init_random_gamma(self, seed: int = 0) -> None:
+ gen = torch.Generator(device="cpu").manual_seed(seed)
+ g = 1.0 + 0.1 * torch.randn(KV_DIM, generator=gen, dtype=torch.float32)
+ self.init_reference_gamma(g)
+
+ def init_reference_gamma(self, gamma: torch.Tensor) -> None:
+ assert gamma.numel() == KV_DIM
+ g = gamma.float().reshape(KV_DIM)
+ self.gamma = g.to(self.device)
+
+ def golden_forward(self, kv: torch.Tensor) -> torch.Tensor:
+ raise RuntimeError("golden_forward is not available in release builds")
diff --git a/tilert/models/glm_5_2_rocm/ops/rmsnorm_projq_wqb.py b/tilert/models/glm_5_2_rocm/ops/rmsnorm_projq_wqb.py
new file mode 100644
index 0000000..e1e13c6
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/rmsnorm_projq_wqb.py
@@ -0,0 +1,76 @@
+"""GLM-5.2 RmsnormProjQWqb op wrapper: quantizer, weight packer, golden, forward."""
+
+import torch
+
+Q_LORA_RANK = 2048
+NOPE_PER_HEAD = 192
+PE_PER_HEAD = 64
+SCALE_BLK_M = 64
+SCALE_BLK_K = 128
+SCALE_COLS = Q_LORA_RANK // SCALE_BLK_K
+FP8_MAX = 448.0
+EPS = 1e-05
+SUPPORTED_SAMPLES = (1, 2, 4)
+SUPPORTED_HEADS = (8, 10)
+NUM_HEADS = 10
+NOPE_DIM = NUM_HEADS * NOPE_PER_HEAD
+PE_DIM = NUM_HEADS * PE_PER_HEAD
+ROWS = NOPE_DIM + PE_DIM
+SCALE_ROWS = ROWS // SCALE_BLK_M
+_CHUNK = 64
+_N_KC = Q_LORA_RANK // _CHUNK
+
+
+def quantize_fp8_block(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ rows = w.shape[0]
+ assert rows % SCALE_BLK_M == 0 and w.shape[1] == Q_LORA_RANK
+ scale_rows = rows // SCALE_BLK_M
+ wf = w.float()
+ blocks = wf.view(scale_rows, SCALE_BLK_M, SCALE_COLS, SCALE_BLK_K)
+ amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12)
+ scales = amax / FP8_MAX
+ q = (blocks / scales).to(torch.float8_e4m3fn)
+ return (
+ q.reshape(rows, Q_LORA_RANK).contiguous(),
+ scales.view(scale_rows, SCALE_COLS).contiguous(),
+ )
+
+
+def swizzle_weights_contig(w_fp8: torch.Tensor) -> torch.Tensor:
+ rows = w_fp8.shape[0]
+ assert rows % 16 == 0 and w_fp8.shape[1] == Q_LORA_RANK
+ w8 = w_fp8.view(torch.uint8)
+ rg = torch.arange(rows // 16)
+ kc = torch.arange(_N_KC)
+ lane = torch.arange(64)
+ sp = torch.arange(2)
+ i = torch.arange(8)
+ RG, KC, L, SP, II = torch.meshgrid(rg, kc, lane, sp, i, indexing="ij")
+ rows_ix = RG * 16 + L % 16
+ ks = KC * _CHUNK + SP * 32 + L // 16 * 8 + II
+ return w8[rows_ix, ks].reshape(-1).contiguous()
+
+
+class RmsnormProjQWqbGlm5:
+ """Op class: golden and tilert forwards share weights."""
+
+ def __init__(self, device: str = "cuda:0", num_heads: int = NUM_HEADS):
+ self.device = device
+ self.num_heads = num_heads
+ self.nope_dim = num_heads * NOPE_PER_HEAD
+ self.pe_dim = num_heads * PE_PER_HEAD
+ self.rows = self.nope_dim + self.pe_dim
+ self.w_fp8: torch.Tensor | None = None
+ self.scales: torch.Tensor | None = None
+ self.packed: torch.Tensor | None = None
+ self.gamma: torch.Tensor | None = None
+ self.gamma_arg: torch.Tensor | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def init_reference_weights(self, w: torch.Tensor, gamma: torch.Tensor) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def golden_forward(self, q_down: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ raise RuntimeError("golden_forward is not available in release builds")
diff --git a/tilert/models/glm_5_2_rocm/ops/rmsnorm_projq_wqi.py b/tilert/models/glm_5_2_rocm/ops/rmsnorm_projq_wqi.py
new file mode 100644
index 0000000..e938e86
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/rmsnorm_projq_wqi.py
@@ -0,0 +1,64 @@
+"""GLM-5.2 RmsnormProjQWqi op wrapper: quantizer, weight packer, golden, forward."""
+
+import torch
+
+Q_LORA_RANK = 2048
+INDEX_HEADS = 32
+INDEX_DIM = 128
+ROWS = INDEX_HEADS * INDEX_DIM
+SCALE_BLK = 128
+SCALE_COLS = Q_LORA_RANK // SCALE_BLK
+SCALE_ROWS = ROWS // SCALE_BLK
+FP8_MAX = 448.0
+EPS = 1e-05
+SUPPORTED_SAMPLES = (1, 2, 4)
+_CHUNK = 64
+_N_KC = Q_LORA_RANK // _CHUNK
+
+
+def quantize_fp8_block(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ rows = w.shape[0]
+ assert rows % SCALE_BLK == 0 and w.shape[1] == Q_LORA_RANK
+ wf = w.float()
+ blocks = wf.view(rows // SCALE_BLK, SCALE_BLK, SCALE_COLS, SCALE_BLK)
+ amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12)
+ scales = amax / FP8_MAX
+ q = (blocks / scales).to(torch.float8_e4m3fn)
+ return (
+ q.reshape(rows, Q_LORA_RANK).contiguous(),
+ scales.view(rows // SCALE_BLK, SCALE_COLS).contiguous(),
+ )
+
+
+def swizzle_weights_contig(w_fp8: torch.Tensor) -> torch.Tensor:
+ rows = w_fp8.shape[0]
+ assert rows % 16 == 0 and w_fp8.shape[1] == Q_LORA_RANK
+ w8 = w_fp8.view(torch.uint8)
+ rg = torch.arange(rows // 16)
+ kc = torch.arange(_N_KC)
+ lane = torch.arange(64)
+ sp = torch.arange(2)
+ i = torch.arange(8)
+ RG, KC, L, SP, II = torch.meshgrid(rg, kc, lane, sp, i, indexing="ij")
+ return w8[RG * 16 + L % 16, KC * _CHUNK + SP * 32 + L // 16 * 8 + II].reshape(-1).contiguous()
+
+
+class RmsnormProjQWqiGlm5:
+ """Op class: golden and tilert forwards share the same weights."""
+
+ def __init__(self, device: str = "cuda:0"):
+ self.device = device
+ self.w_fp8: torch.Tensor | None = None
+ self.scales: torch.Tensor | None = None
+ self.packed: torch.Tensor | None = None
+ self.gamma: torch.Tensor | None = None
+ self.gamma_arg: torch.Tensor | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def init_reference_weights(self, w: torch.Tensor, gamma: torch.Tensor) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def golden_forward(self, q: torch.Tensor) -> torch.Tensor:
+ raise RuntimeError("golden_forward is not available in release builds")
diff --git a/tilert/models/glm_5_2_rocm/ops/rmsnorm_projx_wqakis.py b/tilert/models/glm_5_2_rocm/ops/rmsnorm_projx_wqakis.py
new file mode 100644
index 0000000..4e850f0
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/rmsnorm_projx_wqakis.py
@@ -0,0 +1,84 @@
+"""GLM-5.2 RmsnormProjXWqakis op wrapper: quantizer, both weight packers, golden, forward."""
+
+import torch
+
+HIDDEN = 6144
+Q_DIM = 2048
+KI_DIM = 128
+IS_DIM = 32
+ROWS = Q_DIM + KI_DIM
+SCALE_BLK = 128
+SCALE_ROWS = ROWS // SCALE_BLK
+SCALE_COLS = HIDDEN // SCALE_BLK
+FP8_MAX = 448.0
+EPS = 1e-05
+SUPPORTED_SAMPLES = (1, 2, 4)
+_CHUNK = 64
+_WIS_CHUNK = 64
+_WIS_ROWS = 8
+
+
+def quantize_fp8_block(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ rows, k = w.shape
+ assert rows % SCALE_BLK == 0 and k == HIDDEN
+ blocks = w.float().view(rows // SCALE_BLK, SCALE_BLK, SCALE_COLS, SCALE_BLK)
+ amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12)
+ scales = amax / FP8_MAX
+ q = (blocks / scales).to(torch.float8_e4m3fn)
+ return (
+ q.reshape(rows, k).contiguous(),
+ scales.view(rows // SCALE_BLK, SCALE_COLS).contiguous(),
+ )
+
+
+def swizzle_weights_contig(w_fp8: torch.Tensor) -> torch.Tensor:
+ rows, k = w_fp8.shape
+ assert rows % 16 == 0 and k % _CHUNK == 0
+ w8 = w_fp8.view(torch.uint8)
+ rg = torch.arange(rows // 16)
+ kc = torch.arange(k // _CHUNK)
+ lane = torch.arange(64)
+ sp = torch.arange(2)
+ i = torch.arange(8)
+ RG, KC, L, SP, II = torch.meshgrid(rg, kc, lane, sp, i, indexing="ij")
+ return w8[RG * 16 + L % 16, KC * _CHUNK + SP * 32 + L // 16 * 8 + II].reshape(-1).contiguous()
+
+
+def swizzle_wis_bf16(w_bf16: torch.Tensor) -> torch.Tensor:
+ rows, k = w_bf16.shape
+ assert rows % _WIS_ROWS == 0 and k % _WIS_CHUNK == 0
+ w16 = w_bf16.to(torch.bfloat16).view(torch.int16)
+ rg = torch.arange(rows // _WIS_ROWS)
+ kc = torch.arange(k // _WIS_CHUNK)
+ lane = torch.arange(64)
+ i = torch.arange(8)
+ RG, KC, L, II = torch.meshgrid(rg, kc, lane, i, indexing="ij")
+ packed = w16[
+ RG * _WIS_ROWS + L % _WIS_ROWS, KC * _WIS_CHUNK + L // 8 % 2 * 32 + L // 16 * 8 + II
+ ]
+ return packed.reshape(-1).contiguous().view(torch.uint8)
+
+
+class RmsnormProjXWqakisGlm5:
+ """Op class: golden and tilert forwards share the same weights."""
+
+ def __init__(self, device: str = "cuda:0"):
+ self.device = device
+ self.w_fp8: torch.Tensor | None = None
+ self.scales: torch.Tensor | None = None
+ self.packed: torch.Tensor | None = None
+ self.wis_bf16: torch.Tensor | None = None
+ self.wis_packed: torch.Tensor | None = None
+ self.gamma: torch.Tensor | None = None
+ self.gamma_arg: torch.Tensor | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def init_reference_weights(
+ self, w: torch.Tensor, wis: torch.Tensor, gamma: torch.Tensor
+ ) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def golden_forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ raise RuntimeError("golden_forward is not available in release builds")
diff --git a/tilert/models/glm_5_2_rocm/ops/rmsnorm_projx_wqkva.py b/tilert/models/glm_5_2_rocm/ops/rmsnorm_projx_wqkva.py
new file mode 100644
index 0000000..47bfda6
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/rmsnorm_projx_wqkva.py
@@ -0,0 +1,72 @@
+"""GLM-5.2 RmsnormProjXWqkva op wrapper: quantizer, weight packer, golden, forward."""
+
+import torch
+
+HIDDEN = 6144
+Q_DIM = 2048
+KV_DIM = 512
+PE_DIM = 64
+ROWS = Q_DIM + KV_DIM + PE_DIM
+SCALE_BLK = 128
+SCALE_ROWS = -(-ROWS // SCALE_BLK)
+SCALE_COLS = HIDDEN // SCALE_BLK
+FP8_MAX = 448.0
+EPS = 1e-05
+SUPPORTED_SAMPLES = (1, 2, 4)
+_CHUNK = 64
+_N_RG = ROWS // 16
+_N_KC = HIDDEN // _CHUNK
+
+
+def quantize_fp8_block(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ assert w.shape == (ROWS, HIDDEN)
+ wf = w.float()
+ padded = torch.zeros(SCALE_ROWS * SCALE_BLK, HIDDEN, dtype=torch.float32)
+ padded[:ROWS] = wf
+ blocks = padded.view(SCALE_ROWS, SCALE_BLK, SCALE_COLS, SCALE_BLK)
+ amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12)
+ scales = amax / FP8_MAX
+ q = (blocks / scales).to(torch.float8_e4m3fn)
+ q_rows = q.reshape(SCALE_ROWS * SCALE_BLK, HIDDEN)[:ROWS].contiguous()
+ return (q_rows, scales.view(SCALE_ROWS, SCALE_COLS).contiguous())
+
+
+def swizzle_weights_contig(w_fp8: torch.Tensor) -> torch.Tensor:
+ assert w_fp8.shape == (ROWS, HIDDEN)
+ w8 = w_fp8.view(torch.uint8)
+ rg = torch.arange(_N_RG)
+ kc = torch.arange(_N_KC)
+ lane = torch.arange(64)
+ sp = torch.arange(2)
+ i = torch.arange(8)
+ RG, KC, L, SP, II = torch.meshgrid(rg, kc, lane, sp, i, indexing="ij")
+ rows = RG * 16 + L % 16
+ ks = KC * _CHUNK + SP * 32 + L // 16 * 8 + II
+ return w8[rows, ks].reshape(-1).contiguous()
+
+
+_N_T8 = ROWS // 8
+_N_KC128 = HIDDEN // 128
+
+
+class RmsnormProjXWqkvaGlm5:
+ """Op class: golden and tilert forwards share weights."""
+
+ def __init__(self, device: str = "cuda:0"):
+ self.device = device
+ self.w_fp8: torch.Tensor | None = None
+ self.scales: torch.Tensor | None = None
+ self.packed: torch.Tensor | None = None
+ self.gamma: torch.Tensor | None = None
+ self.gamma_arg: torch.Tensor | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def init_reference_weights(self, w: torch.Tensor, gamma: torch.Tensor) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def golden_forward(
+ self, hidden: torch.Tensor, cur_pos: torch.Tensor, seq_len: int, pe_cache: torch.Tensor
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ raise RuntimeError("golden_forward is not available in release builds")
diff --git a/tilert/models/glm_5_2_rocm/ops/rotate.py b/tilert/models/glm_5_2_rocm/ops/rotate.py
new file mode 100644
index 0000000..46c1cf2
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/rotate.py
@@ -0,0 +1,100 @@
+"""GLM-5.2 rotate op wrappers: the indexer rotate pair (goldens + forwards)."""
+
+import torch
+
+DIM = 128
+LANES = 16
+PER_LANE = 8
+ROPE_DIM = 64
+SCALE = 0.08838834764831843
+LN_EPS = 1e-06
+INDEX_HEADS = 32
+SUPPORTED_HEADS = (8, 16, 32)
+SUPPORTED_SAMPLES = (1, 2, 4, 8)
+
+
+def rope_interleaved(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+ samples, n, _ = x.shape
+ out = x.clone()
+ rope = x[..., :ROPE_DIM].reshape(samples, n, ROPE_DIM // 2, 2)
+ fr = freqs.reshape(samples, 1, ROPE_DIM // 2, 2)
+ a, b = (rope[..., 0], rope[..., 1])
+ c, d = (fr[..., 0], fr[..., 1])
+ rot = torch.stack([a * c - b * d, a * d + b * c], dim=-1)
+ out[..., :ROPE_DIM] = rot.reshape(samples, n, ROPE_DIM)
+ return out
+
+
+def hadamard128(x: torch.Tensor) -> torch.Tensor:
+ shape = x.shape
+ v = x.reshape(-1, LANES, PER_LANE).clone()
+ for st in range(3):
+ stride = 1 << st
+ for j in range(4):
+ lo = j & stride - 1
+ idx = (j - lo) * 2 + lo
+ a = v[:, :, idx].clone()
+ b = v[:, :, idx + stride].clone()
+ v[:, :, idx] = a + b
+ v[:, :, idx + stride] = a - b
+ lane = torch.arange(LANES)
+ for st in range(4):
+ mask = 1 << st
+ partner = v.index_select(1, (lane ^ mask).to(v.device))
+ sign = torch.where((lane & mask).bool(), torch.tensor(-1.0), torch.tensor(1.0)).to(v.device)
+ v = sign[None, :, None] * v + partner
+ return v.reshape(shape)
+
+
+class RotateGlm5:
+ """RoPE + Hadamard on the indexer queries."""
+
+ def __init__(self, device: str = "cuda:0", num_heads: int = INDEX_HEADS):
+ assert num_heads in SUPPORTED_HEADS
+ self.device = device
+ self.num_heads = num_heads
+
+ def golden_forward_f32(self, iq: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+ samples = iq.shape[0]
+ x = iq.float().cpu().reshape(samples, self.num_heads, DIM)
+ x = rope_interleaved(x, freqs.float().cpu())
+ return hadamard128(x) * SCALE
+
+ def golden_forward(self, iq: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+ raise RuntimeError("golden_forward is not available in release builds")
+
+
+class LayerNormRopeRotateGlm5:
+ """LayerNorm + RoPE + Hadamard into the ki cache row at cur_pos."""
+
+ def __init__(self, device: str = "cuda:0"):
+ self.device = device
+ self.weight: torch.Tensor | None = None
+ self.bias: torch.Tensor | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def rows_golden_f32(self, ki: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+ assert self.weight is not None and self.bias is not None
+ samples = ki.shape[0]
+ x = ki.float().cpu().reshape(samples, 1, DIM)
+ mean = x.mean(dim=-1, keepdim=True)
+ var = ((x - mean) ** 2).mean(dim=-1, keepdim=True)
+ x = (x - mean) * torch.rsqrt(var + LN_EPS)
+ x = x * self.weight.float().cpu()[None, None, :] + self.bias.float().cpu()[None, None, :]
+ x = rope_interleaved(x, freqs.float().cpu())
+ return (hadamard128(x) * SCALE).reshape(samples, DIM)
+
+ def rows_golden(self, ki: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+ return self.rows_golden_f32(ki, freqs).to(torch.bfloat16)
+
+ def golden_forward(
+ self,
+ ki: torch.Tensor,
+ ki_cache: torch.Tensor,
+ freqs: torch.Tensor,
+ cur_pos: torch.Tensor,
+ seq_len: int,
+ ) -> torch.Tensor:
+ raise RuntimeError("golden_forward is not available in release builds")
diff --git a/tilert/models/glm_5_2_rocm/ops/sparse_index.py b/tilert/models/glm_5_2_rocm/ops/sparse_index.py
new file mode 100644
index 0000000..307339f
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/sparse_index.py
@@ -0,0 +1,201 @@
+"""GLM-5.2 sparse index op wrapper: indexer scores + top-2048 selection."""
+
+import numpy as np
+import torch
+
+INDEX_HEADS = 32
+INDEX_DIM = 128
+TOPK = 2048
+RADIX = 256
+HIST_ROW = 520
+HIST_LOCAL = 256 * 256
+HIST_REPL = 2 * 64 * 256
+HIST_FLAGS = 256 + 1
+HIST_WORDS = HIST_ROW + HIST_REPL + HIST_LOCAL + HIST_FLAGS
+HIST_MAX_SAMPLES = 8
+PAIRS = 4096 * 64
+SM_SCALE = 0.17677669529 * 0.08838834764831843
+SUPPORTED_SAMPLES = (1, 2, 4, 8)
+
+
+def logits_stride(cur_pos: torch.Tensor, seq_len: int) -> int:
+ m = int(cur_pos.max().item()) + seq_len
+ return (m + 7) // 8 * 8
+
+
+def key32(v: np.ndarray) -> np.ndarray:
+ b = v.astype(np.float32).view(np.uint32)
+ return np.where(b & np.uint32(2147483648), ~b, b | np.uint32(2147483648))
+
+
+def key8(v: np.ndarray) -> np.ndarray:
+ f = np.asarray(v, dtype=np.float32)
+ h = f.astype(np.float16)
+ over = np.abs(h.astype(np.float32)) > np.abs(f)
+ hb = h.view(np.uint16)
+ hb = np.where(over, hb - np.uint16(1), hb).astype(np.uint16)
+ k = np.where(hb & np.uint16(32768), ~hb, hb | np.uint16(32768)).astype(np.uint16)
+ return (k >> np.uint16(8)).astype(np.uint32)
+
+
+FP8_MAX = 448.0
+FP8_AMAX_EPS = 0.0001
+
+
+def quantize_rows(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ assert x.shape[-1] == INDEX_DIM
+ xf = x.float()
+ amax = xf.abs().amax(dim=-1)
+ scale = amax.clamp_min(FP8_AMAX_EPS) * torch.tensor(
+ 1.0 / FP8_MAX, dtype=torch.float32, device=xf.device
+ )
+ q = (xf / scale[..., None]).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
+ return (q.view(torch.uint8), scale.to(torch.float32))
+
+
+def dequantize_rows(q8: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+ return q8.view(torch.float8_e4m3fn).float() * scale.float()[..., None]
+
+
+class SparseIndexGlm5:
+ """The scores, the histogram workspace and the selection."""
+
+ OP_NAMES = ("glm5_sparse_index_op", "glm5_sparse_index_fp8_op", "glm5_topk_select_op")
+
+ def __init__(self, device: str = "cuda:0", topk: int = TOPK, fp8: bool = False):
+ self.device = device
+ self.topk = topk
+ self.fp8 = fp8
+
+ def alloc(self, samples: int, stride: int) -> dict:
+ dev = self.device
+ return {
+ "logits": torch.zeros(samples, stride, dtype=torch.float32, device=dev),
+ "hist": torch.zeros(HIST_MAX_SAMPLES * HIST_WORDS, dtype=torch.int32, device=dev),
+ "tie_pairs": torch.zeros(samples * 2 * PAIRS, dtype=torch.int32, device=dev),
+ "idx": torch.zeros(samples, self.topk, dtype=torch.int32, device=dev),
+ }
+
+ @staticmethod
+ def logits_golden(
+ iq_rt: torch.Tensor,
+ ki_cache: torch.Tensor,
+ idx_scores: torch.Tensor,
+ cur_pos: torch.Tensor,
+ seq_len: int,
+ stride: int,
+ ) -> torch.Tensor:
+ samples = iq_rt.shape[0]
+ batch = cur_pos.numel()
+ q = iq_rt.float().cpu().reshape(samples, INDEX_HEADS, INDEX_DIM)
+ kc = ki_cache.float().cpu()
+ w = idx_scores.float().cpu()
+ out = torch.zeros(samples, stride, dtype=torch.float32)
+ for b in range(batch):
+ n = int(cur_pos[b]) + seq_len
+ keys = kc[b, :n]
+ for s in range(seq_len):
+ sample = b * seq_len + s
+ sc = torch.relu(q[sample] @ keys.T * SM_SCALE)
+ out[sample, :n] = (sc * w[sample][:, None]).sum(dim=0)
+ return out
+
+ def select_golden(
+ self, logits: torch.Tensor, cur_pos: torch.Tensor, seq_len: int
+ ) -> list[dict]:
+ samples = logits.shape[0]
+ lg = logits.float().cpu().numpy()
+ res = []
+ for sample in range(samples):
+ b, s = (sample // seq_len, sample % seq_len)
+ bound = int(cur_pos[b]) + s + 1
+ keys = key32(lg[sample, :bound])
+ if bound <= self.topk:
+ res.append({"identity": True, "bound": bound})
+ continue
+ kstar = int(np.sort(keys)[::-1][self.topk - 1])
+ res.append(
+ {
+ "identity": False,
+ "bound": bound,
+ "kstar": kstar,
+ "strict": set(np.nonzero(keys > kstar)[0].tolist()),
+ "equal": set(np.nonzero(keys == kstar)[0].tolist()),
+ }
+ )
+ return res
+
+ def tilert_scores(
+ self,
+ iq_rt: torch.Tensor,
+ ki_cache: torch.Tensor,
+ idx_scores: torch.Tensor,
+ cur_pos: torch.Tensor,
+ seq_len: int,
+ ws: dict,
+ ) -> torch.Tensor:
+ torch.ops.tilert.glm5_sparse_index_op(
+ iq_rt, ki_cache, idx_scores, cur_pos, ws["logits"], ws["hist"], seq_len, self.topk
+ )
+ return ws["logits"]
+
+ def tilert_scores_fp8(
+ self,
+ iq_rt8: torch.Tensor,
+ iq_scale: torch.Tensor,
+ ki_cache8: torch.Tensor,
+ ki_scale: torch.Tensor,
+ idx_scores: torch.Tensor,
+ cur_pos: torch.Tensor,
+ seq_len: int,
+ ws: dict,
+ ) -> torch.Tensor:
+ torch.ops.tilert.glm5_sparse_index_fp8_op(
+ iq_rt8,
+ iq_scale,
+ ki_cache8,
+ ki_scale,
+ idx_scores,
+ cur_pos,
+ ws["logits"],
+ ws["hist"],
+ seq_len,
+ self.topk,
+ )
+ return ws["logits"]
+
+ def tilert_select(self, cur_pos: torch.Tensor, seq_len: int, ws: dict) -> torch.Tensor:
+ torch.ops.tilert.glm5_topk_select_op(
+ ws["logits"], ws["hist"], ws["tie_pairs"], cur_pos, ws["idx"], seq_len, self.topk
+ )
+ return ws["idx"]
+
+ def tilert_forward(
+ self,
+ iq_rt: torch.Tensor,
+ ki_cache: torch.Tensor,
+ idx_scores: torch.Tensor,
+ cur_pos: torch.Tensor,
+ seq_len: int,
+ ws: dict | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, dict]:
+ samples = iq_rt.shape[0]
+ if ws is None:
+ ws = self.alloc(samples, logits_stride(cur_pos, seq_len))
+ if self.fp8:
+ q8, qs = quantize_rows(iq_rt.view(samples, INDEX_HEADS, INDEX_DIM))
+ k8, ks = quantize_rows(ki_cache)
+ logits = self.tilert_scores_fp8(
+ q8.view(samples, -1).contiguous(),
+ qs.contiguous(),
+ k8.contiguous(),
+ ks.contiguous(),
+ idx_scores,
+ cur_pos,
+ seq_len,
+ ws,
+ )
+ else:
+ logits = self.tilert_scores(iq_rt, ki_cache, idx_scores, cur_pos, seq_len, ws)
+ idx = self.tilert_select(cur_pos, seq_len, ws)
+ return (logits, idx, ws)
diff --git a/tilert/models/glm_5_2_rocm/ops/sparse_select_s0123.py b/tilert/models/glm_5_2_rocm/ops/sparse_select_s0123.py
new file mode 100644
index 0000000..00bc741
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/sparse_select_s0123.py
@@ -0,0 +1,121 @@
+"""GLM-5.2 fused sparse selector op wrapper (sparse_select_s0123)."""
+
+import torch
+
+from tilert.models.glm_5_2_rocm.ops import rmsnorm_projq_wqi as s1_mod
+from tilert.models.glm_5_2_rocm.ops import rmsnorm_projx_wqakis as s0_mod
+
+Q_DIM = s0_mod.Q_DIM
+KI_DIM = s0_mod.KI_DIM
+IS_DIM = s0_mod.IS_DIM
+IQ_DIM = s1_mod.ROWS
+
+
+class Exchange:
+ """The selector's scratch -- allocate ONCE and share across layers."""
+
+ def __init__(self, samples: int, device: torch.device | str = "cuda"):
+ self.samples = samples
+ self.q_pairs = torch.zeros(samples, Q_DIM // 2, 2, dtype=torch.int32, device=device)
+ self.ki_pairs = torch.zeros(samples, KI_DIM // 2, 2, dtype=torch.int32, device=device)
+ self.iq_pairs = torch.zeros(samples, IQ_DIM // 2, 2, dtype=torch.int32, device=device)
+
+
+def fused_forward(
+ hidden_in: torch.Tensor,
+ s0,
+ s1,
+ s2,
+ s3,
+ rope_freqs: torch.Tensor,
+ cur_pos: torch.Tensor,
+ ki_cache: torch.Tensor,
+ ex: Exchange,
+ *,
+ seq_len: int = 1,
+ tag: int = 1,
+ want_q: bool = False,
+ want_ki: bool = False,
+ want_iq: bool = False,
+ fp8: dict | None = None,
+ want_iq_rt: bool = True,
+) -> tuple[
+ torch.Tensor | None, torch.Tensor | None, torch.Tensor, torch.Tensor | None, torch.Tensor
+]:
+ assert s0.packed is not None and s0.wis_packed is not None
+ assert s1.packed is not None
+ assert s2.weight is not None and s2.bias is not None
+ assert s3.num_heads * KI_DIM == IQ_DIM, "must be the 32-head GPU0 shape"
+ assert (
+ s0.gamma_arg is not None
+ and s0.gamma_arg.numel() > 0
+ and (s1.gamma_arg is not None)
+ and (s1.gamma_arg.numel() > 0)
+ ), "the fused kernel has f32-gamma paths only"
+ samples = hidden_in.size(0)
+ assert samples == ex.samples
+ dev = hidden_in.device
+ bf = torch.bfloat16
+ q = torch.empty(samples, Q_DIM, dtype=bf, device=dev) if want_q else None
+ ki = torch.empty(samples, KI_DIM, dtype=bf, device=dev) if want_ki else None
+ is_ = torch.empty(samples, IS_DIM, dtype=bf, device=dev)
+ iq = torch.empty(samples, IQ_DIM, dtype=bf, device=dev) if want_iq else None
+ if fp8 is not None:
+ extra = (fp8["ki_cache8"], fp8["ki_scale"], fp8["iq_rt8"], fp8["iq_scale"])
+ else:
+ assert want_iq_rt, "only the fp8 arm can drop the bf16 iq_rt"
+ extra = (None, None, None, None)
+ iq_rt = torch.empty(samples if want_iq_rt else 0, IQ_DIM, dtype=bf, device=dev)
+ torch.ops.tilert.glm5_sparse_select_s0123_op(
+ hidden_in,
+ s0.gamma_arg,
+ s0.packed,
+ s0.scales,
+ s0.wis_packed,
+ q,
+ ki,
+ is_,
+ ex.q_pairs,
+ ex.ki_pairs,
+ ex.iq_pairs,
+ s1.gamma_arg,
+ s1.packed,
+ s1.scales,
+ iq,
+ s2.weight,
+ s2.bias,
+ rope_freqs,
+ cur_pos,
+ ki_cache,
+ iq_rt,
+ seq_len,
+ tag,
+ *extra,
+ )
+ return (q, ki, is_, iq, iq_rt if want_iq_rt else None)
+
+
+def golden_forward(
+ hidden_in: torch.Tensor,
+ s0,
+ s1,
+ s2,
+ s3,
+ rope_freqs: torch.Tensor,
+ cur_pos: torch.Tensor,
+ ki_cache: torch.Tensor,
+ *,
+ seq_len: int = 1,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ raise RuntimeError("golden_forward is not available in release builds")
+
+
+def alloc_fp8(
+ samples: int, batch: int, cache_len: int, device: torch.device | str = "cuda"
+) -> dict:
+ return {
+ "ki_cache8": torch.zeros(batch, cache_len, KI_DIM, dtype=torch.uint8, device=device),
+ "ki_scale": torch.zeros(batch, cache_len, dtype=torch.float32, device=device),
+ "iq_rt8": torch.zeros(samples, IQ_DIM, dtype=torch.uint8, device=device),
+ "iq_scale": torch.zeros(samples, IQ_DIM // KI_DIM, dtype=torch.float32, device=device),
+ }
diff --git a/tilert/models/glm_5_2_rocm/ops/top1_allreduce.py b/tilert/models/glm_5_2_rocm/ops/top1_allreduce.py
new file mode 100644
index 0000000..5429783
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/top1_allreduce.py
@@ -0,0 +1,53 @@
+"""GLM-5.2 Top1Allreduce op wrapper: golden + forward."""
+
+import torch
+
+VOCAB_SHARD = 19360
+NUM_PES = 8
+
+
+def sym_bytes(samples: int) -> int:
+ return int(torch.ops.tilert.glm5_top1_allreduce_sym_bytes(samples))
+
+
+def sym_buffer(samples: int, device) -> torch.Tensor:
+ return torch.zeros(sym_bytes(samples), dtype=torch.uint8, device=device)
+
+
+def sym_table(buffers: list[torch.Tensor], device) -> torch.Tensor:
+ ptrs = [int(b.data_ptr()) for b in buffers]
+ return torch.tensor(ptrs, dtype=torch.int64, device=device)
+
+
+def argmax_lowest_idx(logits: torch.Tensor) -> torch.Tensor:
+ vals = logits.max(dim=-1, keepdim=True).values
+ return (logits == vals).int().argmax(dim=-1).to(torch.int32)
+
+
+class Top1AllreduceGlm5:
+ """Greedy token selection over the TP8 vocab shards."""
+
+ OP_NAME = "glm5_top1_allreduce_op"
+
+ def __init__(self, device: str = "cuda:0"):
+ self.device = device
+
+ def local_golden(self, logits: torch.Tensor, mype: int) -> torch.Tensor:
+ return argmax_lowest_idx(logits) + mype * VOCAB_SHARD
+
+ @staticmethod
+ def merged_golden(logits_by_rank: list[torch.Tensor]) -> torch.Tensor:
+ return argmax_lowest_idx(torch.cat(logits_by_rank, dim=-1))
+
+ def tilert_forward(
+ self,
+ logits: torch.Tensor,
+ mype: int = 0,
+ npes: int = 1,
+ flag: int = 1,
+ sym: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ samples = logits.shape[0]
+ token = torch.empty(samples, dtype=torch.int32, device=logits.device)
+ torch.ops.tilert.glm5_top1_allreduce_op(logits, sym, mype, npes, flag, token)
+ return token
diff --git a/tilert/models/glm_5_2_rocm/ops/top_p.py b/tilert/models/glm_5_2_rocm/ops/top_p.py
new file mode 100644
index 0000000..ef1d905
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/top_p.py
@@ -0,0 +1,164 @@
+"""GLM-5.2 top-p op wrapper: golden + forward for distributed sampling."""
+
+import numpy as np
+import torch
+
+VOCAB_SHARD = 19360
+NUM_PES = 8
+TOP_K = 256
+SEND_BYTES = 2080
+_M64 = (1 << 64) - 1
+_SEED_MUL = 19349663
+_POS_MUL = 73856093
+_STEP_MUL = 8589934591
+_COL_MUL = 479001599
+_COIN_MUL = 11400714819323198485
+_COIN_XOR = 13787848793156543929
+NO_DRAFT = -1
+
+
+def send_bytes() -> int:
+ return int(torch.ops.tilert.glm5_top_p_send_bytes())
+
+
+def sym_bytes(samples: int) -> int:
+ return int(torch.ops.tilert.glm5_top_p_sym_bytes(samples))
+
+
+def sym_buffer(samples: int, device) -> torch.Tensor:
+ return torch.zeros(sym_bytes(samples), dtype=torch.uint8, device=device)
+
+
+def sym_table(buffers: list[torch.Tensor], device) -> torch.Tensor:
+ ptrs = [int(b.data_ptr()) for b in buffers]
+ return torch.tensor(ptrs, dtype=torch.int64, device=device)
+
+
+def gumbel_uniform(seed: int, position: int, col: int) -> np.float32:
+ step_seed = seed * _SEED_MUL & _M64 ^ position * _POS_MUL & _M64
+ hashed = step_seed * _STEP_MUL & _M64 ^ col * _COL_MUL & _M64
+ u = np.float32(hashed % (1 << 24)) * np.float32(1.0 / (1 << 24))
+ return np.float32(min(max(u, np.float32(1e-10)), np.float32(1.0 - 1e-10)))
+
+
+def verify_coin(seed: int, position: int) -> np.float32:
+ step_seed = seed * _SEED_MUL & _M64 ^ position * _POS_MUL & _M64
+ hashed = step_seed * _COIN_MUL & _M64 ^ _COIN_XOR
+ u = np.float32(hashed % (1 << 24)) * np.float32(1.0 / (1 << 24))
+ return np.float32(min(max(u, np.float32(1e-10)), np.float32(1.0 - 1e-10)))
+
+
+def _inv_t(temperature: float) -> np.float32:
+ if temperature < 1e-06:
+ return np.float32(1e30)
+ if temperature != 1.0:
+ return np.float32(1.0 / temperature)
+ return np.float32(1.0)
+
+
+def local_scores_golden(
+ logits: torch.Tensor, temperature: float
+) -> tuple[torch.Tensor, torch.Tensor, float, float]:
+ x = logits.float().cpu().numpy().astype(np.float32) * _inv_t(temperature)
+ m = np.float32(x.max())
+ e = np.exp(x - m, dtype=np.float32)
+ order = np.lexsort((np.arange(e.shape[0]), -e))[:TOP_K]
+ return (
+ torch.from_numpy(e[order].copy()),
+ torch.from_numpy(order.astype(np.int32)),
+ float(m),
+ float(e.sum(dtype=np.float32)),
+ )
+
+
+def sample_golden(
+ per_rank: list[tuple[torch.Tensor, torch.Tensor, float, float]],
+ top_p: float,
+ seed: int,
+ position: int,
+ draft: int = NO_DRAFT,
+) -> tuple[int, float]:
+ npes = len(per_rank)
+ gm = np.float32(max((np.float32(m) for _, _, m, _ in per_rank)))
+ scales = [np.float32(np.exp(np.float32(m) - gm)) for _, _, m, _ in per_rank]
+ gl = np.float32(0.0)
+ for (_, _, _, l), s in zip(per_rank, scales):
+ gl = np.float32(gl + s * np.float32(l))
+ vals = np.concatenate(
+ [v.numpy().astype(np.float32) * s for (v, _, _, _), s in zip(per_rank, scales)]
+ )
+ idx = np.concatenate(
+ [i.numpy().astype(np.int64) + p * VOCAB_SHARD for p, (_, i, _, _) in enumerate(per_rank)]
+ )
+ order = np.lexsort((idx, -vals))
+ vals, idx = (vals[order], idx[order])
+ probs = (vals / gl).astype(np.float32)
+ cum = np.cumsum(probs, dtype=np.float32)
+ over = np.nonzero(cum > np.float32(top_p))[0]
+ cutoff = int(over[0]) if over.size else npes * TOP_K - 1
+ accept_col = -1
+ if draft != NO_DRAFT:
+ hit = np.nonzero(idx[: cutoff + 1] == draft)[0]
+ if hit.size:
+ dc = int(hit[0])
+ p_draft = np.float32(probs[dc] / max(cum[cutoff], np.float32(1e-20)))
+ if verify_coin(seed, position) < p_draft:
+ accept_col = dc
+ else:
+ probs = probs.copy()
+ probs[dc] = np.float32(0.0)
+ if accept_col >= 0:
+ return (int(idx[accept_col]), float(probs[accept_col]))
+ best, best_col = (-np.float32(np.finfo(np.float32).max), 0)
+ for col in range(cutoff + 1):
+ u = gumbel_uniform(seed, position, col)
+ g = np.float32(-np.log(-np.log(u)))
+ pert = np.float32(np.log(probs[col] + np.float32(1e-10)) + g)
+ if pert > best:
+ best, best_col = (pert, col)
+ return (int(idx[best_col]), float(probs[best_col]))
+
+
+class TopPGlm5:
+ """Distributed nucleus sampling (single-GPU npes=1 arm included)."""
+
+ def __init__(self, device: str = "cuda:0"):
+ self.device = device
+
+ def tilert_forward(
+ self,
+ logits: torch.Tensor,
+ temperature: float,
+ top_p: float,
+ seeds: torch.Tensor,
+ positions: torch.Tensor,
+ mype: int = 0,
+ npes: int = 1,
+ flag: int = 1,
+ sym: torch.Tensor | None = None,
+ draft_tokens: torch.Tensor | None = None,
+ verify_seq: int = 1,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ samples = logits.shape[0]
+ dev = logits.device
+ send_buf = torch.empty(samples * SEND_BYTES, dtype=torch.uint8, device=dev)
+ token = torch.empty(samples, dtype=torch.int32, device=dev)
+ prob = torch.empty(samples, dtype=torch.float32, device=dev)
+ torch.ops.tilert.glm5_top_p_local_scores_op(
+ logits, temperature, send_buf, sym, mype, npes, flag
+ )
+ torch.ops.tilert.glm5_top_p_sample_op(
+ send_buf,
+ top_p,
+ seeds,
+ positions,
+ sym,
+ mype,
+ npes,
+ flag,
+ token,
+ prob,
+ draft_tokens,
+ verify_seq,
+ )
+ return (token, prob)
diff --git a/tilert/models/glm_5_2_rocm/ops/unprojo_allreduce.py b/tilert/models/glm_5_2_rocm/ops/unprojo_allreduce.py
new file mode 100644
index 0000000..a25f61f
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/unprojo_allreduce.py
@@ -0,0 +1,145 @@
+"""GLM-5.2 UnprojOAllReduce op wrapper: packer, golden, forward."""
+
+import os
+
+import torch
+
+HIDDEN = 6144
+V_HEAD_DIM = 256
+SCALE_BLK = 128
+NUM_PES = 8
+FP8_MAX = 448.0
+SUPPORTED_HEADS = (8, 10, 16)
+PROTO_0, PROTO_1 = (0, 1)
+BLOCKS = 256
+ROWS_PER_BLOCK = 24
+MAX_SEQ = 4
+
+
+def quantize_fp8_block(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ rows, k = w.shape
+ assert rows % SCALE_BLK == 0 and k % SCALE_BLK == 0
+ sr, sc = (rows // SCALE_BLK, k // SCALE_BLK)
+ blocks = w.float().view(sr, SCALE_BLK, sc, SCALE_BLK)
+ amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12)
+ scales = amax / FP8_MAX
+ q = (blocks / scales).to(torch.float8_e4m3fn)
+ return (q.reshape(rows, k).contiguous(), scales.view(sr, sc).contiguous())
+
+
+def swizzle_v2(w_fp8: torch.Tensor) -> torch.Tensor:
+ rows, k = w_fp8.shape
+ assert rows % 8 == 0 and k % (8 * 64) == 0
+ units = k // 8 // 64
+ pairs, odd = (units // 2, units % 2)
+ aligned_k = 8 * pairs * 128
+ w8 = w_fp8.view(torch.uint8)
+ lane = torch.arange(64)
+ row_in_group = (lane >> 5) * 4 + (lane & 3)
+ kslice = (lane >> 2 & 7) * 4
+ i = torch.arange(4)
+ s4 = torch.arange(4)
+
+ def seg(k_base):
+ return k_base[:, None, :, None] + kslice[None, :, None, None] + i[None, None, None, :]
+
+ pieces = []
+ if pairs:
+ kb = torch.arange(8 * pairs)[:, None] * 128 + s4[None, :] * 32
+ pieces.append(seg(kb))
+ if odd:
+ j = torch.arange(4)
+ wave = 2 * j[:, None] + (s4[None, :] >> 1)
+ kb = aligned_k + wave * 64 + (s4[None, :] & 1) * 32
+ pieces.append(seg(kb))
+ kk = torch.cat(pieces, dim=0)
+ kidx = kk.reshape(-1)
+ ridx = row_in_group[None, :, None, None].expand_as(kk).reshape(-1)
+ assert kidx.numel() == 8 * k
+ rg = torch.arange(rows // 8)
+ rows_ix = rg[:, None] * 8 + ridx[None, :]
+ ks = kidx[None, :].expand(rows // 8, -1)
+ return w8[rows_ix, ks].reshape(-1).contiguous()
+
+
+def sym_bytes(samples: int) -> int:
+ return int(torch.ops.tilert.glm5_unprojo_allreduce_sym_bytes(samples))
+
+
+def sym_buffer(samples: int, device) -> torch.Tensor:
+ return torch.zeros(sym_bytes(samples), dtype=torch.uint8, device=device)
+
+
+def sym_table(buffers: list[torch.Tensor], device) -> torch.Tensor:
+ assert len(buffers) == NUM_PES
+ return torch.tensor([b.data_ptr() for b in buffers], dtype=torch.int64, device=device)
+
+
+def enable_peer_access(ndev: int = NUM_PES) -> int:
+ return int(torch.ops.tilert.glm5_enable_peer_access(ndev))
+
+
+class UnprojOAllReduceGlm5:
+ """One rank's output projection: W_o shard + the ops that run on it."""
+
+ def __init__(self, num_heads: int = 10, device: str = "cuda:0") -> None:
+ assert num_heads in SUPPORTED_HEADS
+ self.num_heads = num_heads
+ self.k = num_heads * V_HEAD_DIM
+ self.device = device
+ self.w_fp8: torch.Tensor | None = None
+ self.scales: torch.Tensor | None = None
+ self.packed: dict[int, torch.Tensor] = {}
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def init_reference_weights(self, w: torch.Tensor) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def partial_golden(self, proj_o: torch.Tensor) -> torch.Tensor:
+ assert self.w_fp8 is not None and self.scales is not None
+ x = proj_o.float().cpu()
+ wq = self.w_fp8.float().cpu()
+ scales = self.scales.cpu()
+ acc = torch.zeros(x.shape[0], HIDDEN, dtype=torch.float32)
+ for rb in range(HIDDEN // SCALE_BLK):
+ r0, r1 = (rb * SCALE_BLK, (rb + 1) * SCALE_BLK)
+ part = torch.zeros(x.shape[0], SCALE_BLK, dtype=torch.float32)
+ for kb in range(self.k // SCALE_BLK):
+ k0, k1 = (kb * SCALE_BLK, (kb + 1) * SCALE_BLK)
+ part += x[:, k0:k1] @ wq[r0:r1, k0:k1].T * scales[rb, kb]
+ acc[:, r0:r1] = part
+ return acc.to(torch.bfloat16)
+
+ @staticmethod
+ def reduce_golden(partials: list[torch.Tensor], residual: torch.Tensor | None) -> torch.Tensor:
+ acc = torch.zeros(partials[0].shape, dtype=torch.float32)
+ for p in partials:
+ acc += p.float().cpu()
+ if residual is not None:
+ acc += residual.float().cpu()
+ return acc.to(torch.bfloat16)
+
+ def golden_forward(self, proj_o: torch.Tensor, residual: torch.Tensor | None) -> torch.Tensor:
+ raise RuntimeError("golden_forward is not available in release builds")
+
+ def tilert_forward(
+ self,
+ proj_o: torch.Tensor,
+ residual: torch.Tensor | None = None,
+ proto: int = PROTO_0,
+ sym: torch.Tensor | None = None,
+ mype: int = 0,
+ npes: int = 1,
+ flag: int = 1,
+ out: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ assert self.packed is not None and self.scales is not None
+ os.environ["TILERT_GLM5_AR_PROTO"] = str(proto)
+ if out is None:
+ out = torch.empty(proj_o.shape[0], HIDDEN, dtype=torch.bfloat16, device=proj_o.device)
+ torch.ops.tilert.glm5_unprojo_allreduce_op(
+ proj_o, self.packed, self.scales, residual, sym, mype, npes, flag, out
+ )
+ return out
diff --git a/tilert/models/glm_5_2_rocm/ops/upgate_silu.py b/tilert/models/glm_5_2_rocm/ops/upgate_silu.py
new file mode 100644
index 0000000..20b1964
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/upgate_silu.py
@@ -0,0 +1,150 @@
+"""GLM-5.2 FFN up/gate op wrappers: packers, goldens, forwards."""
+
+import torch
+
+HIDDEN = 6144
+DENSE_INTER = 1536
+MOE_INTER = 256
+TOP_K = 8
+MOE_SLOTS = TOP_K + 1
+NUM_MOE_WEIGHTS = 257
+SCALE_BLK = 128
+FP8_MAX = 448.0
+EPS = 1e-05
+PAIR = 8
+_CHUNK = 64
+DENSE_SAMPLES = (1, 2, 4)
+MOE_SAMPLES = (1, 2, 4, 8)
+
+
+def pair_interleave_rows(inter: int) -> torch.Tensor:
+ p = torch.arange(2 * inter)
+ t, u, v = (p // 16, p % 16 // PAIR, p % PAIR)
+ return u * inter + t * PAIR + v
+
+
+def quantize_fp8_block(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ rows, k = w.shape
+ assert rows % SCALE_BLK == 0 and k % SCALE_BLK == 0
+ sr, sc = (rows // SCALE_BLK, k // SCALE_BLK)
+ blocks = w.float().view(sr, SCALE_BLK, sc, SCALE_BLK)
+ amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12)
+ scales = amax / FP8_MAX
+ q = (blocks / scales).to(torch.float8_e4m3fn)
+ return (q.reshape(rows, k).contiguous(), scales.view(sr, sc).contiguous())
+
+
+def swizzle_pair_interleaved(w_fp8: torch.Tensor, inter: int) -> torch.Tensor:
+ rows, k = w_fp8.shape
+ assert rows == 2 * inter and k % _CHUNK == 0
+ w8 = w_fp8.view(torch.uint8)
+ perm = pair_interleave_rows(inter)
+ t = torch.arange(rows // 16)
+ kc = torch.arange(k // _CHUNK)
+ lane = torch.arange(64)
+ sp = torch.arange(2)
+ i = torch.arange(8)
+ T, KC, L, SP, II = torch.meshgrid(t, kc, lane, sp, i, indexing="ij")
+ rows_ix = perm[T * 16 + L % 16]
+ ks = KC * _CHUNK + SP * 32 + L // 16 * 8 + II
+ return w8[rows_ix, ks].reshape(-1).contiguous()
+
+
+def _silu(x: torch.Tensor) -> torch.Tensor:
+ return x * torch.sigmoid(x)
+
+
+def _blocked_gemv(
+ act: torch.Tensor, wq: torch.Tensor, scales: torch.Tensor, row0: int
+) -> torch.Tensor:
+ rows, k = wq.shape
+ assert rows % SCALE_BLK == 0 and row0 % SCALE_BLK == 0
+ sr0 = row0 // SCALE_BLK
+ nrb = rows // SCALE_BLK
+ row_scale = scales[sr0 : sr0 + nrb].repeat_interleave(SCALE_BLK, dim=0)
+ out = torch.zeros(act.shape[0], rows, dtype=torch.float32)
+ for kb in range(k // SCALE_BLK):
+ k0, k1 = (kb * SCALE_BLK, (kb + 1) * SCALE_BLK)
+ out += act[:, k0:k1] @ wq[:, k0:k1].T * row_scale[None, :, kb]
+ return out
+
+
+class RmsnormUpGateSiluGlm5:
+ """The dense MLP's up/gate half (layers 0-2)."""
+
+ def __init__(self, device: str = "cuda:0") -> None:
+ self.device = device
+ self.inter = DENSE_INTER
+ self.w_fp8: torch.Tensor | None = None
+ self.scales: torch.Tensor | None = None
+ self.packed: torch.Tensor | None = None
+ self.gamma: torch.Tensor | None = None
+ self.gamma_arg: torch.Tensor | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def init_reference_weights(self, w: torch.Tensor, gamma: torch.Tensor) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def golden_forward(self, hidden: torch.Tensor) -> torch.Tensor:
+ raise RuntimeError("golden_forward is not available in release builds")
+
+ def tilert_forward(self, hidden: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+ assert self.packed is not None and self.scales is not None
+ if out is None:
+ out = torch.empty(
+ hidden.shape[0], self.inter, dtype=torch.bfloat16, device=hidden.device
+ )
+ torch.ops.tilert.glm5_rmsnorm_upgate_silu_op(
+ hidden, self.gamma_arg, self.packed, self.scales, out
+ )
+ return out
+
+
+class MoeUpGateSiluGlm5:
+ """MoE slots 2+3: 257 experts' up/gate halves (index 0 = shared)."""
+
+ def __init__(self, device: str = "cuda:0", num_weights: int = NUM_MOE_WEIGHTS):
+ self.device = device
+ self.num_weights = num_weights
+ self.inter = MOE_INTER
+ self.w_fp8: torch.Tensor | None = None
+ self.scales: torch.Tensor | None = None
+ self.packed: torch.Tensor | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def init_reference_weights(self, w: torch.Tensor) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def golden_forward(self, norm_hidden: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
+ raise RuntimeError("golden_forward is not available in release builds")
+
+ def golden_routed_forward(
+ self, norm_hidden: torch.Tensor, indices: torch.Tensor
+ ) -> torch.Tensor:
+ return self.golden_forward(norm_hidden, indices)[:, 1:]
+
+
+class SharedUpGateSiluGlm5:
+ """MoE slot 3 standalone, S x 16-CTA shape: RmsNorm + the SHARED expert's up/gate + SiLU in one launch, reading unproj_o directly."""
+
+ def __init__(self, device: str = "cuda:0") -> None:
+ self.device = device
+ self.moe = MoeUpGateSiluGlm5(device=device, num_weights=1)
+ self.gamma: torch.Tensor | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def golden_norm(self, hidden: torch.Tensor) -> torch.Tensor:
+ assert self.gamma is not None
+ x = hidden.float().cpu()
+ gamma = self.gamma.float().cpu()
+ rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + EPS)
+ return (x * gamma[None, :] * rms).to(torch.bfloat16)
+
+ def golden_forward(self, hidden: torch.Tensor) -> torch.Tensor:
+ raise RuntimeError("golden_forward is not available in release builds")
diff --git a/tilert/models/glm_5_2_rocm/ops/upgate_silu_w8a8.py b/tilert/models/glm_5_2_rocm/ops/upgate_silu_w8a8.py
new file mode 100644
index 0000000..33a0106
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/ops/upgate_silu_w8a8.py
@@ -0,0 +1,60 @@
+"""GLM-5.2 MoE up/gate + SiLU, W8A8: the python twins of the device kernel."""
+
+import torch
+
+from tilert.models.glm_5_2_rocm.ops.upgate_silu import pair_interleave_rows
+
+SEG_K = 128
+
+
+def swizzle_pair_interleaved_k128(w_fp8: torch.Tensor, inter: int) -> torch.Tensor:
+ rows, k = w_fp8.shape
+ assert rows == 2 * inter and k % SEG_K == 0
+ w8 = w_fp8.view(torch.uint8)
+ perm = pair_interleave_rows(inter)
+ t = torch.arange(rows // 16)
+ c = torch.arange(k // SEG_K)
+ h = torch.arange(2)
+ lane = torch.arange(64)
+ i = torch.arange(16)
+ T, C, H, L, II = torch.meshgrid(t, c, h, lane, i, indexing="ij")
+ rows_ix = perm[T * 16 + L % 16]
+ ks = C * SEG_K + L // 16 * 32 + H * 16 + II
+ return w8[rows_ix, ks].reshape(-1).contiguous()
+
+
+FP8_MAX = 448.0
+FP8_AMAX_EPS = 0.0001
+
+
+def quant_std_blocks(x_bf16: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ shape = x_bf16.shape
+ xb = x_bf16.float().reshape(*shape[:-1], shape[-1] // SEG_K, SEG_K)
+ amax = xb.abs().amax(dim=-1)
+ scale = amax.clamp_min(FP8_AMAX_EPS) * torch.tensor(
+ 1.0 / FP8_MAX, dtype=torch.float32, device=xb.device
+ )
+ inv = torch.ones_like(scale) / scale
+ q = (xb * inv.unsqueeze(-1)).to(torch.float8_e4m3fn).reshape(shape)
+ return (q, scale)
+
+
+def quant_act_row(norm_bf16: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ assert norm_bf16.dtype == torch.bfloat16
+ return quant_std_blocks(norm_bf16)
+
+
+def quant_mid_rows(mid_bf16: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ assert mid_bf16.dtype == torch.bfloat16
+ return quant_std_blocks(mid_bf16)
+
+
+__all__ = [
+ "FP8_AMAX_EPS",
+ "FP8_MAX",
+ "SEG_K",
+ "quant_act_row",
+ "quant_mid_rows",
+ "quant_std_blocks",
+ "swizzle_pair_interleaved_k128",
+]
diff --git a/tilert/models/glm_5_2_rocm/weight_converter.py b/tilert/models/glm_5_2_rocm/weight_converter.py
new file mode 100644
index 0000000..13e0ff8
--- /dev/null
+++ b/tilert/models/glm_5_2_rocm/weight_converter.py
@@ -0,0 +1,1780 @@
+"""GLM-5.2 weight converter: HF FP8 checkpoint -> device-sharded TileRT weights."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import shutil
+import time
+import zlib
+from collections.abc import Callable
+from contextlib import ExitStack
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+from tilert import logger
+from tilert.models.glm_5_2_rocm.checkpoint_config import (
+ CONVERTER_VERSION,
+ describe_hf_config,
+ load_hf_config,
+ sha256_file,
+ validate_hf_config,
+)
+from tilert.models.glm_5_2_rocm.model_args import (
+ KIND_DENSE,
+ KIND_MOE_FULL,
+ KIND_MOE_SHARED,
+ ModelArgsGlm52,
+ layer_kind,
+)
+from tilert.models.glm_5_2_rocm.ops.llm_preprocess import make_freqs_cis
+
+__all__ = [
+ "CheckpointReader",
+ "Glm52WeightConverter",
+ "load_rank_params",
+ "random_rank_params",
+ "selftest_swizzles",
+]
+FP8_MAX = 448.0
+BLK = 128
+N_BANK = 257
+TP8_HEADS = 8
+
+
+def _u8(t: torch.Tensor) -> torch.Tensor:
+ return t.view(torch.uint8) if t.dtype == torch.float8_e4m3fn else t
+
+
+def swizzle_fp8_contig8(w8: torch.Tensor) -> torch.Tensor:
+ w8 = _u8(w8)
+ *lead, rows, k = w8.shape
+ assert rows % 16 == 0 and k % 64 == 0
+ v = w8.reshape(*lead, rows // 16, 16, k // 64, 2, 4, 8)
+ n = len(lead)
+ perm = tuple(range(n)) + tuple(x + n for x in (0, 2, 4, 1, 3, 5))
+ return v.permute(perm).reshape(*lead, rows * k).contiguous()
+
+
+def swizzle_bf16_16x32(w: torch.Tensor) -> torch.Tensor:
+ rows, k = w.shape
+ assert rows % 16 == 0 and k % 32 == 0
+ v = w.to(torch.bfloat16).view(torch.int16).reshape(rows // 16, 16, k // 32, 4, 8)
+ return v.permute(0, 2, 3, 1, 4).reshape(-1).contiguous().view(torch.uint8)
+
+
+def swizzle_wis_8x64(w: torch.Tensor) -> torch.Tensor:
+ rows, k = w.shape
+ assert rows % 8 == 0 and k % 64 == 0
+ w16 = w.to(torch.bfloat16).view(torch.int16)
+ rg = torch.arange(rows // 8)
+ kc = torch.arange(k // 64)
+ lane = torch.arange(64)
+ i = torch.arange(8)
+ RG, KC, L, II = torch.meshgrid(rg, kc, lane, i, indexing="ij")
+ packed = w16[RG * 8 + L % 8, KC * 64 + L // 8 % 2 * 32 + L // 16 * 8 + II]
+ return packed.reshape(-1).contiguous().view(torch.uint8)
+
+
+def swizzle_fp8_m4(w8: torch.Tensor) -> torch.Tensor:
+ w8 = _u8(w8)
+ *lead, rows, k = w8.shape
+ assert rows % 8 == 0 and k % 128 == 0
+ v = w8.reshape(*lead, rows // 8, 2, 4, k // 128, 4, 8, 4)
+ n = len(lead)
+ perm = tuple(range(n)) + tuple(x + n for x in (0, 3, 1, 5, 2, 4, 6))
+ return v.permute(perm).reshape(*lead, rows * k).contiguous()
+
+
+def swizzle_fp8_v2(w8: torch.Tensor) -> torch.Tensor:
+ from tilert.models.glm_5_2_rocm.ops.unprojo_allreduce import swizzle_v2
+
+ return swizzle_v2(_u8(w8).view(torch.float8_e4m3fn))
+
+
+def _pack_wo(w8: torch.Tensor) -> torch.Tensor:
+ return swizzle_fp8_v2(w8)
+
+
+def swizzle_bf16_m4(w: torch.Tensor) -> torch.Tensor:
+ rows, k = w.shape
+ assert rows % 8 == 0 and k % 64 == 0
+ v = w.to(torch.bfloat16).view(torch.uint16).reshape(rows // 8, 2, 4, k // 64, 2, 8, 4)
+ return v.permute(0, 3, 1, 5, 2, 4, 6).reshape(-1).contiguous().view(torch.uint8)
+
+
+def pair_interleave(w8: torch.Tensor, inter: int) -> torch.Tensor:
+ *lead, rows, k = w8.shape
+ assert rows == 2 * inter and inter % 8 == 0
+ v = w8.reshape(*lead, 2, inter // 8, 8, k)
+ n = len(lead)
+ perm = tuple(range(n)) + tuple(x + n for x in (1, 0, 2, 3))
+ return v.permute(perm).reshape(*lead, rows, k)
+
+
+def swizzle_pair_interleaved(w8: torch.Tensor, inter: int) -> torch.Tensor:
+ return swizzle_fp8_contig8(pair_interleave(_u8(w8), inter))
+
+
+def quantize_fp8_block(
+ w: torch.Tensor, blk_k: int = BLK, blk_m: int = BLK
+) -> tuple[torch.Tensor, torch.Tensor]:
+ rows, k = w.shape
+ assert rows % blk_m == 0 and k % blk_k == 0
+ blocks = w.float().view(rows // blk_m, blk_m, k // blk_k, blk_k)
+ amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12)
+ scales = amax / FP8_MAX
+ q = (blocks / scales).to(torch.float8_e4m3fn)
+ return (q.reshape(rows, k).contiguous(), scales.view(rows // blk_m, k // blk_k).contiguous())
+
+
+def quantize_fp8_block_padded(
+ w: torch.Tensor, blk_k: int = BLK, blk_m: int = BLK
+) -> tuple[torch.Tensor, torch.Tensor]:
+ rows, k = w.shape
+ if rows % blk_m == 0:
+ return quantize_fp8_block(w, blk_k, blk_m)
+ n_blk = -(-rows // blk_m)
+ padded = torch.zeros(n_blk * blk_m, k, dtype=torch.float32)
+ padded[:rows] = w.float()
+ q, scales = quantize_fp8_block(padded, blk_k, blk_m)
+ return (q[:rows].contiguous(), scales)
+
+
+def dequant_fp8(
+ w8: torch.Tensor, scale_inv: torch.Tensor, blk_m: int = BLK, blk_k: int = BLK
+) -> torch.Tensor:
+ rows, k = w8.shape
+ s = scale_inv.float()
+ s = s.repeat_interleave(blk_m, dim=0)[:rows]
+ s = s.repeat_interleave(blk_k, dim=1)[:, :k]
+ return _u8(w8).view(torch.float8_e4m3fn).float() * s
+
+
+ATTN_SCALE_BLK_M = 64
+M4_SCALE_BLK_K = 64
+M6_SCALE_BLK_K = 128
+
+
+def attn_lossless_slices(
+ qb8: torch.Tensor,
+ qb_s: torch.Tensor,
+ kvb8: torch.Tensor,
+ kvb_s: torch.Tensor,
+ h0: int,
+ hv: int,
+ heads: int,
+ args: ModelArgsGlm52,
+) -> dict[str, torch.Tensor]:
+ nope, rope, vdim, kvr, qlr = (
+ args.qk_nope_head_dim,
+ args.qk_rope_head_dim,
+ args.v_head_dim,
+ args.kv_lora_rank,
+ args.q_lora_rank,
+ )
+ qk = nope + rope
+ kvd = nope + vdim
+ sm = ATTN_SCALE_BLK_M
+ qb8 = _u8(qb8).view(args.n_heads, qk, qlr)
+ kvb8 = _u8(kvb8).view(args.n_heads, kvd, kvr)
+ qb_s, kvb_s = (qb_s.float(), kvb_s.float())
+ assert qb_s.shape == (args.n_heads * qk // BLK, qlr // BLK), qb_s.shape
+ assert kvb_s.shape == (args.n_heads * kvd // BLK, kvr // BLK), kvb_s.shape
+ hh = torch.arange(hv) + h0
+ out: dict[str, torch.Tensor] = {}
+ wqb = torch.zeros(heads, qk, qlr, dtype=torch.uint8)
+ wqb[:hv] = qb8[h0 : h0 + hv]
+ out["wqb"] = (
+ torch.cat(
+ [
+ wqb[:, :nope, :].reshape(heads * nope, qlr),
+ wqb[:, nope:, :].reshape(heads * rope, qlr),
+ ]
+ )
+ .contiguous()
+ .view(torch.float8_e4m3fn)
+ )
+ qs = torch.ones(heads * qk // sm, qlr // BLK, dtype=torch.float32)
+ nope_stripes = nope // sm
+ for si in range(nope_stripes):
+ blk = (hh * qk + si * sm) // BLK
+ qs[torch.arange(hv) * nope_stripes + si] = qb_s[blk]
+ pe_base = heads * nope // sm
+ for si in range(rope // sm):
+ blk = (hh * qk + nope + si * sm) // BLK
+ qs[pe_base + torch.arange(hv) * (rope // sm) + si] = qb_s[blk]
+ out["wqb_scales"] = qs.contiguous()
+ wkvb1 = torch.zeros(heads, kvr, nope, dtype=torch.uint8)
+ wkvb1[:hv] = kvb8[h0 : h0 + hv, :nope, :].transpose(-1, -2)
+ out["wkvb1"] = wkvb1.reshape(heads * kvr, nope).contiguous().view(torch.float8_e4m3fn)
+ k1s = torch.ones(heads * kvr // sm, nope // M4_SCALE_BLK_K, dtype=torch.float32)
+ c_stripes = kvr // sm
+ for kb in range(nope // M4_SCALE_BLK_K):
+ rblk = (hh * kvd + kb * M4_SCALE_BLK_K) // BLK
+ for t in range(c_stripes):
+ cblk = t * sm // BLK
+ k1s[torch.arange(hv) * c_stripes + t, kb] = kvb_s[rblk, cblk]
+ out["wkvb1_scales"] = k1s.contiguous()
+ wkvb2 = torch.zeros(heads, vdim, kvr, dtype=torch.uint8)
+ wkvb2[:hv] = kvb8[h0 : h0 + hv, nope:, :]
+ out["wkvb2"] = wkvb2.reshape(heads * vdim, kvr).contiguous().view(torch.float8_e4m3fn)
+ k2s = torch.ones(heads * vdim // sm, kvr // M6_SCALE_BLK_K, dtype=torch.float32)
+ v_stripes = vdim // sm
+ for si in range(v_stripes):
+ rblk = (hh * kvd + nope + si * sm) // BLK
+ k2s[torch.arange(hv) * v_stripes + si] = kvb_s[rblk]
+ out["wkvb2_scales"] = k2s.contiguous()
+ return out
+
+
+class CheckpointReader:
+ """safetensors index + lazily opened shard handles, with a layer cache."""
+
+ def __init__(self, model_dir: str) -> None:
+ self.model_dir = model_dir
+ idx = os.path.join(model_dir, "model.safetensors.index.json")
+ with open(idx) as f:
+ self.weight_map: dict[str, str] = json.load(f)["weight_map"]
+ self._stack = ExitStack()
+ self._handles: dict[str, object] = {}
+ self._cache: dict[str, torch.Tensor] | None = None
+
+ def begin_layer(self) -> None:
+ self._cache = {}
+
+ def end_layer(self) -> None:
+ self._cache = None
+
+ def get(self, key: str) -> torch.Tensor:
+ if self._cache is not None and key in self._cache:
+ return self._cache[key]
+ shard = self.weight_map[key]
+ if shard not in self._handles:
+ self._handles[shard] = self._stack.enter_context(
+ safe_open(os.path.join(self.model_dir, shard), framework="pt")
+ )
+ t = self._handles[shard].get_tensor(key)
+ if self._cache is not None:
+ self._cache[key] = t
+ return t
+
+ def close(self) -> None:
+ self._stack.close()
+ self._handles.clear()
+ self._cache = None
+
+
+def _gamma(t: torch.Tensor) -> torch.Tensor:
+ return t.float().contiguous()
+
+
+def _head_slice(rank: int, args: ModelArgsGlm52) -> tuple[int, int]:
+ h0 = (rank - 1) * args.local_heads
+ return (h0, min(args.n_heads, h0 + args.local_heads) - h0)
+
+
+def shard_mla_layer(
+ rd: CheckpointReader,
+ pre: str,
+ rank: int,
+ args: ModelArgsGlm52,
+ work_dev: str,
+ tp8: bool = False,
+) -> dict[str, torch.Tensor]:
+ if tp8:
+ h0, hv, heads = (rank * TP8_HEADS, TP8_HEADS, TP8_HEADS)
+ else:
+ h0, hv = _head_slice(rank, args)
+ heads = args.local_heads
+ vdim = args.v_head_dim
+ out: dict[str, torch.Tensor] = {}
+ out["in_gamma"] = _gamma(rd.get(f"{pre}input_layernorm.weight"))
+ out["wqkva"] = torch.cat(
+ [
+ rd.get(f"{pre}self_attn.q_a_proj.weight"),
+ rd.get(f"{pre}self_attn.kv_a_proj_with_mqa.weight"),
+ ]
+ )
+ out["wqkva_scales"] = torch.cat(
+ [
+ rd.get(f"{pre}self_attn.q_a_proj.weight_scale_inv").float(),
+ rd.get(f"{pre}self_attn.kv_a_proj_with_mqa.weight_scale_inv").float(),
+ ]
+ ).contiguous()
+ out["q_gamma"] = _gamma(rd.get(f"{pre}self_attn.q_a_layernorm.weight"))
+ out.update(
+ attn_lossless_slices(
+ rd.get(f"{pre}self_attn.q_b_proj.weight"),
+ rd.get(f"{pre}self_attn.q_b_proj.weight_scale_inv"),
+ rd.get(f"{pre}self_attn.kv_b_proj.weight"),
+ rd.get(f"{pre}self_attn.kv_b_proj.weight_scale_inv"),
+ h0,
+ hv,
+ heads,
+ args,
+ )
+ )
+ out["kv_gamma"] = _gamma(rd.get(f"{pre}self_attn.kv_a_layernorm.weight"))
+ wo8 = rd.get(f"{pre}self_attn.o_proj.weight")
+ wo_si = rd.get(f"{pre}self_attn.o_proj.weight_scale_inv").float()
+ wo = torch.zeros(args.dim, heads * vdim, dtype=torch.uint8)
+ wo[:, : hv * vdim] = _u8(wo8[:, h0 * vdim : (h0 + hv) * vdim])
+ out["wo"] = wo.view(torch.float8_e4m3fn)
+ wos = torch.ones(args.dim // BLK, heads * vdim // BLK, dtype=torch.float32)
+ wos[:, : hv * vdim // BLK] = wo_si[:, h0 * vdim // BLK : (h0 + hv) * vdim // BLK]
+ out["wo_scales"] = wos.contiguous()
+ return out
+
+
+def shard_indexer_layer(
+ rd: CheckpointReader, pre: str, args: ModelArgsGlm52
+) -> dict[str, torch.Tensor]:
+ out: dict[str, torch.Tensor] = {}
+ out["in_gamma"] = _gamma(rd.get(f"{pre}input_layernorm.weight"))
+ out["wqaki"] = torch.cat(
+ [rd.get(f"{pre}self_attn.q_a_proj.weight"), rd.get(f"{pre}self_attn.indexer.wk.weight")]
+ )
+ out["wqaki_scales"] = torch.cat(
+ [
+ rd.get(f"{pre}self_attn.q_a_proj.weight_scale_inv").float(),
+ rd.get(f"{pre}self_attn.indexer.wk.weight_scale_inv").float(),
+ ]
+ ).contiguous()
+ out["wis"] = (
+ rd.get(f"{pre}self_attn.indexer.weights_proj.weight").to(torch.bfloat16).contiguous()
+ )
+ out["q_gamma"] = _gamma(rd.get(f"{pre}self_attn.q_a_layernorm.weight"))
+ out["wqi"] = rd.get(f"{pre}self_attn.indexer.wq_b.weight")
+ out["wqi_scales"] = rd.get(f"{pre}self_attn.indexer.wq_b.weight_scale_inv").float().contiguous()
+ out["knorm_w"] = _gamma(rd.get(f"{pre}self_attn.indexer.k_norm.weight"))
+ out["knorm_b"] = _gamma(rd.get(f"{pre}self_attn.indexer.k_norm.bias"))
+ return out
+
+
+def shard_dense_ffn(
+ rd: CheckpointReader, pre: str, rank: int, args: ModelArgsGlm52
+) -> dict[str, torch.Tensor]:
+ inter = args.dense_inter_shard
+ r0, sb, nb = (rank * inter, rank * inter // BLK, inter // BLK)
+ out: dict[str, torch.Tensor] = {}
+ out["post_gamma"] = _gamma(rd.get(f"{pre}post_attention_layernorm.weight"))
+ out["wug"] = torch.cat(
+ [
+ rd.get(f"{pre}mlp.gate_proj.weight")[r0 : r0 + inter],
+ rd.get(f"{pre}mlp.up_proj.weight")[r0 : r0 + inter],
+ ]
+ )
+ out["wug_scales"] = torch.cat(
+ [
+ rd.get(f"{pre}mlp.gate_proj.weight_scale_inv").float()[sb : sb + nb],
+ rd.get(f"{pre}mlp.up_proj.weight_scale_inv").float()[sb : sb + nb],
+ ]
+ ).contiguous()
+ out["wdown"] = rd.get(f"{pre}mlp.down_proj.weight")[:, r0 : r0 + inter].contiguous()
+ out["wdown_scales"] = (
+ rd.get(f"{pre}mlp.down_proj.weight_scale_inv").float()[:, sb : sb + nb].contiguous()
+ )
+ return out
+
+
+def shard_moe_ffn(
+ rd: CheckpointReader, pre: str, rank: int, args: ModelArgsGlm52
+) -> dict[str, torch.Tensor]:
+ inter = args.moe_inter_shard
+ r0, sb, nb = (rank * inter, rank * inter // BLK, inter // BLK)
+
+ def ep(e: int) -> str:
+ return f"{pre}mlp.shared_experts." if e == 0 else f"{pre}mlp.experts.{e - 1}."
+
+ out: dict[str, torch.Tensor] = {}
+ out["post_gamma"] = _gamma(rd.get(f"{pre}post_attention_layernorm.weight"))
+ out["router"] = rd.get(f"{pre}mlp.gate.weight").to(torch.bfloat16).contiguous()
+ out["moe_bias"] = _gamma(rd.get(f"{pre}mlp.gate.e_score_correction_bias"))
+ ug = torch.empty(N_BANK, 2 * inter, args.dim, dtype=torch.uint8)
+ ug_sc = torch.empty(N_BANK, 2 * nb, args.dim // BLK, dtype=torch.float32)
+ dn = torch.empty(N_BANK, args.dim, inter, dtype=torch.uint8)
+ dn_sc = torch.empty(N_BANK, args.dim // BLK, nb, dtype=torch.float32)
+ for e in range(N_BANK):
+ p = ep(e)
+ ug[e, :inter] = _u8(rd.get(f"{p}gate_proj.weight")[r0 : r0 + inter])
+ ug[e, inter:] = _u8(rd.get(f"{p}up_proj.weight")[r0 : r0 + inter])
+ ug_sc[e, :nb] = rd.get(f"{p}gate_proj.weight_scale_inv").float()[sb : sb + nb]
+ ug_sc[e, nb:] = rd.get(f"{p}up_proj.weight_scale_inv").float()[sb : sb + nb]
+ dn[e] = _u8(rd.get(f"{p}down_proj.weight")[:, r0 : r0 + inter])
+ dn_sc[e] = rd.get(f"{p}down_proj.weight_scale_inv").float()[:, sb : sb + nb]
+ out["moe_ug"] = ug.view(torch.float8_e4m3fn)
+ out["moe_ug_scales"] = ug_sc
+ out["moe_down"] = dn.view(torch.float8_e4m3fn)
+ out["moe_down_scales"] = dn_sc
+ return out
+
+
+def shard_layer(
+ rd: CheckpointReader,
+ i: int,
+ rank: int,
+ args: ModelArgsGlm52,
+ work_dev: str,
+ attn_tp8: bool = False,
+) -> dict[str, torch.Tensor]:
+ pre = f"model.layers.{i}."
+ kind = layer_kind(i)
+ out: dict[str, torch.Tensor] = {}
+ if rank == 0:
+ if kind != KIND_MOE_SHARED:
+ out.update(shard_indexer_layer(rd, pre, args))
+ else:
+ out.update(shard_mla_layer(rd, pre, rank, args, work_dev))
+ if attn_tp8 and kind == KIND_MOE_SHARED:
+ for k, v in shard_mla_layer(rd, pre, rank, args, work_dev, tp8=True).items():
+ out[f"attn_tp8.{k}"] = v
+ if kind == KIND_DENSE:
+ out.update(shard_dense_ffn(rd, pre, rank, args))
+ else:
+ out.update(shard_moe_ffn(rd, pre, rank, args))
+ return out
+
+
+def shard_tail(rd: CheckpointReader, rank: int, args: ModelArgsGlm52) -> dict[str, torch.Tensor]:
+ vs = args.vocab_shard
+ return {
+ "final_gamma": _gamma(rd.get("model.norm.weight")),
+ "head": (
+ rd.get("lm_head.weight")[rank * vs : (rank + 1) * vs].to(torch.bfloat16).contiguous()
+ ),
+ }
+
+
+def shard_mtp(
+ rd: CheckpointReader, rank: int, args: ModelArgsGlm52, work_dev: str
+) -> dict[str, torch.Tensor]:
+ i = args.n_layers
+ pre = f"model.layers.{i}."
+ inter = args.dim // 4
+ half, sl = (1, rank - 4) if rank >= 4 else (0, rank)
+ c0 = half * args.dim + sl * inter
+ eh = rd.get(f"{pre}eh_proj.weight").to(torch.bfloat16)
+ out: dict[str, torch.Tensor] = {
+ "mtp.eh_w": eh[:, c0 : c0 + inter].contiguous(),
+ "mtp.e_gamma": _gamma(rd.get(f"{pre}enorm.weight")),
+ "mtp.h_gamma": _gamma(rd.get(f"{pre}hnorm.weight")),
+ "mtp.head_gamma": _gamma(rd.get(f"{pre}shared_head.norm.weight")),
+ }
+ for k, v in shard_layer(rd, i, rank, args, work_dev).items():
+ out[f"mtp_layer.{k}"] = v
+ return out
+
+
+class _ShardWriter:
+ """Accumulate tensors, flush ~shard_bytes files, write an index at close."""
+
+ def __init__(self, out_dir: str, shard_bytes: int = 8 << 30) -> None:
+ os.makedirs(out_dir, exist_ok=True)
+ self.out_dir = out_dir
+ self.shard_bytes = shard_bytes
+ self.pending: dict[str, torch.Tensor] = {}
+ self.pending_bytes = 0
+ self.files: list[dict[str, torch.Tensor]] = []
+ self.weight_map: dict[str, str] = {}
+ self.n_files = 0
+ self.total = 0
+
+ def add(self, name: str, t: torch.Tensor) -> None:
+ self.pending[name] = t
+ self.pending_bytes += t.numel() * t.element_size()
+ if self.pending_bytes >= self.shard_bytes:
+ self.flush()
+
+ def flush(self) -> None:
+ if not self.pending:
+ return
+ self.n_files += 1
+ fname = f"model-{self.n_files:05d}.safetensors"
+ save_file(self.pending, os.path.join(self.out_dir, fname))
+ for k, v in self.pending.items():
+ self.weight_map[k] = fname
+ self.total += v.numel() * v.element_size()
+ self.pending = {}
+ self.pending_bytes = 0
+
+ def close(self) -> None:
+ self.flush()
+ index = {"metadata": {"total_size": self.total}, "weight_map": self.weight_map}
+ with open(os.path.join(self.out_dir, "model.safetensors.index.json"), "w") as f:
+ json.dump(index, f, indent=1)
+
+
+class Glm52WeightConverter:
+ """Offline stage-1 driver: HF checkpoint -> save_dir/{shared,rank0..7}."""
+
+ _PASSTHROUGH = (
+ "tokenizer.json",
+ "tokenizer_config.json",
+ "special_tokens_map.json",
+ "chat_template.jinja",
+ "config.json",
+ "generation_config.json",
+ )
+
+ def __init__(
+ self,
+ model_dir: str,
+ save_dir: str,
+ args: ModelArgsGlm52 | None = None,
+ n_layers: int | None = None,
+ work_dev: str = "cpu",
+ num_mtp: int | None = None,
+ attn_tp8: bool = True,
+ ) -> None:
+ self.args = args or ModelArgsGlm52()
+ self.model_dir = model_dir
+ self.save_dir = save_dir
+ self.n_layers = self.args.n_layers if n_layers is None else n_layers
+ self.work_dev = work_dev
+ self.num_mtp = self.args.num_mtp if num_mtp is None else num_mtp
+ self.hf_config = load_hf_config(model_dir)
+ validate_hf_config(self.hf_config, self.args, n_layers=self.n_layers, num_mtp=self.num_mtp)
+ self.rd = CheckpointReader(model_dir)
+ self.attn_tp8 = attn_tp8
+
+ def convert(self) -> None:
+ args = self.args
+ writers = [
+ _ShardWriter(os.path.join(self.save_dir, f"rank{r}")) for r in range(args.num_devices)
+ ]
+ for i in range(self.n_layers):
+ t0 = time.time()
+ self.rd.begin_layer()
+ for r in range(args.num_devices):
+ shard = shard_layer(self.rd, i, r, args, self.work_dev, attn_tp8=self.attn_tp8)
+ for k, v in shard.items():
+ writers[r].add(f"layer_{i}.{k}", v)
+ self.rd.end_layer()
+ logger.info("converted layer %d/%d (%.1fs)", i + 1, self.n_layers, time.time() - t0)
+ self.rd.begin_layer()
+ for r in range(args.num_devices):
+ for k, v in shard_tail(self.rd, r, args).items():
+ writers[r].add(k, v)
+ if self.num_mtp > 0:
+ for k, v in shard_mtp(self.rd, r, args, self.work_dev).items():
+ writers[r].add(k, v)
+ writers[r].close()
+ self.rd.end_layer()
+ if self.num_mtp > 0:
+ logger.info("converted the MTP module (layer %d)", args.n_layers)
+ shared = os.path.join(self.save_dir, "shared")
+ os.makedirs(shared, exist_ok=True)
+ save_file(
+ {"embed": self.rd.get("model.embed_tokens.weight").to(torch.bfloat16).contiguous()},
+ os.path.join(shared, "embed.safetensors"),
+ )
+ for name in self._PASSTHROUGH:
+ src = os.path.join(self.model_dir, name)
+ if os.path.isfile(src):
+ shutil.copy2(src, os.path.join(shared, name))
+ with open(os.path.join(self.save_dir, "tilert_meta.json"), "w") as f:
+ json.dump(
+ {
+ "model": "glm_5_2",
+ "n_layers": self.n_layers,
+ "num_mtp": self.num_mtp,
+ "tp": 8,
+ "attn_tp8": self.attn_tp8,
+ "attn_fp8_lossless": True,
+ "converter_version": CONVERTER_VERSION,
+ "source_model_dir": os.path.abspath(self.model_dir),
+ "source_config_sha256": sha256_file(
+ os.path.join(self.model_dir, "config.json")
+ ),
+ "source_index_sha256": sha256_file(
+ os.path.join(self.model_dir, "model.safetensors.index.json")
+ ),
+ "kv_dtype": "bf16",
+ "weight_format": "fp8_e4m3_block128; attn 64-row scale stripes",
+ "hf_config": describe_hf_config(self.hf_config),
+ },
+ f,
+ indent=1,
+ )
+ self.rd.close()
+ logger.info(f"device-sharded weights written to {self.save_dir}")
+
+
+def _write_json_atomic(path: str, obj: dict) -> None:
+ tmp = path + ".tmp"
+ with open(tmp, "w") as f:
+ json.dump(obj, f, indent=1)
+ os.replace(tmp, path)
+
+
+def _provenance_fields(model_dir: str) -> dict[str, str]:
+ return {
+ "source_model_dir": os.path.abspath(model_dir),
+ "source_config_sha256": sha256_file(os.path.join(model_dir, "config.json")),
+ "source_index_sha256": sha256_file(os.path.join(model_dir, "model.safetensors.index.json")),
+ }
+
+
+def _stored_tensor(save_dir: str, rank: int, key: str) -> torch.Tensor:
+ rank_dir = os.path.join(save_dir, f"rank{rank}")
+ with open(os.path.join(rank_dir, "model.safetensors.index.json")) as f:
+ wm = json.load(f)["weight_map"]
+ if key not in wm:
+ raise KeyError(f"rank{rank}: {key} is not in the conversion's index")
+ with safe_open(os.path.join(rank_dir, wm[key]), framework="pt") as f:
+ return f.get_tensor(key)
+
+
+def bind_source_checkpoint(
+ rd: CheckpointReader,
+ model_dir: str,
+ save_dir: str,
+ meta: dict,
+ args: ModelArgsGlm52 | None = None,
+) -> dict[str, str]:
+ args = args or ModelArgsGlm52()
+ prov = _provenance_fields(model_dir)
+ for key in ("source_config_sha256", "source_index_sha256"):
+ if meta.get(key) and meta[key] != prov[key]:
+ raise ValueError(
+ f"{save_dir} was cut from a checkpoint whose {key[7:-7]} differs from {model_dir}'s ({key}: {meta[key][:12]}.. vs {prov[key][:12]}..); refusing to mix checkpoints"
+ )
+ shared_cfg = os.path.join(save_dir, "shared", "config.json")
+ if os.path.isfile(shared_cfg) and sha256_file(shared_cfg) != prov["source_config_sha256"]:
+ raise ValueError(
+ f"{save_dir}/shared/config.json differs from {model_dir}/config.json; refusing to mix checkpoints"
+ )
+
+ def same(a: torch.Tensor, b: torch.Tensor) -> bool:
+ if a.dtype != b.dtype or tuple(a.shape) != tuple(b.shape):
+ return False
+ return torch.equal(a.contiguous(), b.contiguous())
+
+ rd.begin_layer()
+ try:
+ for r in (0, args.num_devices - 1):
+ want = shard_tail(rd, r, args)
+ for k, v in want.items():
+ got = _stored_tensor(save_dir, r, k)
+ if not same(got, v):
+ raise ValueError(
+ f"{save_dir} rank{r} {k} != the slice of {model_dir} ({got.dtype}{tuple(got.shape)} vs {v.dtype}{tuple(v.shape)}): the conversion was cut from a different checkpoint; refusing to mix them"
+ )
+ with safe_open(os.path.join(save_dir, "shared", "embed.safetensors"), framework="pt") as f:
+ emb = f.get_tensor("embed")
+ if not same(emb, rd.get("model.embed_tokens.weight").to(torch.bfloat16)):
+ raise ValueError(
+ f"{save_dir}/shared/embed.safetensors != {model_dir}'s embed_tokens: the conversion was cut from a different checkpoint"
+ )
+ finally:
+ rd.end_layer()
+ logger.info("%s is bound to %s (tail + embed bit-exact)", save_dir, model_dir)
+ return prov
+
+
+def stamp_provenance(model_dir: str, save_dir: str, args: ModelArgsGlm52 | None = None) -> None:
+ args = args or ModelArgsGlm52()
+ meta_path = os.path.join(save_dir, "tilert_meta.json")
+ with open(meta_path) as f:
+ meta = json.load(f)
+ if not meta.get("attn_fp8_lossless"):
+ raise RuntimeError(
+ f"{save_dir} still holds re-quantized attention tensors; run --augment-attn-lossless first (converter format {CONVERTER_VERSION} is the lossless one)"
+ )
+ hf_cfg = load_hf_config(model_dir)
+ validate_hf_config(
+ hf_cfg,
+ args,
+ n_layers=int(meta.get("n_layers", args.n_layers)),
+ num_mtp=int(meta.get("num_mtp", 0)),
+ max_seq_len=1,
+ )
+ rd = CheckpointReader(model_dir)
+ try:
+ prov = bind_source_checkpoint(rd, model_dir, save_dir, meta, args)
+ finally:
+ rd.close()
+ meta.update(prov)
+ meta.setdefault("model", "glm_5_2")
+ meta.setdefault("tp", 8)
+ meta["converter_version"] = CONVERTER_VERSION
+ meta.setdefault("kv_dtype", "bf16")
+ meta.setdefault("weight_format", "fp8_e4m3_block128; attn 64-row scale stripes")
+ meta["hf_config"] = describe_hf_config(hf_cfg)
+ _write_json_atomic(meta_path, meta)
+ logger.info("provenance stamped: %s <- %s", save_dir, model_dir)
+
+
+def augment_attn_tp8(
+ model_dir: str,
+ save_dir: str,
+ args: ModelArgsGlm52 | None = None,
+ work_dev: str = "cpu",
+ store_dir: str | None = None,
+) -> None:
+ args = args or ModelArgsGlm52()
+ meta_path = os.path.join(save_dir, "tilert_meta.json")
+ with open(meta_path) as f:
+ meta = json.load(f)
+ if meta.get("attn_tp8"):
+ logger.info("%s already carries the attn_tp8 set; nothing to do", save_dir)
+ return
+ n_layers = meta.get("n_layers", args.n_layers)
+ shared = [i for i in range(n_layers) if layer_kind(i) == KIND_MOE_SHARED]
+ indexes = []
+ for r in range(args.num_devices):
+ with open(os.path.join(save_dir, f"rank{r}", "model.safetensors.index.json")) as f:
+ indexes.append(json.load(f))
+ rd = CheckpointReader(model_dir)
+ prov = bind_source_checkpoint(rd, model_dir, save_dir, meta, args)
+ pending: list[dict[str, torch.Tensor]] = [{} for _ in range(args.num_devices)]
+ for n, i in enumerate(shared):
+ rd.begin_layer()
+ for r in range(args.num_devices):
+ tp8 = shard_mla_layer(rd, f"model.layers.{i}.", r, args, work_dev, tp8=True)
+ for k, v in tp8.items():
+ pending[r][f"layer_{i}.attn_tp8.{k}"] = v
+ rd.end_layer()
+ logger.info("augmented layer %d (%d/%d shared)", i, n + 1, len(shared))
+ rd.close()
+ fname = "attn-tp8-00001.safetensors"
+ for r in range(args.num_devices):
+ rank_dir = os.path.join(save_dir, f"rank{r}")
+ if store_dir is not None:
+ real_dir = os.path.join(store_dir, f"rank{r}")
+ os.makedirs(real_dir, exist_ok=True)
+ real = os.path.abspath(os.path.join(real_dir, fname))
+ save_file(pending[r], real)
+ link = os.path.join(rank_dir, fname)
+ if os.path.islink(link) or os.path.exists(link):
+ os.remove(link)
+ os.symlink(real, link)
+ else:
+ save_file(pending[r], os.path.join(rank_dir, fname))
+ idx = indexes[r]
+ for k, v in pending[r].items():
+ if k not in idx["weight_map"]:
+ idx["metadata"]["total_size"] += v.numel() * v.element_size()
+ idx["weight_map"][k] = fname
+ _write_json_atomic(os.path.join(rank_dir, "model.safetensors.index.json"), idx)
+ logger.info("rank %d: %d attn_tp8 tensors appended", r, len(pending[r]))
+ meta["attn_tp8"] = True
+ meta.update(prov)
+ _write_json_atomic(meta_path, meta)
+ logger.info("attn_tp8 augment complete: %s", save_dir)
+
+
+ATTN_LOSSLESS_FILE = "attn-lossless-00001.safetensors"
+ATTN_LOSSLESS_KEYS = ("wqb", "wqb_scales", "wkvb1", "wkvb1_scales", "wkvb2", "wkvb2_scales")
+
+
+def _attn_lossless_targets(
+ meta: dict, args: ModelArgsGlm52
+) -> list[tuple[str, str, int, list[int], bool]]:
+ n_layers = meta.get("n_layers", args.n_layers)
+ tp7_ranks = list(range(1, args.num_devices))
+ all_ranks = list(range(args.num_devices))
+ out = []
+ for i in range(n_layers):
+ pre = f"model.layers.{i}."
+ out.append((pre, f"layer_{i}.", i, tp7_ranks, False))
+ if meta.get("attn_tp8") and layer_kind(i) == KIND_MOE_SHARED:
+ out.append((pre, f"layer_{i}.attn_tp8.", i, all_ranks, True))
+ if meta.get("num_mtp", 0) > 0:
+ i = args.n_layers
+ out.append((f"model.layers.{i}.", "mtp_layer.", i, tp7_ranks, False))
+ return out
+
+
+def _rank_geometry(r: int, tp8: bool, args: ModelArgsGlm52) -> tuple[int, int, int]:
+ if tp8:
+ return (r * TP8_HEADS, TP8_HEADS, TP8_HEADS)
+ h0, hv = _head_slice(r, args)
+ return (h0, hv, args.local_heads)
+
+
+def augment_attn_lossless(
+ model_dir: str, save_dir: str, args: ModelArgsGlm52 | None = None, store_dir: str | None = None
+) -> None:
+ args = args or ModelArgsGlm52()
+ meta_path = os.path.join(save_dir, "tilert_meta.json")
+ with open(meta_path) as f:
+ meta = json.load(f)
+ if meta.get("attn_fp8_lossless"):
+ logger.info("%s already holds lossless attention slices; nothing to do", save_dir)
+ return
+ targets = _attn_lossless_targets(meta, args)
+ indexes = []
+ for r in range(args.num_devices):
+ with open(os.path.join(save_dir, f"rank{r}", "model.safetensors.index.json")) as f:
+ indexes.append(json.load(f))
+ rd = CheckpointReader(model_dir)
+ prov = bind_source_checkpoint(rd, model_dir, save_dir, meta, args)
+ pending: list[dict[str, torch.Tensor]] = [{} for _ in range(args.num_devices)]
+ last_layer = None
+ for n, (pre, kp, i, ranks, tp8) in enumerate(targets):
+ if i != last_layer:
+ rd.end_layer()
+ rd.begin_layer()
+ last_layer = i
+ qb8 = rd.get(f"{pre}self_attn.q_b_proj.weight")
+ qb_s = rd.get(f"{pre}self_attn.q_b_proj.weight_scale_inv")
+ kvb8 = rd.get(f"{pre}self_attn.kv_b_proj.weight")
+ kvb_s = rd.get(f"{pre}self_attn.kv_b_proj.weight_scale_inv")
+ for r in ranks:
+ h0, hv, heads = _rank_geometry(r, tp8, args)
+ sl = attn_lossless_slices(qb8, qb_s, kvb8, kvb_s, h0, hv, heads, args)
+ for k in ATTN_LOSSLESS_KEYS:
+ key = f"{kp}{k}"
+ if key not in indexes[r]["weight_map"]:
+ raise KeyError(f"rank{r}: {key} is not in the conversion's index")
+ pending[r][key] = sl[k]
+ logger.info("lossless attention set %d/%d (%s)", n + 1, len(targets), kp)
+ rd.close()
+ for r in range(args.num_devices):
+ rank_dir = os.path.join(save_dir, f"rank{r}")
+ if store_dir is not None:
+ real_dir = os.path.join(store_dir, f"rank{r}")
+ os.makedirs(real_dir, exist_ok=True)
+ real = os.path.abspath(os.path.join(real_dir, ATTN_LOSSLESS_FILE))
+ save_file(pending[r], real)
+ link = os.path.join(rank_dir, ATTN_LOSSLESS_FILE)
+ if os.path.islink(link) or os.path.exists(link):
+ os.remove(link)
+ os.symlink(real, link)
+ else:
+ save_file(pending[r], os.path.join(rank_dir, ATTN_LOSSLESS_FILE))
+ idx = indexes[r]
+ for k in pending[r]:
+ idx["weight_map"][k] = ATTN_LOSSLESS_FILE
+ _write_json_atomic(os.path.join(rank_dir, "model.safetensors.index.json"), idx)
+ logger.info("rank %d: %d attention tensors repointed", r, len(pending[r]))
+ meta["attn_fp8_lossless"] = True
+ meta.update(prov)
+ _write_json_atomic(meta_path, meta)
+ logger.info("attn lossless augment complete: %s", save_dir)
+
+
+def verify_attn_lossless(
+ model_dir: str, save_dir: str, layers: list[int], args: ModelArgsGlm52 | None = None
+) -> None:
+ args = args or ModelArgsGlm52()
+ with open(os.path.join(save_dir, "tilert_meta.json")) as f:
+ meta = json.load(f)
+ if not meta.get("attn_fp8_lossless"):
+ raise RuntimeError(f"{save_dir} has no lossless attention set to verify")
+ nope, rope, vdim, kvr = (
+ args.qk_nope_head_dim,
+ args.qk_rope_head_dim,
+ args.v_head_dim,
+ args.kv_lora_rank,
+ )
+ rd = CheckpointReader(model_dir)
+ loaders: dict[int, tuple[str, dict]] = {}
+
+ def conv(r: int, key: str) -> torch.Tensor:
+ if r not in loaders:
+ rank_dir = os.path.join(save_dir, f"rank{r}")
+ with open(os.path.join(rank_dir, "model.safetensors.index.json")) as f:
+ loaders[r] = (rank_dir, json.load(f)["weight_map"])
+ rank_dir, wm = loaders[r]
+ with safe_open(os.path.join(rank_dir, wm[key]), framework="pt") as f:
+ return f.get_tensor(key)
+
+ n_checked = 0
+ for pre, kp, i, ranks, tp8 in _attn_lossless_targets(meta, args):
+ if i not in layers:
+ continue
+ rd.begin_layer()
+ qb = dequant_fp8(
+ rd.get(f"{pre}self_attn.q_b_proj.weight"),
+ rd.get(f"{pre}self_attn.q_b_proj.weight_scale_inv"),
+ ).view(args.n_heads, nope + rope, -1)
+ kvb = dequant_fp8(
+ rd.get(f"{pre}self_attn.kv_b_proj.weight"),
+ rd.get(f"{pre}self_attn.kv_b_proj.weight_scale_inv"),
+ ).view(args.n_heads, nope + vdim, kvr)
+ for r in ranks:
+ h0, hv, heads = _rank_geometry(r, tp8, args)
+ got = dequant_fp8(conv(r, f"{kp}wqb"), conv(r, f"{kp}wqb_scales"), ATTN_SCALE_BLK_M)
+ ref = torch.zeros(heads, nope + rope, qb.shape[-1])
+ ref[:hv] = qb[h0 : h0 + hv]
+ ref = torch.cat(
+ [ref[:, :nope].reshape(heads * nope, -1), ref[:, nope:].reshape(heads * rope, -1)]
+ )
+ assert torch.equal(got, ref), f"{kp}wqb rank{r} differs from the checkpoint"
+ got = dequant_fp8(
+ conv(r, f"{kp}wkvb1"),
+ conv(r, f"{kp}wkvb1_scales"),
+ ATTN_SCALE_BLK_M,
+ M4_SCALE_BLK_K,
+ )
+ ref = torch.zeros(heads, kvr, nope)
+ ref[:hv] = kvb[h0 : h0 + hv, :nope, :].transpose(-1, -2)
+ assert torch.equal(got, ref.reshape(heads * kvr, nope)), f"{kp}wkvb1 rank{r}"
+ got = dequant_fp8(
+ conv(r, f"{kp}wkvb2"),
+ conv(r, f"{kp}wkvb2_scales"),
+ ATTN_SCALE_BLK_M,
+ M6_SCALE_BLK_K,
+ )
+ ref = torch.zeros(heads, vdim, kvr)
+ ref[:hv] = kvb[h0 : h0 + hv, nope:, :]
+ assert torch.equal(got, ref.reshape(heads * vdim, kvr)), f"{kp}wkvb2 rank{r}"
+ n_checked += 3
+ rd.end_layer()
+ logger.info("verified %s: bit-exact against the checkpoint", kp)
+ rd.close()
+ if n_checked == 0:
+ raise RuntimeError(f"no attention set matched layers {layers}")
+ logger.info("verify_attn_lossless: %d tensors bit-exact", n_checked)
+
+
+Getter = Callable[[str], torch.Tensor]
+
+
+def verify_shards_against_checkpoint(
+ model_dir: str,
+ save_dir: str,
+ layers: list[int] | None = None,
+ args: ModelArgsGlm52 | None = None,
+) -> None:
+ args = args or ModelArgsGlm52()
+ with open(os.path.join(save_dir, "tilert_meta.json")) as f:
+ meta = json.load(f)
+ n_layers = int(meta.get("n_layers", args.n_layers))
+ attn_tp8 = bool(meta.get("attn_tp8", False))
+ rd = CheckpointReader(model_dir)
+ loaders: list[tuple[str, dict[str, str]]] = []
+ for r in range(args.num_devices):
+ rank_dir = os.path.join(save_dir, f"rank{r}")
+ with open(os.path.join(rank_dir, "model.safetensors.index.json")) as f:
+ loaders.append((rank_dir, json.load(f)["weight_map"]))
+ handles: dict[str, object] = {}
+
+ def stored(r: int, key: str) -> torch.Tensor:
+ rank_dir, wm = loaders[r]
+ if key not in wm:
+ raise RuntimeError(f"rank{r}: {key} missing from the index")
+ path = os.path.join(rank_dir, wm[key])
+ if path not in handles:
+ handles[path] = safe_open(path, framework="pt").__enter__()
+ return handles[path].get_tensor(key)
+
+ def same(a: torch.Tensor, b: torch.Tensor) -> bool:
+ if a.dtype != b.dtype or tuple(a.shape) != tuple(b.shape):
+ return False
+ if a.dtype == torch.float8_e4m3fn:
+ a, b = (a.view(torch.uint8), b.view(torch.uint8))
+ return torch.equal(a.contiguous(), b.contiguous())
+
+ def check(r: int, d: dict[str, torch.Tensor], tag: str) -> int:
+ n = 0
+ for k, want in d.items():
+ got = stored(r, k)
+ if not same(got, want):
+ raise RuntimeError(
+ f"rank{r} {k}: stored {got.dtype}{tuple(got.shape)} != checkpoint slice {want.dtype}{tuple(want.shape)} ({tag})"
+ )
+ n += 1
+ return n
+
+ if layers is None:
+ layers = list(range(n_layers))
+ if meta.get("num_mtp", 0) > 0:
+ layers.append(args.n_layers)
+ layers += [-1]
+ total = 0
+ t0 = time.time()
+ for i in layers:
+ for r in range(args.num_devices):
+ if i == -1:
+ total += check(r, shard_tail(rd, r, args), "tail")
+ elif i == args.n_layers:
+ total += check(r, shard_mtp(rd, r, args, "cpu"), "mtp")
+ else:
+ d = shard_layer(rd, i, r, args, "cpu", attn_tp8=attn_tp8)
+ total += check(r, {f"layer_{i}.{k}": v for k, v in d.items()}, f"layer {i}")
+ logger.info(
+ "verify_shards_against_checkpoint: %s ok (%d tensors so far, %.0f s)",
+ "tail" if i == -1 else f"layer {i}",
+ total,
+ time.time() - t0,
+ )
+ if -1 in layers:
+ with safe_open(os.path.join(save_dir, "shared", "embed.safetensors"), framework="pt") as f:
+ emb = f.get_tensor("embed")
+ want = rd.get("model.embed_tokens.weight").to(torch.bfloat16)
+ if not same(emb, want):
+ raise RuntimeError("shared/embed.safetensors != checkpoint embed_tokens")
+ total += 1
+ for h in handles.values():
+ h.__exit__(None, None, None)
+ logger.info(
+ "verify_shards_against_checkpoint %s: %d tensors bit-exact vs %s",
+ save_dir,
+ total,
+ model_dir,
+ )
+
+
+def _st_header(path: str) -> tuple[int, dict]:
+ with open(path, "rb") as f:
+ n = int.from_bytes(f.read(8), "little")
+ hdr = json.loads(f.read(n))
+ return (8 + n, hdr)
+
+
+def prune_unreferenced(
+ save_dir: str,
+ args: ModelArgsGlm52 | None = None,
+ dry_run: bool = False,
+ chunk_bytes: int = 256 << 20,
+) -> None:
+ args = args or ModelArgsGlm52()
+ total_removed = 0
+ for r in range(args.num_devices):
+ rank_dir = os.path.join(save_dir, f"rank{r}")
+ idx_path = os.path.join(rank_dir, "model.safetensors.index.json")
+ with open(idx_path) as f:
+ idx = json.load(f)
+ wm: dict[str, str] = idx["weight_map"]
+ for fname in sorted(os.listdir(rank_dir)):
+ if not fname.endswith(".safetensors"):
+ continue
+ link = os.path.join(rank_dir, fname)
+ real = os.path.realpath(link)
+ base, hdr = _st_header(real)
+ meta_entry = hdr.pop("__metadata__", None)
+ stale = sorted(k for k in hdr if wm.get(k) != fname)
+ if not stale:
+ continue
+ keep = [k for k in hdr if k not in set(stale)]
+ removed = sum(hdr[k]["data_offsets"][1] - hdr[k]["data_offsets"][0] for k in stale)
+ total_removed += removed
+ logger.info(
+ "rank %d %s: %d unreferenced tensors (%.2f GB) to drop, %d kept%s",
+ r,
+ fname,
+ len(stale),
+ removed / 1000000000.0,
+ len(keep),
+ " [dry run]" if dry_run else "",
+ )
+ for k in stale[:6]:
+ logger.info(" - %s", k)
+ if dry_run:
+ continue
+ for k in keep:
+ if wm.get(k) != fname:
+ raise RuntimeError(f"{k} kept but not mapped to {fname}")
+ new_hdr: dict = {}
+ if meta_entry is not None:
+ new_hdr["__metadata__"] = meta_entry
+ off = 0
+ for k in keep:
+ e = hdr[k]
+ n = e["data_offsets"][1] - e["data_offsets"][0]
+ new_hdr[k] = {
+ "dtype": e["dtype"],
+ "shape": e["shape"],
+ "data_offsets": [off, off + n],
+ }
+ off += n
+ hb = json.dumps(new_hdr, separators=(",", ":")).encode()
+ hb += b" " * (-len(hb) % 8)
+ need = 8 + len(hb) + off
+ st = os.statvfs(os.path.dirname(real))
+ if st.f_bavail * st.f_frsize < need + (1 << 30):
+ free_gb = st.f_bavail * st.f_frsize / 1000000000.0
+ raise RuntimeError(
+ f"{os.path.dirname(real)}: {free_gb:.1f} GB free, need {need / 1000000000.0:.1f} GB to rewrite {fname}"
+ )
+ tmp = real + ".prune.tmp"
+ digests: dict[str, bytes] = {}
+ buf = bytearray(chunk_bytes)
+ mv = memoryview(buf)
+ src = os.open(real, os.O_RDONLY)
+ dst = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 384)
+ try:
+ os.write(dst, len(hb).to_bytes(8, "little") + hb)
+ for k in keep:
+ lo, hi = hdr[k]["data_offsets"]
+ h = hashlib.blake2b(digest_size=16)
+ pos = base + lo
+ while pos < base + hi:
+ want = min(chunk_bytes, base + hi - pos)
+ got = os.preadv(src, [mv[:want]], pos)
+ if got <= 0:
+ raise OSError(f"{fname}:{k}: short read at {pos}")
+ h.update(mv[:got])
+ w = 0
+ while w < got:
+ w += os.write(dst, mv[w:got])
+ pos += got
+ digests[k] = h.digest()
+ os.fsync(dst)
+ finally:
+ os.close(src)
+ os.close(dst)
+ nb, nh = _st_header(tmp)
+ nh.pop("__metadata__", None)
+ if list(nh) != keep or os.path.getsize(tmp) != need:
+ raise RuntimeError(f"{tmp}: header/size mismatch after rewrite")
+ src = os.open(tmp, os.O_RDONLY)
+ try:
+ for k in keep:
+ lo, hi = nh[k]["data_offsets"]
+ if nh[k]["shape"] != hdr[k]["shape"] or nh[k]["dtype"] != hdr[k]["dtype"]:
+ raise RuntimeError(f"{tmp}:{k}: entry mismatch")
+ h = hashlib.blake2b(digest_size=16)
+ pos = nb + lo
+ while pos < nb + hi:
+ want = min(chunk_bytes, nb + hi - pos)
+ got = os.preadv(src, [mv[:want]], pos)
+ if got <= 0:
+ raise OSError(f"{tmp}:{k}: short read at {pos}")
+ h.update(mv[:got])
+ pos += got
+ if h.digest() != digests[k]:
+ raise RuntimeError(f"{tmp}:{k}: blake2b mismatch after rewrite")
+ finally:
+ os.close(src)
+ os.replace(tmp, real)
+ dfd = os.open(os.path.dirname(real), os.O_RDONLY)
+ os.fsync(dfd)
+ os.close(dfd)
+ logger.info("rank %d %s: rewritten, %d tensors verified bit-exact", r, fname, len(keep))
+ if not dry_run:
+ tot = 0
+ for fname in sorted(set(wm.values())):
+ _, h = _st_header(os.path.realpath(os.path.join(rank_dir, fname)))
+ h.pop("__metadata__", None)
+ tot += sum(e["data_offsets"][1] - e["data_offsets"][0] for e in h.values())
+ idx.setdefault("metadata", {})["total_size"] = tot
+ _write_json_atomic(idx_path, idx)
+ logger.info(
+ "prune_unreferenced %s: %.2f GB of unreferenced tensors%s",
+ save_dir,
+ total_removed / 1000000000.0,
+ " would be removed" if dry_run else " removed",
+ )
+
+
+def audit_referenced(save_dir: str, args: ModelArgsGlm52 | None = None) -> None:
+ args = args or ModelArgsGlm52()
+ n = 0
+ for r in range(args.num_devices):
+ rank_dir = os.path.join(save_dir, f"rank{r}")
+ with open(os.path.join(rank_dir, "model.safetensors.index.json")) as f:
+ wm = json.load(f)["weight_map"]
+ seen: set[str] = set()
+ for fname in sorted(os.listdir(rank_dir)):
+ if not fname.endswith(".safetensors"):
+ continue
+ _, hdr = _st_header(os.path.realpath(os.path.join(rank_dir, fname)))
+ hdr.pop("__metadata__", None)
+ for k in hdr:
+ if wm.get(k) != fname:
+ raise RuntimeError(f"rank{r}/{fname}: unreferenced tensor {k}")
+ seen.add(k)
+ n += len(hdr)
+ missing = set(wm) - seen
+ if missing:
+ raise RuntimeError(f"rank{r}: {len(missing)} indexed tensors missing from shards")
+ logger.info("audit_referenced %s: %d tensors, all referenced, none stale", save_dir, n)
+
+
+def _pack_attn(g: Getter, p: str, rank: int, kind: int) -> list[torch.Tensor]:
+ if kind == KIND_MOE_SHARED:
+ q = f"{p}attn_tp8."
+ return [
+ g(f"{q}in_gamma"),
+ swizzle_fp8_contig8(g(f"{q}wqkva")),
+ g(f"{q}wqkva_scales"),
+ g(f"{q}q_gamma"),
+ swizzle_fp8_contig8(g(f"{q}wqb")),
+ g(f"{q}wqb_scales"),
+ g(f"{q}kv_gamma"),
+ swizzle_fp8_contig8(g(f"{q}wkvb1")),
+ g(f"{q}wkvb1_scales"),
+ swizzle_fp8_contig8(g(f"{q}wkvb2")),
+ g(f"{q}wkvb2_scales"),
+ _pack_wo(g(f"{q}wo")),
+ g(f"{q}wo_scales"),
+ ]
+ if rank == 0:
+ if kind == KIND_MOE_SHARED:
+ return []
+ return [
+ g(f"{p}in_gamma"),
+ swizzle_fp8_contig8(g(f"{p}wqaki")),
+ g(f"{p}wqaki_scales"),
+ swizzle_wis_8x64(g(f"{p}wis")),
+ g(f"{p}q_gamma"),
+ swizzle_fp8_contig8(g(f"{p}wqi")),
+ g(f"{p}wqi_scales"),
+ g(f"{p}knorm_w"),
+ g(f"{p}knorm_b"),
+ ]
+ return [
+ g(f"{p}in_gamma"),
+ swizzle_fp8_contig8(g(f"{p}wqkva")),
+ g(f"{p}wqkva_scales"),
+ g(f"{p}q_gamma"),
+ swizzle_fp8_contig8(g(f"{p}wqb")),
+ g(f"{p}wqb_scales"),
+ g(f"{p}kv_gamma"),
+ swizzle_fp8_contig8(g(f"{p}wkvb1")),
+ g(f"{p}wkvb1_scales"),
+ swizzle_fp8_contig8(g(f"{p}wkvb2")),
+ g(f"{p}wkvb2_scales"),
+ _pack_wo(g(f"{p}wo")),
+ g(f"{p}wo_scales"),
+ ]
+
+
+def swizzle_fp8_v4(w8: torch.Tensor) -> torch.Tensor:
+ e, dim, inter = w8.shape
+ assert dim % 24 == 0 and inter == 256
+ v = w8.view(torch.uint8).view(e, dim // 24, 6, 4, 8, 2, 16)
+ return v.permute(0, 1, 4, 5, 2, 3, 6).reshape(e, -1).contiguous()
+
+
+def swizzle_pair_interleaved_k128(w8: torch.Tensor, inter: int) -> torch.Tensor:
+ wp = pair_interleave(_u8(w8), inter)
+ *lead, rows, k = wp.shape
+ assert rows % 16 == 0 and k % 128 == 0
+ v = wp.view(*lead, rows // 16, 16, k // 128, 4, 2, 16)
+ n = len(lead)
+ perm = tuple(range(n)) + tuple(x + n for x in (0, 2, 4, 3, 1, 5))
+ return v.permute(perm).reshape(*lead, rows * k).contiguous()
+
+
+def swizzle_fp8_down_k128(w8: torch.Tensor) -> torch.Tensor:
+ e, dim, inter = w8.shape
+ assert dim % 24 == 0 and inter == 256
+ v = w8.view(torch.uint8).view(e, dim // 24, 24, 2, 4, 2, 16)
+ mains = v[:, :, :16].permute(0, 1, 3, 5, 4, 2, 6).reshape(e, dim // 24, 4096)
+ tails = v[:, :, 16:].permute(0, 1, 3, 5, 4, 2, 6).reshape(e, dim // 24, 2048)
+ return torch.cat([mains, tails], dim=-1).reshape(e, -1).contiguous()
+
+
+def _env_flag(name: str, dflt: bool = False) -> bool:
+ v = os.environ.get(name)
+ if v is None or v == "":
+ return dflt
+ if v in ("0", "1"):
+ return v == "1"
+ raise RuntimeError(
+ f'{name}: expected "0" or "1", got {v!r} (the lever vars are a loader/dispatch contract, so a value the two sides could read differently is refused)'
+ )
+
+
+def moe_w8a8_enabled() -> bool:
+ return _env_flag("TILERT_GLM5_MOE_W8A8", dflt=True)
+
+
+def add_moe_arm_cli(ap) -> None:
+ ap.add_argument(
+ "--w8a16moe",
+ action="store_true",
+ help="run the MoE chain on the w8a16 v4 monokernel (fp8 weights, bf16 activations) instead of the default W8A8 one",
+ )
+
+
+def apply_moe_arm(args) -> str:
+ w8a16 = bool(getattr(args, "w8a16moe", False))
+ os.environ["TILERT_GLM5_MOE_W8A8"] = "0" if w8a16 else "1"
+ return "w8a16" if w8a16 else "w8a8"
+
+
+def fp8_ki_enabled() -> bool:
+ return _env_flag("TILERT_GLM5_FP8_KI")
+
+
+def add_index_arm_cli(ap) -> None:
+ ap.add_argument(
+ "--fp8-ki",
+ action="store_true",
+ help="fp8 e4m3 ki cache + iq_rt with per-row scales; stage 1 on the fp8 MFMA kernel (default: bf16)",
+ )
+
+
+def apply_index_arm(args) -> str:
+ on = bool(getattr(args, "fp8_ki", False))
+ os.environ["TILERT_GLM5_FP8_KI"] = "1" if on else "0"
+ return "fp8-ki" if on else "bf16-ki"
+
+
+def fp8_kv_enabled() -> bool:
+ return _env_flag("TILERT_GLM5_FP8_KV")
+
+
+def add_kv_arm_cli(ap) -> None:
+ ap.add_argument(
+ "--fp8-kv",
+ action="store_true",
+ help="fp8 e4m3 kv latent cache with per-128-block scales (528 B rows), Q quantized in-kernel, fp8 MFMA score (default: bf16)",
+ )
+
+
+def apply_kv_arm(args) -> str:
+ on = bool(getattr(args, "fp8_kv", False))
+ os.environ["TILERT_GLM5_FP8_KV"] = "1" if on else "0"
+ return "fp8-kv" if on else "bf16-kv"
+
+
+def _pack_ffn(
+ g: Getter, p: str, kind: int, args: ModelArgsGlm52, moe_w8a8: bool = False
+) -> list[torch.Tensor]:
+ if kind == KIND_DENSE:
+ return [
+ g(f"{p}post_gamma"),
+ swizzle_pair_interleaved(g(f"{p}wug"), args.dense_inter_shard),
+ g(f"{p}wug_scales"),
+ _pack_wo(g(f"{p}wdown")),
+ g(f"{p}wdown_scales"),
+ ]
+ if moe_w8a8:
+ moe_ug = swizzle_pair_interleaved_k128(g(f"{p}moe_ug"), args.moe_inter_shard).reshape(-1)
+ else:
+ moe_ug = swizzle_pair_interleaved(g(f"{p}moe_ug"), args.moe_inter_shard).reshape(-1)
+ if moe_w8a8:
+ moe_down = swizzle_fp8_down_k128(g(f"{p}moe_down")).reshape(-1)
+ else:
+ moe_down = swizzle_fp8_v4(g(f"{p}moe_down")).reshape(-1)
+ return [
+ g(f"{p}post_gamma"),
+ swizzle_bf16_m4(g(f"{p}router")),
+ g(f"{p}moe_bias"),
+ moe_ug,
+ g(f"{p}moe_ug_scales"),
+ moe_down,
+ g(f"{p}moe_down_scales"),
+ ]
+
+
+def pack_rank_params(
+ get: Getter,
+ embed: torch.Tensor,
+ args: ModelArgsGlm52,
+ rank: int,
+ device: str,
+ n_layers: int,
+ num_mtp: int = 0,
+ freqs_cis: torch.Tensor | None = None,
+) -> list[torch.Tensor]:
+ if freqs_cis is None:
+ freqs_cis = make_freqs_cis(args.max_seq_len, theta=args.rope_theta, device="cpu")
+ if tuple(freqs_cis.shape) != (args.max_seq_len, 64) or freqs_cis.dtype != torch.float32:
+ raise ValueError(
+ f"freqs_cis must be [{args.max_seq_len}, 64] f32, got {tuple(freqs_cis.shape)} {freqs_cis.dtype}"
+ )
+
+ def g(name: str) -> torch.Tensor:
+ return get(name).to(device, non_blocking=True)
+
+ params: list[torch.Tensor] = [embed.to(torch.bfloat16).to(device), freqs_cis.to(device)]
+ w8a8 = moe_w8a8_enabled()
+ for i in range(n_layers):
+ kind = layer_kind(i)
+ params += _pack_attn(g, f"layer_{i}.", rank, kind)
+ params += _pack_ffn(g, f"layer_{i}.", kind, args, moe_w8a8=w8a8)
+ params.append(g("final_gamma"))
+ params.append(swizzle_bf16_16x32(g("head")).view(torch.uint8))
+ if num_mtp > 0:
+ params += [
+ swizzle_bf16_16x32(g("mtp.eh_w")).view(torch.uint8),
+ g("mtp.e_gamma"),
+ g("mtp.h_gamma"),
+ g("mtp.head_gamma"),
+ ]
+ params += _pack_attn(g, "mtp_layer.", rank, KIND_MOE_FULL)
+ params += _pack_ffn(g, "mtp_layer.", KIND_MOE_FULL, args, moe_w8a8=w8a8)
+ return params
+
+
+_ST_DTYPE = {
+ "BOOL": torch.bool,
+ "U8": torch.uint8,
+ "I8": torch.int8,
+ "I16": torch.int16,
+ "I32": torch.int32,
+ "I64": torch.int64,
+ "F16": torch.float16,
+ "BF16": torch.bfloat16,
+ "F32": torch.float32,
+ "F64": torch.float64,
+ "F8_E4M3": getattr(torch, "float8_e4m3fn", None),
+ "F8_E5M2": getattr(torch, "float8_e5m2", None),
+}
+
+
+class RankBlobLoader:
+ """Read one converted rank shard with pread() and land it on its GPU."""
+
+ def __init__(self, rank_dir: str, device: str, staging_bytes: int = 128 << 20) -> None:
+ with open(os.path.join(rank_dir, "model.safetensors.index.json")) as f:
+ self._map: dict[str, str] = json.load(f)["weight_map"]
+ self._dir = rank_dir
+ self._dev = device
+ self._fds: dict[str, int] = {}
+ self._hdr: dict[str, tuple[int, dict]] = {}
+ self._buf = [
+ torch.empty(staging_bytes, dtype=torch.uint8, pin_memory=True) for _ in range(2)
+ ]
+ self._mv = [memoryview(b.numpy()) for b in self._buf]
+ self._ev = [torch.cuda.Event() for _ in range(2)]
+ self._slot = 0
+
+ def _shard(self, name: str):
+ if name not in self._hdr:
+ path = os.path.join(self._dir, name)
+ with open(path, "rb") as f:
+ n = int.from_bytes(f.read(8), "little")
+ meta = json.loads(f.read(n))
+ self._hdr[name] = (8 + n, meta)
+ self._fds[name] = os.open(path, os.O_RDONLY)
+ return (self._hdr[name], self._fds[name])
+
+ def get(self, key: str) -> torch.Tensor:
+ (base, meta), fd = self._shard(self._map[key])
+ entry = meta[key]
+ lo, hi = entry["data_offsets"]
+ dt = _ST_DTYPE[entry["dtype"]]
+ if dt is None:
+ raise RuntimeError(f"{key}: this torch has no {entry['dtype']}")
+ nbytes = hi - lo
+ out = torch.empty(nbytes, dtype=torch.uint8, device=self._dev)
+ cap = self._buf[0].numel()
+ off = 0
+ while off < nbytes:
+ s = self._slot
+ self._ev[s].synchronize()
+ want = min(cap, nbytes - off)
+ got = os.preadv(fd, [self._mv[s][:want]], base + lo + off)
+ if got <= 0:
+ raise OSError(f"{key}: short read at {off} of {nbytes}")
+ out[off : off + got].copy_(self._buf[s][:got], non_blocking=True)
+ self._ev[s].record()
+ self._slot ^= 1
+ off += got
+ return out.view(dt).reshape(entry["shape"])
+
+ def close(self) -> None:
+ torch.cuda.synchronize(self._dev)
+ for fd in self._fds.values():
+ os.close(fd)
+ self._fds.clear()
+ self._buf.clear()
+ self._mv.clear()
+
+
+def load_rank_params(
+ save_dir: str,
+ args: ModelArgsGlm52,
+ rank: int,
+ device: str,
+ n_layers: int | None = None,
+ embed: torch.Tensor | None = None,
+ num_mtp: int | None = None,
+ freqs_cis: torch.Tensor | None = None,
+) -> list[torch.Tensor]:
+ n_layers = args.n_layers if n_layers is None else n_layers
+ num_mtp = args.num_mtp if num_mtp is None else num_mtp
+ rd = RankBlobLoader(os.path.join(save_dir, f"rank{rank}"), device)
+ if embed is None:
+ with safe_open(os.path.join(save_dir, "shared", "embed.safetensors"), framework="pt") as f:
+ embed = f.get_tensor("embed")
+ try:
+ return pack_rank_params(
+ rd.get, embed, args, rank, device, n_layers, num_mtp, freqs_cis=freqs_cis
+ )
+ finally:
+ rd.close()
+
+
+def random_rank_params(
+ args: ModelArgsGlm52,
+ rank: int,
+ device: str,
+ n_layers: int,
+ seed: int = 0,
+ num_mtp: int = 0,
+ freqs_cis: torch.Tensor | None = None,
+) -> list[torch.Tensor]:
+ gen = torch.Generator(device=device).manual_seed(seed)
+ heads, dim = (args.local_heads, args.dim)
+ cache: dict[str, torch.Tensor] = {}
+ dk = {"device": device, "generator": gen}
+
+ def rep_dk(name: str) -> dict:
+ h = zlib.crc32(name.encode()) & 4294967295
+ g = torch.Generator(device=device).manual_seed((seed * 1000003 + h) % 2**63)
+ return {"device": device, "generator": g}
+
+ def rep_gm(name: str, n: int) -> torch.Tensor:
+ return 1.0 + 0.1 * torch.randn(n, **rep_dk(name))
+
+ def rep_bf(name: str, rows: int, k: int) -> torch.Tensor:
+ return (torch.randn(rows, k, **rep_dk(name)) * k ** (-0.5)).to(torch.bfloat16)
+
+ def rq(
+ rows: int, k: int, blk_k: int = BLK, blk_m: int = BLK
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ w = torch.randn(rows, k, **dk) * k ** (-0.5)
+ return quantize_fp8_block_padded(w, blk_k, blk_m)
+
+ def gm(n: int) -> torch.Tensor:
+ return 1.0 + 0.1 * torch.randn(n, **dk)
+
+ def bf(rows: int, k: int) -> torch.Tensor:
+ return (torch.randn(rows, k, **dk) * k ** (-0.5)).to(torch.bfloat16)
+
+ def fill_layer(p: str, kind: int) -> None:
+ if rank == 0:
+ if kind != KIND_MOE_SHARED:
+ cache[f"{p}in_gamma"], cache[f"{p}q_gamma"] = (gm(dim), gm(2048))
+ cache[f"{p}wqaki"], cache[f"{p}wqaki_scales"] = rq(2176, dim)
+ cache[f"{p}wis"] = bf(32, dim)
+ cache[f"{p}wqi"], cache[f"{p}wqi_scales"] = rq(4096, 2048)
+ cache[f"{p}knorm_w"] = gm(128)
+ cache[f"{p}knorm_b"] = 0.1 * gm(128)
+ else:
+ cache[f"{p}in_gamma"], cache[f"{p}q_gamma"] = (gm(dim), gm(2048))
+ cache[f"{p}kv_gamma"] = gm(512)
+ cache[f"{p}wqkva"], cache[f"{p}wqkva_scales"] = rq(2624, dim)
+ sm = ATTN_SCALE_BLK_M
+ cache[f"{p}wqb"], cache[f"{p}wqb_scales"] = rq(heads * 256, 2048, BLK, sm)
+ cache[f"{p}wkvb1"], cache[f"{p}wkvb1_scales"] = rq(heads * 512, 192, 64, sm)
+ cache[f"{p}wkvb2"], cache[f"{p}wkvb2_scales"] = rq(heads * 256, 512, BLK, sm)
+ cache[f"{p}wo"], cache[f"{p}wo_scales"] = rq(dim, heads * 256)
+ if kind == KIND_DENSE:
+ cache[f"{p}post_gamma"] = rep_gm(f"{p}post_gamma", dim)
+ cache[f"{p}wug"], cache[f"{p}wug_scales"] = rq(2 * args.dense_inter_shard, dim)
+ cache[f"{p}wdown"], cache[f"{p}wdown_scales"] = rq(dim, args.dense_inter_shard)
+ else:
+ inter = args.moe_inter_shard
+ cache[f"{p}post_gamma"] = rep_gm(f"{p}post_gamma", dim)
+ cache[f"{p}router"] = rep_bf(f"{p}router", 256, dim)
+ cache[f"{p}moe_bias"] = 0.01 * torch.randn(256, **rep_dk(f"{p}moe_bias"))
+ ug8, ugs = rq(N_BANK * 2 * inter, dim)
+ cache[f"{p}moe_ug"] = ug8.view(N_BANK, 2 * inter, dim)
+ cache[f"{p}moe_ug_scales"] = ugs.view(N_BANK, 2 * inter // BLK, dim // BLK)
+ dn8, dns = rq(N_BANK * dim, inter)
+ cache[f"{p}moe_down"] = dn8.view(N_BANK, dim, inter)
+ cache[f"{p}moe_down_scales"] = dns.view(N_BANK, dim // BLK, inter // BLK)
+
+ def fill_attn_tp8(p: str, li: int) -> None:
+ g2 = torch.Generator(device=device).manual_seed(seed * 7919 + 977 * li + 13)
+ dk2 = {"device": device, "generator": g2}
+
+ def rq2(rows: int, k: int, blk_k: int = BLK, blk_m: int = BLK):
+ w = torch.randn(rows, k, **dk2) * k ** (-0.5)
+ return quantize_fp8_block_padded(w, blk_k, blk_m)
+
+ def gm2(n: int) -> torch.Tensor:
+ return 1.0 + 0.1 * torch.randn(n, **dk2)
+
+ q = f"{p}attn_tp8."
+ h8 = TP8_HEADS
+ cache[f"{q}in_gamma"], cache[f"{q}q_gamma"] = (gm2(dim), gm2(2048))
+ cache[f"{q}kv_gamma"] = gm2(512)
+ cache[f"{q}wqkva"], cache[f"{q}wqkva_scales"] = rq2(2624, dim)
+ sm = ATTN_SCALE_BLK_M
+ cache[f"{q}wqb"], cache[f"{q}wqb_scales"] = rq2(h8 * 256, 2048, BLK, sm)
+ cache[f"{q}wkvb1"], cache[f"{q}wkvb1_scales"] = rq2(h8 * 512, 192, 64, sm)
+ cache[f"{q}wkvb2"], cache[f"{q}wkvb2_scales"] = rq2(h8 * 256, 512, BLK, sm)
+ cache[f"{q}wo"], cache[f"{q}wo_scales"] = rq2(dim, h8 * 256)
+
+ for i in range(n_layers):
+ fill_layer(f"layer_{i}.", layer_kind(i))
+ if layer_kind(i) == KIND_MOE_SHARED:
+ fill_attn_tp8(f"layer_{i}.", i)
+ cache["final_gamma"] = rep_gm("final_gamma", dim)
+ cache["head"] = bf(args.vocab_shard, dim)
+ if num_mtp > 0:
+ cache["mtp.eh_w"] = bf(dim, dim // 4)
+ cache["mtp.e_gamma"] = rep_gm("mtp.e_gamma", dim)
+ cache["mtp.h_gamma"] = rep_gm("mtp.h_gamma", dim)
+ cache["mtp.head_gamma"] = rep_gm("mtp.head_gamma", dim)
+ fill_layer("mtp_layer.", KIND_MOE_FULL)
+ embed = torch.randn(args.vocab_size, dim, **rep_dk("embed")).to(torch.bfloat16)
+ return pack_rank_params(
+ lambda k: cache[k], embed, args, rank, device, n_layers, num_mtp, freqs_cis=freqs_cis
+ )
+
+
+def selftest_swizzles() -> None:
+ from tilert.models.glm_5_2_rocm.ops import moe_router
+ from tilert.models.glm_5_2_rocm.ops import rmsnorm_head_proj as h1
+ from tilert.models.glm_5_2_rocm.ops import rmsnorm_projq_wqb as m1
+ from tilert.models.glm_5_2_rocm.ops import rmsnorm_projx_wqakis as s0
+ from tilert.models.glm_5_2_rocm.ops import rmsnorm_projx_wqkva, upgate_silu
+
+ g = torch.Generator().manual_seed(0)
+ w = torch.randn(2624, 6144, generator=g)
+ q8, _ = rmsnorm_projx_wqkva.quantize_fp8_block(w)
+ ref = rmsnorm_projx_wqkva.swizzle_weights_contig(q8)
+ assert torch.equal(ref, swizzle_fp8_contig8(q8)), "fp8 contig-8 mismatch"
+ wr = torch.randn(256, 6144, generator=g).to(torch.bfloat16)
+ assert torch.equal(
+ moe_router.swizzle_router_bf16(wr), swizzle_bf16_16x32(wr)
+ ), "router bf16 mismatch"
+ from tilert.models.glm_5_2_rocm.ops import eh_proj_allreduce as t0
+ from tilert.models.glm_5_2_rocm.ops import moe_down_allreduce as mdown
+
+ assert torch.equal(
+ t0.swizzle_256_bf16(wr).view(torch.uint8), swizzle_bf16_m4(wr)
+ ), "router m4 bf16 mismatch"
+ wd = torch.randn(6144, 256, generator=g)
+ d8, _ = mdown.quantize_fp8_block(wd)
+ assert torch.equal(mdown.swizzle_m4(d8), swizzle_fp8_m4(d8)), "down m4 fp8 mismatch"
+ from tilert.models.glm_5_2_rocm.ops import unprojo_allreduce as m7
+
+ wo = torch.randn(6144, 2560, generator=g)
+ o8, _ = mdown.quantize_fp8_block(wo)
+ assert torch.equal(m7.swizzle_v2(o8), swizzle_fp8_v2(o8)), "wo v2 fp8 mismatch"
+ wh = torch.randn(19360, 6144, generator=g).to(torch.bfloat16)
+ assert torch.equal(h1.swizzle_head_bf16(wh), swizzle_bf16_16x32(wh)), "head bf16 mismatch"
+ wis = torch.randn(32, 6144, generator=g).to(torch.bfloat16)
+ assert torch.equal(s0.swizzle_wis_bf16(wis), swizzle_wis_8x64(wis)), "wis bf16 mismatch"
+ for inter in (256, 1536):
+ wu = torch.randn(2 * inter, 6144, generator=g)
+ u8, _ = upgate_silu.quantize_fp8_block(wu)
+ ref = upgate_silu.swizzle_pair_interleaved(u8, inter)
+ got = swizzle_pair_interleaved(u8, inter)
+ assert torch.equal(ref, got), f"pair-interleave({inter}) mismatch"
+ wq = torch.randn(2560, 2048, generator=g)
+ q8b, sref = quantize_fp8_block(wq, BLK, ATTN_SCALE_BLK_M)
+ q8m, sm = m1.quantize_fp8_block(wq)
+ assert torch.equal(q8b.view(torch.uint8), q8m.view(torch.uint8))
+ assert torch.equal(sref, sm)
+ print("selftest_swizzles: OK")
+
+
+if __name__ == "__main__":
+ import argparse
+
+ ap = argparse.ArgumentParser(
+ description="Convert HF GLM-5.2-FP8 weights to TileRT TP8 device shards."
+ )
+ ap.add_argument("--model_dir", type=str, help="HF checkpoint directory")
+ ap.add_argument("--save_dir", type=str, help="output directory")
+ ap.add_argument("--layers", type=int, default=None, help="convert first N layers")
+ ap.add_argument("--device", type=str, default="cpu", help="work device for requantization")
+ ap.add_argument(
+ "--num_mtp",
+ type=int,
+ default=None,
+ help="convert the MTP module too (default: model_args.num_mtp)",
+ )
+ ap.add_argument(
+ "--augment-attn-tp8",
+ action="store_true",
+ help="append the attn_tp8 set to an EXISTING conversion in save_dir (additive + idempotent; no full re-conversion)",
+ )
+ ap.add_argument(
+ "--store-dir",
+ type=str,
+ default=None,
+ help="augment only: write the new shards here and symlink them into the rank dirs (for a save_dir on a full filesystem)",
+ )
+ ap.add_argument(
+ "--augment-attn-lossless",
+ action="store_true",
+ help="replace the re-quantized attention tensors of an EXISTING conversion in save_dir with lossless checkpoint byte slices (in place, idempotent; no full re-conversion)",
+ )
+ ap.add_argument(
+ "--verify-attn-lossless",
+ type=str,
+ default=None,
+ metavar="LAYERS",
+ help="comma-separated layers: check the conversion's attention tensors dequantize bit-exactly to the checkpoint's",
+ )
+ ap.add_argument(
+ "--verify-shards",
+ type=str,
+ default=None,
+ metavar="LAYERS",
+ help="'all' or comma-separated layers (78 = MTP block, -1 = tail+embed): assert every stored tensor equals the converter's byte slice of the checkpoint (no re-quantization anywhere); needs --model_dir",
+ )
+ ap.add_argument(
+ "--prune-unreferenced",
+ action="store_true",
+ help="rewrite the shards in save_dir dropping every tensor the index no longer maps there (the superseded re-quantized Wq_b/Wkv_b after --augment-attn-lossless); every kept tensor is hash-verified",
+ )
+ ap.add_argument(
+ "--audit-referenced",
+ action="store_true",
+ help="fail unless every tensor in every shard of save_dir is the one its index maps there (no stale bytes)",
+ )
+ ap.add_argument("--dry-run", action="store_true", help="with --prune-unreferenced: report only")
+ ap.add_argument(
+ "--stamp-provenance",
+ action="store_true",
+ help="write the provenance / format metadata the loader requires into an older conversion's tilert_meta.json, after binding save_dir to --model_dir by content (tail + embed bit-exact); idempotent",
+ )
+ ap.add_argument("--selftest", action="store_true", help="run the swizzle selftest")
+ cli = ap.parse_args()
+ if cli.selftest:
+ selftest_swizzles()
+ elif cli.stamp_provenance:
+ assert cli.model_dir and cli.save_dir, "--model_dir and --save_dir required"
+ stamp_provenance(
+ model_dir=os.path.expanduser(cli.model_dir), save_dir=os.path.expanduser(cli.save_dir)
+ )
+ elif cli.verify_shards is not None:
+ assert cli.model_dir and cli.save_dir, "--model_dir and --save_dir required"
+ verify_shards_against_checkpoint(
+ model_dir=os.path.expanduser(cli.model_dir),
+ save_dir=os.path.expanduser(cli.save_dir),
+ layers=(
+ None
+ if cli.verify_shards == "all"
+ else [int(x) for x in cli.verify_shards.split(",")]
+ ),
+ )
+ elif cli.prune_unreferenced:
+ assert cli.save_dir, "--save_dir required"
+ prune_unreferenced(os.path.expanduser(cli.save_dir), dry_run=cli.dry_run)
+ elif cli.audit_referenced:
+ assert cli.save_dir, "--save_dir required"
+ audit_referenced(os.path.expanduser(cli.save_dir))
+ elif cli.augment_attn_lossless:
+ assert cli.model_dir and cli.save_dir, "--model_dir and --save_dir required"
+ augment_attn_lossless(
+ model_dir=os.path.expanduser(cli.model_dir),
+ save_dir=os.path.expanduser(cli.save_dir),
+ store_dir=os.path.expanduser(cli.store_dir) if cli.store_dir else None,
+ )
+ elif cli.verify_attn_lossless is not None:
+ assert cli.model_dir and cli.save_dir, "--model_dir and --save_dir required"
+ verify_attn_lossless(
+ model_dir=os.path.expanduser(cli.model_dir),
+ save_dir=os.path.expanduser(cli.save_dir),
+ layers=[int(x) for x in cli.verify_attn_lossless.split(",")],
+ )
+ elif cli.augment_attn_tp8:
+ assert cli.model_dir and cli.save_dir, "--model_dir and --save_dir required"
+ augment_attn_tp8(
+ model_dir=os.path.expanduser(cli.model_dir),
+ save_dir=os.path.expanduser(cli.save_dir),
+ work_dev=cli.device,
+ store_dir=os.path.expanduser(cli.store_dir) if cli.store_dir else None,
+ )
+ else:
+ assert cli.model_dir and cli.save_dir, "--model_dir and --save_dir required"
+ Glm52WeightConverter(
+ model_dir=os.path.expanduser(cli.model_dir),
+ save_dir=os.path.expanduser(cli.save_dir),
+ n_layers=cli.layers,
+ work_dev=cli.device,
+ num_mtp=cli.num_mtp,
+ attn_tp8=True,
+ ).convert()
diff --git a/tilert/models/misc_rocm/__init__.py b/tilert/models/misc_rocm/__init__.py
new file mode 100644
index 0000000..2ae2839
--- /dev/null
+++ b/tilert/models/misc_rocm/__init__.py
@@ -0,0 +1 @@
+pass
diff --git a/tilert/models/misc_rocm/ops/__init__.py b/tilert/models/misc_rocm/ops/__init__.py
new file mode 100644
index 0000000..2419d2d
--- /dev/null
+++ b/tilert/models/misc_rocm/ops/__init__.py
@@ -0,0 +1,3 @@
+from tilert.models.misc_rocm.ops.rmsnorm_raw import RMSNormRaw, rmsnorm_raw
+
+__all__ = ["RMSNormRaw", "rmsnorm_raw"]
diff --git a/tilert/models/misc_rocm/ops/gemv_w8a16_m4_cdna4.py b/tilert/models/misc_rocm/ops/gemv_w8a16_m4_cdna4.py
new file mode 100644
index 0000000..d8672a5
--- /dev/null
+++ b/tilert/models/misc_rocm/ops/gemv_w8a16_m4_cdna4.py
@@ -0,0 +1,48 @@
+"""gemv_w8a16_m4_cdna4 op wrapper: 4x4-MFMA-tile packing + forward."""
+
+import torch
+
+from tilert.models.misc_rocm.ops.gemv_w8a16_mfma import (
+ HIDDEN,
+ ROWS,
+ SCALE_BLK,
+ SUPPORTED_SEQS,
+ GemvW8A16Mfma,
+)
+
+ROWS_PER_BLOCK = 8
+WAVES = 8
+CHUNK = 128
+_N_RG = ROWS // ROWS_PER_BLOCK
+_N_KC = HIDDEN // CHUNK
+
+
+def swizzle_weights_m4(w_fp8: torch.Tensor) -> torch.Tensor:
+ assert w_fp8.shape == (ROWS, HIDDEN)
+ w8 = w_fp8.view(torch.uint8)
+ rg = torch.arange(_N_RG)
+ kc = torch.arange(_N_KC)
+ lane = torch.arange(64)
+ s = torch.arange(4)
+ i = torch.arange(4)
+ RG, KC, L, S, II = torch.meshgrid(rg, kc, lane, s, i, indexing="ij")
+ rows = RG * ROWS_PER_BLOCK + (L >> 5) * 4 + (L & 3)
+ ks = KC * CHUNK + S * 32 + (L >> 2 & 7) * 4 + II
+ return w8[rows, ks].reshape(-1).contiguous()
+
+
+class GemvW8A16M4Cdna4(GemvW8A16Mfma):
+ """4x4-tile CDNA4 op."""
+
+ def __init__(self, device: str = "cuda:0"):
+ super().__init__(device)
+ self.packed_m4: torch.Tensor | None = None
+
+ def init_reference_weights(self, w: torch.Tensor) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def tilert_forward(self, hidden: torch.Tensor) -> torch.Tensor:
+ assert self.packed_m4 is not None
+ out = torch.empty(hidden.size(0), ROWS, dtype=torch.bfloat16, device=hidden.device)
+ torch.ops.tilert.gemv_w8a16_m4_cdna4_op(hidden, self.packed_m4, self.scales, out)
+ return out
diff --git a/tilert/models/misc_rocm/ops/gemv_w8a16_mfma.py b/tilert/models/misc_rocm/ops/gemv_w8a16_mfma.py
new file mode 100644
index 0000000..40426ac
--- /dev/null
+++ b/tilert/models/misc_rocm/ops/gemv_w8a16_mfma.py
@@ -0,0 +1,74 @@
+"""gemv_w8a16_mfma op wrapper: quantizer, weight swizzler, golden, forward."""
+
+import torch
+
+ROWS = 2048
+HIDDEN = 6144
+SCALE_BLK = 128
+FP8_MAX = 448.0
+SUPPORTED_SEQS = (1, 2, 4)
+_ROWS_PER_BLOCK = 32
+_CHUNK = 64
+
+
+def quantize_fp8_block(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ assert w.shape == (ROWS, HIDDEN)
+ blocks = w.float().view(ROWS // SCALE_BLK, SCALE_BLK, HIDDEN // SCALE_BLK, SCALE_BLK)
+ amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12)
+ scales = amax / FP8_MAX
+ q = (blocks / scales).to(torch.float8_e4m3fn)
+ return (
+ q.view(ROWS // SCALE_BLK, SCALE_BLK, HIDDEN // SCALE_BLK, SCALE_BLK)
+ .permute(0, 1, 2, 3)
+ .reshape(ROWS, HIDDEN)
+ .contiguous(),
+ scales.view(ROWS // SCALE_BLK, HIDDEN // SCALE_BLK).contiguous(),
+ )
+
+
+def swizzle_weights(w_fp8: torch.Tensor) -> torch.Tensor:
+ assert w_fp8.shape == (ROWS, HIDDEN)
+ w8 = w_fp8.view(torch.uint8)
+ b = torch.arange(ROWS // _ROWS_PER_BLOCK)
+ g = torch.arange(2)
+ h = torch.arange(4)
+ c = torch.arange(HIDDEN // 4 // _CHUNK)
+ lane = torch.arange(64)
+ s = torch.arange(4)
+ i = torch.arange(4)
+ B, G, H, C, L, S, II = torch.meshgrid(b, g, h, c, lane, s, i, indexing="ij")
+ rows = B * _ROWS_PER_BLOCK + G * 16 + L % 16
+ ks = H * (HIDDEN // 4) + C * _CHUNK + S * 16 + L // 16 * 4 + II
+ return w8[rows, ks].reshape(-1).contiguous()
+
+
+class GemvW8A16Mfma:
+ """Op class: golden and tilert forwards share weights."""
+
+ def __init__(self, device: str = "cuda:0"):
+ self.device = device
+ self.w_fp8: torch.Tensor | None = None
+ self.scales: torch.Tensor | None = None
+ self.packed: torch.Tensor | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def init_reference_weights(self, w: torch.Tensor) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def dequant(self) -> torch.Tensor:
+ assert self.w_fp8 is not None and self.scales is not None
+ blocks = self.w_fp8.float().view(
+ ROWS // SCALE_BLK, SCALE_BLK, HIDDEN // SCALE_BLK, SCALE_BLK
+ )
+ return (blocks * self.scales[:, None, :, None]).view(ROWS, HIDDEN)
+
+ def golden_forward(self, hidden: torch.Tensor) -> torch.Tensor:
+ raise RuntimeError("golden_forward is not available in release builds")
+
+ def tilert_forward(self, hidden: torch.Tensor) -> torch.Tensor:
+ assert self.packed is not None and self.scales is not None
+ out = torch.empty(hidden.size(0), ROWS, dtype=torch.bfloat16, device=hidden.device)
+ torch.ops.tilert.gemv_w8a16_mfma_op(hidden, self.packed, self.scales, out)
+ return out
diff --git a/tilert/models/misc_rocm/ops/gemv_w8a16_mfma_cdna4.py b/tilert/models/misc_rocm/ops/gemv_w8a16_mfma_cdna4.py
new file mode 100644
index 0000000..fd88137
--- /dev/null
+++ b/tilert/models/misc_rocm/ops/gemv_w8a16_mfma_cdna4.py
@@ -0,0 +1,56 @@
+"""gemv_w8a16_mfma_cdna4 op wrapper: the CDNA4 kernel's two weight packings."""
+
+import torch
+
+from tilert.models.misc_rocm.ops.gemv_w8a16_mfma import (
+ FP8_MAX,
+ HIDDEN,
+ ROWS,
+ SCALE_BLK,
+ SUPPORTED_SEQS,
+ GemvW8A16Mfma,
+ quantize_fp8_block,
+ swizzle_weights,
+)
+
+_CHUNK = 64
+_N_RG = ROWS // 16
+_N_KC = HIDDEN // _CHUNK
+
+
+def swizzle_weights_contig(w_fp8: torch.Tensor) -> torch.Tensor:
+ assert w_fp8.shape == (ROWS, HIDDEN)
+ w8 = w_fp8.view(torch.uint8)
+ rg = torch.arange(_N_RG)
+ kc = torch.arange(_N_KC)
+ lane = torch.arange(64)
+ sp = torch.arange(2)
+ i = torch.arange(8)
+ RG, KC, L, SP, II = torch.meshgrid(rg, kc, lane, sp, i, indexing="ij")
+ rows = RG * 16 + L % 16
+ ks = KC * _CHUNK + SP * 32 + L // 16 * 8 + II
+ return w8[rows, ks].reshape(-1).contiguous()
+
+
+class GemvW8A16MfmaCdna4(GemvW8A16Mfma):
+ """CDNA4 op. Inherits the CDNA2 quantizer/golden so the two kernels are compared on identical numbers; only the packing and the launch differ."""
+
+ def __init__(self, device: str = "cuda:0", contig: bool = True):
+ super().__init__(device)
+ self.contig = contig
+ self.packed_contig: torch.Tensor | None = None
+
+ def init_reference_weights(self, w: torch.Tensor) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def packed_for_kernel(self) -> torch.Tensor:
+ packed = self.packed_contig if self.contig else self.packed
+ assert packed is not None
+ return packed
+
+ def tilert_forward(self, hidden: torch.Tensor) -> torch.Tensor:
+ out = torch.empty(hidden.size(0), ROWS, dtype=torch.bfloat16, device=hidden.device)
+ torch.ops.tilert.gemv_w8a16_mfma_cdna4_op(
+ hidden, self.packed_for_kernel(), self.scales, out
+ )
+ return out
diff --git a/tilert/models/misc_rocm/ops/rmsnorm_projx_wqkv_w8a16_mfma.py b/tilert/models/misc_rocm/ops/rmsnorm_projx_wqkv_w8a16_mfma.py
new file mode 100644
index 0000000..0c1254d
--- /dev/null
+++ b/tilert/models/misc_rocm/ops/rmsnorm_projx_wqkv_w8a16_mfma.py
@@ -0,0 +1,74 @@
+"""rmsnorm_projx_wqkv_w8a16_mfma op wrapper: quantizer, swizzler, golden."""
+
+import torch
+
+HIDDEN = 2048
+Q_DIM = 4096
+KV_DIM = 1024
+ROWS = Q_DIM + KV_DIM
+SCALE_BLK = 128
+FP8_MAX = 448.0
+EPS = 1e-06
+SUPPORTED_SEQS = (1, 2, 4)
+_ROWS_PER_BLOCK = 64
+_GROUPS = _ROWS_PER_BLOCK // 16
+_K_PARTS = 2
+_CHUNK = 64
+_NUM_BLOCKS = ROWS // _ROWS_PER_BLOCK
+_CHUNKS = HIDDEN // _K_PARTS // _CHUNK
+
+
+def quantize_fp8_block(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ assert w.shape == (ROWS, HIDDEN)
+ blocks = w.float().view(ROWS // SCALE_BLK, SCALE_BLK, HIDDEN // SCALE_BLK, SCALE_BLK)
+ amax = blocks.abs().amax(dim=(1, 3), keepdim=True).clamp(min=1e-12)
+ scales = amax / FP8_MAX
+ q = (blocks / scales).to(torch.float8_e4m3fn)
+ return (
+ q.reshape(ROWS, HIDDEN).contiguous(),
+ scales.view(ROWS // SCALE_BLK, HIDDEN // SCALE_BLK).contiguous(),
+ )
+
+
+def swizzle_weights(w_fp8: torch.Tensor) -> torch.Tensor:
+ assert w_fp8.shape == (ROWS, HIDDEN)
+ w8 = w_fp8.view(torch.uint8)
+ v = w8.view(_NUM_BLOCKS, _GROUPS, 16, _K_PARTS, _CHUNKS, 4, 4, 4)
+ return v.permute(0, 1, 3, 4, 6, 2, 5, 7).reshape(-1).contiguous()
+
+
+class RmsNormProjXWqkvW8A16Mfma:
+ """Op class: golden and tilert forwards share weights."""
+
+ def __init__(self, device: str = "cuda:0"):
+ self.device = device
+ self.gamma: torch.Tensor | None = None
+ self.w_fp8: torch.Tensor | None = None
+ self.scales: torch.Tensor | None = None
+ self.packed: torch.Tensor | None = None
+
+ def init_random_weights(self, seed: int = 0) -> None:
+ raise RuntimeError("init_random_weights is not available in release builds")
+
+ def init_reference_weights(self, w: torch.Tensor, gamma: torch.Tensor) -> None:
+ raise RuntimeError("init_reference_weights is not available in release builds")
+
+ def dequant(self) -> torch.Tensor:
+ assert self.w_fp8 is not None and self.scales is not None
+ blocks = self.w_fp8.float().view(
+ ROWS // SCALE_BLK, SCALE_BLK, HIDDEN // SCALE_BLK, SCALE_BLK
+ )
+ return (blocks * self.scales[:, None, :, None]).view(ROWS, HIDDEN)
+
+ def golden_forward(self, hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ raise RuntimeError("golden_forward is not available in release builds")
+
+ def tilert_forward(self, hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+ assert self.packed is not None and self.scales is not None
+ assert self.gamma is not None
+ q_out = torch.empty(hidden.size(0), Q_DIM, dtype=torch.bfloat16, device=hidden.device)
+ kv_out = torch.empty(hidden.size(0), KV_DIM, dtype=torch.bfloat16, device=hidden.device)
+ torch.ops.tilert.rmsnorm_projx_wqkv_w8a16_mfma_op(
+ hidden, self.gamma, self.packed, self.scales, q_out, kv_out
+ )
+ return (q_out, kv_out)
diff --git a/tilert/models/misc_rocm/ops/rmsnorm_raw.py b/tilert/models/misc_rocm/ops/rmsnorm_raw.py
new file mode 100644
index 0000000..afb435f
--- /dev/null
+++ b/tilert/models/misc_rocm/ops/rmsnorm_raw.py
@@ -0,0 +1,28 @@
+"""rmsnorm_raw op wrapper."""
+
+import torch
+
+_EPS_FOR_DIM = {2048: 1e-06, 6144: 1e-05, 8192: 1e-06}
+SUPPORTED_DIMS = tuple(_EPS_FOR_DIM)
+
+
+def rmsnorm_raw(
+ hidden_in: torch.Tensor, gamma: torch.Tensor, hidden_out: torch.Tensor
+) -> torch.Tensor:
+ torch.ops.tilert.rmsnorm_raw_op(hidden_in, gamma, hidden_out)
+ return hidden_out
+
+
+class RMSNormRaw:
+
+ def __init__(self, dim: int):
+ assert dim in SUPPORTED_DIMS, f"dim must be one of {SUPPORTED_DIMS}"
+ self.dim = dim
+ self.eps = _EPS_FOR_DIM[dim]
+
+ def golden_forward(self, hidden_in: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor:
+ raise RuntimeError("golden_forward is not available in release builds")
+
+ def tilert_forward(self, hidden_in: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor:
+ hidden_out = torch.empty_like(hidden_in)
+ return rmsnorm_raw(hidden_in, gamma, hidden_out)
diff --git a/tilert/pd_vllm/__init__.py b/tilert/pd_vllm/__init__.py
index 2266c0c..2ae2839 100644
--- a/tilert/pd_vllm/__init__.py
+++ b/tilert/pd_vllm/__init__.py
@@ -1 +1 @@
-"""vLLM-prefill + TileRT-decode PD disaggregation for GLM-5 and DeepSeek-V3.2."""
+pass
diff --git a/tilert/pd_vllm/capabilities.py b/tilert/pd_vllm/capabilities.py
new file mode 100644
index 0000000..6e1426f
--- /dev/null
+++ b/tilert/pd_vllm/capabilities.py
@@ -0,0 +1,274 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+__all__ = [
+ "CapabilityError",
+ "CapabilityUnavailable",
+ "InvalidParameter",
+ "NodeCapabilities",
+ "PROFILE_FIELD_NAMES",
+ "STATIC_FIELD_NAMES",
+ "TYPED_FIELD_NAMES",
+ "engine_capabilities",
+ "validate_generation_request",
+]
+
+
+class CapabilityError(Exception):
+ error_type = "capability_error"
+ http_status = 500
+
+ def to_payload(self) -> dict[str, str]:
+ return {"error": str(self), "error_type": self.error_type}
+
+
+class InvalidParameter(CapabilityError):
+ error_type = "invalid_parameter"
+ http_status = 400
+
+
+class CapabilityUnavailable(CapabilityError):
+ error_type = "capability_unavailable"
+ http_status = 501
+
+
+@dataclass(frozen=True)
+class NodeCapabilities:
+ penalties: bool = False
+ ignore_eos: bool = False
+
+ def intersect(self, other: NodeCapabilities) -> NodeCapabilities:
+ return NodeCapabilities(
+ penalties=self.penalties and other.penalties,
+ ignore_eos=self.ignore_eos and other.ignore_eos,
+ )
+
+ def to_payload(self) -> dict[str, bool]:
+ return {"penalties": self.penalties, "ignore_eos": self.ignore_eos}
+
+ @classmethod
+ def from_payload(cls, payload: object) -> NodeCapabilities:
+ if not isinstance(payload, dict):
+ return cls()
+ caps = payload.get("capabilities", payload)
+ if not isinstance(caps, dict):
+ return cls()
+ return cls(
+ penalties=caps.get("penalties") is True, ignore_eos=caps.get("ignore_eos") is True
+ )
+
+
+def engine_capabilities(engine: object) -> NodeCapabilities:
+
+ def _ask(name: str) -> bool:
+ probe = getattr(engine, name, None)
+ if not callable(probe):
+ return False
+ try:
+ return bool(probe())
+ except Exception:
+ return False
+
+ return NodeCapabilities(
+ penalties=_ask("supports_penalties"), ignore_eos=_ask("supports_ignore_eos")
+ )
+
+
+_EMPTY, _UNSET, _NUMBER, _FLAG, _COUNT = ("empty", "unset", "number", "flag", "count")
+_STATIC_FIELDS: dict[str, tuple] = {
+ "stop_token_ids": (
+ _EMPTY,
+ None,
+ None,
+ "the decode loop uses the model's own stop set and accepts no per-request ids",
+ ),
+ "min_tokens": (
+ _COUNT,
+ 0,
+ 0,
+ "the decode loop cannot suppress its stop set for a minimum length",
+ ),
+ "frequency_penalty": (
+ _NUMBER,
+ 0.0,
+ None,
+ "the decode sampler implements repetition and presence penalties only",
+ ),
+ "min_p": (_NUMBER, 0.0, None, "the decode sampler implements top-p and top-k only"),
+ "seed": (
+ _UNSET,
+ None,
+ None,
+ "the decode sampler's seed is per-process, so a per-request seed cannot make the reply reproducible",
+ ),
+ "logit_bias": (_EMPTY, None, None, "the decode sampler has no per-request logit bias"),
+ "bad_words": (_EMPTY, None, None, "the decode loop has no bad-words matcher"),
+ "allowed_token_ids": (_EMPTY, None, None, "the decode sampler has no per-request allow list"),
+ "structured_outputs": (
+ _EMPTY,
+ None,
+ None,
+ "use response_format, which this stack translates into a decode-side grammar",
+ ),
+ "n": (
+ _COUNT,
+ 1,
+ 1,
+ "one prefill KV state is transferred per request, so the decode node produces exactly one sequence",
+ ),
+ "best_of": (
+ _COUNT,
+ 1,
+ 1,
+ "the decode node produces exactly one sequence, so there is nothing to select from",
+ ),
+ "use_beam_search": (_FLAG, False, None, "the decode node runs single-sequence AR/MTP decode"),
+ "prompt_logprobs": (
+ _UNSET,
+ None,
+ None,
+ "the prefill instance is asked for one token, so no prompt distribution is collected",
+ ),
+ "logprob_token_ids": (
+ _EMPTY,
+ None,
+ None,
+ "the decode export returns the top candidates, not a caller-chosen vocab subset",
+ ),
+ "skip_special_tokens": (
+ _FLAG,
+ True,
+ None,
+ "the reply is detokenised for the output parser, which always consumes special tokens",
+ ),
+}
+_PROFILE_FIELDS: dict[str, tuple] = {
+ "repetition_penalty": (
+ _NUMBER,
+ 1.0,
+ "penalties",
+ "this model's decode runtime has no penalty pre-pass",
+ ),
+ "presence_penalty": (
+ _NUMBER,
+ 0.0,
+ "penalties",
+ "this model's decode runtime has no penalty pre-pass",
+ ),
+ "ignore_eos": (
+ _FLAG,
+ False,
+ "ignore_eos",
+ "this model's decode loop does not clear its stop set",
+ ),
+}
+_TYPED_FIELDS: dict[str, tuple] = {
+ "temperature": (_NUMBER, 0.0),
+ "top_p": (_NUMBER, 0.0),
+ "top_k": (_COUNT, None),
+}
+STATIC_FIELD_NAMES = tuple(_STATIC_FIELDS)
+PROFILE_FIELD_NAMES = tuple(_PROFILE_FIELDS)
+TYPED_FIELD_NAMES = tuple(_TYPED_FIELDS)
+_GREEDY_TEMPERATURE = 1e-05
+
+
+def _is_greedy(body: dict) -> bool:
+ raw = body.get("temperature")
+ if raw is None:
+ return False
+ try:
+ return _lenient_float("temperature", raw) < _GREEDY_TEMPERATURE
+ except InvalidParameter:
+ return False
+
+
+def _as_number(field: str, value: object) -> float:
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ raise InvalidParameter(f"{field} must be a number, got {type(value).__name__}")
+ return float(value)
+
+
+def _as_flag(field: str, value: object) -> bool:
+ if not isinstance(value, bool):
+ raise InvalidParameter(f"{field} must be a boolean, got {type(value).__name__}")
+ return value
+
+
+def _as_count(field: str, value: object, minimum: int) -> int:
+ if isinstance(value, bool) or not isinstance(value, int):
+ raise InvalidParameter(f"{field} must be an integer, got {type(value).__name__}")
+ if value < minimum:
+ raise InvalidParameter(f"{field} must be >= {minimum}, got {value}")
+ return value
+
+
+def _lenient_float(field: str, raw: object) -> float:
+ if isinstance(raw, bool):
+ raise InvalidParameter(f"{field} must be a number, got bool")
+ try:
+ return float(raw)
+ except (TypeError, ValueError):
+ raise InvalidParameter(f"{field} must be a number, got {raw!r}") from None
+
+
+def _lenient_int(field: str, raw: object) -> int:
+ value = _lenient_float(field, raw)
+ if value != int(value):
+ raise InvalidParameter(f"{field} must be an integer, got {raw!r}")
+ return int(value)
+
+
+def _is_neutral(field: str, value: object, kind: str, neutral, minimum: int | None) -> bool:
+ if value is None:
+ return True
+ if kind == _EMPTY:
+ return not value
+ if kind == _UNSET:
+ return False
+ if kind == _NUMBER:
+ return _as_number(field, value) == neutral
+ if kind == _FLAG:
+ return _as_flag(field, value) is neutral
+ return _as_count(field, value, minimum if minimum is not None else 0) == neutral
+
+
+def validate_generation_request(
+ body: dict, capabilities: NodeCapabilities | None = None, adopted: dict | None = None
+) -> None:
+ for field, (kind, minimum) in _TYPED_FIELDS.items():
+ if body.get(field) is None:
+ continue
+ if kind == _NUMBER:
+ value = _lenient_float(field, body[field])
+ if minimum is not None and value < minimum:
+ raise InvalidParameter(f"{field} must be >= {minimum}, got {value}")
+ else:
+ _lenient_int(field, body[field])
+ greedy = _is_greedy(body)
+ for field, (kind, neutral, minimum, why) in _STATIC_FIELDS.items():
+ if field not in body:
+ continue
+ if _is_neutral(field, body[field], kind, neutral, minimum):
+ continue
+ if field == "seed" and greedy:
+ continue
+ raise CapabilityUnavailable(
+ f"{field} is not supported by the TileRT decode stage: {why}. The vLLM prefill instance would apply it to the first token and the decode node would ignore it for the rest of the reply, so the request is refused instead of served incorrectly."
+ )
+ caps = capabilities or NodeCapabilities()
+ for field, (kind, neutral, attr, why) in _PROFILE_FIELDS.items():
+ if field in body:
+ value, origin = (body[field], "the request")
+ elif adopted is not None and field in adopted:
+ value, origin = (adopted[field], "this deployment's defaults")
+ else:
+ continue
+ if _is_neutral(field, value, kind, neutral, None):
+ continue
+ if getattr(caps, attr):
+ continue
+ raise CapabilityUnavailable(
+ f"{field} (from {origin}) is not supported by the decode node serving this pool: {why}. It would apply to the first token only, so the request is refused instead of served incorrectly."
+ )
diff --git a/tilert/pd_vllm/decode_pool.py b/tilert/pd_vllm/decode_pool.py
new file mode 100644
index 0000000..0cc53a2
--- /dev/null
+++ b/tilert/pd_vllm/decode_pool.py
@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+import logging
+import threading
+import time
+
+import requests
+
+from tilert.pd_vllm.capabilities import NodeCapabilities, engine_capabilities
+
+logger = logging.getLogger("pd_vllm.pool")
+__all__ = ["DecodeNode", "NodeLease", "Pool", "acquire_lease", "cancel_decode"]
+QUEUE_LOG_SECONDS = 0.1
+
+
+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
+ self.caps: NodeCapabilities | None = None
+ self.caps_at: float = 0.0
+
+ @property
+ def http_base(self) -> str:
+ return f"http://{self.host}:{self.http_port}"
+
+
+class Pool:
+ CAPS_TTL_S = 60.0
+ CAPS_TIMEOUT_S = 2.0
+
+ def __init__(self, nodes: list[DecodeNode], queue_timeout: float = 0.0):
+ self.nodes = nodes
+ self.queue_timeout = queue_timeout
+ self._cv = threading.Condition()
+ self._caps_lock = threading.Lock()
+
+ def acquire(self) -> DecodeNode | None:
+ deadline = time.monotonic() + self.queue_timeout
+ with self._cv:
+ while True:
+ for n in self.nodes:
+ if not n.busy:
+ n.busy = True
+ return n
+ remaining = deadline - time.monotonic()
+ if remaining <= 0:
+ return None
+ self._cv.wait(remaining)
+
+ def release(self, node: DecodeNode) -> None:
+ with self._cv:
+ node.busy = False
+ self._cv.notify()
+
+ def _node_caps(self, node: DecodeNode) -> NodeCapabilities:
+ with self._caps_lock:
+ if node.caps is not None and time.time() - node.caps_at < self.CAPS_TTL_S:
+ return node.caps
+ try:
+ r = requests.get(f"{node.http_base}/capabilities", timeout=self.CAPS_TIMEOUT_S)
+ r.raise_for_status()
+ caps = NodeCapabilities.from_payload(r.json())
+ except Exception as e:
+ logger.warning(
+ "capability probe failed for %s (%s); treating every optional field as unsupported",
+ node.http_base,
+ e,
+ )
+ return NodeCapabilities()
+ with self._caps_lock:
+ node.caps, node.caps_at = (caps, time.time())
+ logger.info("capabilities for %s: %s", node.http_base, caps.to_payload())
+ return caps
+
+ def capabilities(self) -> NodeCapabilities:
+ result: NodeCapabilities | None = None
+ for n in self.nodes:
+ caps = self._node_caps(n)
+ result = caps if result is None else result.intersect(caps)
+ return result or NodeCapabilities()
+
+
+def cancel_decode(node, rid: str) -> None:
+ try:
+ requests.post(f"{node.http_base}/pd/cancel", json={"rid": rid}, timeout=5)
+ except Exception:
+ logger.warning("cancel POST failed for %s", rid)
+
+
+class NodeLease:
+
+ def __init__(self, pool: Pool, node: DecodeNode):
+ self.pool = pool
+ self.node = node
+ self.rid: str | None = None
+ self.dispatched = False
+ self._released = False
+
+ @property
+ def released(self) -> bool:
+ """Whether this lease has already gone back to the pool. Read by the
+ backstop that reclaims a lease whose handler never released it, so the
+ normal path is not logged as a leak."""
+ return self._released
+
+ def release(self, *, terminated: bool = False) -> None:
+ if self._released:
+ return
+ self._released = True
+ self.pool.release(self.node)
+ if self.dispatched and (not terminated) and (self.rid is not None):
+ threading.Thread(target=cancel_decode, args=(self.node, self.rid), daemon=True).start()
+
+ def __enter__(self) -> NodeLease:
+ return self
+
+ def __exit__(self, *exc) -> None:
+ self.release()
+
+
+def acquire_lease(pool: Pool) -> tuple[NodeLease | None, float]:
+ t0 = time.monotonic()
+ node = pool.acquire()
+ waited = time.monotonic() - t0
+ if node is None:
+ return (None, waited)
+ if waited >= QUEUE_LOG_SECONDS:
+ logger.info("queued %.1fs for decode node %s", waited, node.host)
+ return (NodeLease(pool, node), waited)
diff --git a/tilert/pd_vllm/decode_response.py b/tilert/pd_vllm/decode_response.py
new file mode 100644
index 0000000..8425b3d
--- /dev/null
+++ b/tilert/pd_vllm/decode_response.py
@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any
+
+__all__ = [
+ "BUSY",
+ "DecodeReader",
+ "OK",
+ "PROPAGATE",
+ "REFUSED",
+ "RETRY",
+ "SERVER_ERROR",
+ "TRUNCATED",
+ "TYPED_ERROR",
+ "UNTYPED_ERROR",
+ "PROPAGATED_ERROR_STATUS",
+ "PROPAGATED_ERROR_TYPES",
+ "terminal_verdict",
+ "decode_refusal",
+ "classify_decode_status",
+]
+RETRY = "retry"
+BUSY = "busy"
+PROPAGATE = "propagate"
+SERVER_ERROR = "server_error"
+
+
+def classify_decode_status(
+ status: int, payload: Any, *, attempts_left: bool, propagated_types: tuple[str, ...]
+) -> str:
+ if status == 429:
+ return RETRY if attempts_left else BUSY
+ if status == 200:
+ raise ValueError("200 is not an error status")
+ if isinstance(payload, dict) and payload.get("error_type") in propagated_types:
+ return PROPAGATE
+ return SERVER_ERROR
+
+
+PROPAGATED_ERROR_STATUS = {
+ "grammar_unsupported": 501,
+ "invalid_grammar": 400,
+ "grammar_violation": 400,
+ "grammar_backend_unavailable": 500,
+ "logprobs_unavailable": 501,
+ "capability_unavailable": 501,
+ "invalid_parameter": 400,
+ "request_cancelled": 499,
+}
+PROPAGATED_ERROR_TYPES = frozenset(PROPAGATED_ERROR_STATUS)
+OK = "ok"
+REFUSED = "refused"
+TRUNCATED = "truncated"
+TYPED_ERROR = "typed_error"
+UNTYPED_ERROR = "untyped_error"
+
+
+def terminal_verdict(reader, *, client_gone: bool = False) -> tuple[str, dict, int]:
+ if reader.refusal is not None:
+ return (REFUSED, reader.refusal, reader.refusal_status)
+ if reader.node_error is not None:
+ error_type = reader.node_error.get("error_type")
+ if error_type in PROPAGATED_ERROR_TYPES:
+ return (TYPED_ERROR, reader.node_error, PROPAGATED_ERROR_STATUS[error_type])
+ return (UNTYPED_ERROR, reader.node_error, 502)
+ if not reader.node_terminated and (not reader.stop_hit) and (not client_gone):
+ return (
+ TRUNCATED,
+ {
+ "error": "decode stream ended without a terminal message",
+ "error_type": "decode_truncated",
+ "rid": reader.rid,
+ },
+ 502,
+ )
+ return (OK, {}, 200)
+
+
+def decode_refusal(verdict: str, status: int, payload: Any, rid: str) -> tuple[dict, int]:
+ if verdict == BUSY:
+ return ({"error": "decode node busy", "error_type": "decode_busy", "rid": rid}, 429)
+ if verdict == PROPAGATE:
+ return (payload, status)
+ if verdict == SERVER_ERROR:
+ return ({"error": "decode call failed", "status": status, "rid": rid}, 502)
+ raise ValueError(f"{verdict} is not a refusal")
+
+
+@dataclass
+class DecodeReader:
+ stream: Any = None
+ logprobs_req: Any = None
+ rid: str = ""
+ token_ids: list[int] = field(default_factory=list)
+ timing: dict = field(default_factory=dict)
+ finish_reason: str = "stop"
+ node_terminated: bool = False
+ refusal: dict | None = None
+ refusal_status: int = 501
+ node_error: dict | None = None
+ _stopped: bool = False
+
+ @property
+ def finished(self) -> bool:
+ return self.node_terminated or self._stopped or self.refusal is not None
+
+ @property
+ def stop_hit(self) -> bool:
+ return self._stopped
+
+ def feed(self, line: str) -> list:
+ if not line:
+ return []
+ return self._message(json.loads(line))
+
+ def feed_blocking(self, body: dict) -> list:
+ self.timing = body.get("timing_ms", {})
+ self._set_finish(self.timing.get("finish_reason", "stop"))
+ self.node_terminated = True
+ lp = body.get("logprobs") or {}
+ return self._tokens({"t": body["token_ids"], **lp})
+
+ def _message(self, msg: dict) -> list:
+ if "t" in msg:
+ return self._tokens(msg)
+ if "done" in msg:
+ self.node_terminated = True
+ self.timing = msg.get("timing_ms", {})
+ self._set_finish(msg.get("finish_reason", "stop"))
+ declared = msg.get("n")
+ got = len(self.stream.token_ids) if self.stream is not None else len(self.token_ids)
+ if declared is not None and declared != got:
+ self.refusal = {
+ "error": f"decode node declared {declared} tokens and sent {got}",
+ "error_type": "decode_truncated",
+ "rid": self.rid,
+ }
+ self.refusal_status = 502
+ return []
+ if "error" in msg:
+ self.node_terminated = True
+ self.node_error = msg
+ return []
+ return []
+
+ def _set_finish(self, reason: str) -> None:
+ self.finish_reason = "stop" if reason == "cancelled" else reason
+
+ def _tokens(self, msg: dict) -> list:
+ ids = msg["t"]
+ seen = len(self.stream.token_ids) if self.stream is not None else len(self.token_ids)
+ if self.logprobs_req is not None and (
+ not _logprobs_line_ok(msg, len(ids), self.logprobs_req, seen)
+ ):
+ self.refusal = {
+ "error": "decode node returned no logprobs",
+ "error_type": "logprobs_unavailable",
+ "rid": self.rid,
+ }
+ return []
+ if self.stream is None:
+ self.token_ids += ids
+ return []
+ out = self.stream.push(ids, msg.get("lp"), msg.get("tp"))
+ if self.stream.stop_reason is not None:
+ self._stopped = True
+ return out
+
+
+def _logprobs_line_ok(payload: dict, n_tokens: int, req, seen: int = 0) -> bool:
+ lp = payload.get("lp")
+ if lp is None or len(lp) != n_tokens:
+ return False
+ if any((value is None for i, value in enumerate(lp) if not (seen == 0 and i == 0))):
+ return False
+ if req.top_n > 0:
+ tp = payload.get("tp")
+ if tp is None or len(tp) != n_tokens:
+ return False
+ return True
diff --git a/tilert/pd_vllm/decode_server.py b/tilert/pd_vllm/decode_server.py
index 4ebd777..cdf0ffe 100644
--- a/tilert/pd_vllm/decode_server.py
+++ b/tilert/pd_vllm/decode_server.py
@@ -1,18 +1,3 @@
-"""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
@@ -20,21 +5,30 @@
import os
import queue as queue_mod
import socket
+import sys
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.capabilities import (
+ CapabilityError,
+ engine_capabilities,
+ validate_generation_request,
+)
+from tilert.pd_vllm.grammar_spec import GrammarError, GrammarViolationError
from tilert.pd_vllm.receive_server import ReceiveServer
logger = logging.getLogger("pd_vllm.decode_server")
+_ABANDON_DRAIN_S = 30.0
+_CANCELLED = object()
-DECODE_POLL_S = max(0.0, float(os.environ.get("TILERT_DECODE_POLL_MS") or "200")) / 1000.0
+class LogprobsUnavailable(Exception):
+ pass
class DecodeBody(BaseModel):
@@ -44,17 +38,34 @@ class DecodeBody(BaseModel):
sampling: dict | None = None
timeout_s: float = 120.0
stream: bool = False
+ grammar_spec: dict | None = None
+ enable_thinking: bool = True
+ top_logprobs: int | None = None
+
+
+DECODE_POLL_S = max(0.0, float(os.environ.get("TILERT_DECODE_POLL_MS") or "200")) / 1000.0
def build_app(server: ReceiveServer, engine) -> FastAPI:
app = FastAPI()
lock = threading.Lock()
- state: dict[str, Any] = {"current_rid": None}
+ state = {"current_rid": None}
@app.get("/health")
def health():
return {"status": "ok"}
+ @app.get("/capabilities")
+ def capabilities():
+ caps = engine_capabilities(engine)
+ payload = caps.to_payload()
+ payload["logprobs"] = bool(getattr(engine, "supports_logprobs", lambda: False)())
+ return {
+ "profile": getattr(server.profile, "name", None),
+ "engine": type(engine).__name__,
+ "capabilities": payload,
+ }
+
@app.get("/decode_status")
def decode_status():
busy = lock.locked()
@@ -62,15 +73,9 @@ def decode_status():
@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:
+ if rid and rid == state["current_rid"] and (ev is not None):
ev.set()
logger.info("cancel requested for %s", rid)
return {"cancelled": rid}
@@ -84,18 +89,47 @@ def _cleanup():
engine.reset()
except Exception:
logger.exception("engine reset failed")
- server.release()
+ if state["current_rid"]:
+ server.release(state["current_rid"])
state["current_rid"] = None
state["cancel_event"] = None
lock.release()
+ def _drain_own_kv(rid: str, timeout_s: float, cancel=None):
+ deadline = time.time() + timeout_s
+ while True:
+ if cancel is not None and cancel.is_set():
+ return _CANCELLED
+ remaining = deadline - time.time()
+ if remaining <= 0:
+ return None
+ try:
+ cand = server.completed.get(timeout=min(remaining, 0.5))
+ except queue_mod.Empty:
+ continue
+ if cand.rid == rid:
+ return cand
+ logger.warning("dropping unmatched request %s (waiting for %s)", cand.rid, rid)
+ server.release(cand.rid)
+
+ def _abandon_pending_kv(rid: str, cancel=None) -> None:
+ try:
+ if _drain_own_kv(rid, _ABANDON_DRAIN_S, cancel) is None:
+ logger.warning(
+ "abandoned %s: its KV did not arrive within %.0fs", rid, _ABANDON_DRAIN_S
+ )
+ except Exception:
+ logger.exception("draining KV for abandoned %s failed", rid)
+ finally:
+ _cleanup()
+
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()),
+ " ".join((f"{k}={v}" for k, v in timing.items())),
)
@app.post("/pd/decode")
@@ -105,25 +139,53 @@ def pd_decode(body: DecodeBody):
{"error": "busy", "current_rid": state["current_rid"]}, status_code=429
)
state["current_rid"] = body.rid
+ server.expect(body.rid)
+ cancel = threading.Event()
+ state["cancel_event"] = cancel
t0 = time.time()
- # phase 1: wire wait + convert + inject (common to both modes)
+ if body.top_logprobs is not None and (
+ not getattr(engine, "supports_logprobs", lambda: False)()
+ ):
+ logger.info(
+ "logprobs requested but unsupported by %s (rid=%s)", type(engine).__name__, body.rid
+ )
+ _abandon_pending_kv(body.rid, cancel)
+ return JSONResponse(
+ {
+ "error": f"{type(engine).__name__} does not produce logprobs",
+ "error_type": "logprobs_unavailable",
+ },
+ status_code=501,
+ )
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
+ validate_generation_request(body.sampling or {}, engine_capabilities(engine))
+ except CapabilityError as e:
+ logger.info("sampling rejected for %s: %s (%s)", body.rid, e, e.error_type)
+ _abandon_pending_kv(body.rid, cancel)
+ return JSONResponse(e.to_payload(), status_code=e.http_status)
+ try:
+ grammar_session = engine.prepare_grammar(body.grammar_spec, body.enable_thinking)
+ except GrammarError as e:
+ logger.info("grammar rejected for %s: %s (%s)", body.rid, e, e.error_type)
+ _abandon_pending_kv(body.rid, cancel)
+ return JSONResponse(e.to_payload(), status_code=e.http_status)
+ except Exception as e:
+ logger.exception("grammar prepare failed for %s", body.rid)
+ _abandon_pending_kv(body.rid, cancel)
+ return JSONResponse({"error": str(e)}, status_code=500)
+ try:
+ req = _drain_own_kv(body.rid, body.timeout_s, cancel)
+ if req is _CANCELLED:
+ logger.info("cancelled during KV transfer for %s", body.rid)
+ _cleanup()
+ return JSONResponse(
+ {
+ "error": "cancelled during KV transfer",
+ "error_type": "request_cancelled",
+ "rid": body.rid,
+ },
+ status_code=499,
)
- server.release()
if req is None:
_cleanup()
return JSONResponse(
@@ -140,16 +202,27 @@ def pd_decode(body: DecodeBody):
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
+ want_lp = body.top_logprobs is not None
+ lp_sink: list[tuple[float | None, list]] = []
+
+ def _emit(tok, logprob=None, candidates=None):
+ if want_lp:
+ if logprob is None:
+ if lp_sink:
+ raise LogprobsUnavailable(f"engine emitted token {tok} without a logprob")
+ lp_sink.append((None, []))
+ return tok
+ lp_sink.append((float(logprob), list(candidates or ())))
+ return tok
+
+ def _lp_slice(start: int, count: int) -> dict:
+ rows = lp_sink[start : start + count]
+ return {"lp": [r[0] for r in rows], "tp": [[list(c) for c in r[1]] for r in rows]}
if not body.stream:
try:
@@ -157,7 +230,10 @@ def pd_decode(body: DecodeBody):
first_token_id=body.first_token_id,
max_tokens=body.max_tokens,
sampling=body.sampling,
+ on_token=_emit if want_lp else None,
cancel_event=cancel,
+ grammar_session=grammar_session,
+ **{"top_logprobs": body.top_logprobs} if want_lp else {},
)
timing = {
**pre_timing,
@@ -165,25 +241,37 @@ def pd_decode(body: DecodeBody):
**getattr(engine, "last_stats", {}),
}
_log_reqstat(body, req, len(tokens), timing)
- return {
+ out = {
"rid": body.rid,
"token_ids": tokens,
"seq_len": req.seq_len,
"timing_ms": timing,
}
+ if want_lp:
+ if len(lp_sink) != len(tokens):
+ raise LogprobsUnavailable(
+ f"engine returned {len(lp_sink)} logprob entries for {len(tokens)} tokens"
+ )
+ out["logprobs"] = _lp_slice(0, len(tokens))
+ return out
+ except LogprobsUnavailable as e:
+ logger.info("logprobs unavailable for %s: %s", body.rid, e)
+ return JSONResponse(
+ {"error": str(e), "error_type": "logprobs_unavailable"}, status_code=501
+ )
+ except GrammarViolationError as e:
+ logger.info("grammar violation for %s: %s", body.rid, e)
+ return JSONResponse(e.to_payload(), status_code=e.http_status)
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()
fin: dict = {"loop": None, "ev": None}
def _signal_done() -> None:
- loop, ev = fin["loop"], fin["ev"]
+ loop, ev = (fin["loop"], fin["ev"])
if loop is not None and ev is not None:
loop.call_soon_threadsafe(ev.set)
@@ -193,24 +281,29 @@ def _run():
first_token_id=body.first_token_id,
max_tokens=body.max_tokens,
sampling=body.sampling,
- on_token=q.put,
+ on_token=(lambda *a: q.put(_emit(*a))) if want_lp else q.put,
cancel_event=cancel,
+ grammar_session=grammar_session,
+ **{"top_logprobs": body.top_logprobs} if want_lp else {},
)
q.put(("done", tokens))
_signal_done()
- except Exception as e: # pragma: no cover
+ except GrammarViolationError as e:
+ logger.info("stream grammar violation for %s: %s", body.rid, e)
+ q.put(("error", e.to_payload()))
+ _signal_done()
+ except LogprobsUnavailable as e:
+ logger.info("stream logprobs unavailable for %s: %s", body.rid, e)
+ q.put(("error", {"error": str(e), "error_type": "logprobs_unavailable"}))
+ _signal_done()
+ except Exception as e:
logger.exception("stream decode failed for %s", body.rid)
- q.put(("error", str(e)))
+ q.put(("error", {"error": str(e)}))
_signal_done()
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
@@ -221,6 +314,7 @@ async def _gen():
worker.start()
try:
batch: list[int] = []
+ n_emitted = 0
done_msg = None
last_activity = time.time()
while done_msg is None:
@@ -228,12 +322,16 @@ async def _gen():
first = q.get_nowait()
except queue_mod.Empty:
if time.time() - last_activity > 600:
- yield json.dumps({"error": "decode stalled"}) + "\n"
+ yield (json.dumps({"error": "decode stalled"}) + "\n")
return
await asyncio.sleep(0.001)
continue
if isinstance(first, int):
- yield json.dumps({"t": [first]}) + "\n"
+ line = {"t": [first]}
+ if want_lp:
+ line.update(_lp_slice(0, 1))
+ n_emitted = 1
+ yield (json.dumps(line) + "\n")
else:
done_msg = first
last_activity = time.time()
@@ -252,17 +350,23 @@ async def _gen():
done_msg = item
break
if batch:
- yield json.dumps({"t": batch}) + "\n"
+ line = {"t": batch}
+ if want_lp:
+ line.update(_lp_slice(n_emitted, len(batch)))
+ n_emitted += len(batch)
+ yield (json.dumps(line) + "\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"
+ yield (json.dumps({"error": "decode stalled"}) + "\n")
return
else:
- with contextlib.suppress(TimeoutError):
+ try:
await asyncio.wait_for(fin["ev"].wait(), timeout=DECODE_POLL_S)
+ except TimeoutError:
+ pass
kind, payload = done_msg
if kind == "done":
timing = {
@@ -271,22 +375,22 @@ async def _gen():
**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"
+ 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"
+ yield (json.dumps(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():
@@ -298,8 +402,7 @@ async def _gen():
return app # noqa: R504 (assembled across the function)
-def main() -> None:
- logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
+def build_parser() -> argparse.ArgumentParser:
ap = argparse.ArgumentParser()
ap.add_argument("--engine", choices=["stub", "tilert"], default="stub")
ap.add_argument("--model", default="glm5", help="model profile")
@@ -308,24 +411,40 @@ def main() -> None:
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(
+ "--num-mtp",
+ type=int,
+ choices=(3,),
+ default=3,
+ help="MTP draft depth for speculative decoding; needs --with-mtp",
+ )
ap.add_argument(
"--transport",
choices=["mooncake", "nixl"],
default="mooncake",
- help="RDMA data-plane backend " "(must match prefill's tilert_transport)",
+ 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",
+ help="MLA cache dtype (must match vLLM prefill); MLA-family profiles only",
)
- args = ap.parse_args()
+ ap.add_argument(
+ "--pd-buffer-device",
+ choices=["cuda", "cpu"],
+ default=(os.environ.get("TILERT_PD_BUFFER_DEVICE") or "cuda").lower(),
+ help="where the PD receive buffer lives; 'cpu' = pinned host memory registered with the transport (frees buffer_bytes(max_seq_len) of VRAM, adds one H2D copy per request), 'cuda' = sharded over the visible cards per TILERT_PD_SHARDS",
+ )
+ return ap
+
+def main() -> None:
+ logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
+ args = build_parser().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.
+ num_mtp = profiles.resolve_num_mtp(profile, args.num_mtp, with_mtp=args.with_mtp)
if hasattr(profile, "configure"):
profile.configure(args.kv_cache_dtype)
logger.info(
@@ -334,11 +453,12 @@ def main() -> None:
args.kv_cache_dtype,
profile.layout_version,
)
-
+ if hasattr(profile, "configure_weights") and args.model_weights_dir:
+ profile.configure_weights(args.model_weights_dir)
if args.engine == "stub":
from tilert.pd_vllm.engine_iface import StubEngine
- engine: Any = StubEngine()
+ engine = StubEngine()
else:
logger.info(
"loading TileRT engine (profile=%s, weights=%s)...",
@@ -350,11 +470,17 @@ def main() -> None:
max_seq_len=args.max_seq_len,
with_mtp=args.with_mtp,
ar_steps=8,
+ num_mtp=num_mtp,
+ )
+ logger.info(
+ "TileRT engine ready (cache window %d, num_mtp %d)", engine.max_seq_len, num_mtp
)
- 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
+ profile,
+ max_seq_len=args.max_seq_len,
+ ctrl_port=args.ctrl_port,
+ transport=args.transport,
+ buffer_device="cpu" if args.pd_buffer_device == "cpu" else "cuda:0",
)
app = build_app(server, engine)
logger.info(
@@ -364,10 +490,6 @@ def main() -> None:
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):
diff --git a/tilert/pd_vllm/engine_iface.py b/tilert/pd_vllm/engine_iface.py
index 6ed9133..f94373b 100644
--- a/tilert/pd_vllm/engine_iface.py
+++ b/tilert/pd_vllm/engine_iface.py
@@ -1,39 +1,48 @@
-"""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
+from tilert.pd_vllm.grammar_spec import (
+ GrammarUnsupported,
+ GrammarViolationError,
+ InvalidGrammarError,
+)
+
class PDEngine(Protocol):
+
def inject(self, req: Any) -> None:
- """Restore engine state to 'prefilled seq_len tokens' from req."""
+ pass
+
+ def prepare_grammar(self, grammar_spec: dict | None, enable_thinking: bool = True) -> Any:
+ pass
def decode(
self,
first_token_id: int,
max_tokens: int,
sampling: dict | None,
- on_token: Callable[[int], None] | None = None,
+ on_token: Callable[..., None] | None = None,
cancel_event=None,
+ grammar_session: Any = None,
+ top_logprobs: int | None = None,
) -> list[int]:
- """AR/MTP decode from first_token_id; returns completion ids.
+ pass
+
+ def supports_logprobs(self) -> bool:
+ pass
- 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 supports_penalties(self) -> bool:
+ pass
+
+ def supports_ignore_eos(self) -> bool:
+ pass
def reset(self) -> None:
- """Release per-request state."""
+ pass
class StubEngine:
- """Echo engine for plumbing tests: no GPU, no tilert."""
+ _KNOWN_SPEC_TYPES = ("json_schema", "json_object", "ebnf", "regex", "structural_tag")
def __init__(self, fixed_tokens: tuple[int, ...] = (11, 22, 33)):
self._fixed = fixed_tokens
@@ -43,11 +52,52 @@ def __init__(self, fixed_tokens: tuple[int, ...] = (11, 22, 33)):
def inject(self, req: Any) -> None:
self.injected = req
- def decode(self, first_token_id, max_tokens, sampling, on_token=None, cancel_event=None):
+ def prepare_grammar(self, grammar_spec, enable_thinking=True):
+ if grammar_spec is None:
+ return None
+ if not isinstance(grammar_spec, dict) or "type" not in grammar_spec:
+ raise InvalidGrammarError("grammar spec must be a dict with a 'type'")
+ kind = grammar_spec["type"]
+ if kind == "__backend_missing__":
+ raise GrammarUnsupported("constrained decoding is not supported")
+ if kind not in self._KNOWN_SPEC_TYPES:
+ raise InvalidGrammarError(f"unsupported grammar spec type: {kind!r}")
+ return {"spec": grammar_spec, "enable_thinking": enable_thinking}
+
+ def supports_logprobs(self) -> bool:
+ return True
+
+ def supports_penalties(self) -> bool:
+ return True
+
+ def supports_ignore_eos(self) -> bool:
+ return True
+
+ @staticmethod
+ def fake_logprob(token_id: int) -> float:
+ return -0.5 - 0.25 * (token_id % 4)
+
+ def decode(
+ self,
+ first_token_id,
+ max_tokens,
+ sampling,
+ on_token=None,
+ cancel_event=None,
+ grammar_session=None,
+ top_logprobs=None,
+ ):
+ spec = (grammar_session or {}).get("spec", {})
+ if spec.get("value") == "__violate__":
+ raise GrammarViolationError(f"first token {first_token_id} violates the grammar")
out = ([int(first_token_id)] + list(self._fixed))[:max_tokens]
if on_token:
for t in out:
- on_token(t)
+ if top_logprobs is None:
+ on_token(t)
+ else:
+ cands = [(t + k, self.fake_logprob(t) - 0.5 * k) for k in range(top_logprobs)]
+ on_token(t, self.fake_logprob(t), cands)
self.last_stats = {"finish_reason": "stop"}
return out
diff --git a/tilert/pd_vllm/generation_defaults.py b/tilert/pd_vllm/generation_defaults.py
new file mode 100644
index 0000000..a0ac31d
--- /dev/null
+++ b/tilert/pd_vllm/generation_defaults.py
@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+import json
+import logging
+import os
+from dataclasses import dataclass
+
+from tilert.pd_vllm.sampling import VLLM_DEFAULT_TOP_P
+
+__all__ = [
+ "GUARDED_FIELDS",
+ "NEUTRAL_MIN_P",
+ "NEUTRAL_REPETITION_PENALTY",
+ "GenerationDefaults",
+ "UnsupportedGenerationDefault",
+ "load",
+ "penalties_supported_by",
+]
+logger = logging.getLogger("pd_vllm.generation_defaults")
+VLLM_NEUTRAL_TEMPERATURE = 1.0
+VLLM_NEUTRAL_TOP_K = 0
+GUARDED_FIELDS = {"repetition_penalty": 1.0, "min_p": 0.0}
+NEUTRAL_REPETITION_PENALTY = 1.0
+NEUTRAL_MIN_P = 0.0
+
+
+class UnsupportedGenerationDefault(Exception):
+ pass
+
+
+@dataclass(frozen=True)
+class GenerationDefaults:
+ temperature: float = VLLM_NEUTRAL_TEMPERATURE
+ top_p: float = VLLM_DEFAULT_TOP_P
+ top_k: int = VLLM_NEUTRAL_TOP_K
+ repetition_penalty: float = NEUTRAL_REPETITION_PENALTY
+ source: str = "vllm"
+
+ def resolve(self, body: dict) -> dict:
+ return {
+ "temperature": _as_float("temperature", body.get("temperature"), self.temperature),
+ "top_p": _as_float("top_p", body.get("top_p"), self.top_p),
+ "top_k": _as_int("top_k", body.get("top_k"), self.top_k),
+ "repetition_penalty": _as_float(
+ "repetition_penalty", body.get("repetition_penalty"), self.repetition_penalty
+ ),
+ "min_p": NEUTRAL_MIN_P,
+ }
+
+ def describe(self) -> str:
+ return f"temperature={self.temperature}, top_p={self.top_p}, top_k={self.top_k}, repetition_penalty={self.repetition_penalty} (from {self.source})"
+
+
+def _as_float(field: str, raw, default: float) -> float:
+ if raw is None:
+ return float(default)
+ if isinstance(raw, bool):
+ raise ValueError(f"{field} must be a number, got bool")
+ return float(raw)
+
+
+def _as_int(field: str, raw, default: int) -> int:
+ if raw is None:
+ return int(default)
+ if isinstance(raw, bool):
+ raise ValueError(f"{field} must be an integer, got bool")
+ return int(raw)
+
+
+def _read_config(model_path: str) -> dict:
+ path = os.path.join(model_path, "generation_config.json")
+ if not os.path.isfile(path):
+ logger.info(
+ "no generation_config.json under %s; using vLLM's neutral sampling defaults", model_path
+ )
+ return {}
+ with open(path, encoding="utf-8") as fh:
+ config = json.load(fh)
+ if not isinstance(config, dict):
+ raise UnsupportedGenerationDefault(f"{path} does not contain a JSON object")
+ return config
+
+
+def penalties_supported_by(model: str) -> bool:
+ if not model:
+ return False
+ try:
+ from tilert.pd_vllm.profiles import base as profiles
+
+ profile = profiles.get_profile(model)
+ except Exception as e:
+ logger.warning(
+ "cannot resolve model %r to a profile (%s); treating penalties as unsupported for default resolution",
+ model,
+ e,
+ )
+ return False
+ return bool(getattr(profile, "declares_penalties", False))
+
+
+def _check_guarded(config: dict, source: str, *, model: str, penalties_ok: bool) -> None:
+ for field, neutral in sorted(GUARDED_FIELDS.items()):
+ raw = config.get(field)
+ if raw is None or float(raw) == neutral:
+ continue
+ if field == "repetition_penalty" and penalties_ok:
+ continue
+ if field == "min_p":
+ reason = "no decode runtime implements it on any model"
+ elif model:
+ reason = f"the decode runtime for {model!r} has no penalty pre-pass"
+ else:
+ reason = "the router was not told which model it serves, so it cannot confirm the decode runtime applies it -- pass --model"
+ raise UnsupportedGenerationDefault(
+ f"{source} sets {field}={raw}, but {reason}. The vLLM prefill instance would apply it to the first token while the decode node would not apply it to the rest, and no per-request check can catch it because the client never sent it.\nResolve it explicitly, whichever is true:\n - the value is not wanted: launch with --generation-config vllm, or remove {field} from generation_config.json;\n - the value is wanted: have clients send {field} per request, so the capability gate accepts it on a node that supports it and refuses it on one that does not."
+ )
+
+
+def load(
+ model_path: str = "",
+ source: str = "auto",
+ *,
+ model: str = "",
+ temperature: float | None = None,
+ top_p: float | None = None,
+ top_k: int | None = None,
+ repetition_penalty: float | None = None,
+) -> GenerationDefaults:
+ penalties_ok = penalties_supported_by(model)
+ config: dict = {}
+ if source == "auto" and model_path:
+ config = _read_config(model_path)
+ origin = os.path.join(model_path, "generation_config.json")
+ elif source == "auto":
+ logger.info("no --model-path given; using vLLM's neutral sampling defaults")
+ origin = "vllm neutral defaults"
+ else:
+ origin = "vllm neutral defaults"
+ if config:
+ _check_guarded(config, origin, model=model, penalties_ok=penalties_ok)
+ overrides = {
+ "temperature": temperature,
+ "top_p": top_p,
+ "top_k": top_k,
+ "repetition_penalty": repetition_penalty,
+ }
+ _check_guarded(
+ {k: v for k, v in overrides.items() if v is not None},
+ "command-line overrides",
+ model=model,
+ penalties_ok=penalties_ok,
+ )
+ if not penalties_ok:
+ config = {k: v for k, v in config.items() if k != "repetition_penalty"}
+ resolved = {}
+ for field, neutral in (
+ ("temperature", VLLM_NEUTRAL_TEMPERATURE),
+ ("top_p", VLLM_DEFAULT_TOP_P),
+ ("top_k", VLLM_NEUTRAL_TOP_K),
+ ("repetition_penalty", NEUTRAL_REPETITION_PENALTY),
+ ):
+ if overrides[field] is not None:
+ resolved[field] = overrides[field]
+ elif config.get(field) is not None:
+ resolved[field] = config[field]
+ else:
+ resolved[field] = neutral
+ if any(v is not None for v in overrides.values()):
+ origin = f"{origin} + command-line overrides"
+ defaults = GenerationDefaults(
+ temperature=float(resolved["temperature"]),
+ top_p=float(resolved["top_p"]),
+ top_k=int(resolved["top_k"]),
+ repetition_penalty=float(resolved["repetition_penalty"]),
+ source=origin,
+ )
+ logger.info("sampling defaults for requests that omit a field: %s", defaults.describe())
+ return defaults
diff --git a/tilert/pd_vllm/grammar_spec.py b/tilert/pd_vllm/grammar_spec.py
new file mode 100644
index 0000000..69cc98f
--- /dev/null
+++ b/tilert/pd_vllm/grammar_spec.py
@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+class GrammarError(Exception):
+ error_type = "grammar_error"
+ http_status = 500
+
+ def to_payload(self) -> dict[str, str]:
+ return {"error": str(self), "error_type": self.error_type}
+
+
+class InvalidGrammarError(GrammarError):
+ error_type = "invalid_grammar"
+ http_status = 400
+
+
+class GrammarBackendUnavailable(GrammarError):
+ error_type = "grammar_backend_unavailable"
+ http_status = 500
+
+
+class GrammarViolationError(GrammarError):
+ error_type = "grammar_violation"
+ http_status = 400
+
+
+class GrammarUnsupported(GrammarError):
+ error_type = "grammar_unsupported"
+ http_status = 501
+
+
+_MAX_SCHEMA_DEPTH = 64
+_MAX_SCHEMA_NODES = 1000
+_MAX_ENUM_VALUES = 10000
+
+
+def _walk_schema(node, depth, counts, path):
+ where = path or ""
+ if depth > _MAX_SCHEMA_DEPTH:
+ raise InvalidGrammarError(f"json_schema nests deeper than {_MAX_SCHEMA_DEPTH} at {where}")
+ if isinstance(node, list):
+ for i, item in enumerate(node):
+ _walk_schema(item, depth, counts, f"{path}[{i}]")
+ return
+ if not isinstance(node, dict):
+ return
+ counts["nodes"] += 1
+ if counts["nodes"] > _MAX_SCHEMA_NODES:
+ raise InvalidGrammarError(f"json_schema has over {_MAX_SCHEMA_NODES} subschemas")
+ if isinstance(node.get("enum"), list):
+ counts["enums"] += len(node["enum"])
+ if counts["enums"] > _MAX_ENUM_VALUES:
+ raise InvalidGrammarError(f"json_schema has over {_MAX_ENUM_VALUES} enum values")
+ for key, child in node.items():
+ if key in ("enum", "const", "examples", "default"):
+ continue
+ if isinstance(child, (dict, list)):
+ _walk_schema(child, depth + 1, counts, f"{path}.{key}" if path else key)
+
+
+def validate_grammar_spec(spec: dict[str, Any] | None) -> None:
+ if spec:
+ _walk_schema(spec.get("value"), 0, {"nodes": 0, "enums": 0}, "")
+
+
+_SPEC_TYPES = ("json_schema", "json_object", "ebnf", "regex", "structural_tag")
+
+
+def extract_request_grammar_spec(request: dict[str, Any]) -> dict[str, Any] | None:
+ response_format = request.get("response_format") or {}
+ if not isinstance(response_format, dict):
+ raise InvalidGrammarError("response_format must be an object")
+ rf_type = response_format.get("type")
+ json_schema = None
+ structural_tag = None
+ if rf_type == "json_schema":
+ json_schema = (response_format.get("json_schema") or {}).get("schema")
+ if json_schema is None:
+ raise InvalidGrammarError("response_format json_schema requires json_schema.schema")
+ elif rf_type == "json_object":
+ return {"type": "json_object", "value": None}
+ elif rf_type == "structural_tag":
+ structural_tag = response_format
+ if json_schema is not None:
+ spec = {"type": "json_schema", "value": json_schema}
+ elif request.get("regex") is not None:
+ spec = {"type": "regex", "value": request["regex"]}
+ elif request.get("ebnf") is not None:
+ spec = {"type": "ebnf", "value": request["ebnf"]}
+ elif structural_tag is not None:
+ spec = {"type": "structural_tag", "value": structural_tag}
+ else:
+ return None
+ validate_grammar_spec(spec)
+ return spec
diff --git a/tilert/pd_vllm/logprobs.py b/tilert/pd_vllm/logprobs.py
new file mode 100644
index 0000000..04da4b6
--- /dev/null
+++ b/tilert/pd_vllm/logprobs.py
@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+__all__ = [
+ "GREEDY_TEMPERATURE",
+ "LOGPROB_UNAVAILABLE",
+ "MIN_LOGPROBS_TEMPERATURE",
+ "TOP_LOGPROBS_MAX",
+ "LogprobsRequest",
+ "LogprobsUnsupported",
+ "build_logprobs",
+ "resolve_logprobs_request",
+]
+TOP_LOGPROBS_MAX = 5
+GREEDY_TEMPERATURE = 1e-05
+MIN_LOGPROBS_TEMPERATURE = 0.2
+LOGPROB_UNAVAILABLE = -9999.0
+
+
+class LogprobsUnsupported(Exception):
+ error_type = "invalid_logprobs"
+ http_status = 400
+
+ def to_payload(self) -> dict[str, str]:
+ return {"error": str(self), "error_type": self.error_type}
+
+
+@dataclass(frozen=True)
+class LogprobsRequest:
+ top_n: int
+
+
+def resolve_logprobs_request(body: dict) -> LogprobsRequest | None:
+ enabled = body.get("logprobs")
+ if enabled is not None and (not isinstance(enabled, bool)):
+ raise LogprobsUnsupported(f"logprobs must be a boolean, got {type(enabled).__name__}")
+ raw = body.get("top_logprobs")
+ top_n = 0
+ if raw is not None:
+ if isinstance(raw, bool) or not isinstance(raw, int):
+ raise LogprobsUnsupported(f"top_logprobs must be an integer, got {type(raw).__name__}")
+ if not 0 <= raw <= TOP_LOGPROBS_MAX:
+ raise LogprobsUnsupported(f"top_logprobs must be in [0, {TOP_LOGPROBS_MAX}], got {raw}")
+ if raw > 0 and (not enabled):
+ raise LogprobsUnsupported("when using top_logprobs, logprobs must be set to true")
+ top_n = raw
+ if not enabled:
+ return None
+ temp = body.get("temperature")
+ if isinstance(temp, (int, float)) and (not isinstance(temp, bool)):
+ t = float(temp)
+ if GREEDY_TEMPERATURE <= t < MIN_LOGPROBS_TEMPERATURE:
+ raise LogprobsUnsupported(
+ f"logprobs require temperature >= {MIN_LOGPROBS_TEMPERATURE} or greedy (temperature < {GREEDY_TEMPERATURE}), got {t}"
+ )
+ return LogprobsRequest(top_n=top_n)
+
+
+def _entry(token_id: int, logprob: float | None, decode_one: Callable[[int], str]) -> dict:
+ text = decode_one(token_id)
+ return {
+ "token": text,
+ "logprob": LOGPROB_UNAVAILABLE if logprob is None else _finite(logprob),
+ "bytes": list(text.encode("utf-8")),
+ }
+
+
+def _finite(value: float) -> float:
+ return LOGPROB_UNAVAILABLE if value == float("-inf") else float(value)
+
+
+def build_logprobs(
+ token_ids: list[int],
+ token_logprobs: list[float | None],
+ top_logprobs: list[list[tuple[int, float]]] | None,
+ req: LogprobsRequest,
+ decode_one: Callable[[int], str],
+) -> dict:
+ if len(token_logprobs) != len(token_ids):
+ raise ValueError(
+ f"token_logprobs has {len(token_logprobs)} entries for {len(token_ids)} tokens"
+ )
+ if top_logprobs is not None and len(top_logprobs) != len(token_ids):
+ raise ValueError(
+ f"top_logprobs has {len(top_logprobs)} entries for {len(token_ids)} tokens"
+ )
+ content = []
+ for i, tid in enumerate(token_ids):
+ item = _entry(tid, token_logprobs[i], decode_one)
+ alts = [] if top_logprobs is None else top_logprobs[i][: req.top_n]
+ item["top_logprobs"] = [_entry(alt_id, alt_lp, decode_one) for alt_id, alt_lp in alts]
+ content.append(item)
+ return {"content": content, "refusal": None}
diff --git a/tilert/pd_vllm/oai_parser.py b/tilert/pd_vllm/oai_parser.py
index 931daa9..1c868a3 100644
--- a/tilert/pd_vllm/oai_parser.py
+++ b/tilert/pd_vllm/oai_parser.py
@@ -1,19 +1,4 @@
-"""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 json
import logging
import uuid
from dataclasses import dataclass, field
@@ -23,14 +8,14 @@
@dataclass
class ToolCall:
- call_id: str
+ id: str
name: str
- arguments: str # JSON string (OpenAI convention)
+ arguments: str
def to_openai(self, index: int) -> dict:
return {
"index": index,
- "id": self.call_id,
+ "id": self.id,
"type": "function",
"function": {"name": self.name, "arguments": self.arguments},
}
@@ -47,34 +32,22 @@ 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"),
-}
+_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)}")
+ raise KeyError(f"unknown parser family {family!r}; 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,
- )
+ from vllm.parser.engine.streaming_parser_engine import StreamingParserEngine
mod_name, cfg_name, conv_name = _FAMILIES[family]
mod = importlib.import_module(mod_name)
@@ -94,14 +67,13 @@ def with_thinking(self, thinking: bool) -> "OaiParser":
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 = [], []
+ reasoning, content = ([], [])
slots: dict[int, dict] = {}
for e in events:
if e.type == ET.REASONING_CHUNK:
@@ -115,32 +87,18 @@ def _reduce(self, events) -> Parsed:
for i in sorted(slots):
name = "".join(slots[i]["name"]).strip()
if not name:
- continue # unnamed fragment (heavy truncation) β drop
+ continue
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
@@ -155,7 +113,6 @@ def feed(self, delta_text: str) -> list[dict]:
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
@@ -200,28 +157,35 @@ def _flush_tool(self, index: int, partial: bool) -> dict | None:
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):
+ def __init__(self, tokenizer, skip_special_tokens: bool = False):
self._tok = tokenizer
+ self._skip = skip_special_tokens
self._ids: list[int] = []
self._emitted = 0
+ self._holding = False
+
+ @property
+ def holding(self) -> bool:
+ return self._holding
+
+ def finish(self) -> str:
+ if not self._holding:
+ return ""
+ text = self._tok.decode(self._ids, skip_special_tokens=self._skip)
+ delta = text[self._emitted :]
+ self._emitted = len(text)
+ self._holding = False
+ return delta # noqa: R504 (self._emitted mutated after delta is computed)
def push(self, ids: list[int]) -> str:
self._ids.extend(ids)
- text = self._tok.decode(self._ids, skip_special_tokens=False)
+ text = self._tok.decode(self._ids, skip_special_tokens=self._skip)
if text.endswith("οΏ½"):
+ self._holding = True
return ""
+ self._holding = False
delta = text[self._emitted :]
self._emitted = len(text)
if len(self._ids) > self._FOLD:
diff --git a/tilert/pd_vllm/openai_params.py b/tilert/pd_vllm/openai_params.py
new file mode 100644
index 0000000..b2fb18a
--- /dev/null
+++ b/tilert/pd_vllm/openai_params.py
@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from pydantic import TypeAdapter, ValidationError
+
+__all__ = ["InvalidOutputLength", "resolve_max_tokens"]
+_DECLARED = ("max_tokens", "max_completion_tokens")
+_PRECEDENCE = ("max_completion_tokens", "max_tokens")
+_INT = TypeAdapter(int)
+
+
+class InvalidOutputLength(ValueError):
+ pass
+
+
+def resolve_max_tokens(body: dict, default: int | None = None) -> int | None:
+ seen: dict[str, int] = {}
+ for name in _DECLARED:
+ raw = body.get(name)
+ if raw is None:
+ continue
+ try:
+ seen[name] = _INT.validate_python(raw)
+ except ValidationError:
+ raise InvalidOutputLength(f"{name} must be an integer, got {raw!r}") from None
+ for name in _PRECEDENCE:
+ if name in seen:
+ if seen[name] < 1:
+ raise InvalidOutputLength(f"{name} must be at least 1, got {seen[name]}")
+ return seen[name]
+ return default
diff --git a/tilert/pd_vllm/pd_router.py b/tilert/pd_vllm/pd_router.py
index 245d444..ac48794 100644
--- a/tilert/pd_vllm/pd_router.py
+++ b/tilert/pd_vllm/pd_router.py
@@ -1,34 +1,9 @@
-"""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 -> wait up to
- --queue-timeout, then 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 contextlib
+import functools
import json
import logging
+import os
import threading
import time
@@ -36,56 +11,83 @@
import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse
-
+from pydantic import TypeAdapter, ValidationError
+
+from tilert.pd_vllm import generation_defaults
+from tilert.pd_vllm.capabilities import (
+ CapabilityError,
+ CapabilityUnavailable,
+ InvalidParameter,
+ NodeCapabilities,
+ validate_generation_request,
+)
+from tilert.pd_vllm.decode_pool import DecodeNode, NodeLease, Pool, acquire_lease, cancel_decode
+from tilert.pd_vllm.decode_response import (
+ BUSY,
+ OK,
+ PROPAGATE,
+ PROPAGATED_ERROR_TYPES,
+ REFUSED,
+ RETRY,
+ SERVER_ERROR,
+ TRUNCATED,
+ TYPED_ERROR,
+ UNTYPED_ERROR,
+ DecodeReader,
+ classify_decode_status,
+ decode_refusal,
+ terminal_verdict,
+)
+from tilert.pd_vllm.generation_defaults import GenerationDefaults, UnsupportedGenerationDefault
+from tilert.pd_vllm.grammar_spec import GrammarError, extract_request_grammar_spec
+from tilert.pd_vllm.logprobs import LogprobsUnsupported, resolve_logprobs_request
+from tilert.pd_vllm.openai_params import resolve_max_tokens
+from tilert.pd_vllm.presentation import SseWriter, blocking_choice, blocking_envelope, collect
+from tilert.pd_vllm.presentation import finish_reason as reply_finish_reason
+from tilert.pd_vllm.presentation import sse_chunk, sse_delta, textless_choice, usage_chunk
+from tilert.pd_vllm.reply import CONTENT, REASONING, TOOL_CALL, ReplyStream, as_logprobs
+from tilert.pd_vllm.request_gate import gate_request
+from tilert.pd_vllm.stop_strings import resolve_stop
from tilert.pd_vllm.wire import derive_rid
logger = logging.getLogger("pd_vllm.router")
-
+_HTTP_TIMEOUT_S = float((os.environ.get("TILERT_PD_HTTP_TIMEOUT_S") or "3600").strip() or 3600)
QUEUE_LOG_SECONDS = 0.1
+_DECODE_MAX_TOKENS_DEFAULT = 256
+_DECODE_BUSY_ATTEMPTS = 2
+_DECODE_BUSY_RETRY_S = 0.5
+_RETRY_LOG = "decode node %s busy for %s, retrying in %.1fs"
+
+
+def _log_refusal(verdict: str, node, rid: str, status: int, payload=None) -> None:
+ if verdict == BUSY:
+ logger.warning(
+ "decode node %s still busy for %s after %d attempts",
+ node.http_base,
+ rid,
+ _DECODE_BUSY_ATTEMPTS,
+ )
+ elif verdict == SERVER_ERROR:
+ logger.error(
+ "decode node %s returned %d for %s: %s", node.http_base, status, rid, str(payload)[:200]
+ )
-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:
- """Decode-node reservation; ``queue_timeout`` > 0 waits instead of failing fast."""
+_DISCONNECT_POLL_S = 0.2
+# How long to wait for a cancelled helper task to actually finish. Bounded on
+# purpose: `await ` is NOT guaranteed to return (see _watch).
+_TASK_SETTLE_S = 5.0
- def __init__(self, nodes: list[DecodeNode], queue_timeout: float = 0.0):
- self.nodes = nodes
- self.queue_timeout = queue_timeout
- self._cv = threading.Condition()
- def acquire(self) -> DecodeNode | None:
- """Reserve a node, or None once ``queue_timeout`` elapses."""
- deadline = time.monotonic() + self.queue_timeout
- with self._cv:
- while True:
- for n in self.nodes:
- if not n.busy:
- n.busy = True
- return n
- remaining = deadline - time.monotonic()
- if remaining <= 0:
- return None
- self._cv.wait(remaining)
+class PrefillClientError(Exception):
- def release(self, node: DecodeNode) -> None:
- with self._cv:
- node.busy = False
- self._cv.notify()
+ def __init__(self, status: int, payload):
+ super().__init__(f"vLLM rejected the request with {status}")
+ self.status = status
+ self.payload = payload
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
@@ -100,30 +102,51 @@ def first_token_from_logprobs(resp: dict, is_chat: bool) -> int:
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"
+ f"cannot extract first token id from logprobs ({tok!r}); launch vLLM 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))
-
-
-# Client fields that must not survive into the prefill request, which is
-# forwarded verbatim apart from the fields we set: stream_options contradicts
-# the stream=False we force (vLLM rejects the pair with a 400 during body
-# parsing), and max_completion_tokens takes precedence over max_tokens, so it
-# would override our max_tokens=1. Streaming clients send both.
-_PREFILL_DROP_FIELDS = ("stream_options", "max_completion_tokens")
-
-
-def build_prefill_body(path: str, body: dict, node: DecodeNode) -> dict:
- """The vLLM request that prefills only and hands the KV state to ``node``.
-
- Lives outside ``build_app`` so the rewrite can be exercised without a
- router process, a vLLM instance or a decode node.
- """
+def _token_id_of(tok: object) -> int | None:
+ if isinstance(tok, str) and tok.startswith("token_id:"):
+ try:
+ return int(tok.split(":", 1)[1])
+ except ValueError:
+ return None
+ return None
+
+
+def first_token_logprob_from_prefill(resp: dict, top_n: int):
+ choice = (resp.get("choices") or [{}])[0]
+ content = (choice.get("logprobs") or {}).get("content") or []
+ if not content:
+ return (None, [])
+ entry = content[0]
+ lp = entry.get("logprob")
+ lp = float(lp) if isinstance(lp, (int, float)) else None
+ cands = []
+ for alt in (entry.get("top_logprobs") or [])[:top_n]:
+ alt_id = _token_id_of(alt.get("token"))
+ alt_lp = alt.get("logprob")
+ if alt_id is not None and isinstance(alt_lp, (int, float)):
+ cands.append((alt_id, float(alt_lp)))
+ return (lp, cands)
+
+
+_PREFILL_DROP_FIELDS = (
+ "stream_options",
+ "max_completion_tokens",
+ "stop",
+ "include_stop_str_in_output",
+)
+
+
+def build_prefill_body(
+ path: str,
+ body: dict,
+ node: DecodeNode,
+ logprobs_req=None,
+ defaults: GenerationDefaults | None = None,
+) -> dict:
prefill_body = dict(body)
prefill_body["max_tokens"] = 1
prefill_body["stream"] = False
@@ -131,9 +154,10 @@ def build_prefill_body(path: str, body: dict, node: DecodeNode) -> dict:
prefill_body.pop(field, None)
if path.endswith("chat/completions"):
prefill_body["logprobs"] = True
- prefill_body["top_logprobs"] = 1
+ prefill_body["top_logprobs"] = max(1, logprobs_req.top_n if logprobs_req is not None else 0)
else:
prefill_body["logprobs"] = 1
+ prefill_body.update((defaults or GenerationDefaults()).resolve(body))
prefill_body["kv_transfer_params"] = {
"tilert_host": node.host,
"tilert_ctrl_port": node.ctrl_port,
@@ -141,14 +165,47 @@ def build_prefill_body(path: str, body: dict, node: DecodeNode) -> dict:
return prefill_body
+def build_usage(prompt_tokens, completion_tokens: int) -> dict:
+ prompt = int(prompt_tokens or 0)
+ completion = int(completion_tokens)
+ return {
+ "prompt_tokens": prompt,
+ "completion_tokens": completion,
+ "total_tokens": prompt + completion,
+ }
+
+
+def should_include_usage(body: dict, force: bool = False) -> bool:
+ if force:
+ return True
+ opts = body.get("stream_options")
+ return bool(isinstance(opts, dict) and opts.get("include_usage"))
+
+
+def _json_or_none(response) -> dict | None:
+ try:
+ return response.json()
+ except ValueError:
+ return None
+
+
class RouterCtx:
- """Immutable per-process context (tokenizer, parser factory, config)."""
- def __init__(self, vllm_url: str, pool: Pool, tokenizer, parser_name: str):
+ def __init__(
+ self,
+ vllm_url: str,
+ pool: Pool,
+ tokenizer,
+ parser_name: str,
+ force_include_usage: bool = False,
+ gen_defaults: GenerationDefaults | None = None,
+ ):
self.vllm_url = vllm_url
self.pool = pool
self.tokenizer = tokenizer
self.parser_name = parser_name
+ self.gen_defaults = gen_defaults or GenerationDefaults()
+ self.force_include_usage = force_include_usage
self._parsers = {}
if parser_name != "none":
if tokenizer is None:
@@ -163,21 +220,42 @@ def parser(self, thinking: bool):
return self._parsers.get(thinking)
+def _tilert_bind_lease_to_task(lease):
+ """Release `lease` when the current ASGI task ends, however it ends.
+
+ Idempotent by way of NodeLease.release, so a handler that releases normally
+ is unaffected; this only fires for the exits that run no `finally` of ours
+ -- a streaming response cancelled while the generator is suspended at a
+ `yield`, which leaves that generator suspended and its `finally` unrun.
+ """
+ import asyncio
+ import logging
+
+ task = asyncio.current_task()
+ if task is None:
+ return
+
+ def _reap(_t, lease=lease):
+ if lease.released:
+ return
+ logging.getLogger("pd_vllm.pool").warning(
+ "request task ended without releasing decode node %s (rid=%s); reclaiming it",
+ lease.node.host,
+ lease.rid,
+ )
+ lease.release()
+
+ task.add_done_callback(_reap)
+
+
def build_app(ctx: RouterCtx) -> FastAPI:
app = FastAPI()
pool = ctx.pool
- def _acquire_node() -> tuple[DecodeNode | None, float]:
- """Reserve a node, and report how long the caller queued for it."""
- t0 = time.monotonic()
- node = pool.acquire()
- waited = time.monotonic() - t0
- if node is not None and waited >= QUEUE_LOG_SECONDS:
- logger.info("queued %.1fs for decode node %s", waited, node.host)
- return node, waited
+ def _parser_active(thinking: bool) -> bool:
+ return ctx.parser(thinking) is not None
def _busy_response(waited: float) -> JSONResponse:
- """429 body, distinguishing a full pool from an exhausted queue timeout."""
detail = (
f"no decode node free after waiting {waited:.1f}s"
if pool.queue_timeout > 0
@@ -193,277 +271,448 @@ def health():
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 = build_prefill_body(path, body, node)
- r = requests.post(f"{ctx.vllm_url}{path}", json=prefill_body, timeout=600)
+ def _prefill(path, body, node, logprobs_req=None):
+ prefill_body = build_prefill_body(path, body, node, logprobs_req, ctx.gen_defaults)
+ r = requests.post(f"{ctx.vllm_url}{path}", json=prefill_body, timeout=_HTTP_TIMEOUT_S)
+ if 400 <= r.status_code < 500:
+ try:
+ payload = r.json()
+ except ValueError:
+ payload = {"error": r.text[:500]}
+ raise PrefillClientError(r.status_code, payload)
r.raise_for_status()
return r.json()
def _sampling_of(body):
- return {k: body[k] for k in ("temperature", "top_p", "top_k", "ignore_eos") if k in body}
-
- def _max_tokens_of(body):
- return int(body.get("max_tokens") or body.get("max_completion_tokens") or 256)
+ forwarded = (
+ "temperature",
+ "top_p",
+ "top_k",
+ "repetition_penalty",
+ "presence_penalty",
+ "ignore_eos",
+ )
+ sampling = {k: body[k] for k in forwarded if k in body}
+ sampling.update(ctx.gen_defaults.resolve(body))
+ return sampling
+
+ def _decode_body(
+ rid, first_token_id, body, grammar_spec, *, stream=False, logprobs_req=None, thinking=True
+ ):
+ payload = {
+ "rid": rid,
+ "first_token_id": first_token_id,
+ "max_tokens": resolve_max_tokens(body, _DECODE_MAX_TOKENS_DEFAULT),
+ "sampling": _sampling_of(body),
+ }
+ if stream:
+ payload["stream"] = True
+ if grammar_spec is not None:
+ payload["grammar_spec"] = grammar_spec
+ payload["enable_thinking"] = thinking
+ if logprobs_req is not None:
+ payload["top_logprobs"] = logprobs_req.top_n
+ return payload
- # ββ non-streaming ββββββββββββββββββββββββββββββββββββββββββββββββββββ
def _handle(path: str, body: dict):
- is_chat = path.endswith("chat/completions")
- node, waited = _acquire_node()
- if node is None:
+ try:
+ req = gate_request(path, body, tokenizer=ctx.tokenizer, parser_active=_parser_active)
+ validate_generation_request(body, pool.capabilities(), ctx.gen_defaults.resolve({}))
+ except (CapabilityError, GrammarError, LogprobsUnsupported) as e:
+ return JSONResponse(e.to_payload(), status_code=e.http_status)
+ is_chat, stop, include_stop = (req.is_chat, req.stop, req.include_stop)
+ logprobs_req, grammar_spec = (req.logprobs_req, req.grammar_spec)
+ lease, waited = acquire_lease(pool)
+ if lease is None:
return _busy_response(waited)
+ node = lease.node
t0 = time.time()
+ created = int(t0)
+ reader = None
try:
- prefill = _prefill(path, body, node)
+ prefill = _prefill(path, body, node, logprobs_req)
t_prefill = time.time()
- rid = derive_rid(prefill["id"])
+ rid = lease.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,
+ first_lp = (
+ first_token_logprob_from_prefill(prefill, logprobs_req.top_n)
+ if logprobs_req is not None
+ else None
)
- 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
+ parser = ctx.parser(req.thinking) if is_chat else None
+ asm = (
+ ReplyStream(
+ ctx.tokenizer,
+ stop=stop,
+ include_stop_in_output=include_stop,
+ parser_session=parser.stream() if parser else None,
+ logprobs_req=logprobs_req,
+ first_token_logprob=first_lp,
)
- if is_chat:
- choice["message"] = {"role": "assistant", "content": text}
+ if ctx.tokenizer is not None
+ else None
+ )
+ want_stream = bool(stop)
+ reader = DecodeReader(stream=asm, logprobs_req=logprobs_req, rid=rid)
+ decode_payload = _decode_body(
+ rid,
+ first_token_id,
+ body,
+ grammar_spec,
+ stream=want_stream,
+ logprobs_req=logprobs_req,
+ thinking=req.thinking,
+ )
+ dr = None
+ for attempt in range(1, _DECODE_BUSY_ATTEMPTS + 1):
+ lease.dispatched = True
+ dr = requests.post(
+ f"{node.http_base}/pd/decode",
+ json=decode_payload,
+ timeout=_HTTP_TIMEOUT_S,
+ stream=want_stream,
+ )
+ if dr.status_code == 200:
+ break
+ verdict = classify_decode_status(
+ dr.status_code,
+ _json_or_none(dr),
+ attempts_left=attempt < _DECODE_BUSY_ATTEMPTS,
+ propagated_types=PROPAGATED_ERROR_TYPES,
+ )
+ if want_stream:
+ dr.close()
+ if verdict == RETRY:
+ logger.info(_RETRY_LOG, node.http_base, rid, _DECODE_BUSY_RETRY_S)
+ time.sleep(_DECODE_BUSY_RETRY_S)
+ continue
+ _log_refusal(verdict, node, rid, dr.status_code)
+ body, status = decode_refusal(verdict, dr.status_code, _json_or_none(dr), rid)
+ return JSONResponse(body, status_code=status)
+ emissions = []
+ try:
+ if not want_stream:
+ emissions = reader.feed_blocking(dr.json())
else:
- choice["text"] = text
- choice["token_ids"] = token_ids
-
+ for line in dr.iter_lines(decode_unicode=True):
+ emissions += reader.feed(line)
+ if reader.finished:
+ break
+ finally:
+ if want_stream and dr is not None:
+ dr.close()
+ verdict, payload, status = terminal_verdict(reader)
+ if verdict == TRUNCATED:
+ logger.warning(
+ "decode stream for %s ended after %d tokens with no done/error message",
+ rid,
+ len(asm.token_ids) if asm else 0,
+ )
+ timing = reader.timing
+ finish = reader.finish_reason
+ token_ids = reader.token_ids
+ if verdict != OK:
+ return JSONResponse(payload, status_code=status)
+ if asm is not None:
+ emissions += asm.finish()
+ token_ids = asm.token_ids
+ if asm is None:
+ choice, n_completion = textless_choice(
+ is_chat=is_chat, from_node=finish, token_ids=token_ids
+ )
+ else:
+ choice, n_completion = blocking_choice(
+ collect(emissions),
+ is_chat=is_chat,
+ stream=asm,
+ from_node=finish,
+ logprobs_asked=logprobs_req is not None,
+ 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,
- },
- }
+ blocking_envelope(
+ choice,
+ is_chat=is_chat,
+ prefill=prefill,
+ created=created,
+ usage=build_usage(
+ (prefill.get("usage") or {}).get("prompt_tokens"), n_completion
+ ),
+ timing={"prefill": round(1000 * (t_prefill - t0), 1), **timing},
+ )
)
+ except PrefillClientError as e:
+ logger.info("vLLM rejected the request (%d): %s", e.status, str(e.payload)[:200])
+ return JSONResponse(e.payload, status_code=e.status)
except Exception as e:
logger.exception("pd request failed")
return JSONResponse({"error": str(e)}, status_code=502)
finally:
- pool.release(node)
+ lease.release(terminated=reader is not None and reader.node_terminated)
+
+ async def _settle(task, what: str) -> None:
+ """Wait for a cancelled task to finish, but never indefinitely.
+
+ `await ` reads as a formality and is not one: a task
+ can absorb its cancellation and keep running (see _watch), and then
+ this await is where the request stops for good. asyncio.wait does not
+ re-raise what the task raised, so nothing has to be suppressed here --
+ both callers are best effort.
+ """
+ import asyncio
+
+ done, _pending = await asyncio.wait({task}, timeout=_TASK_SETTLE_S)
+ if not done:
+ logger.warning(
+ "%s did not stop within %.0fs after being cancelled; abandoning it",
+ what,
+ _TASK_SETTLE_S,
+ )
+
+ async def _send_watching_client(client, req, request):
+ import asyncio
+
+ send = asyncio.ensure_future(client.send(req, stream=True))
+ # Stops the watcher WITHOUT relying on cancellation; see _watch.
+ stop = asyncio.Event()
+
+ async def _watch():
+ """Poll for a client disconnect until `stop` is set.
+
+ Driven by the event rather than by cancellation, because cancelling
+ this task is not reliable. Starlette runs its disconnect poll inside
+ an anyio CancelScope it has ALREADY cancelled
+ (`cs.cancel(); await self._receive()`), and such a scope absorbs a
+ CancelledError delivered from outside it just as readily as its own.
+ A watch.cancel() landing in that window is swallowed, the loop goes
+ round again, and the `await watch` below never returns, stranding
+ the caller and its decode lease.
+ """
+ while not stop.is_set():
+ if await request.is_disconnected():
+ return
+ with contextlib.suppress(asyncio.TimeoutError, TimeoutError):
+ await asyncio.wait_for(stop.wait(), _DISCONNECT_POLL_S)
+
+ watch = asyncio.ensure_future(_watch())
+ try:
+ await asyncio.wait({send, watch}, return_when=asyncio.FIRST_COMPLETED)
+ if send.done():
+ return send.result()
+ send.cancel()
+ await _settle(send, "decode send")
+ return None
+ finally:
+ # Order matters: `stop` first, so the watcher ends even if the
+ # cancel that follows is swallowed.
+ stop.set()
+ watch.cancel()
+ await _settle(watch, "disconnect watcher")
- # ββ streaming (chat only) ββββββββββββββββββββββββββββββββββββββββββββ
async def _handle_stream(path: str, body: dict, request: Request):
import anyio
from starlette.concurrency import run_in_threadpool
- node = None # reserved inside the try: both handlers below release it
try:
- # Shielded: a disconnect must not strand a reservation mid-acquire.
+ req = gate_request(path, body, tokenizer=ctx.tokenizer, parser_active=_parser_active)
+ caps = await run_in_threadpool(pool.capabilities)
+ validate_generation_request(body, caps, ctx.gen_defaults.resolve({}))
+ except (CapabilityError, GrammarError, LogprobsUnsupported) as e:
+ return JSONResponse(e.to_payload(), status_code=e.http_status)
+ stop, include_stop = (req.stop, req.include_stop)
+ logprobs_req, grammar_spec = (req.logprobs_req, req.grammar_spec)
+ lease = None
+ try:
with anyio.CancelScope(shield=True):
- node, waited = await run_in_threadpool(_acquire_node)
- if node is None:
+ lease, waited = await run_in_threadpool(acquire_lease, pool)
+ if lease is None:
return _busy_response(waited)
- prefill = await run_in_threadpool(_prefill, path, body, node)
- rid = derive_rid(prefill["id"])
+ _tilert_bind_lease_to_task(lease)
+ node = lease.node
+ prefill = await run_in_threadpool(_prefill, path, body, node, logprobs_req)
+ rid = lease.rid = derive_rid(prefill["id"])
first_token_id = first_token_from_logprobs(prefill, True)
+ first_lp = (
+ first_token_logprob_from_prefill(prefill, logprobs_req.top_n)
+ if logprobs_req is not None
+ else None
+ )
+ except PrefillClientError as e:
+ if lease is not None:
+ lease.release()
+ logger.info("vLLM rejected the stream request (%d): %s", e.status, str(e.payload)[:200])
+ return JSONResponse(e.payload, status_code=e.status)
except Exception as e:
- if node is not None:
- pool.release(node)
+ if lease is not None:
+ lease.release()
logger.exception("pd stream request failed before streaming")
return JSONResponse({"error": str(e)}, status_code=502)
except BaseException:
- # CancelledError is not an Exception; without this the node stays busy.
- if node is not None:
- pool.release(node)
+ if lease is not None:
+ lease.release()
raise
+ import asyncio
+
+ import httpx
+ client = httpx.AsyncClient(timeout=httpx.Timeout(_HTTP_TIMEOUT_S, read=_HTTP_TIMEOUT_S))
+ decode_resp = None
+ try:
+ for attempt in range(1, _DECODE_BUSY_ATTEMPTS + 1):
+ lease.dispatched = True
+ decode_resp = await _send_watching_client(
+ client,
+ client.build_request(
+ "POST",
+ f"{node.http_base}/pd/decode",
+ json=_decode_body(
+ rid,
+ first_token_id,
+ body,
+ grammar_spec,
+ stream=True,
+ logprobs_req=logprobs_req,
+ thinking=req.thinking,
+ ),
+ ),
+ request,
+ )
+ if decode_resp is None:
+ logger.info(
+ "client disconnected while the decode node was still holding headers for %s",
+ rid,
+ )
+ return JSONResponse(
+ {
+ "error": "client disconnected",
+ "error_type": "request_cancelled",
+ "rid": rid,
+ },
+ status_code=499,
+ )
+ if decode_resp.status_code != 429 or attempt == _DECODE_BUSY_ATTEMPTS:
+ break
+ await decode_resp.aclose()
+ logger.info(_RETRY_LOG, node.http_base, rid, _DECODE_BUSY_RETRY_S)
+ await asyncio.sleep(_DECODE_BUSY_RETRY_S)
+ if decode_resp.status_code != 200:
+ payload = None
+ if decode_resp.status_code != 429:
+ try:
+ await decode_resp.aread()
+ payload = decode_resp.json()
+ except Exception:
+ payload = None
+ await decode_resp.aclose()
+ verdict = classify_decode_status(
+ decode_resp.status_code,
+ payload,
+ attempts_left=False,
+ propagated_types=PROPAGATED_ERROR_TYPES,
+ )
+ _log_refusal(verdict, node, rid, decode_resp.status_code, payload)
+ body, status = decode_refusal(verdict, decode_resp.status_code, payload, rid)
+ return JSONResponse(body, status_code=status)
+ except Exception as e:
+ logger.exception("opening the decode stream failed for %s", rid)
+ with contextlib.suppress(Exception):
+ if decode_resp is not None:
+ await decode_resp.aclose()
+ return JSONResponse({"error": str(e)}, status_code=502)
+ finally:
+ if decode_resp is None or decode_resp.status_code != 200:
+ lease.release()
+ with contextlib.suppress(Exception):
+ await client.aclose()
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 _usage_chunk(usage: dict) -> str:
- payload = {
- "id": chunk_id,
- "object": "chat.completion.chunk",
- "created": int(time.time()),
- "model": model,
- "choices": [],
- "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)
+ parser = ctx.parser(req.thinking)
+ created = int(time.time())
+ _chunk = functools.partial(sse_chunk, chunk_id=chunk_id, model=model, created=created)
+ _usage_chunk = functools.partial(
+ usage_chunk, chunk_id=chunk_id, model=model, created=created
+ )
async def _gen():
import anyio
- import httpx
- from tilert.pd_vllm.oai_parser import IncrementalDetok
-
- n_tokens = 0
- saw_tool = False
+ asm = ReplyStream(
+ ctx.tokenizer,
+ stop=stop,
+ include_stop_in_output=include_stop,
+ parser_session=parser.stream() if parser else None,
+ logprobs_req=logprobs_req,
+ first_token_logprob=first_lp,
+ )
+ reader = DecodeReader(stream=asm, logprobs_req=logprobs_req, rid=rid)
+ out = SseWriter(asm, _chunk)
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))
- role_sent = False
-
- def _role_once():
- nonlocal role_sent
- role_sent = True
- return _chunk({"role": "assistant"})
-
+ decode_done = False
try:
- 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 with contextlib.aclosing(decode_resp) as resp:
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 not role_sent:
- yield _role_once()
- 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))
- # R508 false positive: https://github.com/afonasev/flake8-return/issues/137
- elif "done" in msg: # noqa: R508
- finish_reason = msg.get("finish_reason", "stop")
- if finish_reason == "cancelled":
- finish_reason = "stop"
- break
- elif "error" in msg:
- if not role_sent:
- yield _role_once()
- yield _chunk({"content": f"\n[decode error: {msg['error']}]"})
- finish_reason = "stop"
+ for frame in out.frames(reader.feed(line)):
+ yield frame
+ if reader.finished:
+ if reader.stop_hit:
+ logger.info("stop string %r ended %s", asm.stop_reason, rid)
break
+ decode_done = reader.node_terminated
+ finish_reason = reader.finish_reason
+ verdict, payload, _status = terminal_verdict(reader, client_gone=client_gone)
+ if verdict == REFUSED:
+ logger.warning(
+ "decode node %s sent an unusable logprobs line for %s", node.http_base, rid
+ )
+ elif verdict == TRUNCATED:
+ logger.warning(
+ "decode stream for %s ended after %d tokens with no done/error message",
+ rid,
+ len(asm.token_ids) if asm else 0,
+ )
+ if verdict in (REFUSED, TRUNCATED, TYPED_ERROR):
+ for _c in out.fail_closed(payload):
+ yield _c
+ return
+ if verdict == UNTYPED_ERROR:
+ for _c in out.flush_held():
+ yield _c
+ yield _chunk({"content": f"\n[decode error: {payload['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"
- if not role_sent:
- yield _role_once()
- yield _chunk({}, finish=finish_reason)
- yield _usage_chunk(
- {
- "prompt_tokens": prompt_tokens,
- "completion_tokens": n_tokens,
- "total_tokens": (prompt_tokens or 0) + n_tokens,
- }
+ return
+ for _c in out.flush_held():
+ yield _c
+ yield _chunk(
+ {},
+ finish=reply_finish_reason(
+ saw_tool=out.saw_tool, from_node=finish_reason, stream=asm
+ ),
+ stop_reason=asm.stop_reason,
)
+ if should_include_usage(body, ctx.force_include_usage):
+ yield _usage_chunk(build_usage(prompt_tokens, asm.completion_tokens))
yield "data: [DONE]\n\n"
- completed_ok = True
- except Exception:
+ except Exception as exc:
logger.exception("stream failed mid-flight for %s", rid)
+ with contextlib.suppress(Exception):
+ for _c in out.fail_closed(
+ {
+ "error": f"stream failed: {exc}",
+ "error_type": "decode_stream_failed",
+ "rid": rid,
+ }
+ ):
+ yield _c
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()
+ lease.release(terminated=decode_done)
with anyio.CancelScope(shield=True):
await client.aclose()
@@ -476,7 +725,6 @@ async def chat(request: Request):
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")
@@ -509,7 +757,35 @@ def main() -> None:
"--parser",
choices=["glm47", "none"],
default="glm47",
- help="output parser (reasoning + tool calls)",
+ help="output parser (reasoning + tool calls); anything but 'none' loads vLLM's parser engine and needs vllm importable",
+ )
+ ap.add_argument("--model", default="", help="model profile")
+ ap.add_argument(
+ "--generation-config",
+ choices=["auto", "vllm"],
+ default="auto",
+ help="where sampling defaults come from, mirroring vLLM's flag of the same name: 'auto' reads the model's generation_config.json under --model-path, 'vllm' ignores it and uses the neutral defaults. Whichever is chosen, the resolved values are sent explicitly to BOTH legs so they cannot disagree.",
+ )
+ ap.add_argument(
+ "--default-temperature",
+ type=float,
+ default=None,
+ help="override the resolved temperature default (vLLM: --override-generation-config)",
+ )
+ ap.add_argument(
+ "--default-top-p", type=float, default=None, help="override the resolved top_p default"
+ )
+ ap.add_argument(
+ "--default-top-k",
+ type=int,
+ default=None,
+ help="override the resolved top_k default; 0 disables the rank cut, as it does in vLLM",
+ )
+ ap.add_argument(
+ "--default-repetition-penalty",
+ type=float,
+ default=None,
+ help="override the resolved repetition_penalty default. Only executable where --model names a family whose decode runtime implements penalties; 1.0 is the runtime no-op.",
)
ap.add_argument(
"--queue-timeout",
@@ -517,13 +793,16 @@ def main() -> None:
default=0.0,
help="seconds to wait for a free decode node before answering 429 (0: fail fast)",
)
+ ap.add_argument(
+ "--force-include-usage",
+ action="store_true",
+ help="emit the trailing usage chunk on every stream, even when the client omits stream_options.include_usage (vLLM: enable_force_include_usage). Off by default; note it makes every stream end with a choices:[] chunk, which some clients cannot read.",
+ )
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
@@ -531,8 +810,26 @@ def main() -> None:
tokenizer = AutoTokenizer.from_pretrained(
args.model_path, trust_remote_code=True
) # nosec B615
-
- ctx = RouterCtx(args.vllm_url, Pool(nodes, args.queue_timeout), tokenizer, args.parser)
+ try:
+ gen_defaults = generation_defaults.load(
+ args.model_path,
+ args.generation_config,
+ model=args.model,
+ temperature=args.default_temperature,
+ top_p=args.default_top_p,
+ top_k=args.default_top_k,
+ repetition_penalty=args.default_repetition_penalty,
+ )
+ except UnsupportedGenerationDefault as e:
+ raise SystemExit(f"cannot serve with these sampling defaults:\n{e}")
+ ctx = RouterCtx(
+ args.vllm_url,
+ Pool(nodes, args.queue_timeout),
+ tokenizer,
+ args.parser,
+ force_include_usage=args.force_include_usage,
+ gen_defaults=gen_defaults,
+ )
app = build_app(ctx)
logger.info(
"router on :%d -> vllm=%s, %d decode node(s), parser=%s",
diff --git a/tilert/pd_vllm/prefill_connector.py b/tilert/pd_vllm/prefill_connector.py
index abd4819..ec7929b 100644
--- a/tilert/pd_vllm/prefill_connector.py
+++ b/tilert/pd_vllm/prefill_connector.py
@@ -1,27 +1,5 @@
-"""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 os
import queue
import threading
from dataclasses import dataclass, field
@@ -37,6 +15,12 @@
from tilert.pd_vllm.wire import derive_rid
logger = logging.getLogger("pd_vllm.connector")
+_SENT = "sent"
+_REJECTED_TRANSIENT = "rejected_transient"
+_REJECTED_PERMANENT = "rejected_permanent"
+_TRANSIENT_REJECTS = frozenset({"busy", "cancelling"})
+_ADMISSION_ATTEMPTS = 5
+_ADMISSION_BACKOFF_S = 0.2
@dataclass
@@ -49,6 +33,7 @@ class _ReqMeta:
tilert_host: str
tilert_ctrl_port: int
sampling: dict | None = None
+ prompt_token_ids: list = field(default_factory=list)
@dataclass
@@ -58,8 +43,6 @@ class TileRTMetadata(KVConnectorMetadata):
@dataclass
class _Pending:
- """Scheduler-side chunked-prefill accumulation."""
-
req_id: str
prompt_token_ids: list
total_tokens: int
@@ -68,44 +51,39 @@ class _Pending:
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._admission_attempts = int(extra.get("tilert_admission_attempts", _ADMISSION_ATTEMPTS))
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._pd_buffer_device = str(
+ extra.get("tilert_pd_buffer_device")
+ or os.environ.get("TILERT_PD_BUFFER_DEVICE")
+ or "cuda"
+ ).lower()
self._pending: dict[str, _Pending] = {}
-
- # worker-side (lazy)
self._kv_caches: dict = {}
- self._reg = None # profile registration (layer map)
+ self._reg = None
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",
+ "TileRTConnector: role=%s profile=%s target=%s:%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
@@ -114,11 +92,17 @@ 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"))
+ claimed = self._claim(extra.get("kv_transfer_params"))
+ if claimed is not None:
+ claimed = dict(claimed)
+ claimed["_wants_prompt_ids"] = wire.wants_prompt_token_ids(
+ {"repetition_penalty": getattr(sp, "repetition_penalty", 1.0)}
+ )
+ return claimed
return None
def get_num_new_matched_tokens(self, request, num_computed_tokens):
- return 0, False
+ return (0, False)
def update_state_after_alloc(self, request, blocks, num_external_tokens):
pass
@@ -126,16 +110,14 @@ def update_state_after_alloc(self, request, blocks, num_external_tokens):
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)
+ continue
token_ids = list(new_req.prompt_token_ids or [])
if not token_ids:
continue
@@ -151,7 +133,6 @@ def build_connector_meta(self, scheduler_output) -> KVConnectorMetadata:
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)
@@ -171,15 +152,14 @@ def build_connector_meta(self, scheduler_output) -> KVConnectorMetadata:
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]),
+ prompt_token_ids=list(token_ids) if params.get("_wants_prompt_ids") else [],
block_ids_per_group=groups,
- tilert_host=host,
+ tilert_host=params.get("tilert_host") or self._default_host,
tilert_ctrl_port=int(params.get("tilert_ctrl_port", self._default_port)),
sampling=params.get("sampling"),
)
@@ -195,17 +175,32 @@ def _emit(self, req_id, token_ids, groups, params) -> _ReqMeta:
def request_finished(self, request, block_ids):
self._pending.pop(getattr(request, "request_id", ""), None)
- return False, 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)
+ self._collect_kv_scales()
+
+ def _collect_kv_scales(self) -> None:
+ if not getattr(self._profile, "mla_rocm_fp8", False):
+ return
+ ctx = self._vllm_config.compilation_config.static_forward_context
+ scales = []
+ for lid, name, _t, _gi in self._reg.mla_layers:
+ layer = ctx.get(name)
+ k = getattr(layer, "_k_scale", None)
+ if k is None:
+ raise RuntimeError(
+ f"MLA layer {name!r} has no _k_scale; the flat ROCm fp8 KV layout "
+ "cannot be dequantised without it"
+ )
+ scales.append(float(k.item()) if hasattr(k, "item") else float(k))
+ self._profile.set_kv_scales(scales)
def _ensure_worker_ready(self) -> None:
if self._transport is not None:
@@ -213,29 +208,81 @@ def _ensure_worker_ready(self) -> None:
import torch
from vllm.distributed import get_tensor_model_parallel_rank
- self._tp_rank = int(get_tensor_model_parallel_rank())
-
+ self._tp_rank = 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}")
+ host = self._pd_buffer_device == "cpu"
+ senders = len(self._profile.sender_ranks)
+ nshards = int((os.environ.get("TILERT_PD_SHARDS") or "8").strip() or 8)
+ nshards = max(1, min(nshards, torch.cuda.device_count() or 1))
+ if host and nshards > 1:
+ logger.info("staging buffer is host-resident: ignoring TILERT_PD_SHARDS=%d", nshards)
+ nshards = 1
+ if senders > 1:
+ if host:
+ raise RuntimeError(
+ "TILERT_PD_SENDERS>1 and a host-resident staging buffer are mutually "
+ "exclusive: multi-sender exists to keep every shard on its own card"
+ )
+ nshards = senders
+ total = self._profile.staging_bytes(self._reg, self._tp_rank, self._max_seq, nshards)
+ if host and total > 4:
+ from tilert.pd_vllm.transport import alloc_pinned_huge
+ self._staging = alloc_pinned_huge(total)
+ logger.info(
+ "staging buffer: %.2f GiB in pinned host DRAM (rank=%d)",
+ total / 1024**3,
+ self._tp_rank,
+ )
+ elif senders > 1 and total > 4:
+ own = torch.zeros(total, dtype=torch.uint8, device=f"cuda:{dev}")
+ self._staging = [own if i == self._tp_rank else None for i in range(nshards)]
+ logger.info(
+ "staging buffer: own shard only, %.2f GB on cuda:%d (rank=%d of %d senders)",
+ total / 1024**3,
+ dev,
+ self._tp_rank,
+ senders,
+ )
+ elif nshards > 1 and total > 4:
+ self._staging = [
+ torch.zeros(total, dtype=torch.uint8, device=f"cuda:{i}") for i in range(nshards)
+ ]
+ logger.info(
+ "staging buffer: %d shards x %.2f GB on cuda:0..%d (rank=%d)",
+ nshards,
+ total / 1024**3,
+ nshards - 1,
+ self._tp_rank,
+ )
+ else:
+ host = False
+ 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()
+ if isinstance(self._staging, list):
+ rails = self._transport.rails()
+ for i, t in enumerate(self._staging):
+ if t is None:
+ continue
+ self._transport.register(
+ t.data_ptr(), t.numel(), i, f"hip:{i % rails}" if rails else None
+ )
+ logger.info("staging shards pinned to %s rails", rails or "auto")
+ else:
+ self._transport.register(self._staging.data_ptr(), total, dev, host=host)
+ 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,
+ total / 1000000.0,
self._profile.name,
)
@@ -253,9 +300,8 @@ def wait_for_save(self):
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)
+ return
for m in metadata.requests:
try:
sections = self._profile.extract(
@@ -264,22 +310,16 @@ def wait_for_save(self):
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)
+ self._send_q.put({"meta": m, "sections": sections, "seq": sections["seq"]})
def get_finished(self, finished_req_ids):
- return None, None
-
- # ββ background send ββ
+ return (None, None)
def _sender_loop(self) -> None:
while True:
job = self._send_q.get()
try:
- self._send(job)
+ self._send_with_retry(job)
except Exception:
logger.exception("send failed for %s", job["meta"].rid)
@@ -287,9 +327,6 @@ 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()
@@ -300,42 +337,104 @@ def _send(self, job: dict) -> None:
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_proto = hello.get("protocol_version", 1)
+ assert (
+ remote_proto == wire.PROTOCOL_VERSION
+ ), f"control-plane protocol mismatch: decode={remote_proto} vs prefill={wire.PROTOCOL_VERSION}; upgrade both ends together"
+ assert (
+ hello.get("layout_version") == self._profile.layout_version
+ ), f"layout version mismatch: {hello.get('layout_version')} vs {self._profile.layout_version}"
+ remote_senders = int(hello.get("senders", 1) or 1)
+ assert remote_senders == len(self._profile.sender_ranks), (
+ f"sender-count mismatch: decode expects {remote_senders} sender rank(s) but prefill has "
+ f"{len(self._profile.sender_ranks)}; set TILERT_PD_SENDERS to the same value on both roles"
)
+ assert (
+ hello.get("transport") == self._transport.name
+ ), f"transport mismatch: decode={hello.get('transport')} 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,
- },
+ msg = {
+ "rid": m.rid,
+ "rank": self._tp_rank,
+ "seq_len": seq,
+ "last_prompt_token": m.last_prompt_token,
+ "sampling": m.sampling,
+ "admission_window_s": self._admission_window(),
+ }
+ if getattr(self._profile, "kv_scales", None) is not None:
+ msg["kv_scales"] = self._profile.kv_scales
+ if self._tp_rank == 0 and m.prompt_token_ids:
+ msg["prompt_token_ids"] = m.prompt_token_ids
+ wire.send_msg(conn, msg)
+ ack = wire.recv_msg(conn)
+ if not ack.get("accepted"):
+ reason = ack.get("error")
+ logger.warning("decode node refused %s rank=%d: %s", m.rid, self._tp_rank, ack)
+ return _REJECTED_TRANSIENT if reason in _TRANSIENT_REJECTS else _REJECTED_PERMANENT
+ if (
+ ack.get("rid") != m.rid
+ or ack.get("rank") != self._tp_rank
+ or (not isinstance(ack.get("generation"), int))
+ ):
+ logger.error(
+ "discarding %s rank=%d: admission does not match the request (%s)",
+ m.rid,
+ self._tp_rank,
+ ack,
+ )
+ return _REJECTED_PERMANENT
+ generation = ack["generation"]
+ base = (
+ [(t.data_ptr() if t is not None else 0) for t in self._staging]
+ if isinstance(self._staging, list)
+ else self._staging.data_ptr()
)
-
- 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})
+ wire.send_msg(conn, wire.done_msg(m.rid, self._tp_rank, generation))
logger.info(
- "sent %s: rank=%d seq=%d %.1f MB in %.1f ms",
+ "sent %s: rank=%d seq=%d gen=%d %.1f MB in %.1f ms",
m.rid,
self._tp_rank,
seq,
- sum(lens) / 1e6,
+ generation,
+ sum(lens) / 1000000.0,
1000 * (_time.time() - t0),
)
+ return _SENT
finally:
conn.close()
+
+ def _admission_window(self) -> float:
+ return _ADMISSION_BACKOFF_S * (2 ** max(0, self._admission_attempts - 1) - 1)
+
+ def _send_with_retry(self, job: dict) -> None:
+ import time as _time
+
+ m = job["meta"]
+ delay = _ADMISSION_BACKOFF_S
+ for attempt in range(1, self._admission_attempts + 1):
+ outcome = self._send(job)
+ if outcome != _REJECTED_TRANSIENT:
+ return
+ if attempt == self._admission_attempts:
+ break
+ logger.info(
+ "retrying admission for %s rank=%d in %.1fs (attempt %d/%d)",
+ m.rid,
+ self._tp_rank,
+ delay,
+ attempt,
+ self._admission_attempts,
+ )
+ _time.sleep(delay)
+ delay *= 2
+ logger.error(
+ "gave up admitting %s rank=%d after %d attempts: its shard was NOT transferred, so the decode node will wait out its kv_transfer_timeout for this request",
+ m.rid,
+ self._tp_rank,
+ self._admission_attempts,
+ )
diff --git a/tilert/pd_vllm/presentation.py b/tilert/pd_vllm/presentation.py
new file mode 100644
index 0000000..848f42c
--- /dev/null
+++ b/tilert/pd_vllm/presentation.py
@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any
+
+from tilert.pd_vllm.reply import CONTENT, REASONING, TOOL_CALL, as_logprobs
+
+__all__ = [
+ "Collected",
+ "SseWriter",
+ "blocking_envelope",
+ "collect",
+ "finish_reason",
+ "sse_chunk",
+ "sse_delta",
+ "usage_chunk",
+]
+
+
+@dataclass
+class Collected:
+ content: str = ""
+ reasoning: str = ""
+ entries: list = field(default_factory=list)
+ tool_calls: list = field(default_factory=list)
+
+
+def collect(emissions) -> Collected:
+ out = Collected()
+ for e in emissions:
+ if e.channel == CONTENT:
+ out.content += e.text
+ out.entries += e.logprobs
+ elif e.channel == REASONING:
+ out.reasoning += e.text
+ else:
+ out.tool_calls.append(e.tool_call)
+ return out
+
+
+def finish_reason(*, saw_tool: bool, from_node: str, stream) -> str:
+ if saw_tool:
+ return "tool_calls"
+ return stream.finish_reason(from_node)
+
+
+def _tool_calls_field(tool_calls: list) -> list:
+ return [
+ {
+ "index": c["index"],
+ "id": c["id"],
+ "type": "function",
+ "function": {"name": c["name"], "arguments": c["arguments"]},
+ }
+ for c in tool_calls
+ ]
+
+
+def blocking_choice(
+ got: Collected,
+ *,
+ is_chat: bool,
+ stream,
+ from_node: str,
+ logprobs_asked: bool,
+ token_ids: list[int],
+) -> tuple[dict, int]:
+ choice: dict = {
+ "index": 0,
+ "finish_reason": finish_reason(
+ saw_tool=bool(got.tool_calls), from_node=from_node, stream=stream
+ ),
+ "logprobs": as_logprobs(got.entries) if logprobs_asked else None,
+ "stop_reason": stream.stop_reason,
+ }
+ if is_chat:
+ msg = {"role": "assistant", "content": got.content}
+ if got.reasoning:
+ msg["reasoning_content"] = got.reasoning
+ if got.tool_calls:
+ msg["tool_calls"] = _tool_calls_field(got.tool_calls)
+ choice["message"] = msg
+ else:
+ choice["text"] = got.content
+ choice["token_ids"] = token_ids
+ return (choice, stream.completion_tokens)
+
+
+def textless_choice(*, is_chat: bool, from_node: str, token_ids: list[int]) -> tuple[dict, int]:
+ choice: dict = {"index": 0, "finish_reason": from_node, "logprobs": None, "stop_reason": None}
+ if is_chat:
+ choice["message"] = {"role": "assistant", "content": None}
+ else:
+ choice["text"] = None
+ choice["token_ids"] = token_ids
+ return (choice, len(token_ids))
+
+
+def blocking_envelope(
+ choice: dict, *, is_chat: bool, prefill: dict, created: int, usage: dict, timing: dict
+) -> dict:
+ return {
+ "id": prefill["id"],
+ "object": "chat.completion" if is_chat else "text_completion",
+ "created": created,
+ "model": prefill.get("model"),
+ "choices": [choice],
+ "usage": usage,
+ "pd_timing_ms": timing,
+ }
+
+
+def sse_delta(e) -> dict:
+ if e.channel == REASONING:
+ return {"reasoning_content": e.text}
+ if e.channel == CONTENT:
+ return {"content": e.text}
+ return {"tool_calls": _tool_calls_field([e.tool_call])}
+
+
+def sse_chunk(
+ delta: dict,
+ *,
+ chunk_id: str,
+ model: Any,
+ created: int,
+ finish: str | None = None,
+ usage: dict | None = None,
+ logprobs: dict | None = None,
+ stop_reason: str | None = None,
+) -> str:
+ choice: dict = {"index": 0, "delta": delta, "finish_reason": finish}
+ if logprobs is not None:
+ choice["logprobs"] = logprobs
+ if finish is not None:
+ choice["stop_reason"] = stop_reason
+ payload = {
+ "id": chunk_id,
+ "object": "chat.completion.chunk",
+ "created": created,
+ "model": model,
+ "choices": [choice],
+ }
+ if usage is not None:
+ payload["usage"] = usage
+ return _frame(payload)
+
+
+def usage_chunk(usage: dict, *, chunk_id: str, model: Any, created: int) -> str:
+ return _frame(
+ {
+ "id": chunk_id,
+ "object": "chat.completion.chunk",
+ "created": created,
+ "model": model,
+ "choices": [],
+ "usage": usage,
+ }
+ )
+
+
+def _frame(payload: dict) -> str:
+ return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
+
+
+class SseWriter:
+
+ def __init__(self, stream, chunk):
+ self._stream = stream
+ self._chunk = chunk
+ self.role_sent = False
+ self.saw_tool = False
+
+ def role_once(self) -> str:
+ self.role_sent = True
+ return self._chunk({"role": "assistant"})
+
+ def emit(self, e) -> str | None:
+ if e.channel == TOOL_CALL:
+ self.saw_tool = True
+ elif not e.text and (not e.logprobs):
+ return None
+ return self._chunk(sse_delta(e), logprobs=as_logprobs(e.logprobs) if e.logprobs else None)
+
+ def frames(self, emissions):
+ for e in emissions:
+ frame = self.emit(e)
+ if frame is None:
+ continue
+ if not self.role_sent:
+ yield self.role_once()
+ yield frame
+
+ def flush_held(self):
+ yield from self.frames(self._stream.finish())
+ if not self.role_sent:
+ yield self.role_once()
+
+ def fail_closed(self, payload: dict):
+ yield from self.flush_held()
+ yield self._chunk({}, finish="stop")
+ yield _frame({"error": payload})
+ yield "data: [DONE]\n\n"
diff --git a/tilert/pd_vllm/profiles/__init__.py b/tilert/pd_vllm/profiles/__init__.py
index fe663bf..2ae2839 100644
--- a/tilert/pd_vllm/profiles/__init__.py
+++ b/tilert/pd_vllm/profiles/__init__.py
@@ -1 +1 @@
-"""Model profiles for the PD data plane."""
+pass
diff --git a/tilert/pd_vllm/profiles/base.py b/tilert/pd_vllm/profiles/base.py
index f1231d7..74aaebc 100644
--- a/tilert/pd_vllm/profiles/base.py
+++ b/tilert/pd_vllm/profiles/base.py
@@ -1,13 +1,3 @@
-"""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
@@ -16,53 +6,38 @@
class ModelProfile(Protocol):
name: str
num_ranks: int
+ layout_version: 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."""
+ pass
- 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 hello_layout(self, base_ptr, max_seq_len: int) -> dict:
+ pass
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)."""
+ pass
- # ββ 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).
- """
+ pass
- def staging_bytes(self, reg: Any, tp_rank: int, max_seq_len: int) -> int:
- """Per-rank staging-buffer size."""
+ def staging_bytes(self, reg: Any, tp_rank: int, max_seq_len: int, nshards: int = 1) -> int:
+ pass
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``.
- """
+ pass
def rdma_plan(
- self, hello: dict, sections: Any, tp_rank: int, seq_len: int, staging_base: int
+ self, hello: dict, sections: Any, tp_rank: int, seq_len: int, staging_base
) -> tuple[list, list, list]:
- """(src_ptrs, dst_ptrs, lengths) for one mooncake batch write."""
+ pass
- # ββ engine (decode node) βββββββββββββββββββββββββββββββββββββββββββββ
def build_engine(
- self, model_weights_dir: str, max_seq_len: int, with_mtp: bool, ar_steps: int
+ self, model_weights_dir: str, max_seq_len: int, with_mtp: bool, ar_steps: int, num_mtp: int
) -> Any:
- """Construct the decode engine adapter (inject/decode/reset)."""
+ pass
_REGISTRY: dict[str, ModelProfile] = {}
@@ -70,6 +45,16 @@ def build_engine(
"glm5": "glm5",
"glm_5": "glm5",
"glm-5": "glm5",
+ "glm5_2": "glm5_2",
+ "glm_5_2": "glm5_2",
+ "glm-5.2": "glm5_2",
+ "glm5.2": "glm5_2",
+ "glm52": "glm5_2",
+ "glm5_3": "glm5_2",
+ "glm_5_3": "glm5_2",
+ "glm-5.3": "glm5_2",
+ "glm5.3": "glm5_2",
+ "glm53": "glm5_2",
"dsv32": "dsv32",
"deepseek_v3_2": "dsv32",
"deepseek-v3.2": "dsv32",
@@ -88,10 +73,26 @@ def get_profile(name: str) -> ModelProfile:
# 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 == "glm5_2":
+ from tilert.pd_vllm.profiles import glm5_2 # 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)}"
- )
+ raise KeyError(f"unknown model profile {name!r}; registered: {sorted(_REGISTRY)}")
return _REGISTRY[canon]
+
+
+DEFAULT_SUPPORTED_NUM_MTP = (3,)
+
+
+def resolve_num_mtp(profile: ModelProfile, requested: int, *, with_mtp: bool) -> int:
+ supported = tuple(getattr(profile, "supported_num_mtp", DEFAULT_SUPPORTED_NUM_MTP))
+ if requested not in supported:
+ raise ValueError(
+ f"profile {profile.name!r} cannot serve num_mtp={requested}: it supports {list(supported)}."
+ )
+ if requested != 3 and (not with_mtp):
+ raise ValueError(
+ f"num_mtp={requested} needs speculative decoding: pass --with-mtp, or drop --num-mtp to run the non-MTP path."
+ )
+ return requested
diff --git a/tilert/pd_vllm/profiles/dsv32.py b/tilert/pd_vllm/profiles/dsv32.py
index 5691772..14f9660 100644
--- a/tilert/pd_vllm/profiles/dsv32.py
+++ b/tilert/pd_vllm/profiles/dsv32.py
@@ -1,22 +1,17 @@
-"""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,
-)
+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)
+NUM_LAYERS = 61
+LAYOUT_VERSION = 11
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
+ tilert.load_backend("deepseek_v3_2")
from tilert.models.deepseek_v3_2.generator import DSAv32Generator
from tilert.models.deepseek_v3_2.model_args import ModelArgs
diff --git a/tilert/pd_vllm/profiles/glm5.py b/tilert/pd_vllm/profiles/glm5.py
index ec77122..016bc7f 100644
--- a/tilert/pd_vllm/profiles/glm5.py
+++ b/tilert/pd_vllm/profiles/glm5.py
@@ -1,19 +1,9 @@
-"""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
import os
from tilert.pd_vllm.profiles import base
-from tilert.pd_vllm.profiles.mla_nsa import (
- MlaNsaEngineAdapter,
- MlaNsaProfile,
-)
+from tilert.pd_vllm.profiles.mla_nsa import MlaNsaEngineAdapter, MlaNsaProfile
_NO_MTP = (os.environ.get("TILERT_PD_NO_MTP") or "0").strip().lower() not in (
"0",
@@ -21,15 +11,13 @@
"no",
"off",
)
-NUM_LAYERS = 78 if _NO_MTP else 79 # 78 main + 1 MTP draft
-LAYOUT_VERSION = 1010 if _NO_MTP else 10 # glm5 wire family
+NUM_LAYERS = 78 if _NO_MTP else 79
+LAYOUT_VERSION = 1010 if _NO_MTP else 10
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
diff --git a/tilert/pd_vllm/profiles/glm5_2.py b/tilert/pd_vllm/profiles/glm5_2.py
new file mode 100644
index 0000000..859ec35
--- /dev/null
+++ b/tilert/pd_vllm/profiles/glm5_2.py
@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from tilert.pd_vllm.profiles import base
+from tilert.pd_vllm.profiles.glm5_rocm_engine import build_rocm_engine
+from tilert.pd_vllm.profiles.mla_nsa import MlaNsaEngineAdapter, MlaNsaProfile
+
+NUM_LAYERS = 79
+LAYOUT_VERSION = 12
+FULL_LAYERS = [L for L in range(78) if max(L - 2, 0) % 4 == 0] + [78]
+
+
+def _build_engine(model_weights_dir, max_seq_len, with_mtp, ar_steps):
+ return build_rocm_engine(model_weights_dir, max_seq_len, with_mtp, ar_steps)
+
+
+base.register(
+ MlaNsaProfile(
+ name="glm5_2",
+ num_layers=NUM_LAYERS,
+ layout_version=LAYOUT_VERSION,
+ engine_factory=_build_engine,
+ ki_layer_ids=FULL_LAYERS,
+ )
+)
diff --git a/tilert/pd_vllm/profiles/glm5_rocm_engine.py b/tilert/pd_vllm/profiles/glm5_rocm_engine.py
new file mode 100644
index 0000000..b816d2e
--- /dev/null
+++ b/tilert/pd_vllm/profiles/glm5_rocm_engine.py
@@ -0,0 +1,422 @@
+from __future__ import annotations
+
+import ctypes
+import logging
+import os
+
+import torch
+
+from tilert.pd_vllm.grammar_spec import GrammarUnsupported
+from tilert.pd_vllm.sampling import resolve_top_p
+
+logger = logging.getLogger("pd_vllm.profile.glm5_rocm")
+_FP8_MAX = 448.0
+_INDEX_HEAD_DIM = 128
+
+_INJECT_MODE = os.environ.get("TILERT_INJECT_MODE", "rccl").strip().lower()
+
+
+class _RcclBroadcast:
+ _NCCL_CHAR = 0
+
+ def __init__(self, npes: int):
+ self.npes = npes
+ self.lib = None
+ self.path = None
+ err = None
+ for cand in (
+ "librccl.so.1",
+ "librccl.so",
+ os.path.join(os.path.dirname(torch.__file__), "lib", "librccl.so.1"),
+ ):
+ try:
+ self.lib = ctypes.CDLL(cand)
+ self.path = cand
+ break
+ except OSError as exc:
+ err = exc
+ if self.lib is None:
+ raise RuntimeError(f"librccl not loadable: {err}")
+ ver = ctypes.c_int()
+ self.lib.ncclGetVersion(ctypes.byref(ver))
+ self.version = ver.value
+ comms = (ctypes.c_void_p * npes)()
+ devs = (ctypes.c_int * npes)(*range(npes))
+ rc = self.lib.ncclCommInitAll(comms, npes, devs)
+ if rc != 0:
+ raise RuntimeError(f"ncclCommInitAll rc={rc}")
+ self.comms = comms
+ self.streams = [torch.cuda.Stream(device=d) for d in range(npes)]
+
+ def run(self, items) -> None:
+ lib = self.lib
+ rc = lib.ncclGroupStart()
+ if rc != 0:
+ raise RuntimeError(f"ncclGroupStart rc={rc}")
+ for root, src_ptr, nbytes, dsts in items:
+ for r in range(self.npes):
+ send = src_ptr if r == root else dsts[r]
+ rc = lib.ncclBroadcast(
+ ctypes.c_void_p(send),
+ ctypes.c_void_p(dsts[r]),
+ ctypes.c_size_t(nbytes),
+ self._NCCL_CHAR,
+ ctypes.c_int(root),
+ self.comms[r],
+ ctypes.c_void_p(self.streams[r].cuda_stream),
+ )
+ if rc != 0:
+ raise RuntimeError(f"ncclBroadcast rc={rc}")
+ rc = lib.ncclGroupEnd()
+ if rc != 0:
+ raise RuntimeError(f"ncclGroupEnd rc={rc}")
+ for st in self.streams:
+ st.synchronize()
+
+
+def is_rocm_torch() -> bool:
+ return getattr(torch.version, "hip", None) is not None
+
+
+def build_rocm_engine(model_weights_dir, max_seq_len, with_mtp, ar_steps, num_mtp=3):
+ import tilert
+
+ if hasattr(tilert, "load_backend"):
+ tilert.load_backend("glm5_2_rocm")
+ from tilert.models.glm_5_2_rocm.generator import Glm52Generator
+ from tilert.models.glm_5_2_rocm.model_args import ModelArgsGlm52
+
+ if with_mtp and num_mtp not in (1, 3):
+ raise ValueError(f"the ROCm GLM-5.2 engine builds MTP at depth 1 or 3, not {num_mtp}")
+ gen = Glm52Generator(
+ model_weights_dir=model_weights_dir,
+ model_args=ModelArgsGlm52(),
+ max_new_tokens=max(max_seq_len - 256, 4096 - 256),
+ use_topp=True,
+ num_mtp=num_mtp if with_mtp else 0,
+ max_seq_len=max_seq_len,
+ )
+ gen.from_pretrained()
+ return RocmGlm52EngineAdapter(gen, with_mtp, ar_steps=ar_steps)
+
+
+class RocmGlm52EngineAdapter:
+
+ def __init__(
+ self,
+ generator,
+ with_mtp: bool,
+ ar_steps: int = 8,
+ *,
+ pure_tp8: bool | None = None,
+ fp8_ki: bool | None = None,
+ ):
+ self.gen = generator
+ self.dl = generator.decode_layer
+ self.with_mtp = bool(with_mtp) and self.dl.num_mtp > 0
+ if with_mtp and (not self.with_mtp):
+ raise ValueError("--with-mtp requested but the engine was built with num_mtp=0")
+ self.max_seq_len = int(self.dl.args.max_seq_len)
+ self.n_layers = int(self.dl.n_layers)
+ self.npes = int(self.dl.npes)
+ self.mtp_seq_len = self.dl.num_mtp + 1
+ self.ar_steps = max(1, min(1024, int(os.environ.get("GLM5_AR_N", str(ar_steps)))))
+ self.stop_ids = {int(t) for t in generator.stop_token_ids}
+ self.last_stats: dict = {}
+ self._ignore_eos = False
+ self._seq_len = 0
+ from tilert.models.glm_5_2_rocm.model_args import full_layer_ordinals
+ from tilert.models.glm_5_2_rocm.weight_converter import fp8_ki_enabled
+
+ self._full_layers = list(full_layer_ordinals(self.n_layers))
+ self._pure_tp8 = True if pure_tp8 is None else pure_tp8
+ self._fp8_ki = fp8_ki_enabled() if fp8_ki is None else fp8_ki
+ self._rot_streams = None
+ self._rccl = None
+ self._inject_mode = _INJECT_MODE if _INJECT_MODE in ("rccl", "rotate", "legacy") else "rccl"
+ if self._inject_mode == "rccl":
+ try:
+ self._rccl = _RcclBroadcast(self.npes)
+ logger.info(
+ "inject mode=rccl (RCCL %d via %s, %d comms)",
+ self._rccl.version,
+ self._rccl.path,
+ self.npes,
+ )
+ except Exception as exc:
+ logger.warning("inject: RCCL unavailable (%s); falling back to rotate", exc)
+ self._inject_mode = "rotate"
+ else:
+ logger.info("inject mode=%s", self._inject_mode)
+
+ def supports_logprobs(self) -> bool:
+ return False
+
+ def supports_penalties(self) -> bool:
+ return False
+
+ def supports_ignore_eos(self) -> bool:
+ return True
+
+ def prepare_grammar(self, grammar_spec, enable_thinking=True):
+ if grammar_spec is None:
+ return None
+ raise GrammarUnsupported("constrained decoding is not supported")
+
+ def _write_ki(self, slot: torch.Tensor, ki: torch.Tensor, seq: int) -> None:
+ if not self._fp8_ki:
+ slot[0, :seq].copy_(ki.to(slot.device, non_blocking=True))
+ return
+ L = self.max_seq_len
+ nbytes = L * _INDEX_HEAD_DIM
+ x = ki.to(slot.device).float()
+ amax = x.abs().amax(dim=-1).clamp_(min=1e-12)
+ scale = amax / _FP8_MAX
+ q = (x / scale.unsqueeze(-1)).clamp_(-_FP8_MAX, _FP8_MAX)
+ q = q.to(torch.float8_e4m3fn)
+ slot[:nbytes].view(torch.float8_e4m3fn).view(L, _INDEX_HEAD_DIM)[:seq].copy_(q)
+ slot[nbytes : nbytes + L * 4].view(torch.float32)[:seq].copy_(scale)
+
+ def _kv_targets(self, use, n_extra) -> dict:
+ out = {}
+ for rank in range(self.npes):
+ if rank == 0 and (not self._pure_tp8):
+ continue
+ n_pairs = self.n_layers + (n_extra if rank != 0 else 0)
+ out[rank] = min(len(use), n_pairs)
+ return out
+
+ def _move_kv(self, use, seq, n_extra) -> None:
+ if self._inject_mode == "rccl":
+ try:
+ self._move_kv_rccl(use, seq, n_extra)
+ return
+ except Exception:
+ logger.exception(
+ "inject: RCCL path failed; falling back to rotate for the rest of this run"
+ )
+ self._inject_mode = "rotate"
+ if self._inject_mode == "rotate":
+ self._move_kv_rotate(use, seq, n_extra)
+ return
+ self._move_kv_legacy(use, seq, n_extra)
+
+ def _move_kv_legacy(self, use, seq, n_extra) -> None:
+ dl = self.dl
+ for rank, upto in self._kv_targets(use, n_extra).items():
+ caches = dl._caches[rank]
+ for lid in range(upto):
+ _ki, kv, pe = use[lid]
+ caches[2 * lid][0, :seq].copy_(kv, non_blocking=True)
+ caches[2 * lid + 1][0, :seq].copy_(pe, non_blocking=True)
+
+ def _move_kv_rotate(self, use, seq, n_extra) -> None:
+ dl = self.dl
+ upto = self._kv_targets(use, n_extra)
+ by_src = {}
+ for lid in range(len(use)):
+ by_src.setdefault(use[lid][1].device.index, []).append(lid)
+ if self._rot_streams is None:
+ width = self.npes if len(by_src) < self.npes else 1
+ self._rot_streams = {
+ s: [torch.cuda.Stream(device=s) for _ in range(width)] for s in by_src
+ }
+ for k in range(self.npes):
+ for s, lids in by_src.items():
+ r = (s + k) % self.npes
+ if r not in upto:
+ continue
+ pool = self._rot_streams.get(s)
+ if not pool:
+ pool = self._rot_streams[s] = [torch.cuda.Stream(device=s)]
+ caches = dl._caches[r]
+ with torch.cuda.device(s), torch.cuda.stream(pool[r % len(pool)]):
+ for lid in lids:
+ if lid >= upto[r]:
+ continue
+ _ki, kv, pe = use[lid]
+ caches[2 * lid][0, :seq].copy_(kv, non_blocking=True)
+ caches[2 * lid + 1][0, :seq].copy_(pe, non_blocking=True)
+ for pool in self._rot_streams.values():
+ for st in pool:
+ st.synchronize()
+
+ def _move_kv_rccl(self, use, seq, n_extra) -> None:
+ if self._rccl is None:
+ raise RuntimeError("RCCL broadcaster not initialised")
+ dl = self.dl
+ upto = self._kv_targets(use, n_extra)
+ if len(upto) != self.npes:
+ raise RuntimeError("rank 0 is excluded from KV injection; RCCL needs every rank")
+ items = []
+ leftovers = []
+ for lid in range(len(use)):
+ if any(lid >= upto[r] for r in range(self.npes)):
+ leftovers.append(lid)
+ continue
+ _ki, kv, pe = use[lid]
+ for ti, src in ((0, kv), (1, pe)):
+ if not src.is_contiguous():
+ raise RuntimeError(f"layer {lid} tensor {ti} is not contiguous")
+ dsts = []
+ for r in range(self.npes):
+ dst = dl._caches[r][2 * lid + ti][0, :seq]
+ if not dst.is_contiguous():
+ raise RuntimeError(f"cache slot {lid}/{ti} on rank {r} is not contiguous")
+ dsts.append(dst.data_ptr())
+ items.append(
+ (src.device.index, src.data_ptr(), src.numel() * src.element_size(), dsts)
+ )
+ self._rccl.run(items)
+ for lid in leftovers:
+ _ki, kv, pe = use[lid]
+ for r in range(self.npes):
+ if lid >= upto[r]:
+ continue
+ caches = dl._caches[r]
+ caches[2 * lid][0, :seq].copy_(kv, non_blocking=True)
+ caches[2 * lid + 1][0, :seq].copy_(pe, non_blocking=True)
+
+ def inject(self, req) -> None:
+ dl = self.dl
+ layers = req.layers
+ seq = int(req.seq_len)
+ n_extra = 1 if dl.num_mtp > 0 else 0
+ if len(layers) not in (self.n_layers, self.n_layers + 1):
+ raise RuntimeError(
+ f"glm5_rocm inject: got {len(layers)} layers, engine has {self.n_layers} (+1 MTP block)"
+ )
+ if seq <= 0 or seq > self.max_seq_len:
+ raise RuntimeError(f"glm5_rocm inject: seq_len {seq} outside (0, {self.max_seq_len}]")
+ if n_extra and len(layers) == self.n_layers:
+ raise RuntimeError(
+ "glm5_rocm inject: engine has an MTP block but the prefill sent no MTP-layer KV (prefill must run with --speculative-config mtp)"
+ )
+ use = layers[: self.n_layers + n_extra]
+ dl.reset_sequence()
+ self._move_kv(use, seq, n_extra)
+ caches0 = dl._caches[0]
+ ki_base = 2 * self.n_layers if self._pure_tp8 else 0
+ ki_layers = self._full_layers + ([self.n_layers] if n_extra else [])
+ for ki_slot, lid in enumerate(ki_layers):
+ if lid >= len(use):
+ break
+ self._write_ki(caches0[ki_base + ki_slot], use[lid][0], seq)
+ torch.cuda.synchronize()
+ dl.set_cur_pos(seq)
+ self._seq_len = seq
+
+ def decode(
+ self,
+ first_token_id,
+ max_tokens,
+ sampling,
+ on_token=None,
+ cancel_event=None,
+ grammar_session=None,
+ top_logprobs=None,
+ ):
+ if grammar_session is not None:
+ raise GrammarUnsupported("constrained decoding is not supported")
+ if top_logprobs:
+ raise NotImplementedError("logprobs are not supported")
+ sampling = sampling or {}
+ rep = float(sampling.get("repetition_penalty", 1.0) or 1.0)
+ pres = float(sampling.get("presence_penalty", 0.0) or 0.0)
+ if rep != 1.0 or pres != 0.0:
+ raise NotImplementedError(
+ "repetition/presence penalties are not supported by this model's decode runtime"
+ )
+ temp = float(sampling.get("temperature", 1.0))
+ if temp < 1e-05:
+ self.dl.update_sampling(False, 1.0, 1.0)
+ else:
+ self.dl.update_sampling(True, temp, resolve_top_p(sampling))
+ self._ignore_eos = bool(sampling.get("ignore_eos"))
+ first = int(first_token_id)
+ budget = min(int(max_tokens), self.max_seq_len - self._seq_len - 1)
+ if budget <= 0:
+ self.last_stats = {"finish_reason": "length"}
+ return [first]
+ stop_ids = set() if self._ignore_eos else self.stop_ids
+ if first in stop_ids:
+ self.last_stats = {"finish_reason": "stop"}
+ return []
+ tokens = [first]
+ if on_token:
+ on_token(first)
+ base_step = self.dl.step_count
+ if self.with_mtp:
+ finish = self._decode_mtp(first, budget, tokens, stop_ids, on_token, cancel_event)
+ else:
+ finish = self._decode_plain(first, budget, tokens, stop_ids, on_token, cancel_event)
+ self.last_stats = {"finish_reason": finish}
+ if self.with_mtp:
+ self.last_stats.update(self._mtp_stats(base_step))
+ return tokens
+
+ def _mtp_stats(self, base_step: int) -> dict:
+ per_step = self.dl.accepted_step_counts(base_step)
+ if not per_step:
+ return {}
+ return {
+ "mtp_verify_calls": len(per_step),
+ "mtp_accept_mean": round(sum(per_step) / len(per_step), 3),
+ }
+
+ def _emit(self, new, tokens, budget, stop_ids, on_token):
+ for tok in new:
+ tok = int(tok)
+ if tok in stop_ids:
+ return "stop"
+ if len(tokens) >= budget:
+ return "length"
+ tokens.append(tok)
+ if on_token:
+ on_token(tok)
+ return "length" if len(tokens) >= budget else None
+
+ def _decode_mtp(self, first, budget, tokens, stop_ids, on_token, cancel_event):
+ dl = self.dl
+ pos_limit = self.max_seq_len
+ mtp_seq = self.mtp_seq_len
+ chain_slack = max(0, dl.num_mtp - 1)
+ dl.seed_draft(first, first)
+ base = dl.accepted_count
+ produced = 0
+ while True:
+ if cancel_event is not None and cancel_event.is_set():
+ return "cancelled"
+ room = pos_limit - (self._seq_len + produced)
+ k = min(self.ar_steps, (room - chain_slack) // mtp_seq)
+ if k < 1:
+ return "length"
+ got = int(dl.mtp_n(k))
+ new = dl.accepted_tokens(base + produced)
+ produced += got
+ verdict = self._emit(new, tokens, budget, stop_ids, on_token)
+ if verdict:
+ return verdict
+
+ def _decode_plain(self, first, budget, tokens, stop_ids, on_token, cancel_event):
+ dl = self.dl
+ pos_limit = self.max_seq_len
+ base = dl.accepted_count
+ dl.step(first)
+ produced = 0
+ while True:
+ new = dl.accepted_tokens(base + produced)
+ produced += len(new)
+ verdict = self._emit(new, tokens, budget, stop_ids, on_token)
+ if verdict:
+ return verdict
+ if cancel_event is not None and cancel_event.is_set():
+ return "cancelled"
+ room = pos_limit - (self._seq_len + 1 + produced)
+ if room < 1:
+ return "length"
+ n = min(8, budget - len(tokens), room)
+ dl.decode_n(n)
+
+ def reset(self) -> None:
+ pass
diff --git a/tilert/pd_vllm/profiles/mla_nsa.py b/tilert/pd_vllm/profiles/mla_nsa.py
index 2b6d656..4a304c0 100644
--- a/tilert/pd_vllm/profiles/mla_nsa.py
+++ b/tilert/pd_vllm/profiles/mla_nsa.py
@@ -1,5 +1,3 @@
-"""Shared MLA + NSA-KI data plane for the DeepSeek-family models."""
-
from __future__ import annotations
import logging
@@ -10,44 +8,43 @@
import torch
from tilert.pd_vllm import wire
+from tilert.pd_vllm.grammar_spec import (
+ GrammarUnsupported,
+ GrammarViolationError,
+ InvalidGrammarError,
+)
+from tilert.pd_vllm.profiles import base
+from tilert.pd_vllm.sampling import resolve_top_k, resolve_top_p
logger = logging.getLogger("pd_vllm.profile.mla_nsa")
-
-_AR_MTP_API = ("show_hands", "ar_accepted_tokens", "ar_num_accepted")
-_AR_PLAIN_API = ("show_hands_no_mtp", "ar_accepted_tokens_no_mtp")
-
-
-def _has_api(dl, names: tuple[str, ...]) -> bool:
- return all(hasattr(dl, n) for n in names)
-
-
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
+KV_QUANT_BLOCK = 128
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
+KI_TILE = 16
+KV_FP8_BYTES = KV_LORA_RANK
+KV_SCALE_BYTES = KV_LORA_RANK // KV_QUANT_BLOCK * 4
+KV_BYTES_FP8 = KV_FP8_BYTES + KV_SCALE_BYTES
+KV_BYTES_BF16 = KV_LORA_RANK * 2
+PE_BPT = QK_ROPE_HEAD_DIM * 2
+MLA_BPT_FP8 = KV_BYTES_FP8 + PE_BPT
+MLA_BPT_BF16 = (KV_LORA_RANK + QK_ROPE_HEAD_DIM) * 2
+KV_BYTES_ROCM_FP8 = KV_LORA_RANK
+PE_BYTES_ROCM_FP8 = QK_ROPE_HEAD_DIM
+MLA_BPT_ROCM_FP8 = KV_LORA_RANK + QK_ROPE_HEAD_DIM
+_VERSION_BF16_OFFSET = 40
+_VERSION_KI_TILED_OFFSET = 100
+_VERSION_ROCM_FP8_OFFSET = 200
+KI_PAGE_BYTES = PAGE_SIZE * INDEX_HEAD_DIM + PAGE_SIZE * INDEX_HEAD_DIM // KI_QUANT_BLOCK * 4
+
+
+def _default_ki_tiled() -> bool:
+ env = (os.environ.get("TILERT_PD_KI_TILED") or "").strip().lower()
+ if env:
+ return env not in ("0", "false", "no", "off")
+ return getattr(torch.version, "hip", None) is not None
def _max_pages(max_seq_len: int) -> int:
@@ -55,19 +52,15 @@ def _max_pages(max_seq_len: int) -> int:
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)
+ 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)
+ H = torch.from_numpy(_h(d).astype("float32")).to(x.device) * d ** (-0.5)
return (x.float() @ H).to(x.dtype)
@@ -78,82 +71,129 @@ class ConvertedRequest:
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
+ layers: list
+
+
+_EXTRACT_MODE = os.environ.get("TILERT_EXTRACT_MODE", "fast").strip().lower()
+_PD_SENDERS = max(
+ 1, min(int((os.environ.get("TILERT_PD_SENDERS") or "1").strip() or 1), wire.NUM_RANKS)
+)
@dataclass
class _Reg:
- mla_layers: list # [(lid, name, kv_t, gi)] sorted
- ki_layers: list # [(lid, name, ki_t, gi)] sorted
+ mla_layers: list
+ ki_layers: list
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
+ sender_ranks = frozenset(range(_PD_SENDERS))
+ declares_penalties = False
def __init__(
- self, name: str, num_layers: int, layout_version: int, engine_factory, mla_fp8: bool = True
+ self,
+ name: str,
+ num_layers: int,
+ layout_version: int,
+ engine_factory,
+ mla_fp8: bool = True,
+ ki_layer_ids: list[int] | None = None,
+ ki_tiled: bool | None = None,
):
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
+ self.mla_fp8 = mla_fp8
+ self.mla_rocm_fp8 = False
+ self.kv_scales: list[float] | None = None
+ self.ki_tiled = _default_ki_tiled() if ki_tiled is None else bool(ki_tiled)
+ self.ki_layer_ids = ki_layer_ids
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"):
+ if d in ("fp8_rocm", "rocm_fp8"):
+ self.mla_fp8 = False
+ self.mla_rocm_fp8 = True
+ elif d in ("fp8_ds_mla", "fp8", "fp8_e4m3"):
self.mla_fp8 = True
+ self.mla_rocm_fp8 = False
elif d in ("bf16", "bfloat16", "auto"):
self.mla_fp8 = False
+ self.mla_rocm_fp8 = False
else:
raise ValueError(
- f"unknown kv_cache_dtype {kv_cache_dtype!r}; " f"want fp8_ds_mla or bf16"
+ f"unknown kv_cache_dtype {kv_cache_dtype!r}; want fp8_ds_mla, fp8_rocm 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)
+ return (
+ self._base_version
+ + (0 if self.mla_fp8 else _VERSION_BF16_OFFSET)
+ + (_VERSION_KI_TILED_OFFSET if self.ki_tiled else 0)
+ + (_VERSION_ROCM_FP8_OFFSET if self.mla_rocm_fp8 else 0)
+ )
@property
def _kv_bpt(self) -> int:
+ if self.mla_rocm_fp8:
+ return KV_BYTES_ROCM_FP8
return KV_BYTES_FP8 if self.mla_fp8 else KV_BYTES_BF16
+ @property
+ def _pe_bpt(self) -> int:
+ return PE_BYTES_ROCM_FP8 if self.mla_rocm_fp8 else PE_BPT
+
@property
def _mla_bpt(self) -> int:
+ if self.mla_rocm_fp8:
+ return MLA_BPT_ROCM_FP8
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
+ return self.num_layers * max_seq_len * self._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}
+ def _shard_layers(self, nshards: int) -> int:
+ return -(-self.num_layers // nshards)
+
+ def shard_bytes(self, max_seq_len: int, nshards: int) -> int:
+ n = self._shard_layers(nshards)
+ return (
+ n * max_seq_len * self._kv_bpt
+ + n * max_seq_len * self._pe_bpt
+ + n * _max_pages(max_seq_len) * KI_PAGE_BYTES
+ )
+
+ def hello_layout(self, base_ptr, max_seq_len: int) -> dict:
+ if not isinstance(base_ptr, (list, tuple)):
+ kv = base_ptr
+ pe = kv + self._kv_plane(max_seq_len)
+ ki = pe + self._pe_plane(max_seq_len)
+ return {"senders": len(self.sender_ranks), "kv_base": kv, "pe_base": pe, "ki_base": ki}
+ n = self._shard_layers(len(base_ptr))
+ kv_plane = n * max_seq_len * self._kv_bpt
+ pe_plane = n * max_seq_len * self._pe_bpt
+ return {
+ "senders": len(self.sender_ranks),
+ "kv_base": [int(p) for p in base_ptr],
+ "pe_base": [int(p) + kv_plane for p in base_ptr],
+ "ki_base": [int(p) + kv_plane + pe_plane for p in base_ptr],
+ "nshards": len(base_ptr),
+ }
@torch.inference_mode()
def convert(self, buffer, base_ptr, max_seq_len, received, num_devices=1):
@@ -162,27 +202,50 @@ def convert(self, buffer, base_ptr, max_seq_len, received, num_devices=1):
pe_base = self._kv_plane(max_seq_len)
ki_base = pe_base + self._pe_plane(max_seq_len)
kv_bpt = self._kv_bpt
+ pe_bpt = self._pe_bpt
layers = []
+ bufs = list(buffer) if isinstance(buffer, (list, tuple)) else [buffer]
+ nsh = len(bufs)
+ host = not bufs[0].is_cuda
+ scales = self.kv_scales if self.mla_rocm_fp8 else None
+ if nsh > 1:
+ lps = self._shard_layers(nsh)
+ pe_base = lps * max_seq_len * self._kv_bpt
+ ki_base = pe_base + lps * max_seq_len * pe_bpt
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
+ sh, li = (lid % nsh, lid // nsh) if nsh > 1 else (0, lid)
+ buf = bufs[sh]
+ ko = li * max_seq_len * kv_bpt
+ kv_raw = buf[ko : ko + seq * kv_bpt]
+ if host:
+ kv_raw = kv_raw.to("cuda:0", non_blocking=True)
+ kv_raw = kv_raw.view(seq, kv_bpt)
+ if self.mla_rocm_fp8:
+ kv = self._dequant_rocm_fp8(
+ kv_raw, seq, KV_LORA_RANK, scales[lid] if scales else 1.0
+ )
+ elif self.mla_fp8:
+ kv = self._dequant_kv(kv_raw, seq)
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()
+ kv = kv_raw.view(torch.bfloat16).view(seq, KV_LORA_RANK).contiguous()
+ po = pe_base + li * max_seq_len * pe_bpt
+ pe_raw = buf[po : po + seq * pe_bpt]
+ if host:
+ pe_raw = pe_raw.to("cuda:0", non_blocking=True)
+ if self.mla_rocm_fp8:
+ pe = self._dequant_rocm_fp8(
+ pe_raw.view(seq, pe_bpt), seq, QK_ROPE_HEAD_DIM, scales[lid] if scales else 1.0
+ )
+ else:
+ pe = pe_raw.view(torch.bfloat16).view(seq, QK_ROPE_HEAD_DIM).contiguous()
+ io = ki_base + li * _max_pages(max_seq_len) * KI_PAGE_BYTES
+ ki_raw = buf[io : io + npages * KI_PAGE_BYTES]
+ if host:
+ ki_raw = ki_raw.to("cuda:0", non_blocking=True)
+ ki_raw = ki_raw.view(npages, KI_PAGE_BYTES)
+ layers.append((self._dequant_ki(ki_raw, seq, self.ki_tiled), kv, pe))
+ for d in sorted({b.device.index for b in bufs if b.is_cuda} or {0}):
+ torch.cuda.synchronize(d)
return ConvertedRequest(
rid=received.rid,
seq_len=seq,
@@ -193,11 +256,16 @@ def convert(self, buffer, base_ptr, max_seq_len, received, num_devices=1):
)
@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].
+ def _dequant_rocm_fp8(
+ raw: torch.Tensor, seq_len: int, width: int, scale: float
+ ) -> torch.Tensor:
+ x = raw.reshape(-1).contiguous().view(torch.float8_e4m3fn).reshape(seq_len, width)
+ if scale == 1.0:
+ return x.to(torch.bfloat16)
+ return (x.float() * scale).to(torch.bfloat16)
- Per-128-block scale: kv[:, b*128:(b+1)*128] *= scale[:, b].
- """
+ @staticmethod
+ def _dequant_kv(kv_raw: torch.Tensor, seq_len: int) -> torch.Tensor:
nblk = KV_LORA_RANK // KV_QUANT_BLOCK
fp8 = (
kv_raw[:, :KV_FP8_BYTES]
@@ -217,14 +285,20 @@ def _dequant_kv(kv_raw: torch.Tensor, seq_len: int) -> torch.Tensor:
return deq.to(torch.bfloat16)
@staticmethod
- def _dequant_ki(ki_raw: torch.Tensor, seq_len: int) -> torch.Tensor:
+ def _dequant_ki(ki_raw: torch.Tensor, seq_len: int, tiled: bool) -> torch.Tensor:
npages = ki_raw.shape[0]
fp8_bytes = PAGE_SIZE * INDEX_HEAD_DIM
+ plane = ki_raw[:, :fp8_bytes]
+ if tiled:
+ plane = (
+ plane.contiguous()
+ .view(npages, PAGE_SIZE // KI_TILE, INDEX_HEAD_DIM // KI_TILE, KI_TILE, KI_TILE)
+ .permute(0, 1, 3, 2, 4)
+ )
ki_fp8 = (
- ki_raw[:, :fp8_bytes]
- .reshape(-1)
+ plane.reshape(npages * PAGE_SIZE, INDEX_HEAD_DIM)[:seq_len]
+ .contiguous()
.view(torch.float8_e4m3fn)
- .reshape(npages * PAGE_SIZE, INDEX_HEAD_DIM)
)
scale = (
ki_raw[:, fp8_bytes:]
@@ -233,10 +307,9 @@ def _dequant_ki(ki_raw: torch.Tensor, seq_len: int) -> torch.Tensor:
.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)
+ deq = (ki_fp8.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 []):
@@ -244,13 +317,13 @@ def classify_layers(self, kv_caches: dict, kv_cache_config) -> _Reg:
group_of[ln] = gi
def lid_of(name):
- m = re.search(r"\.(\d+)\.", name)
+ m = re.search("\\.(\\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
+ mla, ki = ([], [])
+ for name, val in kv_caches.items():
+ t = val[0] if isinstance(val, (tuple, list)) else val
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))
@@ -258,110 +331,285 @@ def lid_of(name):
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:
+ if len(mla) != self.num_layers:
+ raise RuntimeError(
+ f"{self.name} classify: {len(mla)} MLA layers (expected {self.num_layers}); check --speculative-config and the vLLM layer naming"
+ )
+ ki_ids = [x[0] for x in ki]
+ if not ki or ki[0][0] != 0:
+ raise RuntimeError(
+ f"{self.name} classify: KI layer 0 missing (ids={ki_ids}); cannot expand sparse indexer set"
+ )
+ if len(ki) > self.num_layers or ki_ids != sorted(set(ki_ids)):
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"
+ f"{self.name} classify: bad KI layer set {ki_ids} (num_layers={self.num_layers})"
+ )
+ if self.ki_layer_ids is not None:
+ want = [l for l in self.ki_layer_ids if l < self.num_layers]
+ want_no_mtp = [l for l in want if l != self.num_layers - 1]
+ if ki_ids not in (want, want_no_mtp):
+ raise RuntimeError(
+ f"{self.name} classify: KI layer ids {ki_ids} != expected {want} (or {want_no_mtp} without the MTP tail)"
+ )
+ ki_expanded, cur, idx = ([], None, 0)
+ for L in range(self.num_layers):
+ while idx < len(ki) and ki[idx][0] <= L:
+ cur = ki[idx]
+ idx += 1
+ ki_expanded.append((L, cur[1], cur[2], cur[3]))
+ if len(ki) < self.num_layers:
+ logger.info(
+ "%s: sparse KI %d full layers %s expanded to %d (shared layers reuse previous full layer's indexer)",
+ self.name,
+ len(ki),
+ ki_ids,
+ self.num_layers,
)
- # 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)
+ ki = ki_expanded
t0 = mla[0][2]
bpt = t0.shape[-1] * t0.element_size()
if bpt == MLA_BPT_FP8:
- self.mla_fp8 = True
+ self.mla_fp8, self.mla_rocm_fp8 = (True, False)
elif bpt == MLA_BPT_BF16:
- self.mla_fp8 = False
+ self.mla_fp8, self.mla_rocm_fp8 = (False, False)
+ elif bpt == MLA_BPT_ROCM_FP8:
+ self.mla_fp8, self.mla_rocm_fp8 = (False, True)
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)"
+ f"{self.name}: unexpected MLA cache stride {bpt} B/token; expected {MLA_BPT_FP8} (fp8_ds_mla), {MLA_BPT_BF16} (bf16) or {MLA_BPT_ROCM_FP8} (ROCm flat fp8)"
+ )
+ k0 = ki[0][2]
+ ki_page = int(k0[0].numel() * k0.element_size()) if k0.dim() > 1 else 0
+ if ki_page != KI_PAGE_BYTES:
+ raise RuntimeError(
+ f"{self.name}: KI cache page is {ki_page} B, expected {KI_PAGE_BYTES} "
+ f"({PAGE_SIZE} tokens x {INDEX_HEAD_DIM} fp8 + one fp32 scale per token)"
)
+ layout = (
+ "ROCm flat fp8" if self.mla_rocm_fp8 else ("fp8_ds_mla" if self.mla_fp8 else "bf16")
+ )
logger.info(
- "%s registered %d MLA + %d KI layers, MLA cache=%s",
+ "%s registered %d MLA + %d KI layers, MLA cache=%s (%d B/token, layout v%d)",
self.name,
len(mla),
len(ki),
- "fp8_ds_mla" if self.mla_fp8 else "bf16",
+ layout,
+ bpt,
+ self.layout_version,
)
return _Reg(mla_layers=mla, ki_layers=ki)
- def staging_bytes(self, reg, tp_rank, max_seq_len):
+ def set_kv_scales(self, scales) -> None:
+ if scales is None:
+ self.kv_scales = None
+ return
+ vals = [float(v) for v in scales]
+ if len(vals) != self.num_layers:
+ raise RuntimeError(
+ f"{self.name}: got {len(vals)} kv scales for {self.num_layers} layers"
+ )
+ self.kv_scales = None if all(v == 1.0 for v in vals) else vals
+ logger.info(
+ "%s kv scales: %s",
+ self.name,
+ (
+ "all 1.0 (no dequant multiply)"
+ if self.kv_scales is None
+ else f"min {min(vals):.6g} max {max(vals):.6g}"
+ ),
+ )
+
+ def staging_bytes(self, reg, tp_rank, max_seq_len, nshards: int = 1):
if tp_rank not in self.sender_ranks:
return 4
- return self.buffer_bytes(max_seq_len)
+ if _PD_SENDERS > 1:
+ return self.shard_bytes(max_seq_len, _PD_SENDERS)
+ return (
+ self.shard_bytes(max_seq_len, nshards)
+ if nshards > 1
+ else self.buffer_bytes(max_seq_len)
+ )
@torch.inference_mode()
def extract(self, reg: _Reg, m, tp_rank, staging, max_seq_len):
+ global _EXTRACT_MODE
+ if _EXTRACT_MODE == "legacy":
+ return self._extract_legacy(reg, m, tp_rank, staging, max_seq_len)
+ try:
+ return self._extract_fast(reg, m, tp_rank, staging, max_seq_len)
+ except Exception:
+ logger.exception("extract: fast path failed; using legacy for the rest of this run")
+ _EXTRACT_MODE = "legacy"
+ return self._extract_legacy(reg, m, tp_rank, staging, max_seq_len)
+
+ @torch.inference_mode()
+ def _extract_fast(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]
+ slots_cpu = (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)
+ ki_bt_cpu = torch.tensor(ki_ids[:npages], dtype=torch.long)
+ stgs = list(staging) if isinstance(staging, (list, tuple)) else [staging]
+ nsh = len(stgs)
+ if _PD_SENDERS > 1 and nsh != _PD_SENDERS:
+ raise RuntimeError(
+ f"TILERT_PD_SENDERS={_PD_SENDERS} needs the same number of staging shards, got {nsh}"
+ )
+ lps = self._shard_layers(nsh) if nsh > 1 else self.num_layers
+ pe_bpt = self._pe_bpt
+ pe_base = lps * max_seq_len * self._kv_bpt
+ ki_base = pe_base + lps * max_seq_len * pe_bpt
+ kv_bpt, mla_bpt = (self._kv_bpt, self._mla_bpt)
+ idx, tmp, streams = ({}, {}, {})
+
+ def index_for(dev):
+ if dev not in idx:
+ idx[dev] = (slots_cpu.to(dev), ki_bt_cpu.to(dev))
+ return idx[dev]
+
+ def scratch_for(dev, sh):
+ key = (dev.index, sh)
+ if key not in tmp:
+ tmp[key] = (
+ torch.empty(seq, mla_bpt, dtype=torch.uint8, device=dev),
+ torch.empty(npages, KI_PAGE_BYTES, dtype=torch.uint8, device=dev),
+ )
+ return tmp[key]
+
+ def stream_for(dev, sh):
+ key = (dev.index, sh)
+ if key not in streams:
+ streams[key] = torch.cuda.Stream(device=dev)
+ return streams[key]
- 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
+ sh, li = (lid % nsh, lid // nsh) if nsh > 1 else (0, lid)
+ if _PD_SENDERS > 1 and sh != tp_rank:
+ continue
+ stg = stgs[sh]
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())
+ flat = raw.reshape(-1, mla_bpt)
+ ki_t = reg.ki_layers[lid][2]
+ ki_raw = ki_t if ki_t.dtype == torch.uint8 else ki_t.view(torch.uint8)
+ ki_flat = ki_raw.reshape(ki_t.shape[0], -1)
+ if ki_flat.shape[1] != KI_PAGE_BYTES:
+ raise RuntimeError(f"KI page stride {ki_flat.shape[1]} B != {KI_PAGE_BYTES}")
+ slots, ki_bt = index_for(flat.device)
+ ko = li * max_seq_len * kv_bpt
+ po = pe_base + li * max_seq_len * pe_bpt
+ io = ki_base + li * _max_pages(max_seq_len) * KI_PAGE_BYTES
+ dst_kv = stg[ko : ko + seq * kv_bpt].view(seq, kv_bpt)
+ dst_pe = stg[po : po + seq * pe_bpt].view(seq, pe_bpt)
+ dst_ki = stg[io : io + npages * KI_PAGE_BYTES].view(npages, KI_PAGE_BYTES)
+ same = stg.is_cuda and stg.device == flat.device
+ if same:
+ torch.index_select(flat[:, :kv_bpt], 0, slots, out=dst_kv)
+ torch.index_select(flat[:, kv_bpt:], 0, slots, out=dst_pe)
+ torch.index_select(ki_flat, 0, ki_bt, out=dst_ki)
+ else:
+ rows, kirows = scratch_for(flat.device, sh)
+ with torch.cuda.device(flat.device), torch.cuda.stream(stream_for(flat.device, sh)):
+ torch.index_select(flat, 0, slots, out=rows)
+ dst_kv.copy_(rows[:, :kv_bpt], non_blocking=True)
+ dst_pe.copy_(rows[:, kv_bpt:], non_blocking=True)
+ torch.index_select(ki_flat, 0, ki_bt, out=kirows)
+ dst_ki.copy_(kirows, non_blocking=True)
+ for st in streams.values():
+ st.synchronize()
+ torch.cuda.synchronize()
+ return {"seq": seq, "npages": npages, "stage_max": max_seq_len, "stage_shards": nsh}
+ @torch.inference_mode()
+ def _extract_legacy(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)
+ stgs = list(staging) if isinstance(staging, (list, tuple)) else [staging]
+ nsh = len(stgs)
+ lps = self._shard_layers(nsh) if nsh > 1 else self.num_layers
+ pe_bpt = self._pe_bpt
+ pe_base = lps * max_seq_len * self._kv_bpt
+ ki_base = pe_base + lps * max_seq_len * pe_bpt
+ kv_bpt, mla_bpt = (self._kv_bpt, self._mla_bpt)
+ for lid in range(self.num_layers):
+ sh, li = (lid % nsh, lid // nsh) if nsh > 1 else (0, lid)
+ if _PD_SENDERS > 1 and sh != tp_rank:
+ continue
+ stg = stgs[sh]
+ 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)
+ rows = flat[slots.to(flat.device)]
+ kv_merged = rows[:, :kv_bpt].contiguous()
+ pe = rows[:, kv_bpt:].contiguous()
+ ko = li * max_seq_len * kv_bpt
+ po = pe_base + li * max_seq_len * pe_bpt
+ stg[ko : ko + seq * kv_bpt].copy_(kv_merged.flatten())
+ stg[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_(
+ io = ki_base + li * _max_pages(max_seq_len) * KI_PAGE_BYTES
+ stg[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}
+ return {"seq": seq, "npages": npages, "stage_max": max_seq_len, "stage_shards": nsh}
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)
+ srcs, dsts, lens = ([], [], [])
+ sbases = list(base) if isinstance(base, (list, tuple)) else [base]
+ snsh = len(sbases)
+ s_lps = self._shard_layers(snsh) if snsh > 1 else self.num_layers
kv_bpt = self._kv_bpt
- r_kv, r_pe, r_ki = (int(hello["kv_base"]), int(hello["pe_base"]), int(hello["ki_base"]))
+ pe_bpt = self._pe_bpt
+ s_pe = s_lps * stage_max * kv_bpt
+ s_ki = s_pe + s_lps * stage_max * pe_bpt
+ dnsh = int(hello.get("nshards", 1) or 1)
+ r_kv, r_pe, r_ki = (hello["kv_base"], hello["pe_base"], 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)
+ if dnsh > 1:
+ dsh, dli = lid % dnsh, lid // dnsh
+ d_kv, d_pe, d_ki = int(r_kv[dsh]), int(r_pe[dsh]), int(r_ki[dsh])
+ else:
+ dli = lid
+ d_kv, d_pe, d_ki = int(r_kv), int(r_pe), int(r_ki)
+ ssh, sli = (lid % snsh, lid // snsh) if snsh > 1 else (0, lid)
+ if _PD_SENDERS > 1 and ssh != tp_rank:
+ continue
+ sb = int(sbases[ssh])
+ srcs.append(sb + sli * stage_max * kv_bpt)
+ dsts.append(d_kv + dli * 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)
+ srcs.append(sb + s_pe + sli * stage_max * pe_bpt)
+ dsts.append(d_pe + dli * remote_max * pe_bpt)
+ lens.append(seq_len * pe_bpt)
+ srcs.append(sb + s_ki + sli * _max_pages(stage_max) * KI_PAGE_BYTES)
+ dsts.append(d_ki + dli * _max_pages(remote_max) * KI_PAGE_BYTES)
lens.append(npages * KI_PAGE_BYTES)
- return srcs, dsts, lens
+ return (srcs, dsts, lens)
- def build_engine(self, model_weights_dir, max_seq_len, with_mtp, ar_steps):
+ def build_engine(self, model_weights_dir, max_seq_len, with_mtp, ar_steps, num_mtp=3):
+ assert num_mtp in getattr(
+ self, "supported_num_mtp", base.DEFAULT_SUPPORTED_NUM_MTP
+ ), f"{self.name}: num_mtp={num_mtp} is not in this profile's supported set"
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
@@ -377,7 +625,6 @@ def __init__(self, generator, with_mtp: bool):
@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)
@@ -390,16 +637,41 @@ def inject(self, req) -> None:
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):
+ def prepare_grammar(self, grammar_spec, enable_thinking=True):
+ if grammar_spec is None:
+ return None
+ raise GrammarUnsupported("constrained decoding is not supported")
+
+ def supports_penalties(self) -> bool:
+ return False
+
+ def supports_ignore_eos(self) -> bool:
+ return True
+
+ def decode(
+ self,
+ first_token_id,
+ max_tokens,
+ sampling,
+ on_token=None,
+ cancel_event=None,
+ grammar_session=None,
+ ):
sampling = sampling or {}
+ rep = float(sampling.get("repetition_penalty", 1.0) or 1.0)
+ pres = float(sampling.get("presence_penalty", 0.0) or 0.0)
+ if rep != 1.0 or pres != 0.0:
+ raise NotImplementedError(
+ "repetition/presence penalties are not supported by this model's decode runtime"
+ )
temp = float(sampling.get("temperature", 1.0))
- if temp < 1e-5:
+ if temp < 1e-05:
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)),
+ top_p=resolve_top_p(sampling),
+ top_k=resolve_top_k(sampling),
use_topp=True,
)
self._ignore_eos = bool(sampling.get("ignore_eos"))
@@ -408,10 +680,12 @@ def decode(self, first_token_id, max_tokens, sampling, on_token=None, cancel_eve
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)
+ return self._decode_mtp(first_token_id, budget, on_token, cancel_event, grammar_session)
+ return self._decode_standard(
+ first_token_id, budget, on_token, cancel_event, grammar_session
+ )
- def _decode_mtp(self, first_token_id, budget, on_token, cancel_event):
+ def _decode_mtp(self, first_token_id, budget, on_token, cancel_event, grammar_session=None):
dl = self.gen.decode_layer
T = self.mtp_seq_len
stop_ids = set() if self._ignore_eos else self.stop_ids
@@ -422,61 +696,90 @@ def _decode_mtp(self, first_token_id, budget, on_token, cancel_event):
if int(first_token_id) in stop_ids:
self.last_stats = {"finish_reason": "stop"}
return []
+ finished = False
+ if grammar_session is not None:
+ try:
+ if grammar_session.accept(int(first_token_id)) == "terminated":
+ finished = True
+ except RuntimeError as e:
+ raise GrammarViolationError(
+ f"prefill first token {first_token_id} violates the grammar"
+ ) from e
dl.set_prefill_valid_tokens(0)
ar_steps = max(1, min(1024, int(os.environ.get("GLM5_AR_N", "8"))))
- ar_ok = _has_api(dl, _AR_MTP_API)
+ ar_ok = hasattr(dl, "ar_accepted_tokens") and hasattr(dl, "ar_num_accepted")
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)
- if ar_ok:
- if fwd == 0:
+ accepted, finish, fwd = ([], "length", 0)
+ grammar_mask_written = False
+ try:
+ 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)
+ if grammar_session is not None and fwd >= 1 and (not grammar_session.terminated):
+ chain = draft[0, 1:].cpu().tolist()
+ masks = grammar_session.fill_step_masks(chain)
+ if masks is not None:
+ dl.update_grammar_bitmask(masks)
+ grammar_mask_written = True
+ if fwd == 0 or grammar_session is not None or (not ar_ok):
steps = 1
else:
rem = budget - len(tokens)
steps = max(1, min(ar_steps, -(-rem // T)))
- dl.show_hands(draft, steps)
- acc = dl.ar_accepted_tokens(0).cpu()
- num = dl.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()
- else:
- dl.forward(draft)
- n_acc = int(dl.get_num_accepted(0))
- pred = dl.get_predicted_tokens(0).flatten()
- emitted = [int(pred[i].item()) for i in range(n_acc)]
- per_step = [n_acc]
- if fwd == 0:
+ if ar_ok:
+ dl.show_hands(draft, steps)
+ acc = dl.ar_accepted_tokens(0).cpu()[0]
+ num = dl.ar_num_accepted(0).cpu()[0]
+ 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()
+ else:
+ dl.forward(draft)
+ n_acc = dl.get_num_accepted(0)
+ pred = dl.get_predicted_tokens(0).flatten()
+ emitted = [int(pred[i].item()) for i in range(n_acc)]
+ per_step = [n_acc]
+ if fwd == 0:
+ fwd += 1
+ continue
fwd += 1
- continue
- fwd += 1
- offset = 0
- for na in per_step:
- step_emit = emitted[offset : offset + na]
- offset += na
- for tok in step_emit:
- if len(tokens) >= budget:
- break
- tok = int(tok)
- if tok in stop_ids:
- finished = True
- finish = "stop"
+ offset = 0
+ for na in per_step:
+ step_emit = emitted[offset : offset + na]
+ offset += na
+ for tok in step_emit:
+ if len(tokens) >= budget:
+ break
+ tok = int(tok)
+ if tok in stop_ids:
+ finished = True
+ finish = "stop"
+ break
+ tokens.append(tok)
+ if on_token:
+ on_token(tok)
+ if (
+ grammar_session is not None
+ and grammar_session.accept(tok) == "terminated"
+ ):
+ finished = True
+ finish = "stop"
+ break
+ accepted.append(na)
+ if finished or len(tokens) >= budget:
break
- tokens.append(tok)
- if on_token:
- on_token(tok)
- accepted.append(na)
- if finished or len(tokens) >= budget:
- break
- dl.reset_sequence()
+ finally:
+ if grammar_mask_written:
+ dl.reset_grammar_bitmask()
+ dl.reset_sequence()
self.last_stats = {
"finish_reason": finish,
"mtp_accept_mean": round(sum(accepted) / max(1, len(accepted)), 3),
@@ -484,7 +787,9 @@ def _decode_mtp(self, first_token_id, budget, on_token, cancel_event):
}
return tokens
- def _decode_standard(self, first_token_id, budget, on_token, cancel_event):
+ def _decode_standard(
+ self, first_token_id, budget, on_token, cancel_event, grammar_session=None
+ ):
dl = self.gen.decode_layer
stop_ids = set() if self._ignore_eos else self.stop_ids
torch = self._torch
@@ -494,63 +799,81 @@ def _decode_standard(self, first_token_id, budget, on_token, cancel_event):
if int(first_token_id) in stop_ids:
self.last_stats = {"finish_reason": "stop"}
return []
- if not _has_api(dl, _AR_PLAIN_API):
- return self._decode_plain_per_step(tokens, budget, on_token, cancel_event)
- dl.set_prefill_valid_tokens(0, with_mtp=False)
+ finish = "length"
+ finished = False
+ grammar_mask_written = False
+ ar_ok = hasattr(dl, "show_hands_no_mtp") and hasattr(dl, "ar_accepted_tokens_no_mtp")
ar_steps = max(1, min(1024, int(os.environ.get("GLM5_AR_N", "8"))))
- finish, finished = "length", False
- last_tok = int(first_token_id)
- prev = torch.tensor([last_tok], dtype=torch.int32, device="cuda:0")
- while not finished and len(tokens) < budget:
- if cancel_event is not None and cancel_event.is_set():
- finish = "cancelled"
- break
- steps = max(1, min(ar_steps, budget - len(tokens)))
- dl.show_hands_no_mtp(prev, steps)
- acc = dl.ar_accepted_tokens_no_mtp(0).cpu()
- n_tokens = int(acc[0].item())
- emitted = acc[1 : 1 + n_tokens].tolist()
- for tok in emitted:
- if len(tokens) >= budget:
- break
- tok = int(tok)
- if tok in stop_ids:
- finished = True
- finish = "stop"
- break
- tokens.append(tok)
- last_tok = tok
- if on_token:
- on_token(tok)
+ cur = last_tok = None
+ prev = None
+ if ar_ok:
+ dl.set_prefill_valid_tokens(0, with_mtp=False)
+ last_tok = int(first_token_id)
prev = torch.tensor([last_tok], dtype=torch.int32, device="cuda:0")
- dl.reset_sequence()
- self.last_stats = {"finish_reason": finish}
- return tokens
-
- def _decode_plain_per_step(self, tokens, budget, on_token, cancel_event):
- from tilert.models.deepseek_v3_2.temp_var_indices import Idx
-
- dl = self.gen.decode_layer
- stop_ids = set() if self._ignore_eos else self.stop_ids
- torch = self._torch
- finish = "length"
- cur = torch.tensor(int(tokens[0]), 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()
+ else:
+ cur = torch.tensor(int(first_token_id), dtype=torch.long, device="cuda:0")
+ try:
+ if grammar_session is not None:
+ try:
+ if grammar_session.accept(int(first_token_id)) == "terminated":
+ finish, finished = ("stop", True)
+ except RuntimeError as e:
+ raise GrammarViolationError(
+ f"prefill first token {first_token_id} violates the grammar"
+ ) from e
+ while not finished and len(tokens) < budget:
+ if cancel_event is not None and cancel_event.is_set():
+ finish = "cancelled"
+ break
+ if (
+ grammar_session is not None
+ and grammar_session.active
+ and (not grammar_session.terminated)
+ ):
+ masks = grammar_session.fill_step_masks([])
+ if masks is not None:
+ dl.update_grammar_bitmask(masks)
+ grammar_mask_written = True
+ if ar_ok:
+ steps = (
+ 1
+ if grammar_session is not None
+ else max(1, min(ar_steps, budget - len(tokens)))
+ )
+ dl.show_hands_no_mtp(prev, steps)
+ acc = dl.ar_accepted_tokens_no_mtp(0).cpu()[0]
+ n_tokens = int(acc[0].item())
+ emitted = acc[1 : 1 + n_tokens].tolist()
+ else:
+ from tilert.models.deepseek_v3_2.temp_var_indices import Idx
+
+ res = dl.forward(cur)
+ intermediates, *_ = res[0]
+ nxt = intermediates[Idx.TOKEN_OUT][0][0]
+ emitted = [int(nxt.item())]
+ cur = nxt
+ for tok in emitted:
+ if len(tokens) >= budget:
+ break
+ tok = int(tok)
+ if tok in stop_ids:
+ finished = True
+ finish = "stop"
+ break
+ tokens.append(tok)
+ last_tok = tok
+ if on_token:
+ on_token(tok)
+ if grammar_session is not None and grammar_session.accept(tok) == "terminated":
+ finished = True
+ finish = "stop"
+ break
+ if ar_ok:
+ prev = torch.tensor([last_tok], dtype=torch.int32, device="cuda:0")
+ finally:
+ if grammar_mask_written:
+ dl.reset_grammar_bitmask()
+ dl.reset_sequence()
self.last_stats = {"finish_reason": finish}
return tokens
diff --git a/tilert/pd_vllm/receive_server.py b/tilert/pd_vllm/receive_server.py
index 3ad06f2..215ac62 100644
--- a/tilert/pd_vllm/receive_server.py
+++ b/tilert/pd_vllm/receive_server.py
@@ -1,7 +1,5 @@
-"""Decode-side receive server (W4): Mooncake buffer + TCP control plane."""
-
-import contextlib
import logging
+import os
import queue
import socket
import threading
@@ -13,6 +11,11 @@
from tilert.pd_vllm import wire
logger = logging.getLogger("pd_vllm.receive")
+FREE = "free"
+RESERVED = "reserved"
+TRANSFERRING = "transferring"
+COMPLETE = "complete"
+CANCELLING = "cancelling"
@dataclass
@@ -22,12 +25,21 @@ class ReceivedRequest:
last_prompt_token: int
first_token_id: int | None
sampling: dict | None
+ prompt_token_ids: list = field(default_factory=list)
done_ranks: set = field(default_factory=set)
t_first_conn: float = 0.0
t_complete: float = 0.0
+ generation: int = 0
+ state: str = RESERVED
+ active_writers: int = 0
+
+ @property
+ def has_live_writer(self) -> bool:
+ return self.active_writers > 0
class ReceiveServer:
+
def __init__(
self,
profile,
@@ -37,25 +49,60 @@ def __init__(
device: str = "cuda:0",
request_timeout: float = 120.0,
transport: str = "mooncake",
+ buffer_device: str = "cuda:0",
):
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)
+ self.buffer_device = buffer_device
+ self.host_buffer = buffer_device == "cpu"
+ nshards = int((os.environ.get("TILERT_PD_SHARDS") or "8").strip() or 8)
+ nshards = max(1, min(nshards, torch.cuda.device_count() or 1))
+ if self.host_buffer and nshards > 1:
+ logger.info(
+ "receive buffer is host-resident: ignoring TILERT_PD_SHARDS=%d "
+ "(a DRAM region has no per-card pressure to spread)",
+ nshards,
+ )
+ nshards = 1
+ shard_bytes = 0
+ if self.host_buffer:
+ from tilert.pd_vllm.transport import alloc_pinned_huge
- # RDMA transport (mooncake default / nixl), single cuda:0 registration
+ logger.info(
+ "allocating receive buffer: %.2f GiB in pinned host DRAM (profile=%s)",
+ total / 1024**3,
+ profile.name,
+ )
+ self.buffer = alloc_pinned_huge(total)
+ self.base_ptr = self.buffer.data_ptr()
+ elif nshards > 1:
+ shard_bytes = profile.shard_bytes(max_seq_len, nshards)
+ logger.info(
+ "allocating receive buffer: %d shards x %.2f GB on cuda:0..%d (profile=%s)",
+ nshards,
+ shard_bytes / 1024**3,
+ nshards - 1,
+ profile.name,
+ )
+ self.buffer = [
+ torch.zeros(shard_bytes, dtype=torch.uint8, device=f"cuda:{i}")
+ for i in range(nshards)
+ ]
+ self.base_ptr = [b.data_ptr() for b in self.buffer]
+ else:
+ 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)
from tilert.pd_vllm.transport import make_transport
if hostname is None:
@@ -63,22 +110,28 @@ def __init__(
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)
+ if nshards > 1:
+ rails = self._transport.rails()
+ for i, ptr in enumerate(self.base_ptr):
+ self._transport.register(ptr, shard_bytes, i, f"hip:{i % rails}" if rails else None)
+ logger.info("receive buffer shards pinned to %s rails", rails or "auto")
+ else:
+ self._transport.register(self.base_ptr, total, dev_id, host=self.host_buffer)
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._cancelled: dict[str, float] = {}
+ self._generation = 0
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):
+ try:
self._srv.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)
+ except OSError:
+ pass
self._srv.bind(("::", ctrl_port))
self._srv.listen(32)
self._stop = threading.Event()
@@ -88,19 +141,86 @@ def __init__(
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)."""
+ def release(self, rid: str) -> None:
with self._lock:
+ self._tombstone(rid)
+ cur = self._current
+ if cur is None:
+ return
+ if cur.rid != rid:
+ logger.info("release(%s) ignored: the slot now holds %s", rid, cur.rid)
+ return
+ if cur.has_live_writer and cur.state != COMPLETE:
+ cur.state = CANCELLING
+ logger.warning(
+ "release(%s) with %d writer(s) still active: slot stays claimed (cancelling) until they stop",
+ cur.rid,
+ cur.active_writers,
+ )
+ return
self._current = None
+ def _tombstone(self, rid: str) -> None:
+ now = time.time()
+ self._cancelled = {r: t for r, t in self._cancelled.items() if t > now}
+ self._cancelled[rid] = now + self.request_timeout
+
+ def _extend_tombstone(self, rid: str, window_s) -> None:
+ try:
+ window = float(window_s)
+ except (TypeError, ValueError):
+ return
+ if window <= 0:
+ return
+ deadline = time.time() + window
+ if deadline > self._cancelled.get(rid, 0.0):
+ self._cancelled[rid] = deadline
+
+ def expect(self, rid: str) -> None:
+ with self._lock:
+ if self._cancelled.pop(rid, None) is not None:
+ logger.info("request %s re-announced; its tombstone is dropped", rid)
+
+ def _is_tombstoned(self, rid: str) -> bool:
+ deadline = self._cancelled.get(rid)
+ if deadline is None:
+ return False
+ if deadline <= time.time():
+ del self._cancelled[rid]
+ return False
+ return True
+
+ def _next_generation(self) -> int:
+ self._generation += 1
+ return self._generation
+
+ def _reusable(self, cur: ReceivedRequest) -> bool:
+ if cur.state == COMPLETE:
+ return True
+ if cur.has_live_writer:
+ return False
+ if cur.state in (FREE, CANCELLING):
+ return True
+ return time.time() - cur.t_first_conn >= self.request_timeout
+
+ def state_snapshot(self) -> dict:
+ with self._lock:
+ cur = self._current
+ if cur is None:
+ return {"state": FREE, "rid": None, "generation": None, "active_writers": 0}
+ return {
+ "state": cur.state,
+ "rid": cur.rid,
+ "generation": cur.generation,
+ "active_writers": cur.active_writers,
+ }
+
def close(self) -> None:
self._stop.set()
- with contextlib.suppress(OSError):
+ try:
self._srv.close()
-
- # ββ accept / per-connection handling βββββββββββββββββββββββββββββββββ
+ except OSError:
+ pass
def _accept_loop(self) -> None:
while not self._stop.is_set():
@@ -111,14 +231,78 @@ def _accept_loop(self) -> None:
t = threading.Thread(target=self._handle, args=(conn, addr), daemon=True)
t.start()
+ def _admit(self, req: dict, rid: str, rank: int) -> dict:
+ if req.get("kv_scales") and getattr(self.profile, "kv_scales", None) is None:
+ self.profile.set_kv_scales(req["kv_scales"])
+ with self._lock:
+ if self._is_tombstoned(rid):
+ self._extend_tombstone(rid, req.get("admission_window_s"))
+ logger.warning(
+ "refusing %s rank %d for now: no consumer is waiting for this request",
+ rid,
+ rank,
+ )
+ return wire.reject_msg("cancelling", rid=rid, detail="no_consumer")
+ cur = self._current
+ if cur is not None and cur.rid == rid:
+ if cur.state == CANCELLING:
+ return wire.reject_msg("cancelling", rid=rid)
+ if rank in cur.done_ranks:
+ return wire.reject_msg("duplicate_rank", rid=rid, rank=rank)
+ elif cur is not None and (not self._reusable(cur)):
+ logger.warning(
+ "rejecting %s rank %d (busy with %s, state=%s, writers=%d)",
+ rid,
+ rank,
+ cur.rid,
+ cur.state,
+ cur.active_writers,
+ )
+ return wire.reject_msg("busy", busy_rid=cur.rid, busy_state=cur.state)
+ else:
+ 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"),
+ prompt_token_ids=list(req.get("prompt_token_ids") or []),
+ t_first_conn=time.time(),
+ generation=self._next_generation(),
+ state=RESERVED,
+ )
+ logger.info(
+ "request %s: seq_len=%d (generation %d)", rid, cur.seq_len, cur.generation
+ )
+ if not cur.prompt_token_ids and req.get("prompt_token_ids"):
+ cur.prompt_token_ids = list(req["prompt_token_ids"])
+ logger.info(
+ "request %s: prompt bitmap seeded from rank %d (%d ids)",
+ rid,
+ rank,
+ len(cur.prompt_token_ids),
+ )
+ cur.state = TRANSFERRING
+ cur.active_writers += 1
+ return wire.accept_msg(rid, rank, cur.generation)
+
+ def _writer_left(self, rid: str, generation: int) -> None:
+ with self._lock:
+ cur = self._current
+ if cur is None or cur.generation != generation:
+ return
+ cur.active_writers = max(0, cur.active_writers - 1)
+ if cur.state == CANCELLING and cur.active_writers == 0:
+ logger.info("request %s drained; receive slot free", cur.rid)
+ self._current = None
+
def _handle(self, conn: socket.socket, addr) -> None:
+ admitted: tuple[str, int] | None = 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
+ advisory_busy = self._current is not None and (not self._reusable(self._current))
wire.send_msg(
conn,
wire.hello_msg(
@@ -130,43 +314,39 @@ def _handle(self, conn: socket.socket, addr) -> None:
busy=advisory_busy,
),
)
-
req = wire.recv_msg(conn)
- rid, rank = req["rid"], int(req["rank"])
+ 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"})
+ wire.send_msg(
+ conn,
+ wire.reject_msg(
+ "seq_len exceeds max_seq_len", rid=rid, max_seq_len=self.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)
+ reply = self._admit(req, rid, rank)
+ wire.send_msg(conn, reply)
+ if not reply.get("accepted"):
+ return
+ admitted = (rid, reply["generation"])
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:
+ if cur is None or cur.rid != rid or cur.generation != admitted[1]:
+ logger.warning(
+ "ignoring done from %s rank %d: it was admitted to generation %d and the slot has moved on",
+ rid,
+ rank,
+ admitted[1],
+ )
+ return
+ if cur.state == CANCELLING:
+ logger.warning(
+ "ignoring done from %s rank %d: request was abandoned", rid, rank
+ )
return
cur.done_ranks.add(rank)
logger.info(
@@ -177,6 +357,7 @@ def _handle(self, conn: socket.socket, addr) -> None:
len(self.profile.sender_ranks),
)
if cur.done_ranks >= set(self.profile.sender_ranks):
+ cur.state = COMPLETE
cur.t_complete = time.time()
self.completed.put(cur)
logger.info(
@@ -187,4 +368,6 @@ def _handle(self, conn: socket.socket, addr) -> None:
except Exception:
logger.exception("connection from %s failed", addr)
finally:
+ if admitted is not None:
+ self._writer_left(*admitted)
conn.close()
diff --git a/tilert/pd_vllm/reply.py b/tilert/pd_vllm/reply.py
new file mode 100644
index 0000000..45ca151
--- /dev/null
+++ b/tilert/pd_vllm/reply.py
@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass, field
+from typing import Any
+
+from tilert.pd_vllm.logprobs import LogprobsRequest, build_logprobs
+from tilert.pd_vllm.oai_parser import IncrementalDetok
+from tilert.pd_vllm.stop_strings import StopWindow
+
+__all__ = ["CONTENT", "REASONING", "TOOL_CALL", "Emission", "ReplyStream", "as_logprobs"]
+CONTENT = "content"
+REASONING = "reasoning"
+TOOL_CALL = "tool_call"
+
+
+def as_logprobs(entries: list[dict]) -> dict:
+ return {"content": entries, "refusal": None}
+
+
+@dataclass(frozen=True)
+class Emission:
+ channel: str
+ text: str = ""
+ logprobs: list[dict] = field(default_factory=list)
+ tool_call: dict | None = None
+
+
+@dataclass
+class _Entry:
+ token_id: int
+ logprob: float | None
+ candidates: list[tuple[int, float]]
+ ends_at: int
+
+
+class ReplyStream:
+
+ def __init__(
+ self,
+ tokenizer,
+ *,
+ stop: Iterable[str] = (),
+ include_stop_in_output: bool = False,
+ parser_session: Any | None = None,
+ logprobs_req: LogprobsRequest | None = None,
+ first_token_logprob: tuple[float | None, list] | None = None,
+ ) -> None:
+ stop = list(stop)
+ if stop and parser_session is not None and (logprobs_req is not None):
+ raise ValueError(
+ "stop with both an output parser and logprobs is not attributable; the request gate must refuse it"
+ )
+ self._session = parser_session
+ self._logprobs_req = logprobs_req
+ self._first_lp = first_token_logprob
+ self._detok = IncrementalDetok(tokenizer, skip_special_tokens=parser_session is None)
+ self._decode_one: Callable[[int], str] = lambda t: tokenizer.decode(
+ [t], skip_special_tokens=False
+ )
+ self._window = StopWindow(stop, include_stop_in_output)
+ self._ids: list[int] = []
+ self._chars = 0
+ self._due: list[_Entry] = []
+ self._held: list[dict] = []
+ self._finished = False
+
+ @property
+ def token_ids(self) -> list[int]:
+ return list(self._ids)
+
+ @property
+ def completion_tokens(self) -> int:
+ return len(self._ids)
+
+ @property
+ def stop_reason(self) -> str | None:
+ return self._window.stopped
+
+ def finish_reason(self, from_decode: str) -> str:
+ return "stop" if self._window.stopped is not None else from_decode
+
+ def push(
+ self,
+ token_ids: Iterable[int],
+ logprobs: list[float | None] | None = None,
+ candidates: list[list[tuple[int, float]]] | None = None,
+ ) -> list[Emission]:
+ out: list[Emission] = []
+ for i, tid in enumerate(token_ids):
+ if self._window.stopped is not None:
+ break
+ self._ids.append(tid)
+ delta = self._detok.push([tid])
+ self._chars += len(delta)
+ self._queue(tid, logprobs, candidates, i, bool(delta))
+ self._window.push(delta)
+ out += self._route(self._window.take(limit=self._cap()))
+ return out
+
+ def finish(self) -> list[Emission]:
+ if self._finished:
+ return []
+ self._finished = True
+ tail = self._detok.finish()
+ if tail:
+ self._chars += len(tail)
+ for pending in self._due:
+ if pending.ends_at < 0:
+ pending.ends_at = self._chars
+ self._window.push(tail)
+ out = self._route(self._window.take(final=True))
+ if self._session is not None:
+ out += self._events(self._session.finish(), self._collect())
+ return out
+ leftover = self._collect()
+ if leftover:
+ out.append(Emission(CONTENT, "", leftover))
+ return out
+
+ def _queue(self, tid, logprobs, candidates, i, produced_text) -> None:
+ if self._logprobs_req is None:
+ return
+ lp = logprobs[i] if logprobs and i < len(logprobs) else None
+ cands = list(candidates[i]) if candidates and i < len(candidates) else []
+ if len(self._ids) == 1 and lp is None and (self._first_lp is not None):
+ lp, cands = (self._first_lp[0], list(self._first_lp[1] or []))
+ if produced_text:
+ for pending in self._due:
+ if pending.ends_at < 0:
+ pending.ends_at = self._chars
+ self._due.append(_Entry(tid, lp, cands, ends_at=self._chars if produced_text else -1))
+
+ def _cap(self) -> int | None:
+ if self._logprobs_req is None:
+ return None
+ end = self._chars - self._window.hold
+ cap = self._window.visible
+ for pending in self._due:
+ if 0 <= pending.ends_at <= end:
+ cap = pending.ends_at
+ return cap
+
+ def _route(self, text: str) -> list[Emission]:
+ if self._session is None:
+ entries = self._ready()
+ if not text and (not entries):
+ return []
+ return [Emission(CONTENT, text, entries)]
+ entries = self._held + self._ready()
+ if not text:
+ self._held = entries
+ return []
+ self._held = []
+ return self._events(self._session.feed(text), entries)
+
+ def _events(self, events: list[dict], entries: list[dict]) -> list[Emission]:
+ out: list[Emission] = []
+ for ev in events:
+ kind = ev.get("kind")
+ if kind == "tool":
+ out.append(
+ Emission(
+ TOOL_CALL,
+ tool_call={
+ "index": ev["index"],
+ "id": ev["id"],
+ "name": ev["name"],
+ "arguments": ev["arguments"],
+ },
+ )
+ )
+ elif kind == "reasoning":
+ out.append(Emission(REASONING, ev.get("text", "")))
+ else:
+ out.append(Emission(CONTENT, ev.get("text", ""), entries))
+ entries = []
+ return out
+
+ def _ready(self) -> list[dict]:
+ visible = self._window.visible
+ stopped = self._window.stopped is not None
+ n = 0
+ while (
+ n < len(self._due)
+ and self._due[n].ends_at >= 0
+ and (stopped or self._due[n].ends_at <= visible)
+ ):
+ n += 1
+ return self._build(n)
+
+ def _collect(self) -> list[dict]:
+ return self._build(len(self._due))
+
+ def _build(self, n: int) -> list[dict]:
+ if not n:
+ return []
+ taken, self._due = (self._due[:n], self._due[n:])
+ return build_logprobs(
+ [e.token_id for e in taken],
+ [e.logprob for e in taken],
+ [[tuple(c) for c in e.candidates] for e in taken],
+ self._logprobs_req,
+ self._decode_one,
+ )["content"]
diff --git a/tilert/pd_vllm/request_gate.py b/tilert/pd_vllm/request_gate.py
new file mode 100644
index 0000000..bc96e99
--- /dev/null
+++ b/tilert/pd_vllm/request_gate.py
@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from pydantic import TypeAdapter, ValidationError
+
+from tilert.pd_vllm.capabilities import CapabilityUnavailable, InvalidParameter
+from tilert.pd_vllm.grammar_spec import extract_request_grammar_spec
+from tilert.pd_vllm.logprobs import LogprobsRequest, LogprobsUnsupported, resolve_logprobs_request
+from tilert.pd_vllm.openai_params import InvalidOutputLength, resolve_max_tokens
+from tilert.pd_vllm.stop_strings import resolve_stop
+
+__all__ = [
+ "GatedRequest",
+ "gate_request",
+ "refuse_unattributable_logprobs",
+ "require_tokenizer_for_logprobs",
+ "resolve_stop_request",
+]
+
+
+@dataclass(frozen=True)
+class GatedRequest:
+ is_chat: bool
+ thinking: bool
+ stop: list[str] = field(default_factory=list)
+ include_stop: bool = False
+ logprobs_req: LogprobsRequest | None = None
+ grammar_spec: dict | None = None
+
+
+def gate_request(path: str, body: dict, *, tokenizer, parser_active) -> GatedRequest:
+ is_chat = path.endswith("chat/completions")
+ grammar_spec = extract_request_grammar_spec(body)
+ try:
+ resolve_max_tokens(body)
+ except InvalidOutputLength as e:
+ raise InvalidParameter(str(e)) from None
+ logprobs_req = _logprobs_of(path, body)
+ require_tokenizer_for_logprobs(logprobs_req, tokenizer)
+ stop, include_stop = resolve_stop_request(body, tokenizer)
+ thinking = _thinking_enabled(body) if is_chat else True
+ refuse_unattributable_logprobs(logprobs_req, stop, is_chat and parser_active(thinking))
+ return GatedRequest(
+ is_chat=is_chat,
+ thinking=thinking,
+ stop=stop,
+ include_stop=include_stop,
+ logprobs_req=logprobs_req,
+ grammar_spec=grammar_spec,
+ )
+
+
+def _thinking_enabled(body: dict) -> bool:
+ ctk = body.get("chat_template_kwargs")
+ if ctk is None:
+ return True
+ if not isinstance(ctk, dict):
+ raise InvalidParameter(f"chat_template_kwargs must be an object, got {type(ctk).__name__}")
+ return bool(ctk.get("enable_thinking", True))
+
+
+def _logprobs_of(path: str, body: dict):
+ if path.endswith("chat/completions"):
+ return resolve_logprobs_request(body)
+ if body.get("logprobs") is not None:
+ raise LogprobsUnsupported("logprobs is supported on /v1/chat/completions")
+ return None
+
+
+def refuse_unattributable_logprobs(logprobs_req, stop, has_parser) -> None:
+ if logprobs_req is not None and stop and has_parser:
+ raise CapabilityUnavailable(
+ "logprobs cannot be attributed to message.content when stop strings and an output parser are both in play: the parser buffers across the boundary the stop hold-back creates, so some tokens' channel is undecidable. Drop one of stop / logprobs, or use a router started with --parser none."
+ )
+
+
+def require_tokenizer_for_logprobs(logprobs_req, tokenizer) -> None:
+ if logprobs_req is not None and tokenizer is None:
+ raise CapabilityUnavailable(
+ "logprobs need a tokenizer to name each token, and this router has none: start it with --model-path to serve them."
+ )
+
+
+def resolve_stop_request(body: dict, tokenizer) -> tuple[list[str], bool]:
+ try:
+ stop = resolve_stop(body)
+ except ValueError as e:
+ raise InvalidParameter(str(e)) from e
+ include = False
+ if "include_stop_str_in_output" in body:
+ try:
+ include = TypeAdapter(bool).validate_python(body["include_stop_str_in_output"])
+ except ValidationError:
+ raise InvalidParameter(
+ f"include_stop_str_in_output must be a boolean, got {body['include_stop_str_in_output']!r}"
+ ) from None
+ if stop and tokenizer is None:
+ raise CapabilityUnavailable(
+ "stop strings need a tokenizer to match against the reply text, and this router has none: matching is text-level, so the decode node's token-id stop set cannot serve it."
+ )
+ return (stop, include)
diff --git a/tilert/pd_vllm/sampling.py b/tilert/pd_vllm/sampling.py
new file mode 100644
index 0000000..dfb0866
--- /dev/null
+++ b/tilert/pd_vllm/sampling.py
@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+__all__ = [
+ "GREEDY_LOGPROBS_TOP_P",
+ "TOP_K_DISABLED",
+ "VLLM_DEFAULT_TOP_P",
+ "resolve_top_k",
+ "resolve_top_p",
+]
+GREEDY_LOGPROBS_TOP_P = 1e-09
+VLLM_DEFAULT_TOP_P = 1.0
+
+
+def resolve_top_p(sampling: dict, default: float = VLLM_DEFAULT_TOP_P) -> float:
+ raw = sampling.get("top_p")
+ if raw is None:
+ return float(default)
+ if isinstance(raw, bool):
+ raise ValueError("top_p must be a number, got bool")
+ return float(raw)
+
+
+_KERNEL_TOP_K_POOL = 256
+TOP_K_DISABLED = _KERNEL_TOP_K_POOL
+_TOP_K_APPLIED_MAX = _KERNEL_TOP_K_POOL - 1
+
+
+def resolve_top_k(sampling: dict) -> int:
+ raw = sampling.get("top_k")
+ if raw is None:
+ return TOP_K_DISABLED
+ k = int(raw)
+ if k < 1 or k > _TOP_K_APPLIED_MAX:
+ return TOP_K_DISABLED
+ return k
diff --git a/tilert/pd_vllm/stop_strings.py b/tilert/pd_vllm/stop_strings.py
new file mode 100644
index 0000000..c5df6b3
--- /dev/null
+++ b/tilert/pd_vllm/stop_strings.py
@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+import sys
+
+__all__ = ["StopWindow", "check_stop_strings", "resolve_stop"]
+
+
+def resolve_stop(body: dict) -> list[str]:
+ raw = body.get("stop")
+ if raw is None:
+ return []
+ if isinstance(raw, str):
+ raw = [raw]
+ if not isinstance(raw, (list, tuple)):
+ raise ValueError(f"stop must be a string or list of strings, got {type(raw).__name__}")
+ out = []
+ for s in raw:
+ if not isinstance(s, str):
+ raise ValueError(f"stop entries must be strings, got {type(s).__name__}")
+ if not s:
+ raise ValueError("stop cannot contain an empty string")
+ out.append(s)
+ return out
+
+
+def check_stop_strings(
+ output_text: str, new_char_count: int, stop: list[str], include_in_output: bool
+) -> tuple[str, int] | None:
+ if not new_char_count or not stop:
+ return None
+ best_stop_str: str | None = None
+ best_stop_index = 0
+ best_end = sys.maxsize
+ for stop_str in stop:
+ stop_len = len(stop_str)
+ stop_index = output_text.find(stop_str, 1 - new_char_count - stop_len)
+ if stop_index == -1:
+ continue
+ end = stop_index + stop_len
+ if end < best_end:
+ best_stop_str = stop_str
+ best_stop_index = stop_index
+ best_end = end
+ if best_stop_str is None:
+ return None
+ if include_in_output:
+ if best_end >= len(output_text):
+ return (best_stop_str, -1)
+ return (best_stop_str, best_end)
+ return (best_stop_str, best_stop_index)
+
+
+class StopWindow:
+
+ def __init__(self, stop: list[str], include_in_output: bool = False):
+ self.stop = list(stop)
+ self.include_in_output = include_in_output
+ longest = max((len(s) for s in self.stop), default=1) - 1
+ self._hold = 0 if include_in_output else longest
+ self._keep = longest
+ self._slack = 4096 if self.stop else 0
+ self._text = ""
+ self._base = 0
+ self._taken = 0
+ self._stopped: str | None = None
+
+ @property
+ def stopped(self) -> str | None:
+ return self._stopped
+
+ @property
+ def hold(self) -> int:
+ return self._hold
+
+ @property
+ def visible(self) -> int:
+ return self._taken
+
+ def push(self, delta: str) -> None:
+ if self._stopped is not None or not delta:
+ return
+ self._text += delta
+ hit = check_stop_strings(self._text, len(delta), self.stop, self.include_in_output)
+ if hit is None:
+ return
+ self._stopped, truncate_to = hit
+ if truncate_to != -1:
+ self._text = self._text[:truncate_to]
+
+ def take(self, *, final: bool = False, limit: int | None = None) -> str:
+ end = self._base + len(self._text)
+ if not (final or self._stopped is not None):
+ end -= self._hold
+ if limit is not None and limit < end:
+ end = limit
+ if end <= self._taken:
+ return ""
+ out = self._text[self._taken - self._base : end - self._base]
+ self._taken = end
+ self._trim()
+ return out
+
+ def _trim(self) -> None:
+ if len(self._text) <= self._keep + self._hold + self._slack:
+ return
+ cut = self._taken - self._keep - self._base
+ if cut > 0:
+ self._text = self._text[cut:]
+ self._base += cut
diff --git a/tilert/pd_vllm/transport.py b/tilert/pd_vllm/transport.py
index 7d5e779..e9b25fd 100644
--- a/tilert/pd_vllm/transport.py
+++ b/tilert/pd_vllm/transport.py
@@ -1,5 +1,3 @@
-"""Pluggable RDMA transport for the PD data plane: Mooncake (default) or NIXL."""
-
from __future__ import annotations
import base64
@@ -10,14 +8,20 @@ class Transport:
name = "?"
def init(self, host: str) -> None: ...
- def register(self, ptr: int, nbytes: int, dev_id: int) -> None: ...
+
+ def register(
+ self, ptr: int, nbytes: int, dev_id: int, location: str | None = None, host: bool = False
+ ) -> None: ...
+
+ def rails(self) -> int:
+ return 0
+
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:
@@ -29,11 +33,26 @@ def init(self, host: str) -> None:
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])
+ def register(
+ self, ptr: int, nbytes: int, dev_id: int, location: str | None = None, host: bool = False
+ ) -> None:
+ if host:
+ location = None
+ args = ([ptr], [nbytes]) if location is None else ([ptr], [nbytes], location)
+ ret = self.engine.batch_register_memory(*args)
if ret != 0:
raise RuntimeError(f"Mooncake register failed: {ret}")
+ def rails(self) -> int:
+ try:
+ import json
+
+ topo = self.engine.get_local_topology()
+ topo = json.loads(topo) if isinstance(topo, str) else topo
+ return len({h for v in topo.values() for lst in v for h in lst})
+ except Exception:
+ return 0
+
def local_meta(self) -> dict:
return {"session_id": self.session_id}
@@ -44,26 +63,23 @@ def write(self, remote_meta: dict, srcs, dsts, lens) -> None:
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
+ _MAX_POLL = 2000000
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._remotes: dict[bytes, str] = {}
self._dev = 0
+ self._mem_type = "VRAM"
- 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 register(
+ self, ptr: int, nbytes: int, dev_id: int, location: str | None = None, host: bool = False
+ ) -> None:
+ self._dev = 0 if host else dev_id
+ self._mem_type = "DRAM" if host else "VRAM"
+ self._agent.register_memory([(ptr, nbytes, self._dev, "")], self._mem_type)
def local_meta(self) -> dict:
return {
@@ -79,10 +95,10 @@ def write(self, remote_meta: dict, srcs, dsts, lens) -> None:
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"
+ [(int(s), int(n), self._dev) for s, n in zip(srcs, lens)], self._mem_type
)
rd = self._agent.get_xfer_descs(
- [(int(d), int(n), rdev) for d, n in zip(dsts, lens)], "VRAM"
+ [(int(d), int(n), rdev) for d, n in zip(dsts, lens)], self._mem_type
)
h = self._agent.initialize_xfer("WRITE", ld, rd, rname)
try:
@@ -105,5 +121,51 @@ def write(self, remote_meta: dict, srcs, dsts, lens) -> None:
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)}")
+ raise ValueError(f"unknown transport {name!r}; choices: {sorted(_BACKENDS)}")
return _BACKENDS[key]()
+
+
+def alloc_pinned_huge(total: int):
+ import ctypes
+ import mmap
+ import re
+
+ import torch
+
+ MiB = 1 << 20
+ HUGE = 2 * MiB
+ if total % HUGE:
+ total += HUGE - total % HUGE
+ mm = mmap.mmap(-1, total + HUGE, flags=mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS)
+ raw = ctypes.addressof(ctypes.c_char.from_buffer(mm))
+ off = (-raw) % HUGE
+ addr = raw + off
+ libc = ctypes.CDLL(None, use_errno=True)
+ libc.madvise.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
+ if libc.madvise(ctypes.c_void_p(addr), ctypes.c_size_t(total), ctypes.c_int(14)) != 0:
+ raise RuntimeError(f"madvise(MADV_HUGEPAGE) failed errno={ctypes.get_errno()}")
+ if libc.madvise(ctypes.c_void_p(addr), ctypes.c_size_t(total), ctypes.c_int(23)) != 0:
+ for o in range(off, off + total, HUGE):
+ mm[o] = 0
+ huge = 0
+ for blk in re.split(r"\n(?=[0-9a-f]+-[0-9a-f]+ )", open("/proc/self/smaps").read()):
+ m = re.match(r"([0-9a-f]+)-([0-9a-f]+) ", blk)
+ if not m:
+ continue
+ lo, hi = int(m.group(1), 16), int(m.group(2), 16)
+ if hi <= addr or lo >= addr + total:
+ continue
+ h = re.search(r"AnonHugePages:\s+(\d+) kB", blk)
+ huge += int(h.group(1)) * 1024 if h else 0
+ if huge < total:
+ raise RuntimeError(
+ f"PD host buffer is only {huge / 2**30:.2f} of {total / 2**30:.2f} GiB huge-page backed; "
+ "the RDMA MR would fall back to 4 KiB pages and exceed the per-HCA entry budget "
+ "(check /sys/kernel/mm/transparent_hugepage/{enabled,defrag} and free memory)"
+ )
+ rc = torch.cuda.cudart().cudaHostRegister(addr, total, 0)
+ if int(rc) != 0:
+ raise RuntimeError(f"cudaHostRegister/hipHostRegister failed: {rc}")
+ buf = torch.frombuffer(mm, dtype=torch.uint8, count=total, offset=off)
+ buf._pd_mmap = mm
+ return buf
diff --git a/tilert/pd_vllm/wire.py b/tilert/pd_vllm/wire.py
index 284850c..d50a9b0 100644
--- a/tilert/pd_vllm/wire.py
+++ b/tilert/pd_vllm/wire.py
@@ -1,17 +1,14 @@
-"""Shared control-plane protocol for vLLM-prefill -> TileRT-decode PD."""
-
import json
import socket
import struct
MAGIC = "tilert-pd"
-
+PROTOCOL_VERSION = 2
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")
@@ -23,12 +20,19 @@ def local_ip(probe_addr: str | None = None) -> str:
s.close()
-def derive_rid(request_id: str) -> str:
- """Map a vLLM request/response id to the client-visible rid.
+def wants_prompt_token_ids(sampling: dict | None) -> bool:
+ if not sampling:
+ return False
+ rep = sampling.get("repetition_penalty")
+ if rep is None:
+ return False
+ try:
+ return float(rep) != 1.0
+ except (TypeError, ValueError):
+ return False
- Shared by the prefill connector (internal id) and the router (response id)
- so both agree.
- """
+
+def derive_rid(request_id: str) -> str:
rid = request_id
for prefix in ("chatcmpl-", "cmpl-"):
if rid.startswith(prefix):
@@ -38,7 +42,7 @@ def derive_rid(request_id: str) -> str:
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:
+ if len(parts) == 2 and parts[1].isdigit() and (len(parts[1]) <= 3):
rid = parts[0]
return rid
@@ -74,15 +78,9 @@ def hello_msg(
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,
+ "protocol_version": PROTOCOL_VERSION,
"layout_version": layout_version,
"transport": transport,
"max_seq_len": max_seq_len,
@@ -90,3 +88,15 @@ def hello_msg(
**transport_meta,
**layout,
}
+
+
+def accept_msg(rid: str, rank: int, generation: int) -> dict:
+ return {"accepted": True, "rid": rid, "rank": rank, "generation": generation}
+
+
+def reject_msg(reason: str, **extra) -> dict:
+ return {"accepted": False, "error": reason, **extra}
+
+
+def done_msg(rid: str, rank: int, generation: int) -> dict:
+ return {"done": True, "rid": rid, "rank": rank, "generation": generation}
diff --git a/tilert/tilert_init.py b/tilert/tilert_init.py
index d0cd30f..f699978 100644
--- a/tilert/tilert_init.py
+++ b/tilert/tilert_init.py
@@ -1,18 +1,24 @@
-"""Tilert init operation module."""
+"""Runtime warm-up entry points."""
import torch
-__all__ = [
- "tilert_init",
- "tilert_force_init",
-]
+__all__ = ["tilert_init", "tilert_force_init"]
+
+
+def _has_op(name: str) -> bool:
+ return hasattr(torch.ops.tilert, name)
def tilert_init() -> None:
- """Tilert init operation."""
- torch.ops.tilert.tilert_init_op()
+ if _has_op("tilert_init_op"):
+ torch.ops.tilert.tilert_init_op()
+ return
+ torch.zeros(1, device=f"cuda:{torch.cuda.current_device()}")
+ torch.cuda.synchronize()
def tilert_force_init() -> None:
- """Tilert force init operation."""
- torch.ops.tilert.tilert_force_init_op()
+ if _has_op("tilert_force_init_op"):
+ torch.ops.tilert.tilert_force_init_op()
+ return
+ tilert_init()
diff --git a/tilert/utils.py b/tilert/utils.py
index 4cc3b47..26adbdf 100644
--- a/tilert/utils.py
+++ b/tilert/utils.py
@@ -1,4 +1,4 @@
-"""Utility functions for testing."""
+"""Utility helpers used by the model packages."""
from typing import Any
@@ -7,84 +7,41 @@
__all__ = [
"alloc_misc_ws",
"cosine_similarity",
- "relative_l2_error",
"get_profile_log_tensor",
+ "relative_l2_error",
"SLICES_FOR_TILERT_OP",
]
-
SLICES_FOR_TILERT_OP = 1
def get_profile_log_tensor(
- device_index: int = 0,
- device: torch.device | None = None,
- num_max_insts: int = 64,
+ device_index: int = 0, device: torch.device | None = None, num_max_insts: int = 64
) -> torch.Tensor | None:
- """Get a profile log tensor for the given device index.
-
- Returns ``None`` when no CUDA GPUs are visible so the offline
- weight-conversion path can run with ``CUDA_VISIBLE_DEVICES=""``.
-
- Args:
- device_index: The index of the device.
- device: The device to use.
-
- Returns:
- A profile log tensor, or ``None`` if CUDA is unavailable.
- """
if not torch.cuda.is_available():
return None
if device is None:
device = torch.device("cuda", device_index)
-
- props = torch.cuda.get_device_properties(device_index)
+ device = torch.device(device)
+ props = torch.cuda.get_device_properties(device.index)
num_sm = props.multi_processor_count
-
return torch.zeros(
num_max_insts + 1 + SLICES_FOR_TILERT_OP, num_sm, 16, dtype=torch.uint64, device=device
)
-def alloc_misc_ws(
- num_max_insts: int = 64,
- device_id: int = 0,
-) -> torch.Tensor:
- """Allocate a misc workspace tensor.
-
- Args:
- num_max_insts: Maximum number of profiled instructions.
- device_id: CUDA device index to allocate on.
-
- Returns:
- A zeroed int64 tensor of shape (total_rows, num_sm, 16) on the
- requested CUDA device.
- """
+def alloc_misc_ws(num_max_insts: int = 64, device_id: int = 0) -> torch.Tensor:
+ if not hasattr(torch.ops.tilert, "alloc_misc_ws"):
+ raise RuntimeError(
+ "alloc_misc_ws is not registered by this engine build; use get_profile_log_tensor instead"
+ )
return torch.ops.tilert.alloc_misc_ws(num_max_insts, device_id)
def cosine_similarity(gt: torch.Tensor, out: torch.Tensor) -> torch.Tensor:
- """Calculate the cosine similarity.
-
- Args:
- gt: The ground truth tensor.
- out: The output tensor.
-
- Returns:
- The cosine similarity.
- """
return torch.nn.functional.cosine_similarity(
gt.flatten().float(), out.flatten().float(), dim=-1
)
def relative_l2_error(gt: torch.Tensor, out: torch.Tensor) -> Any:
- """Calculate the relative L2 error.
-
- Args:
- gt: The ground truth tensor.
- out: The output tensor.
-
- Returns:
- The relative L2 error.
- """
- return torch.norm(gt - out) / torch.norm(gt)
+ return torch.norm(gt.float() - out.float()) / torch.norm(gt.float())