Release Ciru v4 (4.0.0): integrated runtime and patch notes
Browse filesPublish the qualified combined runtime, matched source provenance, evaluation summaries and additive model-card updates. Existing target/draft/vision tensors remain unchanged.
This view is limited to 50 files because it contains too many changes. See raw diff
- INSTALL.md +28 -9
- README.md +72 -15
- RELEASE.json +137 -295
- RUNTIME-FIXES.md +37 -1
- SOURCE-PROVENANCE.json +159 -0
- bundle/LICENSES/XGRAMMAR-LICENSE +201 -0
- bundle/native/libornith_attention_iu4.so +2 -2
- bundle/native/libornith_dense_g256.so +0 -0
- bundle/native/libornith_dense_g256_n32.so +0 -0
- bundle/native/libornith_persistent_iu4.so +2 -2
- bundle/native/libornith_routed_storage_n32.so +2 -2
- bundle/packaging/serve.sh +1 -1
- bundle/plugin-site/TOOL-INTEGRATION.json +205 -0
- bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/direct_url.json +0 -1
- bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/uv_build.json +0 -1
- bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/uv_cache.json +0 -1
- bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/INSTALLER +0 -0
- bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/METADATA +1 -1
- bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/RECORD +42 -17
- bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/REQUESTED +0 -0
- bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/WHEEL +0 -0
- bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/entry_points.txt +0 -0
- bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/licenses/LICENSE-APACHE-2.0 +0 -0
- bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/top_level.txt +0 -0
- bundle/plugin-site/ornith_g256/__init__.py +6 -0
- bundle/plugin-site/ornith_g256/_vllm_correctness/chat_completion_serving_r03.py +1319 -0
- bundle/plugin-site/ornith_g256/_vllm_correctness/manifest.json +7 -1
- bundle/plugin-site/ornith_g256/_xgrammar_native/build-receipt.json +29 -0
- bundle/plugin-site/ornith_g256/_xgrammar_native/libxgrammar-candidate.so +3 -0
- bundle/plugin-site/ornith_g256/_xgrammar_native/selection.json +4 -0
- bundle/plugin-site/ornith_g256/attention_compact.py +5 -5
- bundle/plugin-site/ornith_g256/attention_iu4.py +20 -9
- bundle/plugin-site/ornith_g256/attention_iu4_persistent.py +2 -2
- bundle/plugin-site/ornith_g256/optimized/__init__.py +1 -0
- bundle/plugin-site/ornith_g256/optimized/compile_scope.py +30 -0
- bundle/plugin-site/ornith_g256/optimized/draft.py +90 -0
- bundle/plugin-site/ornith_g256/optimized/draft_kernel.py +50 -0
- bundle/plugin-site/ornith_g256/optimized/grammar_parallel.py +122 -0
- bundle/plugin-site/ornith_g256/optimized/greedy.py +40 -0
- bundle/plugin-site/ornith_g256/optimized/greedy_kernel.py +40 -0
- bundle/plugin-site/ornith_g256/optimized/launch.py +17 -0
- bundle/plugin-site/ornith_g256/optimized/mixed_qk_sum.py +66 -0
- bundle/plugin-site/ornith_g256/optimized/pipeline/SOURCE-PROOF.json +34 -0
- bundle/plugin-site/ornith_g256/optimized/pipeline/__init__.py +1 -0
- bundle/plugin-site/ornith_g256/optimized/pipeline/adapter.py +172 -0
- bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/causal_conv1d.py +1307 -0
- bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/chunk_o.py +190 -0
- bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/chunk_scaled_dot_kkt.py +176 -0
- bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/solve_tril.py +558 -0
- bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/wy_fast.py +158 -0
INSTALL.md
CHANGED
|
@@ -1,14 +1,33 @@
|
|
| 1 |
# Install and run Ornith1.5 Ciru Halo Agent
|
| 2 |
|
| 3 |
-
## Updating to Ciru
|
| 4 |
|
| 5 |
-
Download the
|
| 6 |
-
|
| 7 |
-
|
| 8 |
-
|
| 9 |
-
do not need reinstalling. Native schemas apply to automatic tool calls by
|
| 10 |
-
default; an explicit `strict: false` opts a function out of schema enforcement.
|
| 11 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 12 |
|
| 13 |
This release includes the target model, trained DFlash2 drafter, native kernels, custom vLLM plugin, and exact vLLM/AITER runtime wheels and source archives. **Use this runtime; stock `pip install vllm` does not provide the custom quantization or serving path.**
|
| 14 |
|
|
@@ -74,12 +93,12 @@ curl http://127.0.0.1:8000/v1/chat/completions \
|
|
| 74 |
-d '{"model":"ciru-halo-agent","messages":[{"role":"user","content":"Write a Python function that merges overlapping intervals."}],"temperature":0.6,"top_p":0.95,"max_tokens":4096,"chat_template_kwargs":{"enable_thinking":false}}'
|
| 75 |
```
|
| 76 |
|
| 77 |
-
The default profile supplies **262,144 tokens of per-request context capacity**, **eight active sequences**, **44 GiB shared KV/state pool**, prefix caching and adaptive speculation. Input and output share the context window; your agent client must reserve output space and compact history before filling it. Eight independent, fully populated 256K histories are not promised. The default server profile
|
| 78 |
|
| 79 |
First startup compiles/loads GPU kernels and creates caches. Wait for `/health` before sending work. Keep `bundle/cache` writable. Change runtime/model locations with `ORNITH_RUNTIME_ROOT`, `ORNITH_MODEL`, and `ORNITH_DRAFT`. Ordinary users do not need to change quantization or draft-policy settings.
|
| 80 |
|
| 81 |
## Build from source
|
| 82 |
|
| 83 |
-
|
| 84 |
|
| 85 |
Pinned runtime: vLLM `0.1.0rc2.dev9+g9255fd9fb9.rocm100` (base `9255fd9fb9fedf4b29d574a8d8bb21d93892cc98` plus supplied cache overlay), AITER `0.1.0rc1`, Python 3.14.3, PyTorch `2.13.0+rocm10.0.0`, ROCm SDK 10.0.0 and Transformers 5.16.1. Preserve included third-party licenses when redistributing.
|
|
|
|
| 1 |
# Install and run Ornith1.5 Ciru Halo Agent
|
| 2 |
|
| 3 |
+
## Updating to Ciru v4 (4.0.0)
|
| 4 |
|
| 5 |
+
Stop the model process before updating. Download the complete plugin, native
|
| 6 |
+
libraries and launchers together; a plugin-only update is insufficient. Existing
|
| 7 |
+
target, DFlash2, vision and source-prefill weights and pinned engine wheels are
|
| 8 |
+
unchanged. Keep a copy of the previous bundle for rollback.
|
|
|
|
|
|
|
| 9 |
|
| 10 |
+
```bash
|
| 11 |
+
uvx --from huggingface_hub hf download \
|
| 12 |
+
jcbtc/Ornith1.5-Ciru-Halo-Agent-vllm-strix-halo \
|
| 13 |
+
--include 'bundle/plugin-site/*' 'bundle/native/*' 'bundle/serve*.sh' \
|
| 14 |
+
'bundle/packaging/*' README.md INSTALL.md RUNTIME-FIXES.md RELEASE.json \
|
| 15 |
+
'evaluation/*' 'runtime/ciru-v4-ornith-source.tar.gz' \
|
| 16 |
+
--local-dir .
|
| 17 |
+
```
|
| 18 |
+
|
| 19 |
+
Once `bundle/plugin-site/ciru_ornith_g256-4.0.0.dist-info/` is present, remove
|
| 20 |
+
older `ciru_ornith_g256-*.dist-info/` directories so only the 4.0.0 metadata
|
| 21 |
+
remains, then restart the usual launcher. Also update any separately copied
|
| 22 |
+
plugin/native directories. Existing `bundle/paths.env` overrides remain valid;
|
| 23 |
+
a fresh download defaults to `installed-runtime` and `bundle/models/`.
|
| 24 |
+
|
| 25 |
+
Strict automatic tool schemas, the original sampler, 262,144-token context,
|
| 26 |
+
eight active requests, 44 GiB shared cache, native vision and adaptive DFlash2
|
| 27 |
+
policy are preserved. Explicit function `strict: false` retains its opt-out.
|
| 28 |
+
For thinking/agent work, omit a short completion cap and leave room in the
|
| 29 |
+
shared prompt/output context. Rollback requires restoring the previous complete
|
| 30 |
+
plugin/native/launcher set and restarting, not mixing release files.
|
| 31 |
|
| 32 |
This release includes the target model, trained DFlash2 drafter, native kernels, custom vLLM plugin, and exact vLLM/AITER runtime wheels and source archives. **Use this runtime; stock `pip install vllm` does not provide the custom quantization or serving path.**
|
| 33 |
|
|
|
|
| 93 |
-d '{"model":"ciru-halo-agent","messages":[{"role":"user","content":"Write a Python function that merges overlapping intervals."}],"temperature":0.6,"top_p":0.95,"max_tokens":4096,"chat_template_kwargs":{"enable_thinking":false}}'
|
| 94 |
```
|
| 95 |
|
| 96 |
+
The default profile supplies **262,144 tokens of per-request context capacity**, **eight active sequences**, **44 GiB shared KV/state pool**, prefix caching and adaptive speculation. Input and output share the context window; your agent client must reserve output space and compact history before filling it. Eight independent, fully populated 256K histories are not promised. The default server profile enables images. To explicitly select the image launcher, run `bash bundle/serve-vision.sh --host 127.0.0.1 --port 8000` after installation. The matching native BF16 vision encoder and projector are already included in `bundle/models/target/protected-00.safetensors`; no separate GGUF mmproj is required. This enables one image per request at a 1,048,576-pixel budget with up to eight active requests. See [the model card](README.md#vision--image-input) and [vision validation limits](VISION.md).
|
| 97 |
|
| 98 |
First startup compiles/loads GPU kernels and creates caches. Wait for `/health` before sending work. Keep `bundle/cache` writable. Change runtime/model locations with `ORNITH_RUNTIME_ROOT`, `ORNITH_MODEL`, and `ORNITH_DRAFT`. Ordinary users do not need to change quantization or draft-policy settings.
|
| 99 |
|
| 100 |
## Build from source
|
| 101 |
|
| 102 |
+
Use the [Ciru v4 source snapshot](runtime/ciru-v4-ornith-source.tar.gz) for the exact combined plugin, all eight native kernel sources and the build script. The [source repository](https://github.com/ciru-ai/ornith-ciru-halo-agent) retains the development history. See its `BUILD.md` for the native rebuild command. Runtime source archives are provided in this Hugging Face repository under `runtime/`; they include the matching vLLM/AITER source, licenses and release overlay notes. The binary installation above is the tested way to assemble the pinned engine; native source rebuilding is separate from retraining or requantizing the model.
|
| 103 |
|
| 104 |
Pinned runtime: vLLM `0.1.0rc2.dev9+g9255fd9fb9.rocm100` (base `9255fd9fb9fedf4b29d574a8d8bb21d93892cc98` plus supplied cache overlay), AITER `0.1.0rc1`, Python 3.14.3, PyTorch `2.13.0+rocm10.0.0`, ROCm SDK 10.0.0 and Transformers 5.16.1. Preserve included third-party licenses when redistributing.
|
README.md
CHANGED
|
@@ -23,10 +23,11 @@ tags:
|
|
| 23 |
- image-text-to-text
|
| 24 |
---
|
| 25 |
|
| 26 |
-
**Ciru
|
| 27 |
-
|
| 28 |
-
tool
|
| 29 |
-
|
|
|
|
| 30 |
|
| 31 |
# Ornith1.5 Ciru Halo Agent (vllm strix halo)
|
| 32 |
|
|
@@ -36,7 +37,62 @@ Image input is enabled using the upstream BF16 vision weights. See [patch notes]
|
|
| 36 |
|
| 37 |
The design starts with the hardware: packed four-bit weights, four-bit activation paths, Strix Halo four-bit matrix instructions, specialized kernels, and adaptive DFlash2 speculative decoding. Prefix caching and a shared memory pool let agents return to long working histories.
|
| 38 |
|
| 39 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 40 |
|
| 41 |
- **178 tok/s single-request decode** on the ten-question coding speed screen, with **166 ms mean time to first token**.
|
| 42 |
- **295 tok/s aggregate at eight concurrent requests**, completing the ten-question batch in **5.52 seconds**.
|
|
@@ -46,11 +102,11 @@ The design starts with the hardware: packed four-bit weights, four-bit activatio
|
|
| 46 |
|
| 47 |
These figures describe specific workloads, rather than a universal generation rate. The benchmark tables below include quality results and the workloads where other builds are faster.
|
| 48 |
|
| 49 |
-
**[Full benchmarks, runner builds, serving settings, and limitations](https://llm.ciru.ai/research/ornith-strix/)** · **[
|
| 50 |
|
| 51 |
-
## September 13 runtime fixes
|
| 52 |
|
| 53 |
-
**
|
| 54 |
|
| 55 |
## Running Ciru Halo Agent
|
| 56 |
|
|
@@ -71,7 +127,7 @@ Read **[Installation and build instructions](INSTALL.md)** first for Linux prere
|
|
| 71 |
|
| 72 |
### Vision / image input
|
| 73 |
|
| 74 |
-
Image input
|
| 75 |
|
| 76 |
```bash
|
| 77 |
bash bundle/serve-vision.sh --host 127.0.0.1 --port 8000
|
|
@@ -143,7 +199,7 @@ The startup JSON reports the effective `ORNITH_C1_POLICY` (`auto`, `k0`, `k7`, o
|
|
| 143 |
|
| 144 |
The pool is shared. **Eight active requests does not mean eight unrelated, fully populated 256K histories fit at once.** Output must fit in the remaining request context; there is no separate promise of 256K output after a 256K input. Preserve the supplied chat template and use the client’s intended thinking setting.
|
| 145 |
|
| 146 |
-
##
|
| 147 |
|
| 148 |
HumanEval 0–9, thinking off, greedy sampling, natural end-of-sequence, cold prompt salts. Each row runs all ten questions. These are **speed and health checks**, not evidence of general coding quality.
|
| 149 |
|
|
@@ -160,7 +216,7 @@ HumanEval 0–9, thinking off, greedy sampling, natural end-of-sequence, cold pr
|
|
| 160 |
|
| 161 |
For comparison, the fresh **Q4_K_XL + DFlash2 (7-token)** C1 run at 256K capacity completed the same ten-question batch in **17.32 s**, at **131.65 tok/s mean request decode** and **0.351 s mean first-token latency**, with 10/10 health. Ciru Halo Agent completed that batch **1.63× faster**. The measured ROCmFP4 C1 MTP4 mode took **19.56 s**, at **113.31 tok/s**. Complete concurrency sweeps, alternative modes, and their health outcomes are on the research page.
|
| 162 |
|
| 163 |
-
##
|
| 164 |
|
| 165 |
Cold input, without a prefix-cache hit. These are backend prefill-counter rates; each request then generated a natural 17-token acknowledgement. The short acknowledgement is not a sustained decode or reasoning benchmark.
|
| 166 |
|
|
@@ -175,7 +231,7 @@ Cold input, without a prefix-cache hit. These are backend prefill-counter rates;
|
|
| 175 |
|
| 176 |
The complete near-256K cold request took **380.83 s**, versus **1,032.06 s for Q4_K_XL** and **1,007.04 s for ROCmFP4**. ROCmFP4’s lower-depth rows come from its retained author-profile sweep; the near-256K point is the fresh cold capture. The research page identifies each source run.
|
| 177 |
|
| 178 |
-
##
|
| 179 |
|
| 180 |
All rows below use **all ten coding questions**, populated histories, and confirmed prefix-cache reuse. Timing includes queueing, suffix prefill, and generation; preparing the shared history is separate.
|
| 181 |
|
|
@@ -190,7 +246,7 @@ All rows below use **all ten coding questions**, populated histories, and confir
|
|
| 190 |
|
| 191 |
The tested Q4_K_XL and ROCmFP4 configurations did not share the warmed prefix successfully across eight request slots. Their **cached C8 comparison remains unavailable**; this does not establish that the runners cannot support it.
|
| 192 |
|
| 193 |
-
##
|
| 194 |
|
| 195 |
The release was evaluated separately from the short coding speed screen. Full EvalScope runs use the native **541 IFEval, 1,319 GSM8K, and 164 HumanEval** tasks. Tool and Hermes scores use their own native graders.
|
| 196 |
|
|
@@ -207,7 +263,7 @@ Full EvalScope used C8, thinking off, greedy sampling, and a 32,768-token respon
|
|
| 207 |
|
| 208 |
The difficult subset and tool suite are single blocks, not repeated confidence estimates. BF16 agreement uses seven held-out documents with eight adjacent prefix positions each: these 56 correlated anchors measure short-prefix fidelity, not task accuracy. The original BF16 source was fixed; the comparison quants’ exact source-weight ancestry is not established.
|
| 209 |
|
| 210 |
-
### Hermes: three complete passes per model and concurrency
|
| 211 |
|
| 212 |
Twenty native agent scenarios per pass, three seeds, 256K capacity, up to 64 turns, and output bounded by remaining context. Scores below are native scores, not percentages of tasks passed.
|
| 213 |
|
|
@@ -219,7 +275,7 @@ Twenty native agent scenarios per pass, three seeds, 256K capacity, up to 64 tur
|
|
| 219 |
|
| 220 |
Ciru’s C8 mean workflow time was **1.49× faster than ROCmFP4**, while Q4 was slightly faster than Ciru in this suite. One interrupted Q4 transport run was retained separately and replaced once; it is excluded from the three complete passes. The research page includes every pass, score range, token timing, and failure count.
|
| 221 |
|
| 222 |
-
## Where the speed advantage changes
|
| 223 |
|
| 224 |
Ciru Halo Agent’s strongest results are speculative coding, prompt ingestion, and shared-history workloads. **It is not the fastest model on every task.** Four ordinary prose tasks with automatic speculation measured **49.90–63.27 tok/s**, compared with approximately **59.5–59.8 tok/s for Q4** and **73.9–74.1 tok/s for ROCmFP4** in the retained no-speculation prose captures. Ciru’s target-only diagnostic measured approximately **56–58 tok/s**. The diagnostic skips draft work; its allocations still include the drafter.
|
| 225 |
|
|
@@ -231,6 +287,7 @@ Ciru Halo Agent’s strongest results are speculative coding, prompt ingestion,
|
|
| 231 |
|
| 232 |
- **[Ornith Team](https://huggingface.co/ornith-ai/Ornith-1.5-35B-A3B)** created the base Ornith model. Ciru’s work is its hardware-specific quantization, kernels, serving integration, and evaluation; we do not claim authorship of the base model’s training.
|
| 233 |
- **[jzinno](https://huggingface.co/jzinno/Ornith-1.5-35B-A3B-DFlash2)** trained the Ornith DFlash2 drafter, initialized from **[z-lab’s Qwen3.5 DFlash](https://huggingface.co/z-lab/Qwen3.5-35B-A3B-DFlash)**. Credit also goes to the **[DFlash](https://arxiv.org/abs/2602.06036)** and **[DFlash 2](https://inco.ai/blog/dflash2/)** authors and the NVIDIA Nemotron dataset contributors whose work supports that drafter.
|
|
|
|
| 234 |
- **[Qwen](https://github.com/QwenLM)** supplied the underlying Qwen model architecture used by this Ornith checkpoint.
|
| 235 |
- **[vLLM](https://github.com/vllm-project/vllm)**, **[AMD ROCm](https://github.com/ROCm)**, **[AITER](https://github.com/ROCm/aiter)**, **[Composable Kernel](https://github.com/ROCm/composable_kernel)**, **[PyTorch](https://github.com/pytorch/pytorch)**, and **[Triton](https://github.com/triton-lang/triton)** provide the runtime and compiler foundations.
|
| 236 |
- Thanks to **[peculiar-ragdoll](https://huggingface.co/peculiar-ragdoll/Unsloth-Ornith-1.5-35B-A3B)**, **[Daniel Han Chen](https://github.com/danielhanchen/llama.cpp)**, **[julianmb](https://huggingface.co/julianmb/Ornith-1.5-35B-A3B-ROCmFP4-GGUF)**, **[HaloFPX](https://github.com/julianmb/halofpx)**, and the llama.cpp/ROCmFP4 community for the comparison builds. Their artifacts are baselines, not components of Ciru’s target weights. The Q4 quant is a community Unsloth-style release, not an official Unsloth upload.
|
|
|
|
| 23 |
- image-text-to-text
|
| 24 |
---
|
| 25 |
|
| 26 |
+
**Ciru v4 · package 4.0.0 — September 17, 2026.** This continues the Ciru v3
|
| 27 |
+
release line with the combined prefill, attention, draft-path, prefix-cache
|
| 28 |
+
and native tool-runtime improvements. Target, drafter and vision weights
|
| 29 |
+
are unchanged. Update both native libraries and the complete plugin before
|
| 30 |
+
restarting. See [patch notes](RUNTIME-FIXES.md) and [update instructions](INSTALL.md).
|
| 31 |
|
| 32 |
# Ornith1.5 Ciru Halo Agent (vllm strix halo)
|
| 33 |
|
|
|
|
| 37 |
|
| 38 |
The design starts with the hardware: packed four-bit weights, four-bit activation paths, Strix Halo four-bit matrix instructions, specialized kernels, and adaptive DFlash2 speculative decoding. Prefix caching and a shared memory pool let agents return to long working histories.
|
| 39 |
|
| 40 |
+
## Ciru v4 results
|
| 41 |
+
|
| 42 |
+
**Ciru v4 accelerates long-context ingestion and returning-agent workflows.**
|
| 43 |
+
The new measurements below cover the complete combined runtime. Existing
|
| 44 |
+
benchmark sections retain their **Ciru v3** version labels and comparison models.
|
| 45 |
+
|
| 46 |
+
| Workload | Previous published Ciru | Ciru v4 | Change |
|
| 47 |
+
| --- | ---: | ---: | ---: |
|
| 48 |
+
| Cold prefill, 65,536 input tokens | 1,286.66 tok/s | 1,624.19 tok/s | +26.23% |
|
| 49 |
+
| Cold prefill, 253,952 input tokens | 667.61 tok/s | 870.11 tok/s | +30.33% |
|
| 50 |
+
| Cached 63K history, C1 mean decode | 123.34 tok/s | 147.69 tok/s | +19.74% |
|
| 51 |
+
| Cached near-256K history, C1 mean decode | 94.00 tok/s | 118.81 tok/s | +26.40% |
|
| 52 |
+
| Cached near-256K history, C8 mean decode | 11.10 tok/s | 19.86 tok/s | +78.93% |
|
| 53 |
+
| Short coding batch, C8 aggregate | 294.89 tok/s | 304.80 tok/s | +3.36% |
|
| 54 |
+
|
| 55 |
+
These final measurements used Ciru; the accepted published controls used Sozo
|
| 56 |
+
and their original runtime. They compare complete products, including those
|
| 57 |
+
host/runtime differences. They do not isolate a kernel's causal contribution.
|
| 58 |
+
The 97 measured requests ended naturally without API errors; **90/90 coding
|
| 59 |
+
outputs passed both native base and extended checks**. Cold-prefill requests
|
| 60 |
+
produced a natural 17-token acknowledgement, not sustained reasoning output.
|
| 61 |
+
|
| 62 |
+
**Tradeoffs:** short C2/C4 aggregate throughput fell 6.22%/8.84%; C4 mean decode
|
| 63 |
+
fell 2.98%. Cached near-256K C8 decode rose 78.93%, but mean first-token latency
|
| 64 |
+
increased from 18.03 s to 25.29 s. Its complete batch still fell from 53.64 s
|
| 65 |
+
to 49.20 s. See the [complete final speed comparison](evaluation/CIRU-V4-PERFORMANCE.md).
|
| 66 |
+
|
| 67 |
+
A separate same-host C1 HE0–9 test of the combined numerical parent measured
|
| 68 |
+
decode gains of 19.81% and 16.15% in two comparisons, pooled +17.97%. Both were
|
| 69 |
+
positive, but their spread exceeded the preset uncertainty rule. HF passed
|
| 70 |
+
38/40 and final 36/40; all failures were the known HE1 case also seen at BF16.
|
| 71 |
+
This is a familiar speed/health panel, not a full HumanEval score.
|
| 72 |
+
|
| 73 |
+
### BF16 fidelity
|
| 74 |
+
|
| 75 |
+
Short-prefix probability arrays were byte-identical to the previous HF model
|
| 76 |
+
across 3,072 positions. On the reused 32K and 64K panels, each with 1,024 scored
|
| 77 |
+
positions, the changes were mixed:
|
| 78 |
+
|
| 79 |
+
| Metric | 32K: HF → final | 64K: HF → final |
|
| 80 |
+
| --- | ---: | ---: |
|
| 81 |
+
| Panel perplexity, lower is better | 7.96626 → 8.00191 | 5.12747 → 5.09702 |
|
| 82 |
+
| KL distance from BF16, lower is closer | 0.222280 → 0.212698 | 0.116232 → 0.119387 |
|
| 83 |
+
| BF16 top-token agreement | 80.9570% → 80.7617% | 88.2813% → 88.0859% |
|
| 84 |
+
| BF16 top token in model top five | 97.3633% → 97.8516% | 98.9258% → 99.0234% |
|
| 85 |
+
| Top-five overlap | 74.6289% → 74.9609% | 81.5625% → 81.6602% |
|
| 86 |
+
|
| 87 |
+
At 32K, 38 top tokens were recovered and 40 lost; at 64K, 28 were recovered
|
| 88 |
+
and 30 lost. One diagnosed 32K confident error was a real top-token recovery;
|
| 89 |
+
another only crossed the probability threshold. These panels were reused during
|
| 90 |
+
diagnosis and are not fresh held-out quality evidence. The measurements belong
|
| 91 |
+
to the combined numerical parent, whose code and native libraries were preserved
|
| 92 |
+
when the final R03/R04 tool changes were added. See [32K fidelity](evaluation/CIRU-V4-FIDELITY32.md),
|
| 93 |
+
[64K fidelity](evaluation/CIRU-V4-FIDELITY64.md) and [composition proof](evaluation/CIRU-V4-COMPOSITION.json).
|
| 94 |
+
|
| 95 |
+
**Historical Ciru v3 measurements on AMD Ryzen AI Max+ 395 / Radeon 8060S (gfx1151):**
|
| 96 |
|
| 97 |
- **178 tok/s single-request decode** on the ten-question coding speed screen, with **166 ms mean time to first token**.
|
| 98 |
- **295 tok/s aggregate at eight concurrent requests**, completing the ten-question batch in **5.52 seconds**.
|
|
|
|
| 102 |
|
| 103 |
These figures describe specific workloads, rather than a universal generation rate. The benchmark tables below include quality results and the workloads where other builds are faster.
|
| 104 |
|
| 105 |
+
**[Full benchmarks, runner builds, serving settings, and limitations](https://llm.ciru.ai/research/ornith-strix/)** · **[Ciru v4 source snapshot](runtime/ciru-v4-ornith-source.tar.gz)** · **[Source repository](https://github.com/ciru-ai/ornith-ciru-halo-agent)** · **[Credits](CREDITS.md)**
|
| 106 |
|
| 107 |
+
## Earlier September 13 runtime fixes
|
| 108 |
|
| 109 |
+
**Historical plugin-only update; Ciru v4 requires the broader update in [INSTALL.md](INSTALL.md).** This release fixes a reproduced graph-replay cache-corruption crash and rejects malformed or ambiguous tool calls with explicit API errors. The weights, kernels and adaptive DFlash2 settings are unchanged. [Update instructions, root cause and validation scope](RUNTIME-FIXES.md).
|
| 110 |
|
| 111 |
## Running Ciru Halo Agent
|
| 112 |
|
|
|
|
| 127 |
|
| 128 |
### Vision / image input
|
| 129 |
|
| 130 |
+
Image input remains enabled by default in Ciru v4, as introduced in runtime 1.0.2. The explicit image launcher is also available:
|
| 131 |
|
| 132 |
```bash
|
| 133 |
bash bundle/serve-vision.sh --host 127.0.0.1 --port 8000
|
|
|
|
| 199 |
|
| 200 |
The pool is shared. **Eight active requests does not mean eight unrelated, fully populated 256K histories fit at once.** Output must fit in the remaining request context; there is no separate promise of 256K output after a 256K input. Preserve the supplied chat template and use the client’s intended thinking setting.
|
| 201 |
|
| 202 |
+
## Historical v3 coding speed and concurrency
|
| 203 |
|
| 204 |
HumanEval 0–9, thinking off, greedy sampling, natural end-of-sequence, cold prompt salts. Each row runs all ten questions. These are **speed and health checks**, not evidence of general coding quality.
|
| 205 |
|
|
|
|
| 216 |
|
| 217 |
For comparison, the fresh **Q4_K_XL + DFlash2 (7-token)** C1 run at 256K capacity completed the same ten-question batch in **17.32 s**, at **131.65 tok/s mean request decode** and **0.351 s mean first-token latency**, with 10/10 health. Ciru Halo Agent completed that batch **1.63× faster**. The measured ROCmFP4 C1 MTP4 mode took **19.56 s**, at **113.31 tok/s**. Complete concurrency sweeps, alternative modes, and their health outcomes are on the research page.
|
| 218 |
|
| 219 |
+
## Historical v3 prompt ingestion
|
| 220 |
|
| 221 |
Cold input, without a prefix-cache hit. These are backend prefill-counter rates; each request then generated a natural 17-token acknowledgement. The short acknowledgement is not a sustained decode or reasoning benchmark.
|
| 222 |
|
|
|
|
| 231 |
|
| 232 |
The complete near-256K cold request took **380.83 s**, versus **1,032.06 s for Q4_K_XL** and **1,007.04 s for ROCmFP4**. ROCmFP4’s lower-depth rows come from its retained author-profile sweep; the near-256K point is the fresh cold capture. The research page identifies each source run.
|
| 233 |
|
| 234 |
+
## Historical v3 returning agents and large histories
|
| 235 |
|
| 236 |
All rows below use **all ten coding questions**, populated histories, and confirmed prefix-cache reuse. Timing includes queueing, suffix prefill, and generation; preparing the shared history is separate.
|
| 237 |
|
|
|
|
| 246 |
|
| 247 |
The tested Q4_K_XL and ROCmFP4 configurations did not share the warmed prefix successfully across eight request slots. Their **cached C8 comparison remains unavailable**; this does not establish that the runners cannot support it.
|
| 248 |
|
| 249 |
+
## Historical v3 quality and actual agent work
|
| 250 |
|
| 251 |
The release was evaluated separately from the short coding speed screen. Full EvalScope runs use the native **541 IFEval, 1,319 GSM8K, and 164 HumanEval** tasks. Tool and Hermes scores use their own native graders.
|
| 252 |
|
|
|
|
| 263 |
|
| 264 |
The difficult subset and tool suite are single blocks, not repeated confidence estimates. BF16 agreement uses seven held-out documents with eight adjacent prefix positions each: these 56 correlated anchors measure short-prefix fidelity, not task accuracy. The original BF16 source was fixed; the comparison quants’ exact source-weight ancestry is not established.
|
| 265 |
|
| 266 |
+
### Historical v3 Hermes: three complete passes per model and concurrency
|
| 267 |
|
| 268 |
Twenty native agent scenarios per pass, three seeds, 256K capacity, up to 64 turns, and output bounded by remaining context. Scores below are native scores, not percentages of tasks passed.
|
| 269 |
|
|
|
|
| 275 |
|
| 276 |
Ciru’s C8 mean workflow time was **1.49× faster than ROCmFP4**, while Q4 was slightly faster than Ciru in this suite. One interrupted Q4 transport run was retained separately and replaced once; it is excluded from the three complete passes. The research page includes every pass, score range, token timing, and failure count.
|
| 277 |
|
| 278 |
+
## Where the historical v3 speed advantage changes
|
| 279 |
|
| 280 |
Ciru Halo Agent’s strongest results are speculative coding, prompt ingestion, and shared-history workloads. **It is not the fastest model on every task.** Four ordinary prose tasks with automatic speculation measured **49.90–63.27 tok/s**, compared with approximately **59.5–59.8 tok/s for Q4** and **73.9–74.1 tok/s for ROCmFP4** in the retained no-speculation prose captures. Ciru’s target-only diagnostic measured approximately **56–58 tok/s**. The diagnostic skips draft work; its allocations still include the drafter.
|
| 281 |
|
|
|
|
| 287 |
|
| 288 |
- **[Ornith Team](https://huggingface.co/ornith-ai/Ornith-1.5-35B-A3B)** created the base Ornith model. Ciru’s work is its hardware-specific quantization, kernels, serving integration, and evaluation; we do not claim authorship of the base model’s training.
|
| 289 |
- **[jzinno](https://huggingface.co/jzinno/Ornith-1.5-35B-A3B-DFlash2)** trained the Ornith DFlash2 drafter, initialized from **[z-lab’s Qwen3.5 DFlash](https://huggingface.co/z-lab/Qwen3.5-35B-A3B-DFlash)**. Credit also goes to the **[DFlash](https://arxiv.org/abs/2602.06036)** and **[DFlash 2](https://inco.ai/blog/dflash2/)** authors and the NVIDIA Nemotron dataset contributors whose work supports that drafter.
|
| 290 |
+
- **[XGrammar](https://github.com/mlc-ai/xgrammar)** supplies the native grammar implementation; its license and pinned Ciru patch are included in the v4 source snapshot.
|
| 291 |
- **[Qwen](https://github.com/QwenLM)** supplied the underlying Qwen model architecture used by this Ornith checkpoint.
|
| 292 |
- **[vLLM](https://github.com/vllm-project/vllm)**, **[AMD ROCm](https://github.com/ROCm)**, **[AITER](https://github.com/ROCm/aiter)**, **[Composable Kernel](https://github.com/ROCm/composable_kernel)**, **[PyTorch](https://github.com/pytorch/pytorch)**, and **[Triton](https://github.com/triton-lang/triton)** provide the runtime and compiler foundations.
|
| 293 |
- Thanks to **[peculiar-ragdoll](https://huggingface.co/peculiar-ragdoll/Unsloth-Ornith-1.5-35B-A3B)**, **[Daniel Han Chen](https://github.com/danielhanchen/llama.cpp)**, **[julianmb](https://huggingface.co/julianmb/Ornith-1.5-35B-A3B-ROCmFP4-GGUF)**, **[HaloFPX](https://github.com/julianmb/halofpx)**, and the llama.cpp/ROCmFP4 community for the comparison builds. Their artifacts are baselines, not components of Ciru’s target weights. The Q4 quant is a community Unsloth-style release, not an official Unsloth upload.
|
RELEASE.json
CHANGED
|
@@ -1,299 +1,141 @@
|
|
| 1 |
{
|
| 2 |
-
"
|
| 3 |
-
"
|
| 4 |
-
"
|
| 5 |
-
"
|
| 6 |
-
"
|
| 7 |
-
"
|
| 8 |
-
"
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
"
|
| 15 |
-
"
|
| 16 |
-
"
|
| 17 |
-
"
|
| 18 |
-
"
|
| 19 |
-
"
|
| 20 |
-
"
|
| 21 |
-
|
| 22 |
-
|
| 23 |
-
|
| 24 |
-
|
| 25 |
-
"runtime/shell.nix": 416,
|
| 26 |
-
"runtime/ciru-halo-agent-aiter-source.tar.gz": 55734195,
|
| 27 |
-
"runtime/ciru-halo-agent-vllm-source.tar.gz": 40520111,
|
| 28 |
-
"LICENSES/ornith-upstream-model-card.md": 29870,
|
| 29 |
-
"LICENSES/dflash2-upstream-model-card.md": 7675,
|
| 30 |
-
"LICENSES/dflash2-9b4852c05fd0-LICENSE": 11358,
|
| 31 |
-
"LICENSES/dflash2-main-LICENSE": 11358,
|
| 32 |
-
"LICENSES/sources.json": 1365,
|
| 33 |
-
"LICENSES/flash-linear-attention-MIT.txt": 1085,
|
| 34 |
-
"LICENSES/Apache-2.0.txt": 11358,
|
| 35 |
-
"runtime/aiter-jit-gfx1151/module_aiter_core.so": 567024,
|
| 36 |
-
"runtime/wheels/amd_aiter-0.1.0rc1-cp314-cp314-linux_x86_64.whl": 67412360,
|
| 37 |
-
"runtime/wheels/vllm-0.1.0rc2.dev9+g9255fd9fb9.rocm100-cp314-cp314-linux_x86_64.whl": 44318322,
|
| 38 |
-
"bundle/LICENSES/LICENSE-APACHE-2.0": 11358,
|
| 39 |
-
"bundle/models/source-prefill.safetensors": 1761622768,
|
| 40 |
-
"bundle/native/libornith_attention_iu4.so": 186640,
|
| 41 |
-
"bundle/native/libornith_dense_g256.so": 69824,
|
| 42 |
-
"bundle/native/libornith_dense_g256_n32.so": 76848,
|
| 43 |
-
"bundle/native/libornith_head_i8_tile.so": 128952,
|
| 44 |
-
"bundle/native/libornith_persistent_iu4.so": 103792,
|
| 45 |
-
"bundle/native/libornith_routed_direct.so": 156392,
|
| 46 |
-
"bundle/native/libornith_routed_n32.so": 151928,
|
| 47 |
-
"bundle/native/libornith_routed_storage_n32.so": 158536,
|
| 48 |
-
"bundle/packaging/serve.sh": 1474,
|
| 49 |
-
"bundle/cache/aiter/module_aiter_core.so": 567024,
|
| 50 |
-
"bundle/plugin-site/ornith_g256/__init__.py": 1368,
|
| 51 |
-
"bundle/plugin-site/ornith_g256/adaptive_c1.py": 14985,
|
| 52 |
-
"bundle/plugin-site/ornith_g256/attention.py": 772,
|
| 53 |
-
"bundle/plugin-site/ornith_g256/attention_compact.py": 10412,
|
| 54 |
-
"bundle/plugin-site/ornith_g256/attention_fast.py": 12274,
|
| 55 |
-
"bundle/plugin-site/ornith_g256/attention_folded.py": 9570,
|
| 56 |
-
"bundle/plugin-site/ornith_g256/attention_iu4.py": 2575,
|
| 57 |
-
"bundle/plugin-site/ornith_g256/attention_iu4_persistent.py": 12418,
|
| 58 |
-
"bundle/plugin-site/ornith_g256/attention_mixed.py": 4138,
|
| 59 |
-
"bundle/plugin-site/ornith_g256/attention_partition.py": 8642,
|
| 60 |
-
"bundle/plugin-site/ornith_g256/attention_storage.py": 4117,
|
| 61 |
-
"bundle/plugin-site/ornith_g256/attention_tile.py": 9641,
|
| 62 |
-
"bundle/plugin-site/ornith_g256/attention_verify.py": 3755,
|
| 63 |
-
"bundle/plugin-site/ornith_g256/attention_window.py": 3612,
|
| 64 |
-
"bundle/plugin-site/ornith_g256/cache_full1120.py": 3017,
|
| 65 |
-
"bundle/plugin-site/ornith_g256/column_backend.py": 5752,
|
| 66 |
-
"bundle/plugin-site/ornith_g256/column_kernel.py": 9422,
|
| 67 |
-
"bundle/plugin-site/ornith_g256/config.py": 2869,
|
| 68 |
-
"bundle/plugin-site/ornith_g256/dense_n32.py": 2449,
|
| 69 |
-
"bundle/plugin-site/ornith_g256/dense_source.py": 1818,
|
| 70 |
-
"bundle/plugin-site/ornith_g256/dflash_conv_boundary.py": 5508,
|
| 71 |
-
"bundle/plugin-site/ornith_g256/dflash_spec.py": 1393,
|
| 72 |
-
"bundle/plugin-site/ornith_g256/dynamic_graphs.py": 11779,
|
| 73 |
-
"bundle/plugin-site/ornith_g256/gdn_compact.py": 5202,
|
| 74 |
-
"bundle/plugin-site/ornith_g256/gdn_compact_kernel.py": 10042,
|
| 75 |
-
"bundle/plugin-site/ornith_g256/gdn_spec.py": 5696,
|
| 76 |
-
"bundle/plugin-site/ornith_g256/launch.py": 17783,
|
| 77 |
-
"bundle/plugin-site/ornith_g256/lifecycle.py": 8269,
|
| 78 |
-
"bundle/plugin-site/ornith_g256/loader.py": 2933,
|
| 79 |
-
"bundle/plugin-site/ornith_g256/method.py": 10049,
|
| 80 |
-
"bundle/plugin-site/ornith_g256/moe_base.py": 3683,
|
| 81 |
-
"bundle/plugin-site/ornith_g256/native.py": 12196,
|
| 82 |
-
"bundle/plugin-site/ornith_g256/no_spec.py": 2618,
|
| 83 |
-
"bundle/plugin-site/ornith_g256/phase_dispatch.py": 2184,
|
| 84 |
-
"bundle/plugin-site/ornith_g256/prefill_draft.py": 5630,
|
| 85 |
-
"bundle/plugin-site/ornith_g256/prefix_cache.py": 15261,
|
| 86 |
-
"bundle/plugin-site/ornith_g256/runtime.py": 6009,
|
| 87 |
-
"bundle/plugin-site/ornith_g256/worker.py": 13579,
|
| 88 |
-
"bundle/plugin-site/ornith_g256/worker_base.py": 4601,
|
| 89 |
-
"bundle/models/draft/LICENSE": 11358,
|
| 90 |
-
"bundle/models/draft/README.md": 7675,
|
| 91 |
-
"bundle/models/draft/config.json": 1311,
|
| 92 |
-
"bundle/models/draft/download-source.json": 1294,
|
| 93 |
-
"bundle/models/draft/evaluation.json": 8078,
|
| 94 |
-
"bundle/models/draft/manifest.json": 1599,
|
| 95 |
-
"bundle/models/draft/model.safetensors": 1052513488,
|
| 96 |
-
"bundle/models/target/.ciru-source-verified.json": 13078,
|
| 97 |
-
"bundle/models/target/chat_template.jinja": 7536,
|
| 98 |
-
"bundle/models/target/config.json": 3438,
|
| 99 |
-
"bundle/models/target/configuration.json": 58,
|
| 100 |
-
"bundle/models/target/g256-layer-00.safetensors": 434161416,
|
| 101 |
-
"bundle/models/target/g256-layer-01.safetensors": 434161416,
|
| 102 |
-
"bundle/models/target/g256-layer-02.safetensors": 434161416,
|
| 103 |
-
"bundle/models/target/g256-layer-03.safetensors": 430917376,
|
| 104 |
-
"bundle/models/target/g256-layer-04.safetensors": 434161416,
|
| 105 |
-
"bundle/models/target/g256-layer-05.safetensors": 434161416,
|
| 106 |
-
"bundle/models/target/g256-layer-06.safetensors": 434161416,
|
| 107 |
-
"bundle/models/target/g256-layer-07.safetensors": 430917376,
|
| 108 |
-
"bundle/models/target/g256-layer-08.safetensors": 434161416,
|
| 109 |
-
"bundle/models/target/g256-layer-09.safetensors": 434161416,
|
| 110 |
-
"bundle/models/target/g256-layer-10.safetensors": 434161432,
|
| 111 |
-
"bundle/models/target/g256-layer-11.safetensors": 430917392,
|
| 112 |
-
"bundle/models/target/g256-layer-12.safetensors": 434161432,
|
| 113 |
-
"bundle/models/target/g256-layer-13.safetensors": 434161432,
|
| 114 |
-
"bundle/models/target/g256-layer-14.safetensors": 434161432,
|
| 115 |
-
"bundle/models/target/g256-layer-15.safetensors": 430917392,
|
| 116 |
-
"bundle/models/target/g256-layer-16.safetensors": 434161432,
|
| 117 |
-
"bundle/models/target/g256-layer-17.safetensors": 434161432,
|
| 118 |
-
"bundle/models/target/g256-layer-18.safetensors": 434161432,
|
| 119 |
-
"bundle/models/target/g256-layer-19.safetensors": 430917392,
|
| 120 |
-
"bundle/models/target/g256-layer-20.safetensors": 434161432,
|
| 121 |
-
"bundle/models/target/g256-layer-21.safetensors": 434161432,
|
| 122 |
-
"bundle/models/target/g256-layer-22.safetensors": 434161432,
|
| 123 |
-
"bundle/models/target/g256-layer-23.safetensors": 430917392,
|
| 124 |
-
"bundle/models/target/g256-layer-24.safetensors": 434161432,
|
| 125 |
-
"bundle/models/target/g256-layer-25.safetensors": 434161432,
|
| 126 |
-
"bundle/models/target/g256-layer-26.safetensors": 434161432,
|
| 127 |
-
"bundle/models/target/g256-layer-27.safetensors": 430917392,
|
| 128 |
-
"bundle/models/target/g256-layer-28.safetensors": 434161432,
|
| 129 |
-
"bundle/models/target/g256-layer-29.safetensors": 434161432,
|
| 130 |
-
"bundle/models/target/g256-layer-30.safetensors": 434161432,
|
| 131 |
-
"bundle/models/target/g256-layer-31.safetensors": 430917392,
|
| 132 |
-
"bundle/models/target/g256-layer-32.safetensors": 434161432,
|
| 133 |
-
"bundle/models/target/g256-layer-33.safetensors": 434161432,
|
| 134 |
-
"bundle/models/target/g256-layer-34.safetensors": 434161432,
|
| 135 |
-
"bundle/models/target/g256-layer-35.safetensors": 430917392,
|
| 136 |
-
"bundle/models/target/g256-layer-36.safetensors": 434161432,
|
| 137 |
-
"bundle/models/target/g256-layer-37.safetensors": 434161432,
|
| 138 |
-
"bundle/models/target/g256-layer-38.safetensors": 434161432,
|
| 139 |
-
"bundle/models/target/g256-layer-39.safetensors": 430917392,
|
| 140 |
-
"bundle/models/target/generation_config.json": 202,
|
| 141 |
-
"bundle/models/target/head-w8.safetensors": 516505840,
|
| 142 |
-
"bundle/models/target/merges.txt": 3353259,
|
| 143 |
-
"bundle/models/target/model.safetensors.index.json": 171704,
|
| 144 |
-
"bundle/models/target/preprocessor_config.json": 390,
|
| 145 |
-
"bundle/models/target/processor_config.json": 1191,
|
| 146 |
-
"bundle/models/target/protected-00.safetensors": 1913075328,
|
| 147 |
-
"bundle/models/target/protected-01.safetensors": 3842944,
|
| 148 |
-
"bundle/models/target/protected-02.safetensors": 2784544,
|
| 149 |
-
"bundle/models/target/protected-03.safetensors": 3838744,
|
| 150 |
-
"bundle/models/target/protected-04.safetensors": 3842968,
|
| 151 |
-
"bundle/models/target/protected-05.safetensors": 2784560,
|
| 152 |
-
"bundle/models/target/protected-06.safetensors": 3838768,
|
| 153 |
-
"bundle/models/target/protected-07.safetensors": 3842968,
|
| 154 |
-
"bundle/models/target/protected-08.safetensors": 2784560,
|
| 155 |
-
"bundle/models/target/protected-09.safetensors": 3838768,
|
| 156 |
-
"bundle/models/target/protected-10.safetensors": 3842968,
|
| 157 |
-
"bundle/models/target/protected-11.safetensors": 2784560,
|
| 158 |
-
"bundle/models/target/protected-12.safetensors": 3838768,
|
| 159 |
-
"bundle/models/target/protected-13.safetensors": 3842968,
|
| 160 |
-
"bundle/models/target/protected-14.safetensors": 2784560,
|
| 161 |
-
"bundle/models/target/protected-15.safetensors": 1690442880,
|
| 162 |
-
"bundle/models/target/quantize_config.json": 107,
|
| 163 |
-
"bundle/models/target/tokenizer.json": 12807982,
|
| 164 |
-
"bundle/models/target/tokenizer_config.json": 16718,
|
| 165 |
-
"bundle/models/target/video_preprocessor_config.json": 385,
|
| 166 |
-
"bundle/models/target/vocab.json": 6722759,
|
| 167 |
-
"assets/ciru-halo-agent.png": 2407299,
|
| 168 |
-
"VISION.md": 3729,
|
| 169 |
-
"bundle/serve-vision.sh": 252,
|
| 170 |
-
"bundle/plugin-site/ornith_g256/graph_phase_guard.py": 6648,
|
| 171 |
-
"bundle/plugin-site/ornith_g256/parser_contract.py": 4203,
|
| 172 |
-
"bundle/plugin-site/ornith_g256/strict_qwen/__init__.py": 81,
|
| 173 |
-
"bundle/plugin-site/ornith_g256/strict_qwen/qwen3.py": 11658,
|
| 174 |
-
"bundle/plugin-site/ornith_g256/strict_qwen/qwen3_contract.py": 8719,
|
| 175 |
-
"RUNTIME-FIXES.md": 6843,
|
| 176 |
-
"ACCURACY-PATCH-1.0.1.json": 7126,
|
| 177 |
-
"bundle/plugin-site/ornith_g256/runtime_correctness.py": 2151,
|
| 178 |
-
"bundle/plugin-site/ornith_g256/_vllm_correctness/qwen_gdn_linear_attn.py": 78334,
|
| 179 |
-
"bundle/plugin-site/ornith_g256/_vllm_correctness/causal_conv1d.py": 51774,
|
| 180 |
-
"bundle/plugin-site/ornith_g256/_vllm_correctness/gdn_attn.py": 26224,
|
| 181 |
-
"bundle/plugin-site/ornith_g256/_vllm_correctness/gpu_model_runner.py": 345523,
|
| 182 |
-
"bundle/plugin-site/ornith_g256/_vllm_correctness/__init__.py": 79,
|
| 183 |
-
"bundle/plugin-site/ornith_g256/_vllm_correctness/manifest.json": 1607,
|
| 184 |
-
"TOOL-SCHEMA-PATCH-1.0.2.json": 7236,
|
| 185 |
-
"bundle/plugin-site/ornith_g256/native_tool_schema.py": 1897,
|
| 186 |
-
"bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/INSTALLER": 2,
|
| 187 |
-
"bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/METADATA": 254,
|
| 188 |
-
"bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/RECORD": 5595,
|
| 189 |
-
"bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/REQUESTED": 0,
|
| 190 |
-
"bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/uv_build.json": 2,
|
| 191 |
-
"bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/top_level.txt": 12,
|
| 192 |
-
"bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/direct_url.json": 105,
|
| 193 |
-
"bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/uv_cache.json": 137,
|
| 194 |
-
"bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/WHEEL": 91,
|
| 195 |
-
"bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/entry_points.txt": 121,
|
| 196 |
-
"bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/licenses/LICENSE-APACHE-2.0": 11358
|
| 197 |
},
|
| 198 |
-
"
|
| 199 |
-
"
|
| 200 |
-
"
|
| 201 |
-
"
|
| 202 |
-
"
|
| 203 |
-
"
|
| 204 |
-
"
|
| 205 |
-
"
|
| 206 |
-
"
|
| 207 |
-
|
| 208 |
-
|
| 209 |
-
"
|
| 210 |
-
"
|
| 211 |
-
"
|
| 212 |
-
"
|
| 213 |
-
|
| 214 |
-
|
| 215 |
-
|
| 216 |
-
"
|
| 217 |
-
"
|
| 218 |
-
"
|
| 219 |
-
"
|
| 220 |
-
"
|
| 221 |
-
"
|
| 222 |
-
"
|
| 223 |
-
"
|
| 224 |
-
"
|
| 225 |
-
|
| 226 |
-
|
| 227 |
-
"
|
| 228 |
-
"
|
| 229 |
-
|
| 230 |
-
|
| 231 |
-
|
| 232 |
-
|
| 233 |
-
"
|
| 234 |
-
"
|
| 235 |
-
|
| 236 |
-
|
| 237 |
-
"
|
| 238 |
-
"
|
| 239 |
-
"
|
| 240 |
-
"
|
| 241 |
-
"
|
| 242 |
-
"
|
| 243 |
-
|
| 244 |
-
|
| 245 |
-
|
| 246 |
-
|
| 247 |
-
|
| 248 |
-
|
| 249 |
-
|
| 250 |
-
|
| 251 |
-
|
| 252 |
-
|
| 253 |
-
|
| 254 |
-
"
|
| 255 |
-
|
| 256 |
-
|
| 257 |
-
|
| 258 |
-
|
| 259 |
-
|
| 260 |
-
|
| 261 |
-
"
|
| 262 |
-
|
| 263 |
-
|
| 264 |
-
|
| 265 |
-
|
| 266 |
-
"
|
| 267 |
-
"
|
| 268 |
-
"
|
| 269 |
-
"
|
| 270 |
-
"
|
| 271 |
-
|
| 272 |
-
|
| 273 |
-
|
| 274 |
-
|
| 275 |
-
"
|
| 276 |
-
|
| 277 |
-
|
| 278 |
-
|
| 279 |
-
|
| 280 |
-
"
|
| 281 |
-
"
|
| 282 |
-
"
|
| 283 |
-
"
|
| 284 |
-
"
|
| 285 |
-
"
|
| 286 |
-
"
|
| 287 |
-
|
| 288 |
-
|
| 289 |
-
"
|
| 290 |
-
"
|
| 291 |
-
"
|
| 292 |
-
"
|
| 293 |
-
"
|
| 294 |
-
"
|
| 295 |
-
"
|
| 296 |
-
"
|
| 297 |
-
"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 298 |
}
|
| 299 |
}
|
|
|
|
| 1 |
{
|
| 2 |
+
"release": "Ciru v4",
|
| 3 |
+
"version": "4.0.0",
|
| 4 |
+
"runtime_version": "4.0.0",
|
| 5 |
+
"tool_component_version": "1.0.3",
|
| 6 |
+
"model": "Ornith1.5",
|
| 7 |
+
"status": "QUALIFIED_FOR_PUBLICATION",
|
| 8 |
+
"publication_ready": true,
|
| 9 |
+
"qualified_build_sha256": "c970d175624455ca30dd3efbc6f52225f6ad32595764232050342a23ced99d94",
|
| 10 |
+
"hf_parent_revision": "f6f9d19e120cc42a8680b1c1d99470e24a83a4fe",
|
| 11 |
+
"weights_modified": false,
|
| 12 |
+
"pinned_engine_wheels_modified": false,
|
| 13 |
+
"validation": {
|
| 14 |
+
"runtime_build": "c970d175624455ca30dd3efbc6f52225f6ad32595764232050342a23ced99d94",
|
| 15 |
+
"native_performance_requests": 97,
|
| 16 |
+
"coding_health_passes": 90,
|
| 17 |
+
"coding_health_tasks": 90,
|
| 18 |
+
"native_tool_calls": 18,
|
| 19 |
+
"complete_agent_episodes": 3,
|
| 20 |
+
"reports": [
|
| 21 |
+
"evaluation/CIRU-V4-PERFORMANCE.md",
|
| 22 |
+
"evaluation/CIRU-V4-FIDELITY32.md",
|
| 23 |
+
"evaluation/CIRU-V4-FIDELITY64.md"
|
| 24 |
+
]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 25 |
},
|
| 26 |
+
"payload_sha256": {
|
| 27 |
+
".gitattributes": "e8738eab56c30a9c4e6b144e2314425c6cf04e66d46b3a3a44e4d1da65e18ea4",
|
| 28 |
+
"INSTALL.md": "2878687e04890121e3bbbf280944f33f53b0c373b921d85bc4493b1c35552bca",
|
| 29 |
+
"README.md": "fa570a0fe074ce2978953a4840852749c77682051f58b2c00181e82c6ff86aef",
|
| 30 |
+
"RUNTIME-FIXES.md": "9c4fe97a10478429a405e853623d2c0281fe3350227f0505fd4b7ce69d94caf7",
|
| 31 |
+
"SOURCE-PROVENANCE.json": "dde1551f66d56c18610f541a18b27aabb7932a31c90d27a4390d9d39b7f82776",
|
| 32 |
+
"assets/ciru-halo-agent.png": "ddbb4819065a652a6feb20ce760c8292b546800cd2e48d01c560809ad5b58765",
|
| 33 |
+
"bundle/LICENSES/XGRAMMAR-LICENSE": "c71d239df91726fc519c6eb72d318ec65820627232b2f796219e87dcf35d0ab4",
|
| 34 |
+
"bundle/native/libornith_attention_iu4.so": "3d29cf3616c69f9fe2b89f8d3564d479dda682fdc06572886e7b60faccc9a53a",
|
| 35 |
+
"bundle/native/libornith_dense_g256.so": "dc9bbf04f02f951fb444feaea2f612671867aad44929a30fdfadef8e5d5513a9",
|
| 36 |
+
"bundle/native/libornith_dense_g256_n32.so": "db3e265fca6daf51fca412579160db7f98d81cd295fffd54a0c7edc3f5fd45da",
|
| 37 |
+
"bundle/native/libornith_head_i8_tile.so": "2c03de5ed015a052b407436c7e42b50fa33ef76cb46dccab9b226546ca1d1e09",
|
| 38 |
+
"bundle/native/libornith_persistent_iu4.so": "c70ba14d5aebf67b761ea54cfc9a38f8bbedf74cfca8b19050048acda9c8a06c",
|
| 39 |
+
"bundle/native/libornith_routed_direct.so": "5b3ca6826547e0ff1c417e733229e104721cec403ccd8f88f7fcdb1726d6f4ed",
|
| 40 |
+
"bundle/native/libornith_routed_n32.so": "48415e71f6c85e3aefcecf99d72607bb05900244b4aaa9a7ba2c58378d92f5fb",
|
| 41 |
+
"bundle/native/libornith_routed_storage_n32.so": "2002bf6f82a637f5842f7c5f09a09f3b7a0db55edc4b3e35181c8f9a12ffb4b1",
|
| 42 |
+
"bundle/packaging/serve.sh": "78ddf77caf1208aeacbf9fa6ae164a026c3ef1da1bdd43015cdd733a1f1b3b64",
|
| 43 |
+
"bundle/plugin-site/TOOL-INTEGRATION.json": "bc116f818a397f3f883e9209d4848999a5be51f99689ff2a4063c7f221a464d9",
|
| 44 |
+
"bundle/plugin-site/ciru_ornith_g256-4.0.0.dist-info/INSTALLER": "e6184ce10e266134fdcfa401e8f1a95005bcd4f18d16b62b757323e2833fe9a9",
|
| 45 |
+
"bundle/plugin-site/ciru_ornith_g256-4.0.0.dist-info/METADATA": "c0368af5f84a4158453b069c59491aaafe84386bfa5bdaa3901a8a381d648781",
|
| 46 |
+
"bundle/plugin-site/ciru_ornith_g256-4.0.0.dist-info/RECORD": "591c8a6330b025ab1a91b112d6e6ebc7b7160291d9a1f85a0c3835d398b30449",
|
| 47 |
+
"bundle/plugin-site/ciru_ornith_g256-4.0.0.dist-info/REQUESTED": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
|
| 48 |
+
"bundle/plugin-site/ciru_ornith_g256-4.0.0.dist-info/WHEEL": "4a63b1614ee9ccd281a804af409ec38d7dd719417dda5ac684c6f747afe28aa2",
|
| 49 |
+
"bundle/plugin-site/ciru_ornith_g256-4.0.0.dist-info/entry_points.txt": "1127b5c28a5d92219faf90cd5997127a69603ebdeebd1aded4ec87925f3c42a5",
|
| 50 |
+
"bundle/plugin-site/ciru_ornith_g256-4.0.0.dist-info/licenses/LICENSE-APACHE-2.0": "cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30",
|
| 51 |
+
"bundle/plugin-site/ciru_ornith_g256-4.0.0.dist-info/top_level.txt": "11d08c2d79e7f2d0cd05447726190cbce8bc844c3a22e85815811c7dd85e9315",
|
| 52 |
+
"bundle/plugin-site/ornith_g256/__init__.py": "60abbb51a475487e3abcb9e5b34d8836b4d4b36e0bced062d2382dc7b1076cf6",
|
| 53 |
+
"bundle/plugin-site/ornith_g256/_vllm_correctness/__init__.py": "963097106407a87b62b275041516144de15713dee65108accc29e27036ffc059",
|
| 54 |
+
"bundle/plugin-site/ornith_g256/_vllm_correctness/causal_conv1d.py": "230f3e554f5f1d609a37022578c84e42d5c71d46e684101278a7e21af4f7d418",
|
| 55 |
+
"bundle/plugin-site/ornith_g256/_vllm_correctness/chat_completion_serving_r03.py": "5720371f5bcb351e340a3e8d9a0235ec01ff27ed93e4f35542e31b83a055ef45",
|
| 56 |
+
"bundle/plugin-site/ornith_g256/_vllm_correctness/gdn_attn.py": "a3c361d502e8bca2aa46a0baed5ec753a3c62e3d027a298e3ed4692794f9f6ec",
|
| 57 |
+
"bundle/plugin-site/ornith_g256/_vllm_correctness/gpu_model_runner.py": "ef0009acaa67a1a5e066984f34bf619b06e85220ff953a3d9511376f0146dccf",
|
| 58 |
+
"bundle/plugin-site/ornith_g256/_vllm_correctness/manifest.json": "bdef4ca7151a7a5e0410dd3e8bcd2e8a85f125938ec104451cb8777267c3c4d9",
|
| 59 |
+
"bundle/plugin-site/ornith_g256/_vllm_correctness/qwen_gdn_linear_attn.py": "3e58fdddc7a0b42ae3d9b5eabb375a993a4110bad499c37846bde466901dbeb3",
|
| 60 |
+
"bundle/plugin-site/ornith_g256/_xgrammar_native/build-receipt.json": "3bdf77657e5157c35938b66dfb46266608ade3b599a62ce74abb3f807d579e84",
|
| 61 |
+
"bundle/plugin-site/ornith_g256/_xgrammar_native/libxgrammar-candidate.so": "92ab36bfae20c6c4df3f5321ca375204799ac833252921d9d8f6a206491d35f7",
|
| 62 |
+
"bundle/plugin-site/ornith_g256/_xgrammar_native/selection.json": "7cb93b3fe902de6d604bf96f79fdc3efbd8d612e494da838f6a8f0d731eb89ae",
|
| 63 |
+
"bundle/plugin-site/ornith_g256/adaptive_c1.py": "8e48fd874fbdf65183c20bcc7c748d2562dbcafa97e004f18a2818b98980a809",
|
| 64 |
+
"bundle/plugin-site/ornith_g256/attention.py": "2fdd3c0ed84dcb052a9e8f9df4068ddf80a969931bea0a295e022d36b753cda1",
|
| 65 |
+
"bundle/plugin-site/ornith_g256/attention_compact.py": "50df61a02a3890a2825b205ed18c7d18a2b002333d382c572440b58ca9b2b9fe",
|
| 66 |
+
"bundle/plugin-site/ornith_g256/attention_fast.py": "27b1a133cadb87ff3a63c027049e1da81e4112fd5d5260a4440fde85491f838a",
|
| 67 |
+
"bundle/plugin-site/ornith_g256/attention_folded.py": "73dac8203032bce6ceedf213d7eefe5987c921768ae237129e232fefd8d9bba3",
|
| 68 |
+
"bundle/plugin-site/ornith_g256/attention_iu4.py": "15c5341fc848c43ac956546e563cf8ab7f7649f05a674fb0153b0e2d76438cae",
|
| 69 |
+
"bundle/plugin-site/ornith_g256/attention_iu4_persistent.py": "cef74ae423974063212637a49cd64a5a10b9077f3a5e2287746722749a1e6297",
|
| 70 |
+
"bundle/plugin-site/ornith_g256/attention_mixed.py": "cbb530697f2555be6532b7d4d3401f3a33c4c58982659e88892720d21691110a",
|
| 71 |
+
"bundle/plugin-site/ornith_g256/attention_partition.py": "41bfb44dd0d0c46a2c0a725983155018e051cd89edb8527f1195a092bef7ffdd",
|
| 72 |
+
"bundle/plugin-site/ornith_g256/attention_storage.py": "3ab813e221835370bdb503b2ab7d48f46294441b1dd9c652c5dc976438bdb6e2",
|
| 73 |
+
"bundle/plugin-site/ornith_g256/attention_tile.py": "f2bf007132bc712c8e61e036ea8a94b54f18a813d9bdd334eee46436a3ca1fea",
|
| 74 |
+
"bundle/plugin-site/ornith_g256/attention_verify.py": "3a329599694265f4badecc2f9b709b280d9960b967d7192f3959ad9e0d8c6077",
|
| 75 |
+
"bundle/plugin-site/ornith_g256/attention_window.py": "78534f9a1b0beae9ecf80cc5b2a613e4fe4827036c8df0acc7e4dfd956aae1c5",
|
| 76 |
+
"bundle/plugin-site/ornith_g256/cache_full1120.py": "3d4240a749d9671c0bf9b3937200e1cd8d4abfe6699d8c1d6857b188c2ef8a8c",
|
| 77 |
+
"bundle/plugin-site/ornith_g256/column_backend.py": "a035b839369e733b05f68ce83df0d022cb62455016fccf30c2978127ab3c5403",
|
| 78 |
+
"bundle/plugin-site/ornith_g256/column_kernel.py": "27fb053d970184732a46212026afd3152ec07a0dfc3dc4c5cc5cc0d539f17c81",
|
| 79 |
+
"bundle/plugin-site/ornith_g256/config.py": "23210d8b625ef375638e2d6043a92a26181a5b5f1c537d9d4b34c91b7109f36f",
|
| 80 |
+
"bundle/plugin-site/ornith_g256/dense_n32.py": "119fb7d25e9d28b6e520b5285c10de86c7d8ec2fe29a2399e9c086173d765903",
|
| 81 |
+
"bundle/plugin-site/ornith_g256/dense_source.py": "cdc7dd6a13c2439cd405596069c301fe28cb95ff32278f1bafe8118fe357d48e",
|
| 82 |
+
"bundle/plugin-site/ornith_g256/dflash_conv_boundary.py": "a5f46458194cade5ddcbfbb51af30c662d12b9086fc459ecbeb17463339febb1",
|
| 83 |
+
"bundle/plugin-site/ornith_g256/dflash_spec.py": "c948b99eab22d2af5c5fc0ca9f764a0855a3d86aca0609362bf5474a9beff1b3",
|
| 84 |
+
"bundle/plugin-site/ornith_g256/dynamic_graphs.py": "6f51f0dd740f6546829e46bbbe840df2780ce8c49978658eb3a05092554ff90c",
|
| 85 |
+
"bundle/plugin-site/ornith_g256/gdn_compact.py": "8b5c062d10f386d78bc9926ab0576e56359401e5d1621c397a2922e2b426ccab",
|
| 86 |
+
"bundle/plugin-site/ornith_g256/gdn_compact_kernel.py": "0eaf34a2eccef2dabbc9e4f017e6ce84317a28efdcbabecf92b3b4c4c9429a55",
|
| 87 |
+
"bundle/plugin-site/ornith_g256/gdn_spec.py": "0931b1b3601d935bb6acdfcb65d307dae4105645cd387bd9b256dd7a525376a3",
|
| 88 |
+
"bundle/plugin-site/ornith_g256/graph_phase_guard.py": "6cabaf7d4e3be04bbac7a642544d38948de2e07ea7ca38f5927f91bf3ad39163",
|
| 89 |
+
"bundle/plugin-site/ornith_g256/launch.py": "c2854b86e2f97dd2ea0356b4b333922f62b9b3d1b3e71f7993d803c6d0957479",
|
| 90 |
+
"bundle/plugin-site/ornith_g256/lifecycle.py": "3285552177cb8ce8a0149a01f8168d2bc44a739a9395ca84e15acccbb20a0b73",
|
| 91 |
+
"bundle/plugin-site/ornith_g256/loader.py": "a853c6d014214a7adf298548f9175b6e18dd05338917aec16ebcb4fc3c3ecb13",
|
| 92 |
+
"bundle/plugin-site/ornith_g256/method.py": "fc301af72efd234430dbfb839d56e74328457b5b83e14a044aeb9554d8acc707",
|
| 93 |
+
"bundle/plugin-site/ornith_g256/moe_base.py": "26fb517415f2497859f2e3ff1aed8f31e7c56a4aacccfb34c6638e65d7df07c4",
|
| 94 |
+
"bundle/plugin-site/ornith_g256/native.py": "49cce0d67987b9659feb556aafb247423490c2153b7d77b750f055a4fdcda331",
|
| 95 |
+
"bundle/plugin-site/ornith_g256/native_tool_schema.py": "910060689d72b47ad4a7bb1e51ef301e41ff0be1beb9cc6b3df01c423d2edef5",
|
| 96 |
+
"bundle/plugin-site/ornith_g256/no_spec.py": "561106b75328fe0e161fa1243aca4cf50c5049d8e785d22c90aacb3e1e1a0104",
|
| 97 |
+
"bundle/plugin-site/ornith_g256/optimized/__init__.py": "0297f5d476afdddb0556f4f236f95e514db6b9b96788df672a4a0bc0074fd835",
|
| 98 |
+
"bundle/plugin-site/ornith_g256/optimized/compile_scope.py": "18971a73b3a5ad80e6212de7f7e419a1dd00eb7380a5e6c8559f5ac6f5ec8d4d",
|
| 99 |
+
"bundle/plugin-site/ornith_g256/optimized/draft.py": "e801301cf03d69998758d35fb5463754ccc30a18e030ed2b988b26cef43f429c",
|
| 100 |
+
"bundle/plugin-site/ornith_g256/optimized/draft_kernel.py": "3d6da6ef0c144505f888cbd35ba50b6b95ad8e48c376ab73297746bff228919e",
|
| 101 |
+
"bundle/plugin-site/ornith_g256/optimized/grammar_parallel.py": "e9c3608438712521980d7fcd60f58862d0b0a2876620183f304c64a03817edea",
|
| 102 |
+
"bundle/plugin-site/ornith_g256/optimized/greedy.py": "386443ad8729f8fe21ec5fef84ead3acfbd4d6c5dd3b4c1e50a7995ed1d1761f",
|
| 103 |
+
"bundle/plugin-site/ornith_g256/optimized/greedy_kernel.py": "2ebe2a25d1011a45eea7eda3c0964ca1eb2638df7d60f54db02c91435550dc42",
|
| 104 |
+
"bundle/plugin-site/ornith_g256/optimized/launch.py": "c1d42602154674916a94b9907b7746fac3ea48cc505fa69babf116f0f127e03d",
|
| 105 |
+
"bundle/plugin-site/ornith_g256/optimized/mixed_qk_sum.py": "17b18769ae84c2e9707a0257181e5c314e9774fc3b4796f8e6787a8302cb5e1f",
|
| 106 |
+
"bundle/plugin-site/ornith_g256/optimized/pipeline/SOURCE-PROOF.json": "5939bc8a53a831715d05d0fa0290beaa0cb00eda1597fe3bf1f1cfc53699fd0d",
|
| 107 |
+
"bundle/plugin-site/ornith_g256/optimized/pipeline/__init__.py": "0297f5d476afdddb0556f4f236f95e514db6b9b96788df672a4a0bc0074fd835",
|
| 108 |
+
"bundle/plugin-site/ornith_g256/optimized/pipeline/adapter.py": "6af97a8d5e0d79b7dfc32f2f0fe5f5c53b5c1c50d03542ff99c64cae7f3909e1",
|
| 109 |
+
"bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/causal_conv1d.py": "254c540f77b01694696a1ab535abd0e5a23d54e7211e202cf9beead04face16b",
|
| 110 |
+
"bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/chunk_o.py": "a93f951048958265b4d1c3f9e67341329196e7a57525433934215fde4b16e61b",
|
| 111 |
+
"bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/chunk_scaled_dot_kkt.py": "71735d2b1c4a19eecbb6058708481817bb610426c48f96a5f646784afcf0c65c",
|
| 112 |
+
"bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/solve_tril.py": "bdd04b69c0f91562048d25834bd53d298dc1f1bd068b9494af88776b2bf4b1d1",
|
| 113 |
+
"bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/wy_fast.py": "4cc0be519bfcf3812d810a4bebb3dd45fb9b13a8d0ab62d0ad463a1a2e36f61f",
|
| 114 |
+
"bundle/plugin-site/ornith_g256/optimized/postconv.py": "c58ee5d2b8418fbf69eda5d9ee4f342b11cc012ab77ae25b6ce2b2aea51aa50f",
|
| 115 |
+
"bundle/plugin-site/ornith_g256/optimized/prefix.py": "c695a899ab12b58c3b89b6ecc276e82efa3f0833a319911a8528cc4e237710cd",
|
| 116 |
+
"bundle/plugin-site/ornith_g256/optimized/qk.py": "d86f1a94086668ac6402eaec7a19c0ad8e55dc8e6377f8d4a9eeab6aa96f9664",
|
| 117 |
+
"bundle/plugin-site/ornith_g256/optimized/qk_schedule.py": "66038781f0667486fba458ef3118ca89ad9640dc2c2bd05cf8576440f0f178a2",
|
| 118 |
+
"bundle/plugin-site/ornith_g256/optimized/worker.py": "38579b767235e7fcac3f2e6b0507438f155c134e1e85f3d1c205a43bece95512",
|
| 119 |
+
"bundle/plugin-site/ornith_g256/parser_contract.py": "2cebab403dc882569f3ce67d12798df14cebd5c162c72f82becd37bcdd7e2a28",
|
| 120 |
+
"bundle/plugin-site/ornith_g256/phase_dispatch.py": "1a4ebdf2af709e6ca519e0a08f6a537602df178ca575332045617631640f8137",
|
| 121 |
+
"bundle/plugin-site/ornith_g256/prefill_draft.py": "d79bb29a124639543a64e198908e2eaef4cbcf492f2e5e14debcfc9c867816f5",
|
| 122 |
+
"bundle/plugin-site/ornith_g256/prefix_cache.py": "65fb972e285e37741f53b221d2f6380ec464d34a39f36bb66a82dcaacab727d0",
|
| 123 |
+
"bundle/plugin-site/ornith_g256/runtime.py": "db43b1c3766720e0dbee4e8617aee673a533c30d36fcc1fbf43919717d8ab149",
|
| 124 |
+
"bundle/plugin-site/ornith_g256/runtime_correctness.py": "08f41773e623a4f99cf8ab19e434d7e33bac9cfa5eb05a9dc0f4511aab8ba4bc",
|
| 125 |
+
"bundle/plugin-site/ornith_g256/strict_qwen/__init__.py": "ee8a5055249f1413fc22f46421ab77b5485aa39b199c68eb5b56c75dcc40922d",
|
| 126 |
+
"bundle/plugin-site/ornith_g256/strict_qwen/qwen3.py": "33c653447a144b72257e508ccd92610db6d30578e28f98c69c99281773dead1c",
|
| 127 |
+
"bundle/plugin-site/ornith_g256/strict_qwen/qwen3_contract.py": "29c083659c54c81144fdb6c8329c75aaa62cf69d1c45fa7da8bdeeeca4856eaf",
|
| 128 |
+
"bundle/plugin-site/ornith_g256/worker.py": "acc27240c21f3478ca452ddf04244df4ee770efaac38192ef64d328ef22701d5",
|
| 129 |
+
"bundle/plugin-site/ornith_g256/worker_base.py": "a00ef1295172e655c8b3c7d790cfd0316114a9492e483dedfdea036d3b7bf868",
|
| 130 |
+
"bundle/plugin-site/ornith_g256/xgrammar_native.py": "d19c546e124cb11311e4cc31f5833dc663e2789d45090091c0f8fec8a9626623",
|
| 131 |
+
"bundle/serve-vision.sh": "21dd61d76c9d82762480f1f8c0364f9692f1b01c8405c79347c1e2c35d931cd6",
|
| 132 |
+
"bundle/serve.sh": "2950b201e971628425547f13ba68b4a628f230ace612eb1bfe7c9b7b8ead5aea",
|
| 133 |
+
"evaluation/CIRU-V4-C1-SPEED.md": "19cf7335f60272e13efa4ff94bf7ec2ff79c61ad4d694911bbc6e779e852f790",
|
| 134 |
+
"evaluation/CIRU-V4-COMPOSITION.json": "ed41fe291530d9c7f9dbda1c9db2b066fc6c1678b81787091597699ae6e62a68",
|
| 135 |
+
"evaluation/CIRU-V4-FIDELITY32.md": "9858cc15e6b1a767cdaf772587b35abaf1053f5cd39acc18080e425b676cbff9",
|
| 136 |
+
"evaluation/CIRU-V4-FIDELITY64.md": "149143f2c030d138d4e2c5e4d042401b0b6c47f0f8d7fb8eaefdbd796607c8a7",
|
| 137 |
+
"evaluation/CIRU-V4-PERFORMANCE.md": "0e953a5167ec8b7d9037c73be018e0c433bd15e23e28472aac3ace17b99cf310",
|
| 138 |
+
"evaluation/CIRU-V4-SHORT-FIDELITY.md": "84de24ead0949eb28b9697a8f19620312df249456aeddd117cdb3ff434704f1c",
|
| 139 |
+
"runtime/ciru-v4-ornith-source.tar.gz": "e029f28d4bb89c801437834d2fa0109c7a3c906580ec31e39c053579d6859da5"
|
| 140 |
}
|
| 141 |
}
|
RUNTIME-FIXES.md
CHANGED
|
@@ -1,5 +1,41 @@
|
|
| 1 |
# Ciru runtime patch notes
|
| 2 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3 |
## Runtime 1.0.2 — September 15, 2026
|
| 4 |
|
| 5 |
- Enable native vLLM/XGrammar schema constraints for automatic tool calls when
|
|
@@ -43,7 +79,7 @@ all long-generation quality issues are resolved.
|
|
| 43 |
|
| 44 |
This update prevents a reproduced cache-corruption crash and makes malformed tool output fail explicitly. It updates the serving plugin; the released weights, native libraries, runtime wheels, sampler, context pool and adaptive DFlash2 policy are unchanged.
|
| 45 |
|
| 46 |
-
##
|
| 47 |
|
| 48 |
Stop your server before replacing its plugin files. Run this from the directory containing `bundle/` and `runtime/`, then restart your usual launcher:
|
| 49 |
|
|
|
|
| 1 |
# Ciru runtime patch notes
|
| 2 |
|
| 3 |
+
## Ciru v4 (4.0.0) — September 17, 2026
|
| 4 |
+
|
| 5 |
+
Continues the Ciru v3 model release line. The earlier 1.0.x identifiers below
|
| 6 |
+
describe runtime patches; R03/R04 component identity 1.0.3 remains provenance,
|
| 7 |
+
while this complete public release is 4.0.0.
|
| 8 |
+
|
| 9 |
+
- Integrate the retained output32 expert tile and all six qualified GDN
|
| 10 |
+
traversal changes, including the older Apodex campaign work.
|
| 11 |
+
- Use Q4480 prefill chunks with the required wide-capacity dense/expert
|
| 12 |
+
libraries and compiler adapters, plus the qualified prefix-cache repair.
|
| 13 |
+
- Include the persistent-attention correctness repair, exact Q/K fusion,
|
| 14 |
+
fused attention prefill, draft convolution/residual/RMSNorm repair and
|
| 15 |
+
greedy draft-path fusion. Preserve the default adaptive DFlash policy.
|
| 16 |
+
- Integrate strict-grammar scheduling, the native XGrammar finite-repetition
|
| 17 |
+
and shared-memory accounting changes, and continuous usage reporting during
|
| 18 |
+
buffered tool streams. Preserve natural finish reasons and explicit errors
|
| 19 |
+
for malformed output.
|
| 20 |
+
- Preserve all target, draft and BF16 vision weights, the native chat template,
|
| 21 |
+
sampler, context/cache limits and previous correctness overlays.
|
| 22 |
+
- Make the public launcher select the complete optimized worker and use
|
| 23 |
+
relative installation/model defaults when no paths.env override exists.
|
| 24 |
+
|
| 25 |
+
Update **both** native libraries and the complete plugin/launchers using
|
| 26 |
+
[INSTALL.md](INSTALL.md). The pinned vLLM/AITER wheels are retained.
|
| 27 |
+
|
| 28 |
+
Final native performance: cold prefill +26.23% at64K and +30.33% near256K;
|
| 29 |
+
cached C1 decode +19.74% at63K and +26.40% near256K. Short C2/C4 aggregate
|
| 30 |
+
throughput regressed, and near256K C8 first-token latency increased. Numerical
|
| 31 |
+
fidelity results and the complete speed comparison are reported on the model card.
|
| 32 |
+
All individual historical gains overlap and must not be added.
|
| 33 |
+
|
| 34 |
+
The final native performance block completed97 measured requests with90/90
|
| 35 |
+
coding health checks. Native tool integration passed18 exact live calls and
|
| 36 |
+
preserved identical333-token streams across usage modes. See the [Ciru v4 results](README.md#ciru-v4-results).
|
| 37 |
+
|
| 38 |
+
|
| 39 |
## Runtime 1.0.2 — September 15, 2026
|
| 40 |
|
| 41 |
- Enable native vLLM/XGrammar schema constraints for automatic tool calls when
|
|
|
|
| 79 |
|
| 80 |
This update prevents a reproduced cache-corruption crash and makes malformed tool output fail explicitly. It updates the serving plugin; the released weights, native libraries, runtime wheels, sampler, context pool and adaptive DFlash2 policy are unchanged.
|
| 81 |
|
| 82 |
+
## Historical September 13 update instructions
|
| 83 |
|
| 84 |
Stop your server before replacing its plugin files. Run this from the directory containing `bundle/` and `runtime/`, then restart your usual launcher:
|
| 85 |
|
SOURCE-PROVENANCE.json
ADDED
|
@@ -0,0 +1,159 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"release": "Ciru v4",
|
| 3 |
+
"package_version": "4.0.0",
|
| 4 |
+
"qualified_build_sha256": "c970d175624455ca30dd3efbc6f52225f6ad32595764232050342a23ced99d94",
|
| 5 |
+
"native": [
|
| 6 |
+
{
|
| 7 |
+
"source": "kernels/dense/dense_g256_consumer.cpp",
|
| 8 |
+
"library": "libornith_dense_g256.so",
|
| 9 |
+
"source_sha256": "28fe913a6fd5b3e97ae106563a3fe843cf55ba84413a989a7e414ec3357a8150",
|
| 10 |
+
"measured_library_sha256": "dc9bbf04f02f951fb444feaea2f612671867aad44929a30fdfadef8e5d5513a9"
|
| 11 |
+
},
|
| 12 |
+
{
|
| 13 |
+
"source": "kernels/routed/routed_direct_consumer.cpp",
|
| 14 |
+
"library": "libornith_routed_direct.so",
|
| 15 |
+
"source_sha256": "aea3b68893e7e2a8b70ddb1f83a24c444cae6ff8b995ff9794dc545da447a0ff",
|
| 16 |
+
"measured_library_sha256": "5b3ca6826547e0ff1c417e733229e104721cec403ccd8f88f7fcdb1726d6f4ed"
|
| 17 |
+
},
|
| 18 |
+
{
|
| 19 |
+
"source": "kernels/head/head_i8_tile_consumer.cpp",
|
| 20 |
+
"library": "libornith_head_i8_tile.so",
|
| 21 |
+
"source_sha256": "0c14089990d13b63a61fcaa598de2a27d34cae2624ab3c62dd6095e0cb552296",
|
| 22 |
+
"measured_library_sha256": "2c03de5ed015a052b407436c7e42b50fa33ef76cb46dccab9b226546ca1d1e09"
|
| 23 |
+
},
|
| 24 |
+
{
|
| 25 |
+
"source": "kernels/routed_n32/routed_n32_consumer.cpp",
|
| 26 |
+
"library": "libornith_routed_n32.so",
|
| 27 |
+
"source_sha256": "f045389d4cb804fe43d34425ce6c02e817a89bd8cbdb453547f20e36f060be52",
|
| 28 |
+
"measured_library_sha256": "48415e71f6c85e3aefcecf99d72607bb05900244b4aaa9a7ba2c58378d92f5fb"
|
| 29 |
+
},
|
| 30 |
+
{
|
| 31 |
+
"source": "kernels/routed_storage/routed_storage_n32_consumer.cpp",
|
| 32 |
+
"library": "libornith_routed_storage_n32.so",
|
| 33 |
+
"source_sha256": "171dacf56f553dd586dc0bb559b60108c225e93787044c5ee5118c898a5382ea",
|
| 34 |
+
"measured_library_sha256": "2002bf6f82a637f5842f7c5f09a09f3b7a0db55edc4b3e35181c8f9a12ffb4b1"
|
| 35 |
+
},
|
| 36 |
+
{
|
| 37 |
+
"source": "kernels/attention/prefill.cpp",
|
| 38 |
+
"library": "libornith_attention_iu4.so",
|
| 39 |
+
"source_sha256": "7d0a21917fbeff75811d7a7ed9df9438852762c28914eb61ec3b1ba31fc66dae",
|
| 40 |
+
"measured_library_sha256": "3d29cf3616c69f9fe2b89f8d3564d479dda682fdc06572886e7b60faccc9a53a"
|
| 41 |
+
},
|
| 42 |
+
{
|
| 43 |
+
"source": "kernels/attention/persistent.cpp",
|
| 44 |
+
"library": "libornith_persistent_iu4.so",
|
| 45 |
+
"source_sha256": "b430bee407b2c00ee7c2e75dcc39b8dab504767219755c24259f42327975b055",
|
| 46 |
+
"measured_library_sha256": "c70ba14d5aebf67b761ea54cfc9a38f8bbedf74cfca8b19050048acda9c8a06c"
|
| 47 |
+
},
|
| 48 |
+
{
|
| 49 |
+
"source": "kernels/dense_n32/dense_g256_consumer.cpp",
|
| 50 |
+
"library": "libornith_dense_g256_n32.so",
|
| 51 |
+
"source_sha256": "92ff49dc74cdc54a54cbbc2e618ca312f7b17ba502e8c1e83fa2a53570c4567f",
|
| 52 |
+
"measured_library_sha256": "db3e265fca6daf51fca412579160db7f98d81cd295fffd54a0c7edc3f5fd45da"
|
| 53 |
+
}
|
| 54 |
+
],
|
| 55 |
+
"native_source_files": {
|
| 56 |
+
"kernels/attention/persistent.cpp": "b430bee407b2c00ee7c2e75dcc39b8dab504767219755c24259f42327975b055",
|
| 57 |
+
"kernels/attention/prefill.cpp": "7d0a21917fbeff75811d7a7ed9df9438852762c28914eb61ec3b1ba31fc66dae",
|
| 58 |
+
"kernels/dense/dense_g256_api.h": "487320b066d2b033982a02cd8e94b3e2075dbec9f7bdb2f7c53b895a060813a2",
|
| 59 |
+
"kernels/dense/dense_g256_consumer.cpp": "28fe913a6fd5b3e97ae106563a3fe843cf55ba84413a989a7e414ec3357a8150",
|
| 60 |
+
"kernels/dense_n32/dense_g256_api.h": "487320b066d2b033982a02cd8e94b3e2075dbec9f7bdb2f7c53b895a060813a2",
|
| 61 |
+
"kernels/dense_n32/dense_g256_consumer.cpp": "92ff49dc74cdc54a54cbbc2e618ca312f7b17ba502e8c1e83fa2a53570c4567f",
|
| 62 |
+
"kernels/head/head_i8_tile_api.h": "ad11ace0f655b8a883d8888f07f214cb75ba5ff12471ec50a54bd2601f572463",
|
| 63 |
+
"kernels/head/head_i8_tile_consumer.cpp": "0c14089990d13b63a61fcaa598de2a27d34cae2624ab3c62dd6095e0cb552296",
|
| 64 |
+
"kernels/routed/portable_exp.h": "d118a773477b611a0f0eb53c1b403479d0c2d35b02e9970ba3c24bb882da8781",
|
| 65 |
+
"kernels/routed/routed_direct_api.h": "ac2c0076a7db7570f91eed658c8b4646218d98d2dbdb39e84d07c6284ab271d0",
|
| 66 |
+
"kernels/routed/routed_direct_consumer.cpp": "aea3b68893e7e2a8b70ddb1f83a24c444cae6ff8b995ff9794dc545da447a0ff",
|
| 67 |
+
"kernels/routed/routed_direct_prefill.h": "0c6a9be68211d3c0fd7c404a7cccaba93a58e996d4326a1e29122c90abe3b5f8",
|
| 68 |
+
"kernels/routed_n32/portable_exp.h": "d118a773477b611a0f0eb53c1b403479d0c2d35b02e9970ba3c24bb882da8781",
|
| 69 |
+
"kernels/routed_n32/routed_direct_api.h": "ac2c0076a7db7570f91eed658c8b4646218d98d2dbdb39e84d07c6284ab271d0",
|
| 70 |
+
"kernels/routed_n32/routed_direct_prefill.h": "0c6a9be68211d3c0fd7c404a7cccaba93a58e996d4326a1e29122c90abe3b5f8",
|
| 71 |
+
"kernels/routed_n32/routed_n32_consumer.cpp": "f045389d4cb804fe43d34425ce6c02e817a89bd8cbdb453547f20e36f060be52",
|
| 72 |
+
"kernels/routed_storage/portable_exp.h": "d118a773477b611a0f0eb53c1b403479d0c2d35b02e9970ba3c24bb882da8781",
|
| 73 |
+
"kernels/routed_storage/routed_direct_api.h": "ac2c0076a7db7570f91eed658c8b4646218d98d2dbdb39e84d07c6284ab271d0",
|
| 74 |
+
"kernels/routed_storage/routed_storage_n32_consumer.cpp": "171dacf56f553dd586dc0bb559b60108c225e93787044c5ee5118c898a5382ea",
|
| 75 |
+
"kernels/routed_storage/routed_storage_n32_prefill.h": "15e5f92d7c3f06a05010bb997ee07556007122858c04332ceb794a4c3d1a2bf9"
|
| 76 |
+
},
|
| 77 |
+
"plugin_sha256": {
|
| 78 |
+
"plugin-site/ornith_g256/__init__.py": "60abbb51a475487e3abcb9e5b34d8836b4d4b36e0bced062d2382dc7b1076cf6",
|
| 79 |
+
"plugin-site/ornith_g256/_vllm_correctness/__init__.py": "963097106407a87b62b275041516144de15713dee65108accc29e27036ffc059",
|
| 80 |
+
"plugin-site/ornith_g256/_vllm_correctness/causal_conv1d.py": "230f3e554f5f1d609a37022578c84e42d5c71d46e684101278a7e21af4f7d418",
|
| 81 |
+
"plugin-site/ornith_g256/_vllm_correctness/chat_completion_serving_r03.py": "5720371f5bcb351e340a3e8d9a0235ec01ff27ed93e4f35542e31b83a055ef45",
|
| 82 |
+
"plugin-site/ornith_g256/_vllm_correctness/gdn_attn.py": "a3c361d502e8bca2aa46a0baed5ec753a3c62e3d027a298e3ed4692794f9f6ec",
|
| 83 |
+
"plugin-site/ornith_g256/_vllm_correctness/gpu_model_runner.py": "ef0009acaa67a1a5e066984f34bf619b06e85220ff953a3d9511376f0146dccf",
|
| 84 |
+
"plugin-site/ornith_g256/_vllm_correctness/manifest.json": "bdef4ca7151a7a5e0410dd3e8bcd2e8a85f125938ec104451cb8777267c3c4d9",
|
| 85 |
+
"plugin-site/ornith_g256/_vllm_correctness/qwen_gdn_linear_attn.py": "3e58fdddc7a0b42ae3d9b5eabb375a993a4110bad499c37846bde466901dbeb3",
|
| 86 |
+
"plugin-site/ornith_g256/_xgrammar_native/build-receipt.json": "3bdf77657e5157c35938b66dfb46266608ade3b599a62ce74abb3f807d579e84",
|
| 87 |
+
"plugin-site/ornith_g256/_xgrammar_native/selection.json": "7cb93b3fe902de6d604bf96f79fdc3efbd8d612e494da838f6a8f0d731eb89ae",
|
| 88 |
+
"plugin-site/ornith_g256/adaptive_c1.py": "8e48fd874fbdf65183c20bcc7c748d2562dbcafa97e004f18a2818b98980a809",
|
| 89 |
+
"plugin-site/ornith_g256/attention.py": "2fdd3c0ed84dcb052a9e8f9df4068ddf80a969931bea0a295e022d36b753cda1",
|
| 90 |
+
"plugin-site/ornith_g256/attention_compact.py": "50df61a02a3890a2825b205ed18c7d18a2b002333d382c572440b58ca9b2b9fe",
|
| 91 |
+
"plugin-site/ornith_g256/attention_fast.py": "27b1a133cadb87ff3a63c027049e1da81e4112fd5d5260a4440fde85491f838a",
|
| 92 |
+
"plugin-site/ornith_g256/attention_folded.py": "73dac8203032bce6ceedf213d7eefe5987c921768ae237129e232fefd8d9bba3",
|
| 93 |
+
"plugin-site/ornith_g256/attention_iu4.py": "15c5341fc848c43ac956546e563cf8ab7f7649f05a674fb0153b0e2d76438cae",
|
| 94 |
+
"plugin-site/ornith_g256/attention_iu4_persistent.py": "cef74ae423974063212637a49cd64a5a10b9077f3a5e2287746722749a1e6297",
|
| 95 |
+
"plugin-site/ornith_g256/attention_mixed.py": "cbb530697f2555be6532b7d4d3401f3a33c4c58982659e88892720d21691110a",
|
| 96 |
+
"plugin-site/ornith_g256/attention_partition.py": "41bfb44dd0d0c46a2c0a725983155018e051cd89edb8527f1195a092bef7ffdd",
|
| 97 |
+
"plugin-site/ornith_g256/attention_storage.py": "3ab813e221835370bdb503b2ab7d48f46294441b1dd9c652c5dc976438bdb6e2",
|
| 98 |
+
"plugin-site/ornith_g256/attention_tile.py": "f2bf007132bc712c8e61e036ea8a94b54f18a813d9bdd334eee46436a3ca1fea",
|
| 99 |
+
"plugin-site/ornith_g256/attention_verify.py": "3a329599694265f4badecc2f9b709b280d9960b967d7192f3959ad9e0d8c6077",
|
| 100 |
+
"plugin-site/ornith_g256/attention_window.py": "78534f9a1b0beae9ecf80cc5b2a613e4fe4827036c8df0acc7e4dfd956aae1c5",
|
| 101 |
+
"plugin-site/ornith_g256/cache_full1120.py": "3d4240a749d9671c0bf9b3937200e1cd8d4abfe6699d8c1d6857b188c2ef8a8c",
|
| 102 |
+
"plugin-site/ornith_g256/column_backend.py": "a035b839369e733b05f68ce83df0d022cb62455016fccf30c2978127ab3c5403",
|
| 103 |
+
"plugin-site/ornith_g256/column_kernel.py": "27fb053d970184732a46212026afd3152ec07a0dfc3dc4c5cc5cc0d539f17c81",
|
| 104 |
+
"plugin-site/ornith_g256/config.py": "23210d8b625ef375638e2d6043a92a26181a5b5f1c537d9d4b34c91b7109f36f",
|
| 105 |
+
"plugin-site/ornith_g256/dense_n32.py": "119fb7d25e9d28b6e520b5285c10de86c7d8ec2fe29a2399e9c086173d765903",
|
| 106 |
+
"plugin-site/ornith_g256/dense_source.py": "cdc7dd6a13c2439cd405596069c301fe28cb95ff32278f1bafe8118fe357d48e",
|
| 107 |
+
"plugin-site/ornith_g256/dflash_conv_boundary.py": "a5f46458194cade5ddcbfbb51af30c662d12b9086fc459ecbeb17463339febb1",
|
| 108 |
+
"plugin-site/ornith_g256/dflash_spec.py": "c948b99eab22d2af5c5fc0ca9f764a0855a3d86aca0609362bf5474a9beff1b3",
|
| 109 |
+
"plugin-site/ornith_g256/dynamic_graphs.py": "6f51f0dd740f6546829e46bbbe840df2780ce8c49978658eb3a05092554ff90c",
|
| 110 |
+
"plugin-site/ornith_g256/gdn_compact.py": "8b5c062d10f386d78bc9926ab0576e56359401e5d1621c397a2922e2b426ccab",
|
| 111 |
+
"plugin-site/ornith_g256/gdn_compact_kernel.py": "0eaf34a2eccef2dabbc9e4f017e6ce84317a28efdcbabecf92b3b4c4c9429a55",
|
| 112 |
+
"plugin-site/ornith_g256/gdn_spec.py": "0931b1b3601d935bb6acdfcb65d307dae4105645cd387bd9b256dd7a525376a3",
|
| 113 |
+
"plugin-site/ornith_g256/graph_phase_guard.py": "6cabaf7d4e3be04bbac7a642544d38948de2e07ea7ca38f5927f91bf3ad39163",
|
| 114 |
+
"plugin-site/ornith_g256/launch.py": "c2854b86e2f97dd2ea0356b4b333922f62b9b3d1b3e71f7993d803c6d0957479",
|
| 115 |
+
"plugin-site/ornith_g256/lifecycle.py": "3285552177cb8ce8a0149a01f8168d2bc44a739a9395ca84e15acccbb20a0b73",
|
| 116 |
+
"plugin-site/ornith_g256/loader.py": "a853c6d014214a7adf298548f9175b6e18dd05338917aec16ebcb4fc3c3ecb13",
|
| 117 |
+
"plugin-site/ornith_g256/method.py": "fc301af72efd234430dbfb839d56e74328457b5b83e14a044aeb9554d8acc707",
|
| 118 |
+
"plugin-site/ornith_g256/moe_base.py": "26fb517415f2497859f2e3ff1aed8f31e7c56a4aacccfb34c6638e65d7df07c4",
|
| 119 |
+
"plugin-site/ornith_g256/native.py": "49cce0d67987b9659feb556aafb247423490c2153b7d77b750f055a4fdcda331",
|
| 120 |
+
"plugin-site/ornith_g256/native_tool_schema.py": "910060689d72b47ad4a7bb1e51ef301e41ff0be1beb9cc6b3df01c423d2edef5",
|
| 121 |
+
"plugin-site/ornith_g256/no_spec.py": "561106b75328fe0e161fa1243aca4cf50c5049d8e785d22c90aacb3e1e1a0104",
|
| 122 |
+
"plugin-site/ornith_g256/optimized/__init__.py": "0297f5d476afdddb0556f4f236f95e514db6b9b96788df672a4a0bc0074fd835",
|
| 123 |
+
"plugin-site/ornith_g256/optimized/compile_scope.py": "18971a73b3a5ad80e6212de7f7e419a1dd00eb7380a5e6c8559f5ac6f5ec8d4d",
|
| 124 |
+
"plugin-site/ornith_g256/optimized/draft.py": "e801301cf03d69998758d35fb5463754ccc30a18e030ed2b988b26cef43f429c",
|
| 125 |
+
"plugin-site/ornith_g256/optimized/draft_kernel.py": "3d6da6ef0c144505f888cbd35ba50b6b95ad8e48c376ab73297746bff228919e",
|
| 126 |
+
"plugin-site/ornith_g256/optimized/grammar_parallel.py": "e9c3608438712521980d7fcd60f58862d0b0a2876620183f304c64a03817edea",
|
| 127 |
+
"plugin-site/ornith_g256/optimized/greedy.py": "386443ad8729f8fe21ec5fef84ead3acfbd4d6c5dd3b4c1e50a7995ed1d1761f",
|
| 128 |
+
"plugin-site/ornith_g256/optimized/greedy_kernel.py": "2ebe2a25d1011a45eea7eda3c0964ca1eb2638df7d60f54db02c91435550dc42",
|
| 129 |
+
"plugin-site/ornith_g256/optimized/launch.py": "c1d42602154674916a94b9907b7746fac3ea48cc505fa69babf116f0f127e03d",
|
| 130 |
+
"plugin-site/ornith_g256/optimized/mixed_qk_sum.py": "17b18769ae84c2e9707a0257181e5c314e9774fc3b4796f8e6787a8302cb5e1f",
|
| 131 |
+
"plugin-site/ornith_g256/optimized/pipeline/SOURCE-PROOF.json": "5939bc8a53a831715d05d0fa0290beaa0cb00eda1597fe3bf1f1cfc53699fd0d",
|
| 132 |
+
"plugin-site/ornith_g256/optimized/pipeline/__init__.py": "0297f5d476afdddb0556f4f236f95e514db6b9b96788df672a4a0bc0074fd835",
|
| 133 |
+
"plugin-site/ornith_g256/optimized/pipeline/adapter.py": "6af97a8d5e0d79b7dfc32f2f0fe5f5c53b5c1c50d03542ff99c64cae7f3909e1",
|
| 134 |
+
"plugin-site/ornith_g256/optimized/pipeline/candidate/causal_conv1d.py": "254c540f77b01694696a1ab535abd0e5a23d54e7211e202cf9beead04face16b",
|
| 135 |
+
"plugin-site/ornith_g256/optimized/pipeline/candidate/chunk_o.py": "a93f951048958265b4d1c3f9e67341329196e7a57525433934215fde4b16e61b",
|
| 136 |
+
"plugin-site/ornith_g256/optimized/pipeline/candidate/chunk_scaled_dot_kkt.py": "71735d2b1c4a19eecbb6058708481817bb610426c48f96a5f646784afcf0c65c",
|
| 137 |
+
"plugin-site/ornith_g256/optimized/pipeline/candidate/solve_tril.py": "bdd04b69c0f91562048d25834bd53d298dc1f1bd068b9494af88776b2bf4b1d1",
|
| 138 |
+
"plugin-site/ornith_g256/optimized/pipeline/candidate/wy_fast.py": "4cc0be519bfcf3812d810a4bebb3dd45fb9b13a8d0ab62d0ad463a1a2e36f61f",
|
| 139 |
+
"plugin-site/ornith_g256/optimized/postconv.py": "c58ee5d2b8418fbf69eda5d9ee4f342b11cc012ab77ae25b6ce2b2aea51aa50f",
|
| 140 |
+
"plugin-site/ornith_g256/optimized/prefix.py": "c695a899ab12b58c3b89b6ecc276e82efa3f0833a319911a8528cc4e237710cd",
|
| 141 |
+
"plugin-site/ornith_g256/optimized/qk.py": "d86f1a94086668ac6402eaec7a19c0ad8e55dc8e6377f8d4a9eeab6aa96f9664",
|
| 142 |
+
"plugin-site/ornith_g256/optimized/qk_schedule.py": "66038781f0667486fba458ef3118ca89ad9640dc2c2bd05cf8576440f0f178a2",
|
| 143 |
+
"plugin-site/ornith_g256/optimized/worker.py": "38579b767235e7fcac3f2e6b0507438f155c134e1e85f3d1c205a43bece95512",
|
| 144 |
+
"plugin-site/ornith_g256/parser_contract.py": "2cebab403dc882569f3ce67d12798df14cebd5c162c72f82becd37bcdd7e2a28",
|
| 145 |
+
"plugin-site/ornith_g256/phase_dispatch.py": "1a4ebdf2af709e6ca519e0a08f6a537602df178ca575332045617631640f8137",
|
| 146 |
+
"plugin-site/ornith_g256/prefill_draft.py": "d79bb29a124639543a64e198908e2eaef4cbcf492f2e5e14debcfc9c867816f5",
|
| 147 |
+
"plugin-site/ornith_g256/prefix_cache.py": "65fb972e285e37741f53b221d2f6380ec464d34a39f36bb66a82dcaacab727d0",
|
| 148 |
+
"plugin-site/ornith_g256/runtime.py": "db43b1c3766720e0dbee4e8617aee673a533c30d36fcc1fbf43919717d8ab149",
|
| 149 |
+
"plugin-site/ornith_g256/runtime_correctness.py": "08f41773e623a4f99cf8ab19e434d7e33bac9cfa5eb05a9dc0f4511aab8ba4bc",
|
| 150 |
+
"plugin-site/ornith_g256/strict_qwen/__init__.py": "ee8a5055249f1413fc22f46421ab77b5485aa39b199c68eb5b56c75dcc40922d",
|
| 151 |
+
"plugin-site/ornith_g256/strict_qwen/qwen3.py": "33c653447a144b72257e508ccd92610db6d30578e28f98c69c99281773dead1c",
|
| 152 |
+
"plugin-site/ornith_g256/strict_qwen/qwen3_contract.py": "29c083659c54c81144fdb6c8329c75aaa62cf69d1c45fa7da8bdeeeca4856eaf",
|
| 153 |
+
"plugin-site/ornith_g256/worker.py": "acc27240c21f3478ca452ddf04244df4ee770efaac38192ef64d328ef22701d5",
|
| 154 |
+
"plugin-site/ornith_g256/worker_base.py": "a00ef1295172e655c8b3c7d790cfd0316114a9492e483dedfdea036d3b7bf868",
|
| 155 |
+
"plugin-site/ornith_g256/xgrammar_native.py": "d19c546e124cb11311e4cc31f5833dc663e2789d45090091c0f8fec8a9626623"
|
| 156 |
+
},
|
| 157 |
+
"runtime_native_grammar": "92ab36bfae20c6c4df3f5321ca375204799ac833252921d9d8f6a206491d35f7",
|
| 158 |
+
"note": "Same measured model libraries and plugin code; R04 source pinned by patches/xgrammar/SOURCE.json. Portable recipe includes output32 compile definition. Rebuilding with another toolchain need not produce identical binaries."
|
| 159 |
+
}
|
bundle/LICENSES/XGRAMMAR-LICENSE
ADDED
|
@@ -0,0 +1,201 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
Apache License
|
| 2 |
+
Version 2.0, January 2004
|
| 3 |
+
http://www.apache.org/licenses/
|
| 4 |
+
|
| 5 |
+
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
| 6 |
+
|
| 7 |
+
1. Definitions.
|
| 8 |
+
|
| 9 |
+
"License" shall mean the terms and conditions for use, reproduction,
|
| 10 |
+
and distribution as defined by Sections 1 through 9 of this document.
|
| 11 |
+
|
| 12 |
+
"Licensor" shall mean the copyright owner or entity authorized by
|
| 13 |
+
the copyright owner that is granting the License.
|
| 14 |
+
|
| 15 |
+
"Legal Entity" shall mean the union of the acting entity and all
|
| 16 |
+
other entities that control, are controlled by, or are under common
|
| 17 |
+
control with that entity. For the purposes of this definition,
|
| 18 |
+
"control" means (i) the power, direct or indirect, to cause the
|
| 19 |
+
direction or management of such entity, whether by contract or
|
| 20 |
+
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
| 21 |
+
outstanding shares, or (iii) beneficial ownership of such entity.
|
| 22 |
+
|
| 23 |
+
"You" (or "Your") shall mean an individual or Legal Entity
|
| 24 |
+
exercising permissions granted by this License.
|
| 25 |
+
|
| 26 |
+
"Source" form shall mean the preferred form for making modifications,
|
| 27 |
+
including but not limited to software source code, documentation
|
| 28 |
+
source, and configuration files.
|
| 29 |
+
|
| 30 |
+
"Object" form shall mean any form resulting from mechanical
|
| 31 |
+
transformation or translation of a Source form, including but
|
| 32 |
+
not limited to compiled object code, generated documentation,
|
| 33 |
+
and conversions to other media types.
|
| 34 |
+
|
| 35 |
+
"Work" shall mean the work of authorship, whether in Source or
|
| 36 |
+
Object form, made available under the License, as indicated by a
|
| 37 |
+
copyright notice that is included in or attached to the work
|
| 38 |
+
(an example is provided in the Appendix below).
|
| 39 |
+
|
| 40 |
+
"Derivative Works" shall mean any work, whether in Source or Object
|
| 41 |
+
form, that is based on (or derived from) the Work and for which the
|
| 42 |
+
editorial revisions, annotations, elaborations, or other modifications
|
| 43 |
+
represent, as a whole, an original work of authorship. For the purposes
|
| 44 |
+
of this License, Derivative Works shall not include works that remain
|
| 45 |
+
separable from, or merely link (or bind by name) to the interfaces of,
|
| 46 |
+
the Work and Derivative Works thereof.
|
| 47 |
+
|
| 48 |
+
"Contribution" shall mean any work of authorship, including
|
| 49 |
+
the original version of the Work and any modifications or additions
|
| 50 |
+
to that Work or Derivative Works thereof, that is intentionally
|
| 51 |
+
submitted to Licensor for inclusion in the Work by the copyright owner
|
| 52 |
+
or by an individual or Legal Entity authorized to submit on behalf of
|
| 53 |
+
the copyright owner. For the purposes of this definition, "submitted"
|
| 54 |
+
means any form of electronic, verbal, or written communication sent
|
| 55 |
+
to the Licensor or its representatives, including but not limited to
|
| 56 |
+
communication on electronic mailing lists, source code control systems,
|
| 57 |
+
and issue tracking systems that are managed by, or on behalf of, the
|
| 58 |
+
Licensor for the purpose of discussing and improving the Work, but
|
| 59 |
+
excluding communication that is conspicuously marked or otherwise
|
| 60 |
+
designated in writing by the copyright owner as "Not a Contribution."
|
| 61 |
+
|
| 62 |
+
"Contributor" shall mean Licensor and any individual or Legal Entity
|
| 63 |
+
on behalf of whom a Contribution has been received by Licensor and
|
| 64 |
+
subsequently incorporated within the Work.
|
| 65 |
+
|
| 66 |
+
2. Grant of Copyright License. Subject to the terms and conditions of
|
| 67 |
+
this License, each Contributor hereby grants to You a perpetual,
|
| 68 |
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
| 69 |
+
copyright license to reproduce, prepare Derivative Works of,
|
| 70 |
+
publicly display, publicly perform, sublicense, and distribute the
|
| 71 |
+
Work and such Derivative Works in Source or Object form.
|
| 72 |
+
|
| 73 |
+
3. Grant of Patent License. Subject to the terms and conditions of
|
| 74 |
+
this License, each Contributor hereby grants to You a perpetual,
|
| 75 |
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
| 76 |
+
(except as stated in this section) patent license to make, have made,
|
| 77 |
+
use, offer to sell, sell, import, and otherwise transfer the Work,
|
| 78 |
+
where such license applies only to those patent claims licensable
|
| 79 |
+
by such Contributor that are necessarily infringed by their
|
| 80 |
+
Contribution(s) alone or by combination of their Contribution(s)
|
| 81 |
+
with the Work to which such Contribution(s) was submitted. If You
|
| 82 |
+
institute patent litigation against any entity (including a
|
| 83 |
+
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
| 84 |
+
or a Contribution incorporated within the Work constitutes direct
|
| 85 |
+
or contributory patent infringement, then any patent licenses
|
| 86 |
+
granted to You under this License for that Work shall terminate
|
| 87 |
+
as of the date such litigation is filed.
|
| 88 |
+
|
| 89 |
+
4. Redistribution. You may reproduce and distribute copies of the
|
| 90 |
+
Work or Derivative Works thereof in any medium, with or without
|
| 91 |
+
modifications, and in Source or Object form, provided that You
|
| 92 |
+
meet the following conditions:
|
| 93 |
+
|
| 94 |
+
(a) You must give any other recipients of the Work or
|
| 95 |
+
Derivative Works a copy of this License; and
|
| 96 |
+
|
| 97 |
+
(b) You must cause any modified files to carry prominent notices
|
| 98 |
+
stating that You changed the files; and
|
| 99 |
+
|
| 100 |
+
(c) You must retain, in the Source form of any Derivative Works
|
| 101 |
+
that You distribute, all copyright, patent, trademark, and
|
| 102 |
+
attribution notices from the Source form of the Work,
|
| 103 |
+
excluding those notices that do not pertain to any part of
|
| 104 |
+
the Derivative Works; and
|
| 105 |
+
|
| 106 |
+
(d) If the Work includes a "NOTICE" text file as part of its
|
| 107 |
+
distribution, then any Derivative Works that You distribute must
|
| 108 |
+
include a readable copy of the attribution notices contained
|
| 109 |
+
within such NOTICE file, excluding those notices that do not
|
| 110 |
+
pertain to any part of the Derivative Works, in at least one
|
| 111 |
+
of the following places: within a NOTICE text file distributed
|
| 112 |
+
as part of the Derivative Works; within the Source form or
|
| 113 |
+
documentation, if provided along with the Derivative Works; or,
|
| 114 |
+
within a display generated by the Derivative Works, if and
|
| 115 |
+
wherever such third-party notices normally appear. The contents
|
| 116 |
+
of the NOTICE file are for informational purposes only and
|
| 117 |
+
do not modify the License. You may add Your own attribution
|
| 118 |
+
notices within Derivative Works that You distribute, alongside
|
| 119 |
+
or as an addendum to the NOTICE text from the Work, provided
|
| 120 |
+
that such additional attribution notices cannot be construed
|
| 121 |
+
as modifying the License.
|
| 122 |
+
|
| 123 |
+
You may add Your own copyright statement to Your modifications and
|
| 124 |
+
may provide additional or different license terms and conditions
|
| 125 |
+
for use, reproduction, or distribution of Your modifications, or
|
| 126 |
+
for any such Derivative Works as a whole, provided Your use,
|
| 127 |
+
reproduction, and distribution of the Work otherwise complies with
|
| 128 |
+
the conditions stated in this License.
|
| 129 |
+
|
| 130 |
+
5. Submission of Contributions. Unless You explicitly state otherwise,
|
| 131 |
+
any Contribution intentionally submitted for inclusion in the Work
|
| 132 |
+
by You to the Licensor shall be under the terms and conditions of
|
| 133 |
+
this License, without any additional terms or conditions.
|
| 134 |
+
Notwithstanding the above, nothing herein shall supersede or modify
|
| 135 |
+
the terms of any separate license agreement you may have executed
|
| 136 |
+
with Licensor regarding such Contributions.
|
| 137 |
+
|
| 138 |
+
6. Trademarks. This License does not grant permission to use the trade
|
| 139 |
+
names, trademarks, service marks, or product names of the Licensor,
|
| 140 |
+
except as required for reasonable and customary use in describing the
|
| 141 |
+
origin of the Work and reproducing the content of the NOTICE file.
|
| 142 |
+
|
| 143 |
+
7. Disclaimer of Warranty. Unless required by applicable law or
|
| 144 |
+
agreed to in writing, Licensor provides the Work (and each
|
| 145 |
+
Contributor provides its Contributions) on an "AS IS" BASIS,
|
| 146 |
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
| 147 |
+
implied, including, without limitation, any warranties or conditions
|
| 148 |
+
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
| 149 |
+
PARTICULAR PURPOSE. You are solely responsible for determining the
|
| 150 |
+
appropriateness of using or redistributing the Work and assume any
|
| 151 |
+
risks associated with Your exercise of permissions under this License.
|
| 152 |
+
|
| 153 |
+
8. Limitation of Liability. In no event and under no legal theory,
|
| 154 |
+
whether in tort (including negligence), contract, or otherwise,
|
| 155 |
+
unless required by applicable law (such as deliberate and grossly
|
| 156 |
+
negligent acts) or agreed to in writing, shall any Contributor be
|
| 157 |
+
liable to You for damages, including any direct, indirect, special,
|
| 158 |
+
incidental, or consequential damages of any character arising as a
|
| 159 |
+
result of this License or out of the use or inability to use the
|
| 160 |
+
Work (including but not limited to damages for loss of goodwill,
|
| 161 |
+
work stoppage, computer failure or malfunction, or any and all
|
| 162 |
+
other commercial damages or losses), even if such Contributor
|
| 163 |
+
has been advised of the possibility of such damages.
|
| 164 |
+
|
| 165 |
+
9. Accepting Warranty or Additional Liability. While redistributing
|
| 166 |
+
the Work or Derivative Works thereof, You may choose to offer,
|
| 167 |
+
and charge a fee for, acceptance of support, warranty, indemnity,
|
| 168 |
+
or other liability obligations and/or rights consistent with this
|
| 169 |
+
License. However, in accepting such obligations, You may act only
|
| 170 |
+
on Your own behalf and on Your sole responsibility, not on behalf
|
| 171 |
+
of any other Contributor, and only if You agree to indemnify,
|
| 172 |
+
defend, and hold each Contributor harmless for any liability
|
| 173 |
+
incurred by, or claims asserted against, such Contributor by reason
|
| 174 |
+
of your accepting any such warranty or additional liability.
|
| 175 |
+
|
| 176 |
+
END OF TERMS AND CONDITIONS
|
| 177 |
+
|
| 178 |
+
APPENDIX: How to apply the Apache License to your work.
|
| 179 |
+
|
| 180 |
+
To apply the Apache License to your work, attach the following
|
| 181 |
+
boilerplate notice, with the fields enclosed by brackets "[]"
|
| 182 |
+
replaced with your own identifying information. (Don't include
|
| 183 |
+
the brackets!) The text should be enclosed in the appropriate
|
| 184 |
+
comment syntax for the file format. We also recommend that a
|
| 185 |
+
file or class name and description of purpose be included on the
|
| 186 |
+
same "printed page" as the copyright notice for easier
|
| 187 |
+
identification within third-party archives.
|
| 188 |
+
|
| 189 |
+
Copyright [yyyy] [name of copyright owner]
|
| 190 |
+
|
| 191 |
+
Licensed under the Apache License, Version 2.0 (the "License");
|
| 192 |
+
you may not use this file except in compliance with the License.
|
| 193 |
+
You may obtain a copy of the License at
|
| 194 |
+
|
| 195 |
+
http://www.apache.org/licenses/LICENSE-2.0
|
| 196 |
+
|
| 197 |
+
Unless required by applicable law or agreed to in writing, software
|
| 198 |
+
distributed under the License is distributed on an "AS IS" BASIS,
|
| 199 |
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 200 |
+
See the License for the specific language governing permissions and
|
| 201 |
+
limitations under the License.
|
bundle/native/libornith_attention_iu4.so
CHANGED
|
@@ -1,3 +1,3 @@
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:
|
| 3 |
-
size
|
|
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:3d29cf3616c69f9fe2b89f8d3564d479dda682fdc06572886e7b60faccc9a53a
|
| 3 |
+
size 252976
|
bundle/native/libornith_dense_g256.so
CHANGED
|
Binary files a/bundle/native/libornith_dense_g256.so and b/bundle/native/libornith_dense_g256.so differ
|
|
|
bundle/native/libornith_dense_g256_n32.so
CHANGED
|
Binary files a/bundle/native/libornith_dense_g256_n32.so and b/bundle/native/libornith_dense_g256_n32.so differ
|
|
|
bundle/native/libornith_persistent_iu4.so
CHANGED
|
@@ -1,3 +1,3 @@
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:
|
| 3 |
-
size
|
|
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:c70ba14d5aebf67b761ea54cfc9a38f8bbedf74cfca8b19050048acda9c8a06c
|
| 3 |
+
size 104904
|
bundle/native/libornith_routed_storage_n32.so
CHANGED
|
@@ -1,3 +1,3 @@
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
-
oid sha256:
|
| 3 |
-
size
|
|
|
|
| 1 |
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:2002bf6f82a637f5842f7c5f09a09f3b7a0db55edc4b3e35181c8f9a12ffb4b1
|
| 3 |
+
size 166536
|
bundle/packaging/serve.sh
CHANGED
|
@@ -28,4 +28,4 @@ export XDG_CACHE_HOME="$cache_directory" AITER_JIT_DIR="$cache_directory/aiter"
|
|
| 28 |
source "$runtime_root/runtime-env.sh"
|
| 29 |
unset VLLM_SOURCE VLLM_VENV
|
| 30 |
export PYTHONPATH="$plugin_site${PYTHONPATH:+:$PYTHONPATH}"
|
| 31 |
-
exec "$runtime_root/venv/bin/python" -m ornith_g256.launch "${launch_args[@]}" --cache-directory "$cache_directory"
|
|
|
|
| 28 |
source "$runtime_root/runtime-env.sh"
|
| 29 |
unset VLLM_SOURCE VLLM_VENV
|
| 30 |
export PYTHONPATH="$plugin_site${PYTHONPATH:+:$PYTHONPATH}"
|
| 31 |
+
exec "$runtime_root/venv/bin/python" -m ornith_g256.optimized.launch "${launch_args[@]}" --cache-directory "$cache_directory"
|
bundle/plugin-site/TOOL-INTEGRATION.json
ADDED
|
@@ -0,0 +1,205 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"status": "STAGED_PENDING_QUALIFICATION",
|
| 3 |
+
"parent": "/srv/ssd/p3700ba/scratch/crown/optimization-campaign-20260916/artifacts/ornith-combined-v3/bundle/plugin-site/ornith_g256",
|
| 4 |
+
"parent_files": {
|
| 5 |
+
"attention_window.py": "78534f9a1b0beae9ecf80cc5b2a613e4fe4827036c8df0acc7e4dfd956aae1c5",
|
| 6 |
+
"dflash_spec.py": "c948b99eab22d2af5c5fc0ca9f764a0855a3d86aca0609362bf5474a9beff1b3",
|
| 7 |
+
"gdn_compact.py": "8b5c062d10f386d78bc9926ab0576e56359401e5d1621c397a2922e2b426ccab",
|
| 8 |
+
"cache_full1120.py": "3d4240a749d9671c0bf9b3937200e1cd8d4abfe6699d8c1d6857b188c2ef8a8c",
|
| 9 |
+
"attention_storage.py": "3ab813e221835370bdb503b2ab7d48f46294441b1dd9c652c5dc976438bdb6e2",
|
| 10 |
+
"prefill_draft.py": "d79bb29a124639543a64e198908e2eaef4cbcf492f2e5e14debcfc9c867816f5",
|
| 11 |
+
"column_kernel.py": "27fb053d970184732a46212026afd3152ec07a0dfc3dc4c5cc5cc0d539f17c81",
|
| 12 |
+
"moe_base.py": "26fb517415f2497859f2e3ff1aed8f31e7c56a4aacccfb34c6638e65d7df07c4",
|
| 13 |
+
"attention_partition.py": "41bfb44dd0d0c46a2c0a725983155018e051cd89edb8527f1195a092bef7ffdd",
|
| 14 |
+
"lifecycle.py": "3285552177cb8ce8a0149a01f8168d2bc44a739a9395ca84e15acccbb20a0b73",
|
| 15 |
+
"attention_verify.py": "3a329599694265f4badecc2f9b709b280d9960b967d7192f3959ad9e0d8c6077",
|
| 16 |
+
"attention_iu4_persistent.py": "cef74ae423974063212637a49cd64a5a10b9077f3a5e2287746722749a1e6297",
|
| 17 |
+
"loader.py": "a853c6d014214a7adf298548f9175b6e18dd05338917aec16ebcb4fc3c3ecb13",
|
| 18 |
+
"attention_mixed.py": "cbb530697f2555be6532b7d4d3401f3a33c4c58982659e88892720d21691110a",
|
| 19 |
+
"graph_phase_guard.py": "6cabaf7d4e3be04bbac7a642544d38948de2e07ea7ca38f5927f91bf3ad39163",
|
| 20 |
+
"phase_dispatch.py": "1a4ebdf2af709e6ca519e0a08f6a537602df178ca575332045617631640f8137",
|
| 21 |
+
"method.py": "fc301af72efd234430dbfb839d56e74328457b5b83e14a044aeb9554d8acc707",
|
| 22 |
+
"gdn_compact_kernel.py": "0eaf34a2eccef2dabbc9e4f017e6ce84317a28efdcbabecf92b3b4c4c9429a55",
|
| 23 |
+
"attention_tile.py": "f2bf007132bc712c8e61e036ea8a94b54f18a813d9bdd334eee46436a3ca1fea",
|
| 24 |
+
"no_spec.py": "561106b75328fe0e161fa1243aca4cf50c5049d8e785d22c90aacb3e1e1a0104",
|
| 25 |
+
"__init__.py": "1f526dde80a76eb467b1fd30e8b2babdf0b677cb2f246217ac3ab2509d8cf288",
|
| 26 |
+
"gdn_spec.py": "0931b1b3601d935bb6acdfcb65d307dae4105645cd387bd9b256dd7a525376a3",
|
| 27 |
+
"column_backend.py": "a035b839369e733b05f68ce83df0d022cb62455016fccf30c2978127ab3c5403",
|
| 28 |
+
"launch.py": "c2854b86e2f97dd2ea0356b4b333922f62b9b3d1b3e71f7993d803c6d0957479",
|
| 29 |
+
"dense_n32.py": "119fb7d25e9d28b6e520b5285c10de86c7d8ec2fe29a2399e9c086173d765903",
|
| 30 |
+
"prefix_cache.py": "65fb972e285e37741f53b221d2f6380ec464d34a39f36bb66a82dcaacab727d0",
|
| 31 |
+
"runtime.py": "db43b1c3766720e0dbee4e8617aee673a533c30d36fcc1fbf43919717d8ab149",
|
| 32 |
+
"parser_contract.py": "2cebab403dc882569f3ce67d12798df14cebd5c162c72f82becd37bcdd7e2a28",
|
| 33 |
+
"attention_fast.py": "27b1a133cadb87ff3a63c027049e1da81e4112fd5d5260a4440fde85491f838a",
|
| 34 |
+
"attention.py": "2fdd3c0ed84dcb052a9e8f9df4068ddf80a969931bea0a295e022d36b753cda1",
|
| 35 |
+
"native_tool_schema.py": "910060689d72b47ad4a7bb1e51ef301e41ff0be1beb9cc6b3df01c423d2edef5",
|
| 36 |
+
"attention_folded.py": "73dac8203032bce6ceedf213d7eefe5987c921768ae237129e232fefd8d9bba3",
|
| 37 |
+
"dynamic_graphs.py": "6f51f0dd740f6546829e46bbbe840df2780ce8c49978658eb3a05092554ff90c",
|
| 38 |
+
"worker.py": "acc27240c21f3478ca452ddf04244df4ee770efaac38192ef64d328ef22701d5",
|
| 39 |
+
"runtime_correctness.py": "08f41773e623a4f99cf8ab19e434d7e33bac9cfa5eb05a9dc0f4511aab8ba4bc",
|
| 40 |
+
"config.py": "23210d8b625ef375638e2d6043a92a26181a5b5f1c537d9d4b34c91b7109f36f",
|
| 41 |
+
"dflash_conv_boundary.py": "a5f46458194cade5ddcbfbb51af30c662d12b9086fc459ecbeb17463339febb1",
|
| 42 |
+
"native.py": "49cce0d67987b9659feb556aafb247423490c2153b7d77b750f055a4fdcda331",
|
| 43 |
+
"attention_compact.py": "50df61a02a3890a2825b205ed18c7d18a2b002333d382c572440b58ca9b2b9fe",
|
| 44 |
+
"worker_base.py": "a00ef1295172e655c8b3c7d790cfd0316114a9492e483dedfdea036d3b7bf868",
|
| 45 |
+
"attention_iu4.py": "15c5341fc848c43ac956546e563cf8ab7f7649f05a674fb0153b0e2d76438cae",
|
| 46 |
+
"dense_source.py": "cdc7dd6a13c2439cd405596069c301fe28cb95ff32278f1bafe8118fe357d48e",
|
| 47 |
+
"adaptive_c1.py": "8e48fd874fbdf65183c20bcc7c748d2562dbcafa97e004f18a2818b98980a809",
|
| 48 |
+
"_vllm_correctness/gpu_model_runner.py": "ef0009acaa67a1a5e066984f34bf619b06e85220ff953a3d9511376f0146dccf",
|
| 49 |
+
"_vllm_correctness/__init__.py": "963097106407a87b62b275041516144de15713dee65108accc29e27036ffc059",
|
| 50 |
+
"_vllm_correctness/manifest.json": "0ddfa489f4f9ead135006e29aa675e7522afa5cee10596c2115e07f13186478c",
|
| 51 |
+
"_vllm_correctness/causal_conv1d.py": "230f3e554f5f1d609a37022578c84e42d5c71d46e684101278a7e21af4f7d418",
|
| 52 |
+
"_vllm_correctness/gdn_attn.py": "a3c361d502e8bca2aa46a0baed5ec753a3c62e3d027a298e3ed4692794f9f6ec",
|
| 53 |
+
"_vllm_correctness/qwen_gdn_linear_attn.py": "3e58fdddc7a0b42ae3d9b5eabb375a993a4110bad499c37846bde466901dbeb3",
|
| 54 |
+
"strict_qwen/__init__.py": "ee8a5055249f1413fc22f46421ab77b5485aa39b199c68eb5b56c75dcc40922d",
|
| 55 |
+
"strict_qwen/qwen3_contract.py": "29c083659c54c81144fdb6c8329c75aaa62cf69d1c45fa7da8bdeeeca4856eaf",
|
| 56 |
+
"strict_qwen/qwen3.py": "33c653447a144b72257e508ccd92610db6d30578e28f98c69c99281773dead1c",
|
| 57 |
+
"optimized/__init__.py": "0297f5d476afdddb0556f4f236f95e514db6b9b96788df672a4a0bc0074fd835",
|
| 58 |
+
"optimized/compile_scope.py": "18971a73b3a5ad80e6212de7f7e419a1dd00eb7380a5e6c8559f5ac6f5ec8d4d",
|
| 59 |
+
"optimized/draft.py": "e801301cf03d69998758d35fb5463754ccc30a18e030ed2b988b26cef43f429c",
|
| 60 |
+
"optimized/draft_kernel.py": "3d6da6ef0c144505f888cbd35ba50b6b95ad8e48c376ab73297746bff228919e",
|
| 61 |
+
"optimized/greedy.py": "386443ad8729f8fe21ec5fef84ead3acfbd4d6c5dd3b4c1e50a7995ed1d1761f",
|
| 62 |
+
"optimized/greedy_kernel.py": "2ebe2a25d1011a45eea7eda3c0964ca1eb2638df7d60f54db02c91435550dc42",
|
| 63 |
+
"optimized/launch.py": "c1d42602154674916a94b9907b7746fac3ea48cc505fa69babf116f0f127e03d",
|
| 64 |
+
"optimized/mixed_qk_sum.py": "17b18769ae84c2e9707a0257181e5c314e9774fc3b4796f8e6787a8302cb5e1f",
|
| 65 |
+
"optimized/postconv.py": "c58ee5d2b8418fbf69eda5d9ee4f342b11cc012ab77ae25b6ce2b2aea51aa50f",
|
| 66 |
+
"optimized/prefix.py": "c695a899ab12b58c3b89b6ecc276e82efa3f0833a319911a8528cc4e237710cd",
|
| 67 |
+
"optimized/qk.py": "d86f1a94086668ac6402eaec7a19c0ad8e55dc8e6377f8d4a9eeab6aa96f9664",
|
| 68 |
+
"optimized/qk_schedule.py": "66038781f0667486fba458ef3118ca89ad9640dc2c2bd05cf8576440f0f178a2",
|
| 69 |
+
"optimized/worker.py": "38579b767235e7fcac3f2e6b0507438f155c134e1e85f3d1c205a43bece95512",
|
| 70 |
+
"optimized/grammar_parallel.py": "e9c3608438712521980d7fcd60f58862d0b0a2876620183f304c64a03817edea",
|
| 71 |
+
"optimized/pipeline/SOURCE-PROOF.json": "5939bc8a53a831715d05d0fa0290beaa0cb00eda1597fe3bf1f1cfc53699fd0d",
|
| 72 |
+
"optimized/pipeline/__init__.py": "0297f5d476afdddb0556f4f236f95e514db6b9b96788df672a4a0bc0074fd835",
|
| 73 |
+
"optimized/pipeline/adapter.py": "6af97a8d5e0d79b7dfc32f2f0fe5f5c53b5c1c50d03542ff99c64cae7f3909e1",
|
| 74 |
+
"optimized/pipeline/candidate/causal_conv1d.py": "254c540f77b01694696a1ab535abd0e5a23d54e7211e202cf9beead04face16b",
|
| 75 |
+
"optimized/pipeline/candidate/chunk_o.py": "a93f951048958265b4d1c3f9e67341329196e7a57525433934215fde4b16e61b",
|
| 76 |
+
"optimized/pipeline/candidate/chunk_scaled_dot_kkt.py": "71735d2b1c4a19eecbb6058708481817bb610426c48f96a5f646784afcf0c65c",
|
| 77 |
+
"optimized/pipeline/candidate/solve_tril.py": "bdd04b69c0f91562048d25834bd53d298dc1f1bd068b9494af88776b2bf4b1d1",
|
| 78 |
+
"optimized/pipeline/candidate/wy_fast.py": "4cc0be519bfcf3812d810a4bebb3dd45fb9b13a8d0ab62d0ad463a1a2e36f61f"
|
| 79 |
+
},
|
| 80 |
+
"changed_existing_files": [
|
| 81 |
+
"__init__.py",
|
| 82 |
+
"_vllm_correctness/manifest.json"
|
| 83 |
+
],
|
| 84 |
+
"files": {
|
| 85 |
+
"ornith_g256/__init__.py": "60abbb51a475487e3abcb9e5b34d8836b4d4b36e0bced062d2382dc7b1076cf6",
|
| 86 |
+
"ornith_g256/_vllm_correctness/__init__.py": "963097106407a87b62b275041516144de15713dee65108accc29e27036ffc059",
|
| 87 |
+
"ornith_g256/_vllm_correctness/causal_conv1d.py": "230f3e554f5f1d609a37022578c84e42d5c71d46e684101278a7e21af4f7d418",
|
| 88 |
+
"ornith_g256/_vllm_correctness/chat_completion_serving_r03.py": "5720371f5bcb351e340a3e8d9a0235ec01ff27ed93e4f35542e31b83a055ef45",
|
| 89 |
+
"ornith_g256/_vllm_correctness/gdn_attn.py": "a3c361d502e8bca2aa46a0baed5ec753a3c62e3d027a298e3ed4692794f9f6ec",
|
| 90 |
+
"ornith_g256/_vllm_correctness/gpu_model_runner.py": "ef0009acaa67a1a5e066984f34bf619b06e85220ff953a3d9511376f0146dccf",
|
| 91 |
+
"ornith_g256/_vllm_correctness/manifest.json": "bdef4ca7151a7a5e0410dd3e8bcd2e8a85f125938ec104451cb8777267c3c4d9",
|
| 92 |
+
"ornith_g256/_vllm_correctness/qwen_gdn_linear_attn.py": "3e58fdddc7a0b42ae3d9b5eabb375a993a4110bad499c37846bde466901dbeb3",
|
| 93 |
+
"ornith_g256/_xgrammar_native/build-receipt.json": "3bdf77657e5157c35938b66dfb46266608ade3b599a62ce74abb3f807d579e84",
|
| 94 |
+
"ornith_g256/_xgrammar_native/libxgrammar-candidate.so": "92ab36bfae20c6c4df3f5321ca375204799ac833252921d9d8f6a206491d35f7",
|
| 95 |
+
"ornith_g256/_xgrammar_native/selection.json": "7cb93b3fe902de6d604bf96f79fdc3efbd8d612e494da838f6a8f0d731eb89ae",
|
| 96 |
+
"ornith_g256/adaptive_c1.py": "8e48fd874fbdf65183c20bcc7c748d2562dbcafa97e004f18a2818b98980a809",
|
| 97 |
+
"ornith_g256/attention.py": "2fdd3c0ed84dcb052a9e8f9df4068ddf80a969931bea0a295e022d36b753cda1",
|
| 98 |
+
"ornith_g256/attention_compact.py": "50df61a02a3890a2825b205ed18c7d18a2b002333d382c572440b58ca9b2b9fe",
|
| 99 |
+
"ornith_g256/attention_fast.py": "27b1a133cadb87ff3a63c027049e1da81e4112fd5d5260a4440fde85491f838a",
|
| 100 |
+
"ornith_g256/attention_folded.py": "73dac8203032bce6ceedf213d7eefe5987c921768ae237129e232fefd8d9bba3",
|
| 101 |
+
"ornith_g256/attention_iu4.py": "15c5341fc848c43ac956546e563cf8ab7f7649f05a674fb0153b0e2d76438cae",
|
| 102 |
+
"ornith_g256/attention_iu4_persistent.py": "cef74ae423974063212637a49cd64a5a10b9077f3a5e2287746722749a1e6297",
|
| 103 |
+
"ornith_g256/attention_mixed.py": "cbb530697f2555be6532b7d4d3401f3a33c4c58982659e88892720d21691110a",
|
| 104 |
+
"ornith_g256/attention_partition.py": "41bfb44dd0d0c46a2c0a725983155018e051cd89edb8527f1195a092bef7ffdd",
|
| 105 |
+
"ornith_g256/attention_storage.py": "3ab813e221835370bdb503b2ab7d48f46294441b1dd9c652c5dc976438bdb6e2",
|
| 106 |
+
"ornith_g256/attention_tile.py": "f2bf007132bc712c8e61e036ea8a94b54f18a813d9bdd334eee46436a3ca1fea",
|
| 107 |
+
"ornith_g256/attention_verify.py": "3a329599694265f4badecc2f9b709b280d9960b967d7192f3959ad9e0d8c6077",
|
| 108 |
+
"ornith_g256/attention_window.py": "78534f9a1b0beae9ecf80cc5b2a613e4fe4827036c8df0acc7e4dfd956aae1c5",
|
| 109 |
+
"ornith_g256/cache_full1120.py": "3d4240a749d9671c0bf9b3937200e1cd8d4abfe6699d8c1d6857b188c2ef8a8c",
|
| 110 |
+
"ornith_g256/column_backend.py": "a035b839369e733b05f68ce83df0d022cb62455016fccf30c2978127ab3c5403",
|
| 111 |
+
"ornith_g256/column_kernel.py": "27fb053d970184732a46212026afd3152ec07a0dfc3dc4c5cc5cc0d539f17c81",
|
| 112 |
+
"ornith_g256/config.py": "23210d8b625ef375638e2d6043a92a26181a5b5f1c537d9d4b34c91b7109f36f",
|
| 113 |
+
"ornith_g256/dense_n32.py": "119fb7d25e9d28b6e520b5285c10de86c7d8ec2fe29a2399e9c086173d765903",
|
| 114 |
+
"ornith_g256/dense_source.py": "cdc7dd6a13c2439cd405596069c301fe28cb95ff32278f1bafe8118fe357d48e",
|
| 115 |
+
"ornith_g256/dflash_conv_boundary.py": "a5f46458194cade5ddcbfbb51af30c662d12b9086fc459ecbeb17463339febb1",
|
| 116 |
+
"ornith_g256/dflash_spec.py": "c948b99eab22d2af5c5fc0ca9f764a0855a3d86aca0609362bf5474a9beff1b3",
|
| 117 |
+
"ornith_g256/dynamic_graphs.py": "6f51f0dd740f6546829e46bbbe840df2780ce8c49978658eb3a05092554ff90c",
|
| 118 |
+
"ornith_g256/gdn_compact.py": "8b5c062d10f386d78bc9926ab0576e56359401e5d1621c397a2922e2b426ccab",
|
| 119 |
+
"ornith_g256/gdn_compact_kernel.py": "0eaf34a2eccef2dabbc9e4f017e6ce84317a28efdcbabecf92b3b4c4c9429a55",
|
| 120 |
+
"ornith_g256/gdn_spec.py": "0931b1b3601d935bb6acdfcb65d307dae4105645cd387bd9b256dd7a525376a3",
|
| 121 |
+
"ornith_g256/graph_phase_guard.py": "6cabaf7d4e3be04bbac7a642544d38948de2e07ea7ca38f5927f91bf3ad39163",
|
| 122 |
+
"ornith_g256/launch.py": "c2854b86e2f97dd2ea0356b4b333922f62b9b3d1b3e71f7993d803c6d0957479",
|
| 123 |
+
"ornith_g256/lifecycle.py": "3285552177cb8ce8a0149a01f8168d2bc44a739a9395ca84e15acccbb20a0b73",
|
| 124 |
+
"ornith_g256/loader.py": "a853c6d014214a7adf298548f9175b6e18dd05338917aec16ebcb4fc3c3ecb13",
|
| 125 |
+
"ornith_g256/method.py": "fc301af72efd234430dbfb839d56e74328457b5b83e14a044aeb9554d8acc707",
|
| 126 |
+
"ornith_g256/moe_base.py": "26fb517415f2497859f2e3ff1aed8f31e7c56a4aacccfb34c6638e65d7df07c4",
|
| 127 |
+
"ornith_g256/native.py": "49cce0d67987b9659feb556aafb247423490c2153b7d77b750f055a4fdcda331",
|
| 128 |
+
"ornith_g256/native_tool_schema.py": "910060689d72b47ad4a7bb1e51ef301e41ff0be1beb9cc6b3df01c423d2edef5",
|
| 129 |
+
"ornith_g256/no_spec.py": "561106b75328fe0e161fa1243aca4cf50c5049d8e785d22c90aacb3e1e1a0104",
|
| 130 |
+
"ornith_g256/optimized/__init__.py": "0297f5d476afdddb0556f4f236f95e514db6b9b96788df672a4a0bc0074fd835",
|
| 131 |
+
"ornith_g256/optimized/compile_scope.py": "18971a73b3a5ad80e6212de7f7e419a1dd00eb7380a5e6c8559f5ac6f5ec8d4d",
|
| 132 |
+
"ornith_g256/optimized/draft.py": "e801301cf03d69998758d35fb5463754ccc30a18e030ed2b988b26cef43f429c",
|
| 133 |
+
"ornith_g256/optimized/draft_kernel.py": "3d6da6ef0c144505f888cbd35ba50b6b95ad8e48c376ab73297746bff228919e",
|
| 134 |
+
"ornith_g256/optimized/grammar_parallel.py": "e9c3608438712521980d7fcd60f58862d0b0a2876620183f304c64a03817edea",
|
| 135 |
+
"ornith_g256/optimized/greedy.py": "386443ad8729f8fe21ec5fef84ead3acfbd4d6c5dd3b4c1e50a7995ed1d1761f",
|
| 136 |
+
"ornith_g256/optimized/greedy_kernel.py": "2ebe2a25d1011a45eea7eda3c0964ca1eb2638df7d60f54db02c91435550dc42",
|
| 137 |
+
"ornith_g256/optimized/launch.py": "c1d42602154674916a94b9907b7746fac3ea48cc505fa69babf116f0f127e03d",
|
| 138 |
+
"ornith_g256/optimized/mixed_qk_sum.py": "17b18769ae84c2e9707a0257181e5c314e9774fc3b4796f8e6787a8302cb5e1f",
|
| 139 |
+
"ornith_g256/optimized/pipeline/SOURCE-PROOF.json": "5939bc8a53a831715d05d0fa0290beaa0cb00eda1597fe3bf1f1cfc53699fd0d",
|
| 140 |
+
"ornith_g256/optimized/pipeline/__init__.py": "0297f5d476afdddb0556f4f236f95e514db6b9b96788df672a4a0bc0074fd835",
|
| 141 |
+
"ornith_g256/optimized/pipeline/adapter.py": "6af97a8d5e0d79b7dfc32f2f0fe5f5c53b5c1c50d03542ff99c64cae7f3909e1",
|
| 142 |
+
"ornith_g256/optimized/pipeline/candidate/causal_conv1d.py": "254c540f77b01694696a1ab535abd0e5a23d54e7211e202cf9beead04face16b",
|
| 143 |
+
"ornith_g256/optimized/pipeline/candidate/chunk_o.py": "a93f951048958265b4d1c3f9e67341329196e7a57525433934215fde4b16e61b",
|
| 144 |
+
"ornith_g256/optimized/pipeline/candidate/chunk_scaled_dot_kkt.py": "71735d2b1c4a19eecbb6058708481817bb610426c48f96a5f646784afcf0c65c",
|
| 145 |
+
"ornith_g256/optimized/pipeline/candidate/solve_tril.py": "bdd04b69c0f91562048d25834bd53d298dc1f1bd068b9494af88776b2bf4b1d1",
|
| 146 |
+
"ornith_g256/optimized/pipeline/candidate/wy_fast.py": "4cc0be519bfcf3812d810a4bebb3dd45fb9b13a8d0ab62d0ad463a1a2e36f61f",
|
| 147 |
+
"ornith_g256/optimized/postconv.py": "c58ee5d2b8418fbf69eda5d9ee4f342b11cc012ab77ae25b6ce2b2aea51aa50f",
|
| 148 |
+
"ornith_g256/optimized/prefix.py": "c695a899ab12b58c3b89b6ecc276e82efa3f0833a319911a8528cc4e237710cd",
|
| 149 |
+
"ornith_g256/optimized/qk.py": "d86f1a94086668ac6402eaec7a19c0ad8e55dc8e6377f8d4a9eeab6aa96f9664",
|
| 150 |
+
"ornith_g256/optimized/qk_schedule.py": "66038781f0667486fba458ef3118ca89ad9640dc2c2bd05cf8576440f0f178a2",
|
| 151 |
+
"ornith_g256/optimized/worker.py": "38579b767235e7fcac3f2e6b0507438f155c134e1e85f3d1c205a43bece95512",
|
| 152 |
+
"ornith_g256/parser_contract.py": "2cebab403dc882569f3ce67d12798df14cebd5c162c72f82becd37bcdd7e2a28",
|
| 153 |
+
"ornith_g256/phase_dispatch.py": "1a4ebdf2af709e6ca519e0a08f6a537602df178ca575332045617631640f8137",
|
| 154 |
+
"ornith_g256/prefill_draft.py": "d79bb29a124639543a64e198908e2eaef4cbcf492f2e5e14debcfc9c867816f5",
|
| 155 |
+
"ornith_g256/prefix_cache.py": "65fb972e285e37741f53b221d2f6380ec464d34a39f36bb66a82dcaacab727d0",
|
| 156 |
+
"ornith_g256/runtime.py": "db43b1c3766720e0dbee4e8617aee673a533c30d36fcc1fbf43919717d8ab149",
|
| 157 |
+
"ornith_g256/runtime_correctness.py": "08f41773e623a4f99cf8ab19e434d7e33bac9cfa5eb05a9dc0f4511aab8ba4bc",
|
| 158 |
+
"ornith_g256/strict_qwen/__init__.py": "ee8a5055249f1413fc22f46421ab77b5485aa39b199c68eb5b56c75dcc40922d",
|
| 159 |
+
"ornith_g256/strict_qwen/qwen3.py": "33c653447a144b72257e508ccd92610db6d30578e28f98c69c99281773dead1c",
|
| 160 |
+
"ornith_g256/strict_qwen/qwen3_contract.py": "29c083659c54c81144fdb6c8329c75aaa62cf69d1c45fa7da8bdeeeca4856eaf",
|
| 161 |
+
"ornith_g256/worker.py": "acc27240c21f3478ca452ddf04244df4ee770efaac38192ef64d328ef22701d5",
|
| 162 |
+
"ornith_g256/worker_base.py": "a00ef1295172e655c8b3c7d790cfd0316114a9492e483dedfdea036d3b7bf868",
|
| 163 |
+
"ornith_g256/xgrammar_native.py": "d19c546e124cb11311e4cc31f5833dc663e2789d45090091c0f8fec8a9626623"
|
| 164 |
+
},
|
| 165 |
+
"r01_sha256": "e9c3608438712521980d7fcd60f58862d0b0a2876620183f304c64a03817edea",
|
| 166 |
+
"r01_existing_preserved": true,
|
| 167 |
+
"r04": {
|
| 168 |
+
"mode": "candidate",
|
| 169 |
+
"library": "libxgrammar-candidate.so",
|
| 170 |
+
"sha256": "92ab36bfae20c6c4df3f5321ca375204799ac833252921d9d8f6a206491d35f7",
|
| 171 |
+
"source_commit": "557becfb64c503ae9c04344b0047661f43f44320",
|
| 172 |
+
"patch_sha256": "2c42719e948809242f230ec534f967eab1babad9be1ba6dda9c59fd706f20d45",
|
| 173 |
+
"sources": {
|
| 174 |
+
"cpp/grammar.cc": "01d1552b37ecaa72da601a3b6d43975a2c3a0de02cdf797e82c243b30eef8c9d",
|
| 175 |
+
"cpp/grammar_functor.cc": "9410c20b665093b9fb67dc30db148121b819d3aee5ff94eb82558bf79b1f8279"
|
| 176 |
+
},
|
| 177 |
+
"configure": [
|
| 178 |
+
"cmake",
|
| 179 |
+
"-S",
|
| 180 |
+
"/srv/ssd/p3700ba/scratch/crown/runtime-tool-integration-20260917/native-build-r2/source-candidate",
|
| 181 |
+
"-B",
|
| 182 |
+
"/srv/ssd/p3700ba/scratch/crown/runtime-tool-integration-20260917/native-build-r2/build-candidate",
|
| 183 |
+
"-G",
|
| 184 |
+
"Ninja",
|
| 185 |
+
"-DCMAKE_BUILD_TYPE=Release",
|
| 186 |
+
"-DXGRAMMAR_BUILD_PYTHON_BINDINGS=ON",
|
| 187 |
+
"-DCMAKE_CXX_FLAGS_RELEASE=-O3 -DNDEBUG -fno-lto",
|
| 188 |
+
"-DCMAKE_CXX_FLAGS_RELWITHDEBINFO=-O3 -g -DNDEBUG -fno-lto",
|
| 189 |
+
"-DPython_EXECUTABLE=/opt/ciru/glm53-iu4/venv/bin/python",
|
| 190 |
+
"-DPython3_EXECUTABLE=/opt/ciru/glm53-iu4/venv/bin/python"
|
| 191 |
+
],
|
| 192 |
+
"build_seconds": 67.22386657004245,
|
| 193 |
+
"compiler": "g++ (GCC) 15.2.0\nCopyright (C) 2025 Free Software Foundation, Inc.\nThis is free software; see the source for copying conditions. There is NO\nwarranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.\n\n",
|
| 194 |
+
"submodules": " 6689d14c203eed390ae7bb64f56a983cfd7dff9c 3rdparty/cpptrace (v0.1-644-g6689d14)\n bbd2f4d32427e548797929af08cfe2a9cbb3cf12 3rdparty/dlpack (v1.0)\n df1544bcee0c7ce35cd5ea0b3eb8cc81855a4140 3rdparty/googletest (release-1.8.0-3498-gdf1544bc)\n"
|
| 195 |
+
},
|
| 196 |
+
"r03": {
|
| 197 |
+
"file": "chat_completion_serving_r03.py",
|
| 198 |
+
"native_path": "entrypoints/openai/chat_completion/serving.py",
|
| 199 |
+
"native_sha256": "9982953285e9df469032a82fffa4095d0e9d86278bede6e2b91d03d02373d182",
|
| 200 |
+
"sha256": "5720371f5bcb351e340a3e8d9a0235ec01ff27ed93e4f35542e31b83a055ef45"
|
| 201 |
+
},
|
| 202 |
+
"model_settings_changed": false,
|
| 203 |
+
"weights_changed": false,
|
| 204 |
+
"strict_policy_changed": false
|
| 205 |
+
}
|
bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/direct_url.json
DELETED
|
@@ -1 +0,0 @@
|
|
| 1 |
-
{"url":"file:///srv/ssd/sn850x/scratch/crown/ornith-prefix-64k-v1/src/runtime/ornith_g256","dir_info":{}}
|
|
|
|
|
|
bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/uv_build.json
DELETED
|
@@ -1 +0,0 @@
|
|
| 1 |
-
{}
|
|
|
|
|
|
bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/uv_cache.json
DELETED
|
@@ -1 +0,0 @@
|
|
| 1 |
-
{"timestamp":{"secs_since_epoch":1788727205,"nanos_since_epoch":456565118},"commit":null,"tags":null,"env":{},"directories":{"src":null}}
|
|
|
|
|
|
bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/INSTALLER
RENAMED
|
File without changes
|
bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/METADATA
RENAMED
|
@@ -1,6 +1,6 @@
|
|
| 1 |
Metadata-Version: 2.4
|
| 2 |
Name: ciru-ornith-g256
|
| 3 |
-
Version:
|
| 4 |
Summary: Self-contained project adapter for Ornith G256 and DFlash2 on Ciru vLLM
|
| 5 |
Author-email: Ciru <ciru@ciru.ai>
|
| 6 |
Requires-Python: >=3.10
|
|
|
|
| 1 |
Metadata-Version: 2.4
|
| 2 |
Name: ciru-ornith-g256
|
| 3 |
+
Version: 4.0.0
|
| 4 |
Summary: Self-contained project adapter for Ornith G256 and DFlash2 on Ciru vLLM
|
| 5 |
Author-email: Ciru <ciru@ciru.ai>
|
| 6 |
Requires-Python: >=3.10
|
bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/RECORD
RENAMED
|
@@ -1,27 +1,29 @@
|
|
| 1 |
-
|
| 2 |
-
ciru_ornith_g256-
|
| 3 |
-
ciru_ornith_g256-
|
| 4 |
-
ciru_ornith_g256-
|
| 5 |
-
ciru_ornith_g256-
|
| 6 |
-
ciru_ornith_g256-
|
| 7 |
-
ciru_ornith_g256-
|
| 8 |
-
ciru_ornith_g256-
|
| 9 |
-
|
| 10 |
-
ciru_ornith_g256-1.0.2.dist-info/uv_cache.json,sha256=L73WKGonFia8yJ8vK90qJmbQxtUhlvCNsq7D4umUJX4,137
|
| 11 |
-
ornith_g256/__init__.py,sha256=tqhBbqjsIA7JBuJINSvgLpqqNdviaoSPeG-jINdHqPM,1368
|
| 12 |
ornith_g256/_vllm_correctness/__init__.py,sha256=ljCXEGQHqHtisnUEFRYUTeFXE97mUQiszCnicDb_wFk,79
|
| 13 |
ornith_g256/_vllm_correctness/causal_conv1d.py,sha256=Iw8-VU9fHWCaNwIleMhOQtXHHUbmhBASeKfiGvT31Bg,51774
|
|
|
|
| 14 |
ornith_g256/_vllm_correctness/gdn_attn.py,sha256=o8Nh1QLovKKqRqC67V7HU6PGLj0CeimOPtRpJ5T59uw,26224
|
| 15 |
ornith_g256/_vllm_correctness/gpu_model_runner.py,sha256=7wAJrKpnoaXgZphPNL9hmwboUiD_lTo9lRE3bwFG3M8,345523
|
| 16 |
-
ornith_g256/_vllm_correctness/manifest.json,sha256=
|
| 17 |
ornith_g256/_vllm_correctness/qwen_gdn_linear_attn.py,sha256=Plj93cegtCrj2bXquzdamTpBELrUmcN4Rr3kZpAdvrM,78334
|
|
|
|
|
|
|
|
|
|
| 18 |
ornith_g256/adaptive_c1.py,sha256=jkj9h0-99lGDwgvMfHSNJWLbyvqX4ATxiigYuYmAqAk,14985
|
| 19 |
ornith_g256/attention.py,sha256=L908DthNywUqno-d9AaN34CpaZMb6gopXgItNrdTzaE,772
|
| 20 |
-
ornith_g256/attention_compact.py,sha256=
|
| 21 |
ornith_g256/attention_fast.py,sha256=J7GhM8rbh_86Y8AnBJ4dqB5BEv1dUmCkRA_ehUkfg4o,12274
|
| 22 |
ornith_g256/attention_folded.py,sha256=c9rIIDAyvObO7fIT1-7-WYfJIXaK4jcSniMv79jZu6M,9570
|
| 23 |
-
ornith_g256/attention_iu4.py,sha256=
|
| 24 |
-
ornith_g256/attention_iu4_persistent.py,sha256=
|
| 25 |
ornith_g256/attention_mixed.py,sha256=y7UwaX8lVb5lMrfU00AfOjPExYmCZZ6IiScg0haREQo,4138
|
| 26 |
ornith_g256/attention_partition.py,sha256=Qb-0TdDQxGosCnJZgxVQGOBRzYntuFJ_EZWgkr73_90,8642
|
| 27 |
ornith_g256/attention_storage.py,sha256=OrgT4iGDU3C9tQOyq31I9GKURBsd2cZSxdyXZDi9tuI,4117
|
|
@@ -49,6 +51,28 @@ ornith_g256/moe_base.py,sha256=JvtRdBXySXhZ8uP_Gu2PMefFakqszPs0xmOOZdffB8Q,3683
|
|
| 49 |
ornith_g256/native.py,sha256=Sczg1nmHuWWf61Vqr7JHQjSQwhU7fXe3UPBVpP3NozE,12196
|
| 50 |
ornith_g256/native_tool_schema.py,sha256=kQBgaJ1ytHrUp7seUe8wHkH_C-G-ucxrPfAcQj0u3vU,1897
|
| 51 |
ornith_g256/no_spec.py,sha256=VhEGt1Mo_g4WH6EkOspM9QxQSdjnhdIskKrLPh4aAQQ,2618
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 52 |
ornith_g256/parser_contract.py,sha256=LOurQD3IglafPOZ9EnmN8Uzr1cFixy-Cvs03vN1-Kig,4203
|
| 53 |
ornith_g256/phase_dispatch.py,sha256=Gk698q9wnmylGeCgj2pTdgLfF4yldTMgRWF2MWQPgTc,2184
|
| 54 |
ornith_g256/prefill_draft.py,sha256=15uymhJGOVQ6ZOGYkI4urvTLz0kvLl4U3rz8nIZ4FvU,5630
|
|
@@ -58,6 +82,7 @@ ornith_g256/runtime_correctness.py,sha256=CPQXc-YjpPmc-KsZ5DTX4zusnPpesFqdwPRRGq
|
|
| 58 |
ornith_g256/strict_qwen/__init__.py,sha256=7opQVSSfFBP8IvRkIat3tUhao5sZnGjrW1bHXcxAki0,81
|
| 59 |
ornith_g256/strict_qwen/qwen3.py,sha256=M8ZTRHoUS3IlflCMzZJhDbbTBXjij5jGnJkoF3PerRw,11658
|
| 60 |
ornith_g256/strict_qwen/qwen3_contract.py,sha256=KcCDZZxUyBFE_bbIMpx1qqYs9p0cRfp9qL3u7KSFbq8,8719
|
| 61 |
-
ornith_g256/worker.py,sha256=
|
| 62 |
ornith_g256/worker_base.py,sha256=oA7xKVFy5lXIs8fXkM_QMWEUqUkuSD3t_eoDbTt7-Gg,4601
|
| 63 |
-
|
|
|
|
|
|
| 1 |
+
TOOL-INTEGRATION.json,sha256=vBFvgYo5fz-IPpIJ1ISJmaW-UfmWif8qQGPH8iGkZNk,18419
|
| 2 |
+
ciru_ornith_g256-4.0.0.dist-info/INSTALLER,sha256=5hhM4Q4mYTT9z6QB6PGpUAW81PGNFrYrdXMj4oM_6ak,2
|
| 3 |
+
ciru_ornith_g256-4.0.0.dist-info/METADATA,sha256=wDaK9fhKQVhFOwacWUkaqv6EOGv6W9qjkBqKOB1kh4E,254
|
| 4 |
+
ciru_ornith_g256-4.0.0.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
| 5 |
+
ciru_ornith_g256-4.0.0.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
|
| 6 |
+
ciru_ornith_g256-4.0.0.dist-info/entry_points.txt,sha256=ESe1wopdkiGfr5DNWZcSemlgPr3uvRre1OyHkl88QqU,121
|
| 7 |
+
ciru_ornith_g256-4.0.0.dist-info/licenses/LICENSE-APACHE-2.0,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
|
| 8 |
+
ciru_ornith_g256-4.0.0.dist-info/top_level.txt,sha256=EdCMLXnn8tDNBUR3JhkMvOi8hEw6IuhYFYEcfdhekxU,12
|
| 9 |
+
ornith_g256/__init__.py,sha256=YKu7UaR1SH46vLnls02INrTUs24LztBi0jgtx7EHbPY,1571
|
|
|
|
|
|
|
| 10 |
ornith_g256/_vllm_correctness/__init__.py,sha256=ljCXEGQHqHtisnUEFRYUTeFXE97mUQiszCnicDb_wFk,79
|
| 11 |
ornith_g256/_vllm_correctness/causal_conv1d.py,sha256=Iw8-VU9fHWCaNwIleMhOQtXHHUbmhBASeKfiGvT31Bg,51774
|
| 12 |
+
ornith_g256/_vllm_correctness/chat_completion_serving_r03.py,sha256=VyA3H1vLNR40Cj6NmgI17AH_J-2T5PNVQuMbg6BV70U,57061
|
| 13 |
ornith_g256/_vllm_correctness/gdn_attn.py,sha256=o8Nh1QLovKKqRqC67V7HU6PGLj0CeimOPtRpJ5T59uw,26224
|
| 14 |
ornith_g256/_vllm_correctness/gpu_model_runner.py,sha256=7wAJrKpnoaXgZphPNL9hmwboUiD_lTo9lRE3bwFG3M8,345523
|
| 15 |
+
ornith_g256/_vllm_correctness/manifest.json,sha256=ve9MpxUael4EEN0-i80uioXxJZOOwQRFHLh3cmfDxNk,1963
|
| 16 |
ornith_g256/_vllm_correctness/qwen_gdn_linear_attn.py,sha256=Plj93cegtCrj2bXquzdamTpBELrUmcN4Rr3kZpAdvrM,78334
|
| 17 |
+
ornith_g256/_xgrammar_native/build-receipt.json,sha256=O993ZX5RV8NZOLZt-0YmZgit47WZpiznSrs_gH1XnoQ,1654
|
| 18 |
+
ornith_g256/_xgrammar_native/libxgrammar-candidate.so,sha256=kqs2v64gxsTfP1MhyjdSBHmayDMlKSHZ2PaiBkkdNfc,60211248
|
| 19 |
+
ornith_g256/_xgrammar_native/selection.json,sha256=fLk7P-kC3m1gS_lvef3D772NYS5JTag49qjw1zHria4,106
|
| 20 |
ornith_g256/adaptive_c1.py,sha256=jkj9h0-99lGDwgvMfHSNJWLbyvqX4ATxiigYuYmAqAk,14985
|
| 21 |
ornith_g256/attention.py,sha256=L908DthNywUqno-d9AaN34CpaZMb6gopXgItNrdTzaE,772
|
| 22 |
+
ornith_g256/attention_compact.py,sha256=UN9hoCo4kKKCWyBe0Yx9GKKwAjM9OCxXJEC1jKmyuf4,10459
|
| 23 |
ornith_g256/attention_fast.py,sha256=J7GhM8rbh_86Y8AnBJ4dqB5BEv1dUmCkRA_ehUkfg4o,12274
|
| 24 |
ornith_g256/attention_folded.py,sha256=c9rIIDAyvObO7fIT1-7-WYfJIXaK4jcSniMv79jZu6M,9570
|
| 25 |
+
ornith_g256/attention_iu4.py,sha256=FcU0H8hIxDrJVlRuVjz4q392SfBaZ0-wFTsOLXZDjK4,3337
|
| 26 |
+
ornith_g256/attention_iu4_persistent.py,sha256=zvdK5COXQGMhJjeknNZKWhC5B386XiKHdGcidJoeYpc,12418
|
| 27 |
ornith_g256/attention_mixed.py,sha256=y7UwaX8lVb5lMrfU00AfOjPExYmCZZ6IiScg0haREQo,4138
|
| 28 |
ornith_g256/attention_partition.py,sha256=Qb-0TdDQxGosCnJZgxVQGOBRzYntuFJ_EZWgkr73_90,8642
|
| 29 |
ornith_g256/attention_storage.py,sha256=OrgT4iGDU3C9tQOyq31I9GKURBsd2cZSxdyXZDi9tuI,4117
|
|
|
|
| 51 |
ornith_g256/native.py,sha256=Sczg1nmHuWWf61Vqr7JHQjSQwhU7fXe3UPBVpP3NozE,12196
|
| 52 |
ornith_g256/native_tool_schema.py,sha256=kQBgaJ1ytHrUp7seUe8wHkH_C-G-ucxrPfAcQj0u3vU,1897
|
| 53 |
ornith_g256/no_spec.py,sha256=VhEGt1Mo_g4WH6EkOspM9QxQSdjnhdIskKrLPh4aAQQ,2618
|
| 54 |
+
ornith_g256/optimized/__init__.py,sha256=Apf11Hav3dsFVvTyNvleUU22ublniN9nKkoLwAdP2DU,45
|
| 55 |
+
ornith_g256/optimized/compile_scope.py,sha256=GJcac7OlrYDmIS3n9-QZod0A63OApebIVZ9axvXsjU0,1334
|
| 56 |
+
ornith_g256/optimized/draft.py,sha256=6AEwHPA9aZmHWNNftUY3VMzDChjgMO0rmIsmzvQ_Qpw,4327
|
| 57 |
+
ornith_g256/optimized/draft_kernel.py,sha256=PW2m7wwURQX4iMvTW6ULa5WtjkjDdqtzKXdGv_IokZ4,2569
|
| 58 |
+
ornith_g256/optimized/grammar_parallel.py,sha256=6cNghDhxJSGYDX_NYPWIYtCwoodmIBg_MExkoDgX7eo,5431
|
| 59 |
+
ornith_g256/optimized/greedy.py,sha256=OGRDrYcp-P4h7F_vhOrTrPvU1sXdO0weUKeZXtHRdh8,2212
|
| 60 |
+
ornith_g256/optimized/greedy_kernel.py,sha256=Lr4qJdEBGkXup-2jwJZMoesmON99YPVNsCyRQ1VQ3EI,1930
|
| 61 |
+
ornith_g256/optimized/launch.py,sha256=wdQmAhVGdJFqlLmQe3dG-sPqSMxQX6abq_EW8PEn4D0,435
|
| 62 |
+
ornith_g256/optimized/mixed_qk_sum.py,sha256=F7GHaa6EwulwegJXGB5cMU6XdPw7R5b45nh6gwLLXh8,2761
|
| 63 |
+
ornith_g256/optimized/pipeline/SOURCE-PROOF.json,sha256=WTm8ilOoMXFdBdD6ApC-qgywDtoVl_478fHPxTaZ_Q0,1709
|
| 64 |
+
ornith_g256/optimized/pipeline/__init__.py,sha256=Apf11Hav3dsFVvTyNvleUU22ublniN9nKkoLwAdP2DU,45
|
| 65 |
+
ornith_g256/optimized/pipeline/adapter.py,sha256=avl6jV4Nebffwy8vD-X1xTtcHFDQNUL_mcZMrn85CeE,9092
|
| 66 |
+
ornith_g256/optimized/pipeline/candidate/causal_conv1d.py,sha256=JUxUD3ewFpRpahq1NavQ5aI9VOchHiAs-b7q0E-s4Ws,51774
|
| 67 |
+
ornith_g256/optimized/pipeline/candidate/chunk_o.py,sha256=qT-VEEiVgmW00cP55nNBMpGW56V1JUM5NCFf3ksW5hs,5392
|
| 68 |
+
ornith_g256/optimized/pipeline/candidate/chunk_scaled_dot_kkt.py,sha256=cXNdKxxKGe7LtgWHCEgYF7thBCbEj5al9kZ4Svzwxlw,5720
|
| 69 |
+
ornith_g256/optimized/pipeline/candidate/solve_tril.py,sha256=vdBLacD5FWIEjSWDS9U9KY3B8b0Gi5SUr4h3ayv0sdE,19567
|
| 70 |
+
ornith_g256/optimized/pipeline/candidate/wy_fast.py,sha256=TMC-UZv884EtgQpL67PdRfubE6jQq2LQrUY6Gi429h8,4424
|
| 71 |
+
ornith_g256/optimized/postconv.py,sha256=xY7l0rhBj79p7aXZ7k80KxHMASq3euJbbOKyrqUapQ8,7935
|
| 72 |
+
ornith_g256/optimized/prefix.py,sha256=xpWomasStYw7ibbswnboLvo_CDOjGZEahSjMTiN3EM0,2353
|
| 73 |
+
ornith_g256/optimized/qk.py,sha256=2G8alAhmaKxkAursehnArY5V3I5jd_jUqe6raqlvlmQ,4957
|
| 74 |
+
ornith_g256/optimized/qk_schedule.py,sha256=ZgOHgfBmdIb7pFjvMRjKia2WQNwsK9Bc-FdkQPDxeKI,5793
|
| 75 |
+
ornith_g256/optimized/worker.py,sha256=OFebdnI15_ysPy5rBQdDjxVcE04ehfPRwgWkO-zpVRI,2988
|
| 76 |
ornith_g256/parser_contract.py,sha256=LOurQD3IglafPOZ9EnmN8Uzr1cFixy-Cvs03vN1-Kig,4203
|
| 77 |
ornith_g256/phase_dispatch.py,sha256=Gk698q9wnmylGeCgj2pTdgLfF4yldTMgRWF2MWQPgTc,2184
|
| 78 |
ornith_g256/prefill_draft.py,sha256=15uymhJGOVQ6ZOGYkI4urvTLz0kvLl4U3rz8nIZ4FvU,5630
|
|
|
|
| 82 |
ornith_g256/strict_qwen/__init__.py,sha256=7opQVSSfFBP8IvRkIat3tUhao5sZnGjrW1bHXcxAki0,81
|
| 83 |
ornith_g256/strict_qwen/qwen3.py,sha256=M8ZTRHoUS3IlflCMzZJhDbbTBXjij5jGnJkoF3PerRw,11658
|
| 84 |
ornith_g256/strict_qwen/qwen3_contract.py,sha256=KcCDZZxUyBFE_bbIMpx1qqYs9p0cRfp9qL3u7KSFbq8,8719
|
| 85 |
+
ornith_g256/worker.py,sha256=rMJyQMIfNHjKRS3fBCRN9O53DvqsOBku9k0yjvInAdU,13579
|
| 86 |
ornith_g256/worker_base.py,sha256=oA7xKVFy5lXIs8fXkM_QMWEUqUkuSD3t_eoDbTt7-Gg,4601
|
| 87 |
+
ornith_g256/xgrammar_native.py,sha256=0ZxUbhJMsRMR5Mwx9YM9xmPieJ1FCQCRwPj-yKliZiM,2059
|
| 88 |
+
ciru_ornith_g256-4.0.0.dist-info/RECORD,,
|
bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/REQUESTED
RENAMED
|
File without changes
|
bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/WHEEL
RENAMED
|
File without changes
|
bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/entry_points.txt
RENAMED
|
File without changes
|
bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/licenses/LICENSE-APACHE-2.0
RENAMED
|
File without changes
|
bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/top_level.txt
RENAMED
|
File without changes
|
bundle/plugin-site/ornith_g256/__init__.py
CHANGED
|
@@ -1,6 +1,9 @@
|
|
| 1 |
"""Ciru G256 prototype; no installed vLLM files are modified."""
|
| 2 |
|
| 3 |
|
|
|
|
|
|
|
|
|
|
| 4 |
from .runtime_correctness import install as _install_correctness
|
| 5 |
_install_correctness()
|
| 6 |
|
|
@@ -30,3 +33,6 @@ def register():
|
|
| 30 |
|
| 31 |
from .adaptive_c1 import install as install_adaptive_c1
|
| 32 |
install_adaptive_c1()
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
"""Ciru G256 prototype; no installed vLLM files are modified."""
|
| 2 |
|
| 3 |
|
| 4 |
+
from .xgrammar_native import install as _install_native_grammar
|
| 5 |
+
_install_native_grammar()
|
| 6 |
+
|
| 7 |
from .runtime_correctness import install as _install_correctness
|
| 8 |
_install_correctness()
|
| 9 |
|
|
|
|
| 33 |
|
| 34 |
from .adaptive_c1 import install as install_adaptive_c1
|
| 35 |
install_adaptive_c1()
|
| 36 |
+
|
| 37 |
+
from .optimized.grammar_parallel import install as install_parallel_grammar
|
| 38 |
+
install_parallel_grammar()
|
bundle/plugin-site/ornith_g256/_vllm_correctness/chat_completion_serving_r03.py
ADDED
|
@@ -0,0 +1,1319 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
| 3 |
+
|
| 4 |
+
import asyncio
|
| 5 |
+
import time
|
| 6 |
+
from collections.abc import AsyncGenerator, AsyncIterator
|
| 7 |
+
from collections.abc import Sequence as GenericSequence
|
| 8 |
+
from http import HTTPStatus
|
| 9 |
+
from typing import Any, Final, cast
|
| 10 |
+
|
| 11 |
+
from fastapi import Request
|
| 12 |
+
|
| 13 |
+
from vllm.engine.protocol import EngineClient
|
| 14 |
+
from vllm.entrypoints.chat_utils import (
|
| 15 |
+
ChatTemplateContentFormatOption,
|
| 16 |
+
ConversationMessage,
|
| 17 |
+
make_tool_call_id,
|
| 18 |
+
)
|
| 19 |
+
from vllm.entrypoints.generate.base.serving import (
|
| 20 |
+
GenerateBaseServing,
|
| 21 |
+
GenerationError,
|
| 22 |
+
build_per_request_timing_metrics,
|
| 23 |
+
build_spec_decoding_metrics,
|
| 24 |
+
clamp_prompt_logprobs,
|
| 25 |
+
format_token_id_placeholder,
|
| 26 |
+
)
|
| 27 |
+
from vllm.entrypoints.openai.chat_completion.protocol import (
|
| 28 |
+
ChatCompletionLogProb,
|
| 29 |
+
ChatCompletionLogProbs,
|
| 30 |
+
ChatCompletionLogProbsContent,
|
| 31 |
+
ChatCompletionNamedToolChoiceParam,
|
| 32 |
+
ChatCompletionRequest,
|
| 33 |
+
ChatCompletionResponse,
|
| 34 |
+
ChatCompletionResponseChoice,
|
| 35 |
+
ChatCompletionResponseStreamChoice,
|
| 36 |
+
ChatCompletionStreamResponse,
|
| 37 |
+
ChatMessage,
|
| 38 |
+
)
|
| 39 |
+
from vllm.entrypoints.openai.engine.protocol import (
|
| 40 |
+
CompletionTokenUsageInfo,
|
| 41 |
+
DeltaMessage,
|
| 42 |
+
ErrorResponse,
|
| 43 |
+
FunctionCall,
|
| 44 |
+
PerRequestMetrics,
|
| 45 |
+
PromptTokenUsageInfo,
|
| 46 |
+
RequestResponseMetadata,
|
| 47 |
+
ToolCall,
|
| 48 |
+
UsageInfo,
|
| 49 |
+
)
|
| 50 |
+
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
|
| 51 |
+
from vllm.entrypoints.serve.utils.api_utils import get_max_tokens, should_include_usage
|
| 52 |
+
from vllm.entrypoints.serve.utils.request_logger import RequestLogger
|
| 53 |
+
from vllm.entrypoints.serve.utils.tool_calls_utils import (
|
| 54 |
+
maybe_filter_parallel_tool_calls,
|
| 55 |
+
)
|
| 56 |
+
from vllm.inputs import EngineInput, MultiModalPlaceholders
|
| 57 |
+
from vllm.logger import init_logger
|
| 58 |
+
from vllm.logprobs import Logprob
|
| 59 |
+
from vllm.outputs import RequestOutput
|
| 60 |
+
from vllm.parser import ParserManager
|
| 61 |
+
from vllm.parser.abstract_parser import Parser
|
| 62 |
+
from vllm.renderers.online_renderer import OnlineRenderer
|
| 63 |
+
from vllm.sampling_params import BeamSearchParams, SamplingParams
|
| 64 |
+
from vllm.tokenizers import TokenizerLike
|
| 65 |
+
from vllm.utils.collection_utils import as_list
|
| 66 |
+
from vllm.utils.serial_utils import numpy2base64
|
| 67 |
+
|
| 68 |
+
logger = init_logger(__name__)
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
def _get_mm_token_counts(engine_input: EngineInput) -> dict[str, int]:
|
| 72 |
+
"""Sum per-modality placeholder tokens from ``mm_placeholders``.
|
| 73 |
+
|
| 74 |
+
Keyed by modality name; ``PlaceholderRange.length`` is the placeholder's
|
| 75 |
+
prompt token span, so each sum matches the placeholder tokens already
|
| 76 |
+
counted in ``usage.prompt_tokens``.
|
| 77 |
+
"""
|
| 78 |
+
mm_placeholders = cast(
|
| 79 |
+
MultiModalPlaceholders | None, engine_input.get("mm_placeholders")
|
| 80 |
+
)
|
| 81 |
+
return {
|
| 82 |
+
modality: sum(p.length for p in ranges)
|
| 83 |
+
for modality, ranges in (mm_placeholders or {}).items()
|
| 84 |
+
if ranges
|
| 85 |
+
}
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
def _make_prompt_tokens_details(
|
| 89 |
+
enable_prompt_tokens_details: bool,
|
| 90 |
+
num_cached_tokens: int | None,
|
| 91 |
+
num_cache_creation_tokens: int | None,
|
| 92 |
+
mm_token_counts: dict[str, int] | None,
|
| 93 |
+
) -> PromptTokenUsageInfo | None:
|
| 94 |
+
"""Build ``prompt_tokens_details`` from cached + multimodal token counts."""
|
| 95 |
+
if not enable_prompt_tokens_details:
|
| 96 |
+
return None
|
| 97 |
+
if (
|
| 98 |
+
num_cached_tokens is None
|
| 99 |
+
and num_cache_creation_tokens is None
|
| 100 |
+
and not mm_token_counts
|
| 101 |
+
):
|
| 102 |
+
return None
|
| 103 |
+
return PromptTokenUsageInfo(
|
| 104 |
+
cached_tokens=num_cached_tokens,
|
| 105 |
+
created_cache_tokens=num_cache_creation_tokens,
|
| 106 |
+
multimodal_tokens=mm_token_counts or None,
|
| 107 |
+
)
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
def _make_completion_tokens_details(
|
| 111 |
+
reasoning_tokens: int,
|
| 112 |
+
) -> CompletionTokenUsageInfo:
|
| 113 |
+
return CompletionTokenUsageInfo(reasoning_tokens=reasoning_tokens)
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
class OpenAIServingChat(GenerateBaseServing):
|
| 117 |
+
def __init__(
|
| 118 |
+
self,
|
| 119 |
+
engine_client: EngineClient,
|
| 120 |
+
models: OpenAIServingModels,
|
| 121 |
+
response_role: str,
|
| 122 |
+
*,
|
| 123 |
+
online_renderer: "OnlineRenderer",
|
| 124 |
+
request_logger: RequestLogger | None,
|
| 125 |
+
chat_template: str | None,
|
| 126 |
+
chat_template_content_format: ChatTemplateContentFormatOption,
|
| 127 |
+
trust_request_chat_template: bool = False,
|
| 128 |
+
return_tokens_as_token_ids: bool = False,
|
| 129 |
+
reasoning_parser: str = "",
|
| 130 |
+
enable_auto_tools: bool = False,
|
| 131 |
+
exclude_tools_when_tool_choice_none: bool = False,
|
| 132 |
+
tool_parser: str | None = None,
|
| 133 |
+
enable_prompt_tokens_details: bool = False,
|
| 134 |
+
enable_force_include_usage: bool = False,
|
| 135 |
+
enable_log_outputs: bool = False,
|
| 136 |
+
enable_log_deltas: bool = True,
|
| 137 |
+
default_chat_template_kwargs: dict[str, Any] | None = None,
|
| 138 |
+
enable_per_request_metrics: bool = False,
|
| 139 |
+
) -> None:
|
| 140 |
+
super().__init__(
|
| 141 |
+
engine_client=engine_client,
|
| 142 |
+
models=models,
|
| 143 |
+
request_logger=request_logger,
|
| 144 |
+
return_tokens_as_token_ids=return_tokens_as_token_ids,
|
| 145 |
+
)
|
| 146 |
+
|
| 147 |
+
self.online_renderer = online_renderer
|
| 148 |
+
self.response_role = response_role
|
| 149 |
+
self.chat_template = chat_template
|
| 150 |
+
self.chat_template_content_format: Final = chat_template_content_format
|
| 151 |
+
self.trust_request_chat_template = trust_request_chat_template
|
| 152 |
+
self.default_chat_template_kwargs = default_chat_template_kwargs or {}
|
| 153 |
+
self.enable_log_outputs = enable_log_outputs
|
| 154 |
+
self.enable_log_deltas = enable_log_deltas
|
| 155 |
+
|
| 156 |
+
self.enable_auto_tools: bool = enable_auto_tools
|
| 157 |
+
self._include_reasoning_tokens_details = bool(reasoning_parser)
|
| 158 |
+
self.parser_cls = ParserManager.get_parser(
|
| 159 |
+
tool_parser_name=tool_parser,
|
| 160 |
+
reasoning_parser_name=reasoning_parser,
|
| 161 |
+
enable_auto_tools=enable_auto_tools,
|
| 162 |
+
model_name=self.model_config.model,
|
| 163 |
+
is_harmony=self.model_config.hf_config.model_type == "gpt_oss",
|
| 164 |
+
)
|
| 165 |
+
self.exclude_tools_when_tool_choice_none = exclude_tools_when_tool_choice_none
|
| 166 |
+
|
| 167 |
+
self.enable_prompt_tokens_details = enable_prompt_tokens_details
|
| 168 |
+
self.enable_force_include_usage = enable_force_include_usage
|
| 169 |
+
self.enable_per_request_metrics = enable_per_request_metrics
|
| 170 |
+
self.default_sampling_params = self.model_config.get_diff_sampling_param()
|
| 171 |
+
mc = self.model_config
|
| 172 |
+
self.override_max_tokens = (
|
| 173 |
+
self.default_sampling_params.get("max_tokens")
|
| 174 |
+
if mc.generation_config not in ("auto", "vllm")
|
| 175 |
+
else getattr(mc, "override_generation_config", {}).get("max_new_tokens")
|
| 176 |
+
)
|
| 177 |
+
# NOTE(woosuk): While OpenAI's chat completion API supports browsing
|
| 178 |
+
# for some models, currently vLLM doesn't support it. Please use the
|
| 179 |
+
# Responses API instead.
|
| 180 |
+
self.supports_browsing = False
|
| 181 |
+
self.browser_tool = None
|
| 182 |
+
# NOTE(woosuk): Chat completion API does not support code interpreter.
|
| 183 |
+
# Please use the Responses API instead.
|
| 184 |
+
self.supports_code_interpreter = False
|
| 185 |
+
self.python_tool = None
|
| 186 |
+
|
| 187 |
+
def _effective_chat_template_kwargs(
|
| 188 |
+
self, request: ChatCompletionRequest
|
| 189 |
+
) -> dict[str, Any]:
|
| 190 |
+
return (
|
| 191 |
+
request.build_chat_params(
|
| 192 |
+
self.chat_template,
|
| 193 |
+
self.chat_template_content_format,
|
| 194 |
+
)
|
| 195 |
+
.with_defaults(self.default_chat_template_kwargs)
|
| 196 |
+
.chat_template_kwargs
|
| 197 |
+
)
|
| 198 |
+
|
| 199 |
+
def _engine_chat_template_kwargs(
|
| 200 |
+
self, chat_template_kwargs: dict[str, Any]
|
| 201 |
+
) -> dict[str, Any]:
|
| 202 |
+
"""Subclass hook to narrow ``chat_template_kwargs`` for the engine.
|
| 203 |
+
|
| 204 |
+
The same dict is used twice: to build the API-server-side parser
|
| 205 |
+
instance, and to populate
|
| 206 |
+
``EngineCoreRequest.reasoning_parser_kwargs`` for the engine-core
|
| 207 |
+
side. The latter crosses ZMQ as msgpack, so a handler that stashes
|
| 208 |
+
request-scoped state only the API-server-side parser needs (values
|
| 209 |
+
msgpack can't encode, or payloads not worth shipping) can drop
|
| 210 |
+
those entries here without affecting the in-process parser.
|
| 211 |
+
|
| 212 |
+
Must not mutate the argument -- the caller still needs the full
|
| 213 |
+
dict. The default forwards it unchanged.
|
| 214 |
+
"""
|
| 215 |
+
return chat_template_kwargs
|
| 216 |
+
|
| 217 |
+
async def render_chat_request(
|
| 218 |
+
self,
|
| 219 |
+
request: ChatCompletionRequest,
|
| 220 |
+
) -> tuple[list[ConversationMessage], list[EngineInput]] | ErrorResponse:
|
| 221 |
+
"""
|
| 222 |
+
Validate the model and preprocess a chat completion request.
|
| 223 |
+
|
| 224 |
+
Delegates preprocessing logic to OnlineRenderer, adding the
|
| 225 |
+
engine-aware checks (LoRA model validation, engine health).
|
| 226 |
+
|
| 227 |
+
Returns:
|
| 228 |
+
A tuple of (conversation, engine_inputs) on success,
|
| 229 |
+
or an ErrorResponse on failure.
|
| 230 |
+
"""
|
| 231 |
+
error_check_ret = await self._check_model(request)
|
| 232 |
+
if error_check_ret is not None:
|
| 233 |
+
logger.error("Error with model %s", error_check_ret)
|
| 234 |
+
return error_check_ret
|
| 235 |
+
|
| 236 |
+
# If the engine is dead, raise the engine's DEAD_ERROR.
|
| 237 |
+
# This is required for the streaming case, where we return a
|
| 238 |
+
# success status before we actually start generating text :).
|
| 239 |
+
if self.engine_client.errored:
|
| 240 |
+
raise self.engine_client.dead_error
|
| 241 |
+
|
| 242 |
+
return await self.online_renderer.render_chat(request)
|
| 243 |
+
|
| 244 |
+
async def create_chat_completion(
|
| 245 |
+
self,
|
| 246 |
+
request: ChatCompletionRequest,
|
| 247 |
+
raw_request: Request | None = None,
|
| 248 |
+
) -> AsyncGenerator[str, None] | ChatCompletionResponse | ErrorResponse:
|
| 249 |
+
"""
|
| 250 |
+
Chat Completion API similar to OpenAI's API.
|
| 251 |
+
|
| 252 |
+
See https://platform.openai.com/docs/api-reference/chat/create
|
| 253 |
+
for the API specification. This API mimics the OpenAI
|
| 254 |
+
Chat Completion API.
|
| 255 |
+
"""
|
| 256 |
+
return await self._with_kv_transfer_rejection_cleanup(
|
| 257 |
+
self._create_chat_completion(request, raw_request), request, raw_request
|
| 258 |
+
)
|
| 259 |
+
|
| 260 |
+
async def _create_chat_completion(
|
| 261 |
+
self,
|
| 262 |
+
request: ChatCompletionRequest,
|
| 263 |
+
raw_request: Request | None = None,
|
| 264 |
+
) -> AsyncGenerator[str, None] | ChatCompletionResponse | ErrorResponse:
|
| 265 |
+
# Streaming response
|
| 266 |
+
tokenizer = self.renderer.tokenizer
|
| 267 |
+
assert tokenizer is not None
|
| 268 |
+
chat_template_kwargs = self._effective_chat_template_kwargs(request)
|
| 269 |
+
parser: Parser | None = None
|
| 270 |
+
if self.parser_cls is not None:
|
| 271 |
+
parser = self.parser_cls(
|
| 272 |
+
tokenizer,
|
| 273 |
+
request.tools,
|
| 274 |
+
chat_template_kwargs=chat_template_kwargs,
|
| 275 |
+
model_config=self.model_config,
|
| 276 |
+
)
|
| 277 |
+
result = await self.render_chat_request(request)
|
| 278 |
+
if isinstance(result, ErrorResponse):
|
| 279 |
+
return result
|
| 280 |
+
|
| 281 |
+
conversation, engine_inputs = result
|
| 282 |
+
|
| 283 |
+
request_id = (
|
| 284 |
+
f"chatcmpl-{self._base_request_id(raw_request, request.request_id)}"
|
| 285 |
+
)
|
| 286 |
+
|
| 287 |
+
request_metadata = RequestResponseMetadata(request_id=request_id)
|
| 288 |
+
if raw_request:
|
| 289 |
+
raw_request.state.request_metadata = request_metadata
|
| 290 |
+
|
| 291 |
+
lora_request = self._maybe_get_adapters(request, supports_default_mm_loras=True)
|
| 292 |
+
|
| 293 |
+
model_name = self.models.model_name(lora_request)
|
| 294 |
+
|
| 295 |
+
# Extract data_parallel_rank from header (router can inject it)
|
| 296 |
+
data_parallel_rank = self._get_data_parallel_rank(raw_request)
|
| 297 |
+
|
| 298 |
+
# Schedule the request and get the result generator.
|
| 299 |
+
max_model_len = self.model_config.max_model_len
|
| 300 |
+
generators: list[AsyncGenerator[RequestOutput, None]] = []
|
| 301 |
+
mm_token_counts: dict[str, int] | None = None
|
| 302 |
+
for i, engine_input in enumerate(engine_inputs):
|
| 303 |
+
prompt_token_ids = self._extract_prompt_components(engine_input).token_ids
|
| 304 |
+
mm_token_counts = _get_mm_token_counts(engine_input)
|
| 305 |
+
|
| 306 |
+
# If we are creating sub requests for multiple prompts, ensure that they
|
| 307 |
+
# have unique request ids.
|
| 308 |
+
sub_request_id = (
|
| 309 |
+
request_id if len(engine_inputs) == 1 else f"{request_id}_{i}"
|
| 310 |
+
)
|
| 311 |
+
|
| 312 |
+
max_tokens = get_max_tokens(
|
| 313 |
+
max_model_len,
|
| 314 |
+
request.max_completion_tokens
|
| 315 |
+
if request.max_completion_tokens is not None
|
| 316 |
+
else request.max_tokens,
|
| 317 |
+
self._extract_prompt_len(engine_input),
|
| 318 |
+
self.default_sampling_params,
|
| 319 |
+
self.override_max_tokens,
|
| 320 |
+
truncate_prompt_tokens=request.truncate_prompt_tokens,
|
| 321 |
+
)
|
| 322 |
+
|
| 323 |
+
sampling_params: SamplingParams | BeamSearchParams
|
| 324 |
+
if request.use_beam_search:
|
| 325 |
+
sampling_params = request.to_beam_search_params(
|
| 326 |
+
max_tokens, self.default_sampling_params
|
| 327 |
+
)
|
| 328 |
+
else:
|
| 329 |
+
sampling_params = request.to_sampling_params(
|
| 330 |
+
max_tokens,
|
| 331 |
+
self.default_sampling_params,
|
| 332 |
+
)
|
| 333 |
+
|
| 334 |
+
self._log_inputs(
|
| 335 |
+
sub_request_id,
|
| 336 |
+
engine_input,
|
| 337 |
+
params=sampling_params,
|
| 338 |
+
lora_request=lora_request,
|
| 339 |
+
)
|
| 340 |
+
|
| 341 |
+
trace_headers = (
|
| 342 |
+
None
|
| 343 |
+
if raw_request is None
|
| 344 |
+
else await self._get_trace_headers(raw_request.headers)
|
| 345 |
+
)
|
| 346 |
+
session_id = self._get_session_id(request, raw_request)
|
| 347 |
+
|
| 348 |
+
if isinstance(sampling_params, BeamSearchParams):
|
| 349 |
+
generator = self.beam_search(
|
| 350 |
+
prompt=engine_input,
|
| 351 |
+
request_id=sub_request_id,
|
| 352 |
+
params=sampling_params,
|
| 353 |
+
lora_request=lora_request,
|
| 354 |
+
trace_headers=trace_headers,
|
| 355 |
+
session_id=session_id,
|
| 356 |
+
)
|
| 357 |
+
else:
|
| 358 |
+
if not request.include_reasoning:
|
| 359 |
+
reasoning_ended = True
|
| 360 |
+
elif request._grammar_from_parser:
|
| 361 |
+
# The Mistral grammar already includes an optional
|
| 362 |
+
# `think?` rule that handles both reasoning and
|
| 363 |
+
# non-reasoning outputs.
|
| 364 |
+
reasoning_ended = True
|
| 365 |
+
elif parser is not None and parser.reasoning_parser is not None:
|
| 366 |
+
reasoning_ended = parser.is_reasoning_end(prompt_token_ids or [])
|
| 367 |
+
else:
|
| 368 |
+
reasoning_ended = None
|
| 369 |
+
|
| 370 |
+
generator = self.engine_client.generate(
|
| 371 |
+
engine_input,
|
| 372 |
+
sampling_params,
|
| 373 |
+
sub_request_id,
|
| 374 |
+
lora_request=lora_request,
|
| 375 |
+
trace_headers=trace_headers,
|
| 376 |
+
priority=self._get_priority(request, raw_request),
|
| 377 |
+
data_parallel_rank=data_parallel_rank,
|
| 378 |
+
session_id=session_id,
|
| 379 |
+
reasoning_ended=reasoning_ended,
|
| 380 |
+
reasoning_parser_kwargs={
|
| 381 |
+
"chat_template_kwargs": self._engine_chat_template_kwargs(
|
| 382 |
+
chat_template_kwargs
|
| 383 |
+
),
|
| 384 |
+
}
|
| 385 |
+
if parser is not None and parser.reasoning_parser is not None
|
| 386 |
+
else None,
|
| 387 |
+
)
|
| 388 |
+
|
| 389 |
+
generators.append(generator)
|
| 390 |
+
|
| 391 |
+
assert len(generators) == 1
|
| 392 |
+
(result_generator,) = generators
|
| 393 |
+
|
| 394 |
+
if request.stream:
|
| 395 |
+
return self.chat_completion_stream_generator(
|
| 396 |
+
request,
|
| 397 |
+
result_generator,
|
| 398 |
+
request_id,
|
| 399 |
+
model_name,
|
| 400 |
+
conversation,
|
| 401 |
+
tokenizer,
|
| 402 |
+
request_metadata,
|
| 403 |
+
chat_template_kwargs=chat_template_kwargs,
|
| 404 |
+
mm_token_counts=mm_token_counts,
|
| 405 |
+
)
|
| 406 |
+
|
| 407 |
+
return await self.chat_completion_full_generator(
|
| 408 |
+
request,
|
| 409 |
+
result_generator,
|
| 410 |
+
request_id,
|
| 411 |
+
model_name,
|
| 412 |
+
conversation,
|
| 413 |
+
tokenizer,
|
| 414 |
+
request_metadata,
|
| 415 |
+
parser=parser,
|
| 416 |
+
mm_token_counts=mm_token_counts,
|
| 417 |
+
)
|
| 418 |
+
|
| 419 |
+
def get_chat_request_role(self, request: ChatCompletionRequest) -> str:
|
| 420 |
+
if request.add_generation_prompt:
|
| 421 |
+
return self.response_role
|
| 422 |
+
return request.messages[-1]["role"]
|
| 423 |
+
|
| 424 |
+
def _create_chat_message(self, *args: Any, **kwargs: Any) -> ChatMessage:
|
| 425 |
+
"""Construct the response :class:`ChatMessage` for the non-streaming path.
|
| 426 |
+
|
| 427 |
+
The full-generator calls this at every construction site so
|
| 428 |
+
subclasses can swap in a specialized :class:`ChatMessage`
|
| 429 |
+
subclass (e.g. :class:`CohereServingChatV2` returning
|
| 430 |
+
:class:`CohereChatMessage`) without duplicating the branchy
|
| 431 |
+
tool-choice / auto-tools logic that decides which fields are
|
| 432 |
+
populated. The default returns a plain :class:`ChatMessage`.
|
| 433 |
+
"""
|
| 434 |
+
return ChatMessage(*args, **kwargs)
|
| 435 |
+
|
| 436 |
+
def _finalize_response_message(
|
| 437 |
+
self,
|
| 438 |
+
message: ChatMessage,
|
| 439 |
+
*,
|
| 440 |
+
parser: Parser | None,
|
| 441 |
+
) -> ChatMessage:
|
| 442 |
+
"""Subclass hook to enrich a fully-constructed :class:`ChatMessage`.
|
| 443 |
+
|
| 444 |
+
Default is a no-op. Subclasses that need to surface parser-side
|
| 445 |
+
extras (e.g. :class:`CohereServingChatV2` reading grounding
|
| 446 |
+
citations off the reasoning parser and populating
|
| 447 |
+
:class:`CohereChatMessage.citations`) override this to inspect
|
| 448 |
+
``parser`` and mutate/replace ``message``.
|
| 449 |
+
"""
|
| 450 |
+
return message
|
| 451 |
+
|
| 452 |
+
async def chat_completion_stream_generator(
|
| 453 |
+
self,
|
| 454 |
+
request: ChatCompletionRequest,
|
| 455 |
+
result_generator: AsyncIterator[RequestOutput],
|
| 456 |
+
request_id: str,
|
| 457 |
+
model_name: str,
|
| 458 |
+
conversation: list[ConversationMessage],
|
| 459 |
+
tokenizer: TokenizerLike,
|
| 460 |
+
request_metadata: RequestResponseMetadata,
|
| 461 |
+
chat_template_kwargs: dict[str, Any] | None = None,
|
| 462 |
+
mm_token_counts: dict[str, int] | None = None,
|
| 463 |
+
) -> AsyncGenerator[str, None]:
|
| 464 |
+
created_time = int(time.time())
|
| 465 |
+
chunk_object_type: Final = "chat.completion.chunk"
|
| 466 |
+
first_iteration = True
|
| 467 |
+
|
| 468 |
+
# Send response for each token for each request.n (index)
|
| 469 |
+
num_choices = 1 if request.n is None else request.n
|
| 470 |
+
previous_num_tokens = [0] * num_choices
|
| 471 |
+
# TODO: Remove once all reasoning parsers use the Parser Engine.
|
| 472 |
+
generated_token_ids: list[list[int]] = [[] for _ in range(num_choices)]
|
| 473 |
+
previous_reasoning_tokens = [0] * num_choices
|
| 474 |
+
finish_reason_sent = [False] * num_choices
|
| 475 |
+
num_prompt_tokens = 0
|
| 476 |
+
num_cached_tokens = None
|
| 477 |
+
num_cache_creation_tokens = None
|
| 478 |
+
tools_streamed = [False] * num_choices
|
| 479 |
+
|
| 480 |
+
if isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam):
|
| 481 |
+
tool_choice_function_name = request.tool_choice.function.name
|
| 482 |
+
else:
|
| 483 |
+
tool_choice_function_name = None
|
| 484 |
+
|
| 485 |
+
previous_texts = [""] * num_choices
|
| 486 |
+
|
| 487 |
+
try:
|
| 488 |
+
if self.parser_cls is not None:
|
| 489 |
+
if tokenizer is None:
|
| 490 |
+
raise ValueError(
|
| 491 |
+
"Tokenizer not available when `skip_tokenizer_init=True`"
|
| 492 |
+
)
|
| 493 |
+
parsers: list[Parser | None] = [
|
| 494 |
+
self.parser_cls(
|
| 495 |
+
tokenizer,
|
| 496 |
+
request.tools,
|
| 497 |
+
chat_template_kwargs=chat_template_kwargs,
|
| 498 |
+
model_config=self.model_config,
|
| 499 |
+
)
|
| 500 |
+
for _ in range(num_choices)
|
| 501 |
+
]
|
| 502 |
+
else:
|
| 503 |
+
parsers = [None] * num_choices
|
| 504 |
+
except Exception as e:
|
| 505 |
+
logger.exception("Error in parser creation.")
|
| 506 |
+
data = self.create_streaming_error_response(e)
|
| 507 |
+
yield f"data: {data}\n\n"
|
| 508 |
+
yield "data: [DONE]\n\n"
|
| 509 |
+
return
|
| 510 |
+
|
| 511 |
+
stream_options = request.stream_options
|
| 512 |
+
include_usage, include_continuous_usage = should_include_usage(
|
| 513 |
+
stream_options, self.enable_force_include_usage
|
| 514 |
+
)
|
| 515 |
+
|
| 516 |
+
last_res: RequestOutput | None = None
|
| 517 |
+
try:
|
| 518 |
+
async for res in result_generator:
|
| 519 |
+
last_res = res
|
| 520 |
+
if res.prompt_token_ids is not None:
|
| 521 |
+
num_prompt_tokens = len(res.prompt_token_ids)
|
| 522 |
+
if res.encoder_prompt_token_ids is not None:
|
| 523 |
+
num_prompt_tokens += len(res.encoder_prompt_token_ids)
|
| 524 |
+
|
| 525 |
+
# We need to do it here, because if there are exceptions in
|
| 526 |
+
# the result_generator, it needs to be sent as the FIRST
|
| 527 |
+
# response (by the try...catch).
|
| 528 |
+
if first_iteration:
|
| 529 |
+
num_cached_tokens = res.num_cached_tokens
|
| 530 |
+
num_cache_creation_tokens = res.num_cache_creation_tokens
|
| 531 |
+
# Send first response for each request.n (index) with
|
| 532 |
+
# the role
|
| 533 |
+
role = self.get_chat_request_role(request)
|
| 534 |
+
|
| 535 |
+
# ``res.prompt`` is the rendered chat-templated prompt
|
| 536 |
+
prompt_text = res.prompt if request.return_prompt_text else None
|
| 537 |
+
|
| 538 |
+
# NOTE num_choices defaults to 1 so this usually executes
|
| 539 |
+
# once per request
|
| 540 |
+
for i in range(num_choices):
|
| 541 |
+
choice_data = ChatCompletionResponseStreamChoice(
|
| 542 |
+
index=i,
|
| 543 |
+
delta=DeltaMessage(
|
| 544 |
+
role=role,
|
| 545 |
+
content="",
|
| 546 |
+
),
|
| 547 |
+
logprobs=None,
|
| 548 |
+
finish_reason=None,
|
| 549 |
+
)
|
| 550 |
+
|
| 551 |
+
# return prompt_token_ids at the first chunk ever
|
| 552 |
+
chunk = ChatCompletionStreamResponse(
|
| 553 |
+
id=request_id,
|
| 554 |
+
object=chunk_object_type,
|
| 555 |
+
created=created_time,
|
| 556 |
+
choices=[choice_data],
|
| 557 |
+
model=model_name,
|
| 558 |
+
prompt_token_ids=(
|
| 559 |
+
res.prompt_token_ids
|
| 560 |
+
if request.return_token_ids
|
| 561 |
+
else None
|
| 562 |
+
),
|
| 563 |
+
prompt_text=prompt_text,
|
| 564 |
+
)
|
| 565 |
+
|
| 566 |
+
# if continuous usage stats are requested, add it
|
| 567 |
+
if include_continuous_usage:
|
| 568 |
+
chunk.usage = UsageInfo(
|
| 569 |
+
prompt_tokens=num_prompt_tokens,
|
| 570 |
+
completion_tokens=0,
|
| 571 |
+
total_tokens=num_prompt_tokens,
|
| 572 |
+
completion_tokens_details=(
|
| 573 |
+
_make_completion_tokens_details(0)
|
| 574 |
+
if self._include_reasoning_tokens_details
|
| 575 |
+
else None
|
| 576 |
+
),
|
| 577 |
+
)
|
| 578 |
+
|
| 579 |
+
data = chunk.model_dump_json(exclude_unset=True)
|
| 580 |
+
yield f"data: {data}\n\n"
|
| 581 |
+
|
| 582 |
+
# Send response to echo the input portion of the
|
| 583 |
+
# last message
|
| 584 |
+
if request.echo:
|
| 585 |
+
last_msg_content: str | list[dict[str, str]] = ""
|
| 586 |
+
if (
|
| 587 |
+
conversation
|
| 588 |
+
and "content" in conversation[-1]
|
| 589 |
+
and conversation[-1].get("role") == role
|
| 590 |
+
):
|
| 591 |
+
last_msg_content = conversation[-1]["content"] or ""
|
| 592 |
+
|
| 593 |
+
if last_msg_content:
|
| 594 |
+
for i in range(num_choices):
|
| 595 |
+
choice_data = ChatCompletionResponseStreamChoice(
|
| 596 |
+
index=i,
|
| 597 |
+
delta=DeltaMessage(content=last_msg_content),
|
| 598 |
+
logprobs=None,
|
| 599 |
+
finish_reason=None,
|
| 600 |
+
)
|
| 601 |
+
chunk = ChatCompletionStreamResponse(
|
| 602 |
+
id=request_id,
|
| 603 |
+
object=chunk_object_type,
|
| 604 |
+
created=created_time,
|
| 605 |
+
choices=[choice_data],
|
| 606 |
+
model=model_name,
|
| 607 |
+
)
|
| 608 |
+
if include_continuous_usage:
|
| 609 |
+
chunk.usage = UsageInfo(
|
| 610 |
+
prompt_tokens=num_prompt_tokens,
|
| 611 |
+
completion_tokens=0,
|
| 612 |
+
total_tokens=num_prompt_tokens,
|
| 613 |
+
completion_tokens_details=(
|
| 614 |
+
_make_completion_tokens_details(0)
|
| 615 |
+
if self._include_reasoning_tokens_details
|
| 616 |
+
else None
|
| 617 |
+
),
|
| 618 |
+
)
|
| 619 |
+
|
| 620 |
+
data = chunk.model_dump_json(exclude_unset=True)
|
| 621 |
+
yield f"data: {data}\n\n"
|
| 622 |
+
first_iteration = False
|
| 623 |
+
|
| 624 |
+
for output in res.outputs:
|
| 625 |
+
i = output.index
|
| 626 |
+
parser = parsers[i]
|
| 627 |
+
if finish_reason_sent[i]:
|
| 628 |
+
continue
|
| 629 |
+
|
| 630 |
+
if request.logprobs and (
|
| 631 |
+
request.top_logprobs is not None or request.logprob_token_ids
|
| 632 |
+
):
|
| 633 |
+
assert output.logprobs is not None, "Did not output logprobs"
|
| 634 |
+
logprobs = self._create_chat_logprobs(
|
| 635 |
+
token_ids=output.token_ids,
|
| 636 |
+
top_logprobs=output.logprobs,
|
| 637 |
+
tokenizer=tokenizer,
|
| 638 |
+
num_output_top_logprobs=request.top_logprobs,
|
| 639 |
+
logprob_token_ids=request.logprob_token_ids,
|
| 640 |
+
return_as_token_id=request.return_tokens_as_token_ids,
|
| 641 |
+
)
|
| 642 |
+
else:
|
| 643 |
+
logprobs = None
|
| 644 |
+
|
| 645 |
+
delta_text = output.text
|
| 646 |
+
|
| 647 |
+
if (
|
| 648 |
+
not delta_text
|
| 649 |
+
and not output.token_ids
|
| 650 |
+
and not previous_num_tokens[i]
|
| 651 |
+
):
|
| 652 |
+
# Chunked prefill case, don't return empty chunks
|
| 653 |
+
continue
|
| 654 |
+
|
| 655 |
+
delta_message: DeltaMessage | None
|
| 656 |
+
|
| 657 |
+
if parser is not None:
|
| 658 |
+
delta_message = parser.parse_delta(
|
| 659 |
+
delta_text=delta_text,
|
| 660 |
+
delta_token_ids=as_list(output.token_ids),
|
| 661 |
+
request=request,
|
| 662 |
+
prompt_token_ids=res.prompt_token_ids,
|
| 663 |
+
finished=output.finish_reason is not None,
|
| 664 |
+
)
|
| 665 |
+
if delta_message is not None and delta_message.tool_calls:
|
| 666 |
+
tools_streamed[i] = True
|
| 667 |
+
|
| 668 |
+
# handle streaming just a content delta (no parsers)
|
| 669 |
+
else:
|
| 670 |
+
delta_message = DeltaMessage(content=delta_text)
|
| 671 |
+
|
| 672 |
+
previous_texts[i] += delta_text
|
| 673 |
+
|
| 674 |
+
# set the previous values for the next iteration
|
| 675 |
+
previous_num_tokens[i] += len(output.token_ids)
|
| 676 |
+
if parser is not None:
|
| 677 |
+
generated_token_ids[i].extend(output.token_ids)
|
| 678 |
+
previous_reasoning_tokens[i] = parser.count_reasoning_tokens(
|
| 679 |
+
tuple(generated_token_ids[i])
|
| 680 |
+
)
|
| 681 |
+
|
| 682 |
+
# if the message delta is None (e.g. because it was a
|
| 683 |
+
# "control token" for tool calls or the parser otherwise
|
| 684 |
+
# wasn't ready to send a token, then
|
| 685 |
+
# get the next token without streaming a chunk
|
| 686 |
+
# When reasoning is hidden, suppress per-token
|
| 687 |
+
# metadata (logprobs, token_ids) on every chunk to
|
| 688 |
+
# prevent leaking reasoning tokens through decoded
|
| 689 |
+
# token text in logprob entries or raw token IDs.
|
| 690 |
+
hide_stream_metadata = (
|
| 691 |
+
not request.include_reasoning and parser is not None
|
| 692 |
+
)
|
| 693 |
+
if hide_stream_metadata:
|
| 694 |
+
logprobs = None
|
| 695 |
+
|
| 696 |
+
if delta_message is None:
|
| 697 |
+
# NOTE: If return_token_ids is enabled, we still need to
|
| 698 |
+
# send a chunk with token_ids even if delta_message is None
|
| 699 |
+
# to ensure all tokens are included in the response
|
| 700 |
+
# A parser may buffer a whole JSON array/object.
|
| 701 |
+
# Real token progress is independent of argument deltas.
|
| 702 |
+
# Preserve requested usage without exposing partial calls
|
| 703 |
+
# or hidden reasoning metadata. Empty prefill stays quiet.
|
| 704 |
+
if (
|
| 705 |
+
output.finish_reason is None
|
| 706 |
+
and not (include_continuous_usage and output.token_ids)
|
| 707 |
+
and (not request.return_token_ids or hide_stream_metadata)
|
| 708 |
+
):
|
| 709 |
+
continue
|
| 710 |
+
delta_message = DeltaMessage()
|
| 711 |
+
|
| 712 |
+
# Log streaming delta if output logging is enabled
|
| 713 |
+
if self.enable_log_outputs and self.request_logger:
|
| 714 |
+
delta_content_parts = []
|
| 715 |
+
if delta_message.content:
|
| 716 |
+
delta_content_parts.append(delta_message.content)
|
| 717 |
+
if delta_message.reasoning:
|
| 718 |
+
reasoning = delta_message.reasoning
|
| 719 |
+
delta_content_parts.append(f"[reasoning: {reasoning}]")
|
| 720 |
+
if delta_message.tool_calls:
|
| 721 |
+
tool_args = "".join(
|
| 722 |
+
tc.function.arguments
|
| 723 |
+
for tc in delta_message.tool_calls
|
| 724 |
+
if tc.function and tc.function.arguments
|
| 725 |
+
)
|
| 726 |
+
if tool_args:
|
| 727 |
+
delta_content_parts.append(f"[tool_calls: {tool_args}]")
|
| 728 |
+
|
| 729 |
+
if delta_content_parts and self.enable_log_deltas:
|
| 730 |
+
delta_content = " ".join(delta_content_parts)
|
| 731 |
+
self.request_logger.log_outputs(
|
| 732 |
+
request_id=request_id,
|
| 733 |
+
outputs=delta_content,
|
| 734 |
+
output_token_ids=as_list(output.token_ids),
|
| 735 |
+
finish_reason=output.finish_reason,
|
| 736 |
+
is_streaming=True,
|
| 737 |
+
delta=True,
|
| 738 |
+
)
|
| 739 |
+
|
| 740 |
+
include_token_ids = (
|
| 741 |
+
request.return_token_ids and not hide_stream_metadata
|
| 742 |
+
)
|
| 743 |
+
|
| 744 |
+
if output.finish_reason is None:
|
| 745 |
+
# Send token-by-token response for each request.n
|
| 746 |
+
choice_data = ChatCompletionResponseStreamChoice(
|
| 747 |
+
index=i,
|
| 748 |
+
delta=delta_message,
|
| 749 |
+
logprobs=logprobs,
|
| 750 |
+
finish_reason=None,
|
| 751 |
+
token_ids=(
|
| 752 |
+
as_list(output.token_ids) if include_token_ids else None
|
| 753 |
+
),
|
| 754 |
+
)
|
| 755 |
+
|
| 756 |
+
# if the model is finished generating
|
| 757 |
+
else:
|
| 758 |
+
# check for error finish reason and abort streaming
|
| 759 |
+
# finish_reason='error' indicates a retryable error
|
| 760 |
+
self._raise_if_error(output.finish_reason, request_id)
|
| 761 |
+
|
| 762 |
+
# Send the finish response for each request.n only once
|
| 763 |
+
# In OpenAI's API, when a tool is called, the
|
| 764 |
+
# finish_reason is:
|
| 765 |
+
# "tool_calls" for "auto" or "required" tool calls,
|
| 766 |
+
# and "stop" for named tool calls.
|
| 767 |
+
if (
|
| 768 |
+
tools_streamed[i]
|
| 769 |
+
and not tool_choice_function_name
|
| 770 |
+
and output.finish_reason == "stop"
|
| 771 |
+
):
|
| 772 |
+
finish_reason_ = "tool_calls"
|
| 773 |
+
else:
|
| 774 |
+
finish_reason_ = (
|
| 775 |
+
output.finish_reason if output.finish_reason else "stop"
|
| 776 |
+
)
|
| 777 |
+
choice_data = ChatCompletionResponseStreamChoice(
|
| 778 |
+
index=i,
|
| 779 |
+
delta=delta_message,
|
| 780 |
+
logprobs=logprobs,
|
| 781 |
+
finish_reason=finish_reason_,
|
| 782 |
+
stop_reason=output.stop_reason,
|
| 783 |
+
token_ids=(
|
| 784 |
+
as_list(output.token_ids) if include_token_ids else None
|
| 785 |
+
),
|
| 786 |
+
)
|
| 787 |
+
|
| 788 |
+
finish_reason_sent[i] = True
|
| 789 |
+
|
| 790 |
+
choice_data = maybe_filter_parallel_tool_calls(choice_data, request)
|
| 791 |
+
chunk = ChatCompletionStreamResponse(
|
| 792 |
+
id=request_id,
|
| 793 |
+
object=chunk_object_type,
|
| 794 |
+
created=created_time,
|
| 795 |
+
choices=[choice_data],
|
| 796 |
+
model=model_name,
|
| 797 |
+
)
|
| 798 |
+
# Stamp the fingerprint on terminal chunks only (those with
|
| 799 |
+
# finish_reason set). When ``include_usage`` is on, the
|
| 800 |
+
# trailing usage chunk below overrides this as the true
|
| 801 |
+
# final message.
|
| 802 |
+
if (
|
| 803 |
+
not include_usage
|
| 804 |
+
and self.system_fingerprint is not None
|
| 805 |
+
and choice_data.finish_reason is not None
|
| 806 |
+
):
|
| 807 |
+
chunk.system_fingerprint = self.system_fingerprint
|
| 808 |
+
|
| 809 |
+
# handle usage stats if requested & if continuous
|
| 810 |
+
if include_continuous_usage:
|
| 811 |
+
completion_tokens = previous_num_tokens[i]
|
| 812 |
+
chunk.usage = UsageInfo(
|
| 813 |
+
prompt_tokens=num_prompt_tokens,
|
| 814 |
+
completion_tokens=completion_tokens,
|
| 815 |
+
total_tokens=num_prompt_tokens + completion_tokens,
|
| 816 |
+
completion_tokens_details=(
|
| 817 |
+
_make_completion_tokens_details(
|
| 818 |
+
previous_reasoning_tokens[i]
|
| 819 |
+
)
|
| 820 |
+
if self._include_reasoning_tokens_details
|
| 821 |
+
else None
|
| 822 |
+
),
|
| 823 |
+
)
|
| 824 |
+
|
| 825 |
+
data = chunk.model_dump_json(exclude_unset=True)
|
| 826 |
+
yield f"data: {data}\n\n"
|
| 827 |
+
|
| 828 |
+
# once the final token is handled, if stream_options.include_usage
|
| 829 |
+
# is sent, send the usage
|
| 830 |
+
if include_usage:
|
| 831 |
+
completion_tokens = sum(previous_num_tokens)
|
| 832 |
+
final_usage = UsageInfo(
|
| 833 |
+
prompt_tokens=num_prompt_tokens,
|
| 834 |
+
completion_tokens=completion_tokens,
|
| 835 |
+
total_tokens=num_prompt_tokens + completion_tokens,
|
| 836 |
+
completion_tokens_details=_make_completion_tokens_details(
|
| 837 |
+
sum(previous_reasoning_tokens)
|
| 838 |
+
)
|
| 839 |
+
if self._include_reasoning_tokens_details
|
| 840 |
+
else None,
|
| 841 |
+
)
|
| 842 |
+
final_usage.prompt_tokens_details = _make_prompt_tokens_details(
|
| 843 |
+
self.enable_prompt_tokens_details,
|
| 844 |
+
num_cached_tokens,
|
| 845 |
+
num_cache_creation_tokens,
|
| 846 |
+
mm_token_counts,
|
| 847 |
+
)
|
| 848 |
+
|
| 849 |
+
# In streaming, metrics ride on this final usage chunk, which is
|
| 850 |
+
# only emitted when usage reporting is enabled (i.e.
|
| 851 |
+
# ``stream_options.include_usage=true`` or
|
| 852 |
+
# ``--enable-force-include-usage``).
|
| 853 |
+
stream_per_request_metrics: PerRequestMetrics | None = None
|
| 854 |
+
# See note in chat_completion_full_generator: suppress for n>1.
|
| 855 |
+
if (request.n or 1) == 1:
|
| 856 |
+
if self.enable_per_request_metrics:
|
| 857 |
+
last_metrics = (
|
| 858 |
+
last_res.metrics if last_res is not None else None
|
| 859 |
+
)
|
| 860 |
+
stream_per_request_metrics = build_per_request_timing_metrics(
|
| 861 |
+
last_metrics, completion_tokens
|
| 862 |
+
)
|
| 863 |
+
spec_stats = build_spec_decoding_metrics(last_res)
|
| 864 |
+
if spec_stats is not None:
|
| 865 |
+
if stream_per_request_metrics is None:
|
| 866 |
+
stream_per_request_metrics = PerRequestMetrics()
|
| 867 |
+
stream_per_request_metrics.speculative_decoding = spec_stats
|
| 868 |
+
|
| 869 |
+
final_usage_chunk = ChatCompletionStreamResponse(
|
| 870 |
+
id=request_id,
|
| 871 |
+
object=chunk_object_type,
|
| 872 |
+
created=created_time,
|
| 873 |
+
choices=[],
|
| 874 |
+
model=model_name,
|
| 875 |
+
usage=final_usage,
|
| 876 |
+
system_fingerprint=self.system_fingerprint,
|
| 877 |
+
metrics=stream_per_request_metrics,
|
| 878 |
+
)
|
| 879 |
+
final_usage_data = final_usage_chunk.model_dump_json(
|
| 880 |
+
exclude_unset=True, exclude_none=True
|
| 881 |
+
)
|
| 882 |
+
yield f"data: {final_usage_data}\n\n"
|
| 883 |
+
|
| 884 |
+
# report to FastAPI middleware aggregate usage across all choices
|
| 885 |
+
num_completion_tokens = sum(previous_num_tokens)
|
| 886 |
+
request_metadata.final_usage_info = UsageInfo(
|
| 887 |
+
prompt_tokens=num_prompt_tokens,
|
| 888 |
+
completion_tokens=num_completion_tokens,
|
| 889 |
+
total_tokens=num_prompt_tokens + num_completion_tokens,
|
| 890 |
+
completion_tokens_details=_make_completion_tokens_details(
|
| 891 |
+
sum(previous_reasoning_tokens)
|
| 892 |
+
)
|
| 893 |
+
if self._include_reasoning_tokens_details
|
| 894 |
+
else None,
|
| 895 |
+
)
|
| 896 |
+
|
| 897 |
+
# Log complete streaming response if output logging is enabled
|
| 898 |
+
if self.enable_log_outputs and self.request_logger:
|
| 899 |
+
# Log the complete response for each choice
|
| 900 |
+
for i in range(num_choices):
|
| 901 |
+
full_text = (
|
| 902 |
+
previous_texts[i]
|
| 903 |
+
if previous_texts and i < len(previous_texts)
|
| 904 |
+
else f"<streaming_complete: {previous_num_tokens[i]} tokens>"
|
| 905 |
+
)
|
| 906 |
+
self.request_logger.log_outputs(
|
| 907 |
+
request_id=request_id,
|
| 908 |
+
outputs=full_text,
|
| 909 |
+
output_token_ids=None, # Consider also logging all token IDs
|
| 910 |
+
finish_reason="streaming_complete",
|
| 911 |
+
is_streaming=True,
|
| 912 |
+
delta=False,
|
| 913 |
+
)
|
| 914 |
+
|
| 915 |
+
except GenerationError as e:
|
| 916 |
+
yield f"data: {self._convert_generation_error_to_streaming_response(e)}\n\n"
|
| 917 |
+
except Exception as e:
|
| 918 |
+
logger.exception("Error in chat completion stream generator.")
|
| 919 |
+
data = self.create_streaming_error_response(e)
|
| 920 |
+
yield f"data: {data}\n\n"
|
| 921 |
+
# Send the final done message after all response.n are finished
|
| 922 |
+
yield "data: [DONE]\n\n"
|
| 923 |
+
|
| 924 |
+
async def chat_completion_full_generator(
|
| 925 |
+
self,
|
| 926 |
+
request: ChatCompletionRequest,
|
| 927 |
+
result_generator: AsyncIterator[RequestOutput],
|
| 928 |
+
request_id: str,
|
| 929 |
+
model_name: str,
|
| 930 |
+
conversation: list[ConversationMessage],
|
| 931 |
+
tokenizer: TokenizerLike,
|
| 932 |
+
request_metadata: RequestResponseMetadata,
|
| 933 |
+
parser: Parser | None = None,
|
| 934 |
+
mm_token_counts: dict[str, int] | None = None,
|
| 935 |
+
) -> ErrorResponse | ChatCompletionResponse:
|
| 936 |
+
created_time = int(time.time())
|
| 937 |
+
final_res: RequestOutput | None = None
|
| 938 |
+
|
| 939 |
+
try:
|
| 940 |
+
async for res in result_generator:
|
| 941 |
+
final_res = res
|
| 942 |
+
except asyncio.CancelledError:
|
| 943 |
+
return self.create_error_response("Client disconnected")
|
| 944 |
+
|
| 945 |
+
if final_res is None:
|
| 946 |
+
return self.create_error_response(
|
| 947 |
+
"No output received from the engine.",
|
| 948 |
+
err_type="InternalServerError",
|
| 949 |
+
status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
|
| 950 |
+
)
|
| 951 |
+
|
| 952 |
+
choices: list[ChatCompletionResponseChoice] = []
|
| 953 |
+
total_reasoning_tokens = 0
|
| 954 |
+
|
| 955 |
+
role = self.get_chat_request_role(request)
|
| 956 |
+
tool_parser_cls = (
|
| 957 |
+
self.parser_cls.tool_parser_cls if self.parser_cls is not None else None
|
| 958 |
+
)
|
| 959 |
+
for output in final_res.outputs:
|
| 960 |
+
# check for error finish reason and raise GenerationError
|
| 961 |
+
# finish_reason='error' indicates a retryable request-level internal error
|
| 962 |
+
self._raise_if_error(output.finish_reason, request_id)
|
| 963 |
+
token_ids = output.token_ids
|
| 964 |
+
out_logprobs = output.logprobs
|
| 965 |
+
|
| 966 |
+
if request.logprobs and (
|
| 967 |
+
request.top_logprobs is not None or request.logprob_token_ids
|
| 968 |
+
):
|
| 969 |
+
assert out_logprobs is not None, "Did not output logprobs"
|
| 970 |
+
logprobs = self._create_chat_logprobs(
|
| 971 |
+
token_ids=token_ids,
|
| 972 |
+
top_logprobs=out_logprobs,
|
| 973 |
+
num_output_top_logprobs=request.top_logprobs,
|
| 974 |
+
logprob_token_ids=request.logprob_token_ids,
|
| 975 |
+
tokenizer=tokenizer,
|
| 976 |
+
return_as_token_id=request.return_tokens_as_token_ids,
|
| 977 |
+
)
|
| 978 |
+
else:
|
| 979 |
+
logprobs = None
|
| 980 |
+
|
| 981 |
+
if parser is not None:
|
| 982 |
+
reasoning, content, tool_calls = parser.parse(
|
| 983 |
+
output.text,
|
| 984 |
+
request,
|
| 985 |
+
enable_auto_tools=self.enable_auto_tools,
|
| 986 |
+
model_output_token_ids=token_ids,
|
| 987 |
+
)
|
| 988 |
+
suppress_metadata = not request.include_reasoning and parser is not None
|
| 989 |
+
if not request.include_reasoning:
|
| 990 |
+
reasoning = None
|
| 991 |
+
total_reasoning_tokens += parser.count_reasoning_tokens(token_ids)
|
| 992 |
+
if suppress_metadata:
|
| 993 |
+
logprobs = None
|
| 994 |
+
else:
|
| 995 |
+
reasoning = None
|
| 996 |
+
content = output.text
|
| 997 |
+
tool_calls = []
|
| 998 |
+
suppress_metadata = False
|
| 999 |
+
|
| 1000 |
+
auto_tools_called = False
|
| 1001 |
+
is_named_tool_choice = (
|
| 1002 |
+
request.tool_choice is not None
|
| 1003 |
+
and type(request.tool_choice) is ChatCompletionNamedToolChoiceParam
|
| 1004 |
+
)
|
| 1005 |
+
is_required_tool_choice = request.tool_choice == "required"
|
| 1006 |
+
|
| 1007 |
+
# All six construction sites route through ``self._create_chat_message``
|
| 1008 |
+
# so subclasses can swap in a specialized :class:`ChatMessage`
|
| 1009 |
+
# (e.g. the Cohere v2 handler's ``CohereChatMessage``) without
|
| 1010 |
+
# having to duplicate this branch logic.
|
| 1011 |
+
if (not self.enable_auto_tools or not tool_parser_cls) and (
|
| 1012 |
+
not is_named_tool_choice and not is_required_tool_choice
|
| 1013 |
+
):
|
| 1014 |
+
message = self._create_chat_message(
|
| 1015 |
+
role=role, reasoning=reasoning, content=content
|
| 1016 |
+
)
|
| 1017 |
+
|
| 1018 |
+
elif is_named_tool_choice or is_required_tool_choice:
|
| 1019 |
+
message = self._create_chat_message(
|
| 1020 |
+
role=role,
|
| 1021 |
+
reasoning=reasoning,
|
| 1022 |
+
content=content or "",
|
| 1023 |
+
tool_calls=[
|
| 1024 |
+
ToolCall(id=tc.id or make_tool_call_id(), function=tc)
|
| 1025 |
+
for tc in (tool_calls or [])
|
| 1026 |
+
],
|
| 1027 |
+
)
|
| 1028 |
+
|
| 1029 |
+
# if the request doesn't use tool choice
|
| 1030 |
+
# OR specifies to not use a tool
|
| 1031 |
+
elif not request.tool_choice or request.tool_choice == "none":
|
| 1032 |
+
message = self._create_chat_message(
|
| 1033 |
+
role=role, reasoning=reasoning, content=content
|
| 1034 |
+
)
|
| 1035 |
+
|
| 1036 |
+
# handle when there are tools and tool choice is auto
|
| 1037 |
+
elif (
|
| 1038 |
+
request.tools
|
| 1039 |
+
and (request.tool_choice == "auto" or request.tool_choice is None)
|
| 1040 |
+
and self.enable_auto_tools
|
| 1041 |
+
and tool_parser_cls
|
| 1042 |
+
):
|
| 1043 |
+
auto_tools_called = tool_calls is not None and len(tool_calls) > 0
|
| 1044 |
+
if tool_calls:
|
| 1045 |
+
message = self._create_chat_message(
|
| 1046 |
+
role=role,
|
| 1047 |
+
reasoning=reasoning,
|
| 1048 |
+
content=content,
|
| 1049 |
+
tool_calls=[
|
| 1050 |
+
ToolCall(id=tc.id or make_tool_call_id(), function=tc)
|
| 1051 |
+
for tc in tool_calls
|
| 1052 |
+
],
|
| 1053 |
+
)
|
| 1054 |
+
|
| 1055 |
+
else:
|
| 1056 |
+
message = self._create_chat_message(
|
| 1057 |
+
role=role,
|
| 1058 |
+
reasoning=reasoning,
|
| 1059 |
+
content=content,
|
| 1060 |
+
)
|
| 1061 |
+
|
| 1062 |
+
# undetermined case that is still important to handle
|
| 1063 |
+
else:
|
| 1064 |
+
logger.error(
|
| 1065 |
+
"Error in chat_completion_full_generator - cannot determine"
|
| 1066 |
+
" if tools should be extracted. Returning a standard chat "
|
| 1067 |
+
"completion."
|
| 1068 |
+
)
|
| 1069 |
+
message = self._create_chat_message(
|
| 1070 |
+
role=role, reasoning=reasoning, content=content
|
| 1071 |
+
)
|
| 1072 |
+
|
| 1073 |
+
# Subclass hook: enrich the constructed message with any
|
| 1074 |
+
# parser-side extras that don't fit through the plain
|
| 1075 |
+
# ``(reasoning, content, tool_calls)`` tuple. Base is a no-op;
|
| 1076 |
+
# citation-aware handlers use this to surface grounding
|
| 1077 |
+
# metadata cached on the reasoning parser.
|
| 1078 |
+
message = self._finalize_response_message(message, parser=parser)
|
| 1079 |
+
|
| 1080 |
+
# In OpenAI's API, when a tool is called, the finish_reason is:
|
| 1081 |
+
# "tool_calls" for "auto" or "required" tool calls,
|
| 1082 |
+
# and "stop" for named tool calls.
|
| 1083 |
+
is_finish_reason_tool_calls = (
|
| 1084 |
+
auto_tools_called and output.finish_reason == "stop"
|
| 1085 |
+
) or (
|
| 1086 |
+
request.tool_choice
|
| 1087 |
+
and request.tool_choice == "required"
|
| 1088 |
+
and output.finish_reason == "stop"
|
| 1089 |
+
)
|
| 1090 |
+
|
| 1091 |
+
routed_experts_b64 = (
|
| 1092 |
+
numpy2base64(output.routed_experts)
|
| 1093 |
+
if output.routed_experts is not None
|
| 1094 |
+
else None
|
| 1095 |
+
)
|
| 1096 |
+
|
| 1097 |
+
choice_data = ChatCompletionResponseChoice(
|
| 1098 |
+
index=output.index,
|
| 1099 |
+
message=message,
|
| 1100 |
+
logprobs=logprobs,
|
| 1101 |
+
finish_reason="tool_calls"
|
| 1102 |
+
if is_finish_reason_tool_calls
|
| 1103 |
+
else output.finish_reason
|
| 1104 |
+
if output.finish_reason
|
| 1105 |
+
else "stop",
|
| 1106 |
+
stop_reason=output.stop_reason,
|
| 1107 |
+
token_ids=(
|
| 1108 |
+
as_list(output.token_ids)
|
| 1109 |
+
if request.return_token_ids and not suppress_metadata
|
| 1110 |
+
else None
|
| 1111 |
+
),
|
| 1112 |
+
routed_experts=routed_experts_b64,
|
| 1113 |
+
)
|
| 1114 |
+
choice_data = maybe_filter_parallel_tool_calls(choice_data, request)
|
| 1115 |
+
|
| 1116 |
+
choices.append(choice_data)
|
| 1117 |
+
|
| 1118 |
+
if request.echo:
|
| 1119 |
+
last_msg_content: str | list[dict[str, str]] = ""
|
| 1120 |
+
if (
|
| 1121 |
+
conversation
|
| 1122 |
+
and "content" in conversation[-1]
|
| 1123 |
+
and conversation[-1].get("role") == role
|
| 1124 |
+
):
|
| 1125 |
+
last_msg_content = conversation[-1]["content"] or ""
|
| 1126 |
+
if isinstance(last_msg_content, list):
|
| 1127 |
+
last_msg_content = "\n".join(msg["text"] for msg in last_msg_content)
|
| 1128 |
+
|
| 1129 |
+
for choice in choices:
|
| 1130 |
+
full_message = last_msg_content + (choice.message.content or "")
|
| 1131 |
+
choice.message.content = full_message
|
| 1132 |
+
|
| 1133 |
+
assert final_res.prompt_token_ids is not None
|
| 1134 |
+
num_prompt_tokens = len(final_res.prompt_token_ids)
|
| 1135 |
+
if final_res.encoder_prompt_token_ids is not None:
|
| 1136 |
+
num_prompt_tokens += len(final_res.encoder_prompt_token_ids)
|
| 1137 |
+
num_generated_tokens = sum(
|
| 1138 |
+
len(output.token_ids) for output in final_res.outputs
|
| 1139 |
+
)
|
| 1140 |
+
usage = UsageInfo(
|
| 1141 |
+
prompt_tokens=num_prompt_tokens,
|
| 1142 |
+
completion_tokens=num_generated_tokens,
|
| 1143 |
+
total_tokens=num_prompt_tokens + num_generated_tokens,
|
| 1144 |
+
completion_tokens_details=_make_completion_tokens_details(
|
| 1145 |
+
total_reasoning_tokens
|
| 1146 |
+
)
|
| 1147 |
+
if self._include_reasoning_tokens_details
|
| 1148 |
+
else None,
|
| 1149 |
+
)
|
| 1150 |
+
usage.prompt_tokens_details = _make_prompt_tokens_details(
|
| 1151 |
+
self.enable_prompt_tokens_details,
|
| 1152 |
+
final_res.num_cached_tokens,
|
| 1153 |
+
final_res.num_cache_creation_tokens,
|
| 1154 |
+
mm_token_counts,
|
| 1155 |
+
)
|
| 1156 |
+
|
| 1157 |
+
request_metadata.final_usage_info = usage
|
| 1158 |
+
|
| 1159 |
+
per_request_metrics: PerRequestMetrics | None = None
|
| 1160 |
+
# Per-request metrics (timing + spec-decode acceptance) describe a single
|
| 1161 |
+
# generation stream. For n>1 the stats belong to only one of the n
|
| 1162 |
+
# sequences, so they cannot be attributed to the request; suppress.
|
| 1163 |
+
if (request.n or 1) == 1:
|
| 1164 |
+
if self.enable_per_request_metrics:
|
| 1165 |
+
per_request_metrics = build_per_request_timing_metrics(
|
| 1166 |
+
final_res.metrics, num_generated_tokens
|
| 1167 |
+
)
|
| 1168 |
+
spec_stats = build_spec_decoding_metrics(final_res)
|
| 1169 |
+
if spec_stats is not None:
|
| 1170 |
+
if per_request_metrics is None:
|
| 1171 |
+
per_request_metrics = PerRequestMetrics()
|
| 1172 |
+
per_request_metrics.speculative_decoding = spec_stats
|
| 1173 |
+
|
| 1174 |
+
# ``final_res.prompt`` is the rendered chat-templated prompt text
|
| 1175 |
+
prompt_text = final_res.prompt if request.return_prompt_text else None
|
| 1176 |
+
|
| 1177 |
+
response = ChatCompletionResponse(
|
| 1178 |
+
id=request_id,
|
| 1179 |
+
created=created_time,
|
| 1180 |
+
model=model_name,
|
| 1181 |
+
choices=choices,
|
| 1182 |
+
usage=usage,
|
| 1183 |
+
system_fingerprint=self.system_fingerprint,
|
| 1184 |
+
prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs),
|
| 1185 |
+
prompt_token_ids=(
|
| 1186 |
+
final_res.prompt_token_ids if request.return_token_ids else None
|
| 1187 |
+
),
|
| 1188 |
+
prompt_text=prompt_text,
|
| 1189 |
+
kv_transfer_params=final_res.kv_transfer_params,
|
| 1190 |
+
ec_transfer_params=final_res.ec_transfer_params,
|
| 1191 |
+
metrics=per_request_metrics,
|
| 1192 |
+
)
|
| 1193 |
+
|
| 1194 |
+
# Log complete response if output logging is enabled
|
| 1195 |
+
if self.enable_log_outputs and self.request_logger:
|
| 1196 |
+
for choice in choices:
|
| 1197 |
+
output_text = ""
|
| 1198 |
+
if choice.message.content:
|
| 1199 |
+
output_text = choice.message.content
|
| 1200 |
+
elif choice.message.tool_calls:
|
| 1201 |
+
# For tool calls, log the function name and arguments
|
| 1202 |
+
tool_call_descriptions = []
|
| 1203 |
+
for tc in choice.message.tool_calls: # type: ignore
|
| 1204 |
+
function_call: FunctionCall = tc.function # type: ignore
|
| 1205 |
+
tool_call_descriptions.append(
|
| 1206 |
+
f"{function_call.name}({function_call.arguments})"
|
| 1207 |
+
)
|
| 1208 |
+
tool_calls_str = ", ".join(tool_call_descriptions)
|
| 1209 |
+
output_text = f"[tool_calls: {tool_calls_str}]"
|
| 1210 |
+
|
| 1211 |
+
if output_text:
|
| 1212 |
+
# Get the corresponding output token IDs
|
| 1213 |
+
output_token_ids = None
|
| 1214 |
+
if choice.index < len(final_res.outputs):
|
| 1215 |
+
output_token_ids = final_res.outputs[choice.index].token_ids
|
| 1216 |
+
|
| 1217 |
+
self.request_logger.log_outputs(
|
| 1218 |
+
request_id=request_id,
|
| 1219 |
+
outputs=output_text,
|
| 1220 |
+
output_token_ids=output_token_ids,
|
| 1221 |
+
finish_reason=choice.finish_reason,
|
| 1222 |
+
is_streaming=False,
|
| 1223 |
+
delta=False,
|
| 1224 |
+
)
|
| 1225 |
+
|
| 1226 |
+
return response
|
| 1227 |
+
|
| 1228 |
+
def _get_top_logprobs(
|
| 1229 |
+
self,
|
| 1230 |
+
logprobs: dict[int, Logprob],
|
| 1231 |
+
top_logprobs: int | None,
|
| 1232 |
+
tokenizer: TokenizerLike | None,
|
| 1233 |
+
should_return_as_token_id: bool,
|
| 1234 |
+
return_all: bool = False,
|
| 1235 |
+
) -> list[ChatCompletionLogProb]:
|
| 1236 |
+
return [
|
| 1237 |
+
ChatCompletionLogProb(
|
| 1238 |
+
token=(
|
| 1239 |
+
token := self._get_decoded_token(
|
| 1240 |
+
p[1],
|
| 1241 |
+
p[0],
|
| 1242 |
+
tokenizer,
|
| 1243 |
+
return_as_token_id=should_return_as_token_id,
|
| 1244 |
+
)
|
| 1245 |
+
),
|
| 1246 |
+
logprob=max(p[1].logprob, -9999.0),
|
| 1247 |
+
bytes=list(token.encode("utf-8", errors="replace")),
|
| 1248 |
+
)
|
| 1249 |
+
for i, p in enumerate(logprobs.items())
|
| 1250 |
+
if return_all
|
| 1251 |
+
or top_logprobs == -1
|
| 1252 |
+
or (top_logprobs is not None and i < top_logprobs)
|
| 1253 |
+
]
|
| 1254 |
+
|
| 1255 |
+
def _create_chat_logprobs(
|
| 1256 |
+
self,
|
| 1257 |
+
token_ids: GenericSequence[int],
|
| 1258 |
+
top_logprobs: GenericSequence[dict[int, Logprob] | None],
|
| 1259 |
+
tokenizer: TokenizerLike | None,
|
| 1260 |
+
num_output_top_logprobs: int | None = None,
|
| 1261 |
+
logprob_token_ids: list[int] | None = None,
|
| 1262 |
+
return_as_token_id: bool | None = None,
|
| 1263 |
+
) -> ChatCompletionLogProbs:
|
| 1264 |
+
"""Create OpenAI-style logprobs."""
|
| 1265 |
+
logprobs_content: list[ChatCompletionLogProbsContent] = []
|
| 1266 |
+
|
| 1267 |
+
should_return_as_token_id = (
|
| 1268 |
+
return_as_token_id
|
| 1269 |
+
if return_as_token_id is not None
|
| 1270 |
+
else self.return_tokens_as_token_ids
|
| 1271 |
+
)
|
| 1272 |
+
for i, token_id in enumerate(token_ids):
|
| 1273 |
+
step_top_logprobs = top_logprobs[i]
|
| 1274 |
+
if step_top_logprobs is None or step_top_logprobs.get(token_id) is None:
|
| 1275 |
+
if should_return_as_token_id:
|
| 1276 |
+
token = format_token_id_placeholder(token_id)
|
| 1277 |
+
else:
|
| 1278 |
+
if tokenizer is None:
|
| 1279 |
+
raise ValueError(
|
| 1280 |
+
"Unable to get tokenizer because `skip_tokenizer_init=True`"
|
| 1281 |
+
)
|
| 1282 |
+
|
| 1283 |
+
token = tokenizer.decode(token_id)
|
| 1284 |
+
|
| 1285 |
+
logprobs_content.append(
|
| 1286 |
+
ChatCompletionLogProbsContent(
|
| 1287 |
+
token=token,
|
| 1288 |
+
bytes=list(token.encode("utf-8", errors="replace")),
|
| 1289 |
+
)
|
| 1290 |
+
)
|
| 1291 |
+
else:
|
| 1292 |
+
step_token = step_top_logprobs[token_id]
|
| 1293 |
+
step_decoded = step_token.decoded_token
|
| 1294 |
+
|
| 1295 |
+
logprobs_content.append(
|
| 1296 |
+
ChatCompletionLogProbsContent(
|
| 1297 |
+
token=self._get_decoded_token(
|
| 1298 |
+
step_token,
|
| 1299 |
+
token_id,
|
| 1300 |
+
tokenizer,
|
| 1301 |
+
should_return_as_token_id,
|
| 1302 |
+
),
|
| 1303 |
+
logprob=max(step_token.logprob, -9999.0),
|
| 1304 |
+
bytes=(
|
| 1305 |
+
None
|
| 1306 |
+
if step_decoded is None
|
| 1307 |
+
else list(step_decoded.encode("utf-8", errors="replace"))
|
| 1308 |
+
),
|
| 1309 |
+
top_logprobs=self._get_top_logprobs(
|
| 1310 |
+
step_top_logprobs,
|
| 1311 |
+
num_output_top_logprobs,
|
| 1312 |
+
tokenizer,
|
| 1313 |
+
should_return_as_token_id,
|
| 1314 |
+
return_all=bool(logprob_token_ids),
|
| 1315 |
+
),
|
| 1316 |
+
)
|
| 1317 |
+
)
|
| 1318 |
+
|
| 1319 |
+
return ChatCompletionLogProbs(content=logprobs_content)
|
bundle/plugin-site/ornith_g256/_vllm_correctness/manifest.json
CHANGED
|
@@ -1,5 +1,5 @@
|
|
| 1 |
{
|
| 2 |
-
"runtime_version": "1.0.
|
| 3 |
"modules": {
|
| 4 |
"vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn": {
|
| 5 |
"file": "qwen_gdn_linear_attn.py",
|
|
@@ -24,6 +24,12 @@
|
|
| 24 |
"native_path": "v1/worker/gpu_model_runner.py",
|
| 25 |
"native_sha256": "4706fcf4b85158173d88e3b19ffb0f6a280fcb60554cfdac861a01a0cab4d4df",
|
| 26 |
"sha256": "ef0009acaa67a1a5e066984f34bf619b06e85220ff953a3d9511376f0146dccf"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 27 |
}
|
| 28 |
},
|
| 29 |
"upstream": [
|
|
|
|
| 1 |
{
|
| 2 |
+
"runtime_version": "1.0.3",
|
| 3 |
"modules": {
|
| 4 |
"vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn": {
|
| 5 |
"file": "qwen_gdn_linear_attn.py",
|
|
|
|
| 24 |
"native_path": "v1/worker/gpu_model_runner.py",
|
| 25 |
"native_sha256": "4706fcf4b85158173d88e3b19ffb0f6a280fcb60554cfdac861a01a0cab4d4df",
|
| 26 |
"sha256": "ef0009acaa67a1a5e066984f34bf619b06e85220ff953a3d9511376f0146dccf"
|
| 27 |
+
},
|
| 28 |
+
"vllm.entrypoints.openai.chat_completion.serving": {
|
| 29 |
+
"file": "chat_completion_serving_r03.py",
|
| 30 |
+
"native_path": "entrypoints/openai/chat_completion/serving.py",
|
| 31 |
+
"native_sha256": "9982953285e9df469032a82fffa4095d0e9d86278bede6e2b91d03d02373d182",
|
| 32 |
+
"sha256": "5720371f5bcb351e340a3e8d9a0235ec01ff27ed93e4f35542e31b83a055ef45"
|
| 33 |
}
|
| 34 |
},
|
| 35 |
"upstream": [
|
bundle/plugin-site/ornith_g256/_xgrammar_native/build-receipt.json
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"mode": "candidate",
|
| 3 |
+
"library": "libxgrammar-candidate.so",
|
| 4 |
+
"sha256": "92ab36bfae20c6c4df3f5321ca375204799ac833252921d9d8f6a206491d35f7",
|
| 5 |
+
"source_commit": "557becfb64c503ae9c04344b0047661f43f44320",
|
| 6 |
+
"patch_sha256": "2c42719e948809242f230ec534f967eab1babad9be1ba6dda9c59fd706f20d45",
|
| 7 |
+
"sources": {
|
| 8 |
+
"cpp/grammar.cc": "01d1552b37ecaa72da601a3b6d43975a2c3a0de02cdf797e82c243b30eef8c9d",
|
| 9 |
+
"cpp/grammar_functor.cc": "9410c20b665093b9fb67dc30db148121b819d3aee5ff94eb82558bf79b1f8279"
|
| 10 |
+
},
|
| 11 |
+
"configure": [
|
| 12 |
+
"cmake",
|
| 13 |
+
"-S",
|
| 14 |
+
"/srv/ssd/p3700ba/scratch/crown/runtime-tool-integration-20260917/native-build-r2/source-candidate",
|
| 15 |
+
"-B",
|
| 16 |
+
"/srv/ssd/p3700ba/scratch/crown/runtime-tool-integration-20260917/native-build-r2/build-candidate",
|
| 17 |
+
"-G",
|
| 18 |
+
"Ninja",
|
| 19 |
+
"-DCMAKE_BUILD_TYPE=Release",
|
| 20 |
+
"-DXGRAMMAR_BUILD_PYTHON_BINDINGS=ON",
|
| 21 |
+
"-DCMAKE_CXX_FLAGS_RELEASE=-O3 -DNDEBUG -fno-lto",
|
| 22 |
+
"-DCMAKE_CXX_FLAGS_RELWITHDEBINFO=-O3 -g -DNDEBUG -fno-lto",
|
| 23 |
+
"-DPython_EXECUTABLE=/opt/ciru/glm53-iu4/venv/bin/python",
|
| 24 |
+
"-DPython3_EXECUTABLE=/opt/ciru/glm53-iu4/venv/bin/python"
|
| 25 |
+
],
|
| 26 |
+
"build_seconds": 67.22386657004245,
|
| 27 |
+
"compiler": "g++ (GCC) 15.2.0\nCopyright (C) 2025 Free Software Foundation, Inc.\nThis is free software; see the source for copying conditions. There is NO\nwarranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.\n\n",
|
| 28 |
+
"submodules": " 6689d14c203eed390ae7bb64f56a983cfd7dff9c 3rdparty/cpptrace (v0.1-644-g6689d14)\n bbd2f4d32427e548797929af08cfe2a9cbb3cf12 3rdparty/dlpack (v1.0)\n df1544bcee0c7ce35cd5ea0b3eb8cc81855a4140 3rdparty/googletest (release-1.8.0-3498-gdf1544bc)\n"
|
| 29 |
+
}
|
bundle/plugin-site/ornith_g256/_xgrammar_native/libxgrammar-candidate.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:92ab36bfae20c6c4df3f5321ca375204799ac833252921d9d8f6a206491d35f7
|
| 3 |
+
size 60211248
|
bundle/plugin-site/ornith_g256/_xgrammar_native/selection.json
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"mode": "candidate",
|
| 3 |
+
"sha256": "92ab36bfae20c6c4df3f5321ca375204799ac833252921d9d8f6a206491d35f7"
|
| 4 |
+
}
|
bundle/plugin-site/ornith_g256/attention_compact.py
CHANGED
|
@@ -107,9 +107,9 @@ def configure(*, max_num_seqs, max_model_len, max_num_batched_tokens, device,
|
|
| 107 |
raise RuntimeError('Compact attention arena was configured differently')
|
| 108 |
return _arena['bytes']
|
| 109 |
if (not 1 <= max_num_seqs <= 8 or not 1 <= max_model_len <= 262144
|
| 110 |
-
or not 1 <= max_num_batched_tokens <=
|
| 111 |
or torch.cuda.is_current_stream_capturing()):
|
| 112 |
-
raise ValueError('Configure target compact prefill before capture, within C8/256K/
|
| 113 |
_prefill = _graph_safe_prefill()
|
| 114 |
key = torch.empty((max_num_seqs * max_model_len, 2, 256),
|
| 115 |
dtype=torch.bfloat16, device=device)
|
|
@@ -177,14 +177,14 @@ def try_forward(query, key_cache, value_cache, output, block_table,
|
|
| 177 |
num_warps=8, num_stages=1)
|
| 178 |
# Graph-padded rows are not real requests and receive defined zero output.
|
| 179 |
output.zero_()
|
| 180 |
-
if (_arena['iu4'] is not None and requests == 1 and max_query_len
|
| 181 |
and max_seq_len >= 4096 and max_seq_len % 32 == 0
|
| 182 |
and max_query_len <= rows and sm_scale == 256**-.5):
|
| 183 |
_arena['iu4'].forward(query, key, value, output, query_start_loc,
|
| 184 |
-
seq_lens, max_seq_len)
|
| 185 |
logger.info_once('Ornith C1 IU4 prefill active: BF16 cache gather + '
|
| 186 |
'normalized Q/K H256 and P/V H32, signed IU4 QK/PV; '
|
| 187 |
-
'
|
| 188 |
return True
|
| 189 |
_prefill(q=query, k=key, v=value, o=output, softmax_lse=lse,
|
| 190 |
sd_mask=None, sm_scale=sm_scale, alibi_slopes=None, causal=True,
|
|
|
|
| 107 |
raise RuntimeError('Compact attention arena was configured differently')
|
| 108 |
return _arena['bytes']
|
| 109 |
if (not 1 <= max_num_seqs <= 8 or not 1 <= max_model_len <= 262144
|
| 110 |
+
or not 1 <= max_num_batched_tokens <= 8192
|
| 111 |
or torch.cuda.is_current_stream_capturing()):
|
| 112 |
+
raise ValueError('Configure target compact prefill before capture, within C8/256K/8192 tokens')
|
| 113 |
_prefill = _graph_safe_prefill()
|
| 114 |
key = torch.empty((max_num_seqs * max_model_len, 2, 256),
|
| 115 |
dtype=torch.bfloat16, device=device)
|
|
|
|
| 177 |
num_warps=8, num_stages=1)
|
| 178 |
# Graph-padded rows are not real requests and receive defined zero output.
|
| 179 |
output.zero_()
|
| 180 |
+
if (_arena['iu4'] is not None and requests == 1 and max_query_len in (1120, 4480)
|
| 181 |
and max_seq_len >= 4096 and max_seq_len % 32 == 0
|
| 182 |
and max_query_len <= rows and sm_scale == 256**-.5):
|
| 183 |
_arena['iu4'].forward(query, key, value, output, query_start_loc,
|
| 184 |
+
seq_lens, max_seq_len, query_rows=max_query_len)
|
| 185 |
logger.info_once('Ornith C1 IU4 prefill active: BF16 cache gather + '
|
| 186 |
'normalized Q/K H256 and P/V H32, signed IU4 QK/PV; '
|
| 187 |
+
'Q%d and aligned K>=4096, existing BF16 fallback elsewhere', max_query_len)
|
| 188 |
return True
|
| 189 |
_prefill(q=query, k=key, v=value, o=output, softmax_lse=lse,
|
| 190 |
sd_mask=None, sm_scale=sm_scale, alibi_slopes=None, causal=True,
|
bundle/plugin-site/ornith_g256/attention_iu4.py
CHANGED
|
@@ -11,31 +11,42 @@ class NativeAttention:
|
|
| 11 |
raise ValueError("IU4 prefill capacity must be1120..262144")
|
| 12 |
self.max_keys = 65536 if max_model_len <= 65536 else 262144
|
| 13 |
groups = self.max_keys // 32
|
|
|
|
| 14 |
self.library = ctypes.CDLL(str(library_path))
|
| 15 |
self.library.iu4_prepare.argtypes = [ctypes.c_void_p] * 9 + [ctypes.c_int] * 3 + [ctypes.c_void_p]
|
| 16 |
self.library.iu4_prepare.restype = ctypes.c_int
|
| 17 |
self.library.iu4_attention.argtypes = [ctypes.c_void_p] * 7 + [ctypes.c_int] + [ctypes.c_void_p]
|
| 18 |
self.library.iu4_attention.restype = ctypes.c_int
|
| 19 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 20 |
((2, groups, 32, 32), torch.int32), ((2, self.max_keys), torch.float16),
|
| 21 |
((2, groups, 4, 256), torch.int32), ((2, groups, 256), torch.float16)]
|
| 22 |
self.packed = [torch.empty(shape, dtype=dtype, device=device) for shape, dtype in specs]
|
| 23 |
self.pointers = [tensor.data_ptr() for tensor in self.packed]
|
| 24 |
-
self.output = torch.empty((
|
| 25 |
self.bytes = self.output.numel() * self.output.element_size() + sum(tensor.numel() * tensor.element_size() for tensor in self.packed)
|
| 26 |
|
| 27 |
-
def forward(self, query, key, value, output, starts, lengths, max_keys):
|
| 28 |
-
#
|
| 29 |
-
#
|
| 30 |
# No allocation, tensor copy to CPU, synchronization or cache mutation here.
|
| 31 |
-
if not 1120
|
|
|
|
|
|
|
| 32 |
raise ValueError("IU4 prefill requires aligned keys within reserved capacity")
|
|
|
|
|
|
|
|
|
|
|
|
|
| 33 |
stream = torch.cuda.current_stream().cuda_stream
|
| 34 |
-
rc =
|
| 35 |
*self.pointers, max_keys, query.stride(0), query.stride(1), stream)
|
| 36 |
if rc:
|
| 37 |
raise RuntimeError(f'IU4 attention preparation launch failed: HIP {rc}')
|
| 38 |
-
rc =
|
| 39 |
if rc:
|
| 40 |
raise RuntimeError(f'IU4 attention launch failed: HIP {rc}')
|
| 41 |
-
output[:
|
|
|
|
| 11 |
raise ValueError("IU4 prefill capacity must be1120..262144")
|
| 12 |
self.max_keys = 65536 if max_model_len <= 65536 else 262144
|
| 13 |
groups = self.max_keys // 32
|
| 14 |
+
self.query_capacity = 4480
|
| 15 |
self.library = ctypes.CDLL(str(library_path))
|
| 16 |
self.library.iu4_prepare.argtypes = [ctypes.c_void_p] * 9 + [ctypes.c_int] * 3 + [ctypes.c_void_p]
|
| 17 |
self.library.iu4_prepare.restype = ctypes.c_int
|
| 18 |
self.library.iu4_attention.argtypes = [ctypes.c_void_p] * 7 + [ctypes.c_int] + [ctypes.c_void_p]
|
| 19 |
self.library.iu4_attention.restype = ctypes.c_int
|
| 20 |
+
self.library.iu4_prepare_q4480.argtypes = self.library.iu4_prepare.argtypes
|
| 21 |
+
self.library.iu4_prepare_q4480.restype = ctypes.c_int
|
| 22 |
+
self.library.iu4_attention_q4480.argtypes = self.library.iu4_attention.argtypes
|
| 23 |
+
self.library.iu4_attention_q4480.restype = ctypes.c_int
|
| 24 |
+
specs = [((16, self.query_capacity, 32), torch.int32), ((16, self.query_capacity), torch.float16),
|
| 25 |
((2, groups, 32, 32), torch.int32), ((2, self.max_keys), torch.float16),
|
| 26 |
((2, groups, 4, 256), torch.int32), ((2, groups, 256), torch.float16)]
|
| 27 |
self.packed = [torch.empty(shape, dtype=dtype, device=device) for shape, dtype in specs]
|
| 28 |
self.pointers = [tensor.data_ptr() for tensor in self.packed]
|
| 29 |
+
self.output = torch.empty((self.query_capacity, 16, 256), dtype=torch.bfloat16, device=device)
|
| 30 |
self.bytes = self.output.numel() * self.output.element_size() + sum(tensor.numel() * tensor.element_size() for tensor in self.packed)
|
| 31 |
|
| 32 |
+
def forward(self, query, key, value, output, starts, lengths, max_keys, *, query_rows=1120):
|
| 33 |
+
# CPU query metadata selects one of two fixed-Q consumers. The shared
|
| 34 |
+
# arena is capacity-sized; each native call packs with its own Q stride.
|
| 35 |
# No allocation, tensor copy to CPU, synchronization or cache mutation here.
|
| 36 |
+
if (query_rows not in (1120, 4480) or query.shape[0] < query_rows
|
| 37 |
+
or output.shape[0] < query_rows
|
| 38 |
+
or not query_rows <= max_keys <= self.max_keys or max_keys % 32):
|
| 39 |
raise ValueError("IU4 prefill requires aligned keys within reserved capacity")
|
| 40 |
+
prepare = (self.library.iu4_prepare_q4480 if query_rows == 4480
|
| 41 |
+
else self.library.iu4_prepare)
|
| 42 |
+
attend = (self.library.iu4_attention_q4480 if query_rows == 4480
|
| 43 |
+
else self.library.iu4_attention)
|
| 44 |
stream = torch.cuda.current_stream().cuda_stream
|
| 45 |
+
rc = prepare(query.data_ptr(), key.data_ptr(), value.data_ptr(),
|
| 46 |
*self.pointers, max_keys, query.stride(0), query.stride(1), stream)
|
| 47 |
if rc:
|
| 48 |
raise RuntimeError(f'IU4 attention preparation launch failed: HIP {rc}')
|
| 49 |
+
rc = attend(*self.pointers, self.output.data_ptr(), max_keys, stream)
|
| 50 |
if rc:
|
| 51 |
raise RuntimeError(f'IU4 attention launch failed: HIP {rc}')
|
| 52 |
+
output[:query_rows].copy_(self.output[:query_rows])
|
bundle/plugin-site/ornith_g256/attention_iu4_persistent.py
CHANGED
|
@@ -99,7 +99,7 @@ class State:
|
|
| 99 |
raise RuntimeError(f'Expected ten target layers aliasing five banks, got {targets}/{len(self.banks)}')
|
| 100 |
self.device = next(iter(self.banks.values())).cache.device
|
| 101 |
device = self.device
|
| 102 |
-
self.groups = torch.empty(
|
| 103 |
self.contexts = torch.empty(8, dtype=torch.int32, device=device)
|
| 104 |
self.counts = torch.empty(8, dtype=torch.int32, device=device)
|
| 105 |
self.owners = torch.empty(64, dtype=torch.int32, device=device)
|
|
@@ -122,7 +122,7 @@ class State:
|
|
| 122 |
bank = self.banks.get(cache.data_ptr())
|
| 123 |
if bank is None:
|
| 124 |
return
|
| 125 |
-
if slots.numel() >
|
| 126 |
raise RuntimeError('Unexpected persistent IU4 slot capacity/type')
|
| 127 |
self.call('update_cache', [cache, slots, self.groups, bank.busy,
|
| 128 |
bank.valid, *bank.packed], slots.numel())
|
|
|
|
| 99 |
raise RuntimeError(f'Expected ten target layers aliasing five banks, got {targets}/{len(self.banks)}')
|
| 100 |
self.device = next(iter(self.banks.values())).cache.device
|
| 101 |
device = self.device
|
| 102 |
+
self.groups = torch.empty(8192, dtype=torch.int32, device=device)
|
| 103 |
self.contexts = torch.empty(8, dtype=torch.int32, device=device)
|
| 104 |
self.counts = torch.empty(8, dtype=torch.int32, device=device)
|
| 105 |
self.owners = torch.empty(64, dtype=torch.int32, device=device)
|
|
|
|
| 122 |
bank = self.banks.get(cache.data_ptr())
|
| 123 |
if bank is None:
|
| 124 |
return
|
| 125 |
+
if slots.numel() > 8192 or slots.dtype != torch.int64:
|
| 126 |
raise RuntimeError('Unexpected persistent IU4 slot capacity/type')
|
| 127 |
self.call('update_cache', [cache, slots, self.groups, bank.busy,
|
| 128 |
bank.valid, *bank.packed], slots.numel())
|
bundle/plugin-site/ornith_g256/optimized/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
"""Retained Ornith runtime optimizations."""
|
bundle/plugin-site/ornith_g256/optimized/compile_scope.py
ADDED
|
@@ -0,0 +1,30 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Keep model-only fusion hooks outside compiler autotuning microbenchmarks."""
|
| 2 |
+
import json
|
| 3 |
+
import threading
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
def install(original_run, output):
|
| 7 |
+
from torch._inductor.runtime.triton_heuristics import CachingAutotuner
|
| 8 |
+
from torch._inductor.codegen.wrapper import PythonWrapperCodegen
|
| 9 |
+
selected_run=CachingAutotuner.run
|
| 10 |
+
original_tuning=PythonWrapperCodegen.generate_and_run_autotune_block
|
| 11 |
+
local=threading.local()
|
| 12 |
+
state=dict(autotune_blocks=0,original_autotune_calls=0,active=False)
|
| 13 |
+
out=Path(output)/'compile-scope.json'
|
| 14 |
+
def persist():out.write_text(json.dumps(state,indent=2)+'\n')
|
| 15 |
+
def tuning(wrapper,*args,**kwargs):
|
| 16 |
+
previous=getattr(local,'autotuning',False)
|
| 17 |
+
local.autotuning=True;state['active']=True
|
| 18 |
+
state['autotune_blocks']+=1
|
| 19 |
+
try:return original_tuning(wrapper,*args,**kwargs)
|
| 20 |
+
finally:
|
| 21 |
+
local.autotuning=previous;state['active']=previous;persist()
|
| 22 |
+
def run(kernel,*args,**kwargs):
|
| 23 |
+
if getattr(local,'autotuning',False):
|
| 24 |
+
state['original_autotune_calls']+=1
|
| 25 |
+
return original_run(kernel,*args,**kwargs)
|
| 26 |
+
return selected_run(kernel,*args,**kwargs)
|
| 27 |
+
PythonWrapperCodegen.generate_and_run_autotune_block=tuning
|
| 28 |
+
CachingAutotuner.run=run
|
| 29 |
+
persist()
|
| 30 |
+
return lambda:dict(state)
|
bundle/plugin-site/ornith_g256/optimized/draft.py
ADDED
|
@@ -0,0 +1,90 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Isolated draft residual/RMSNorm repair; public target and sampler preserved."""
|
| 2 |
+
import ast
|
| 3 |
+
import hashlib
|
| 4 |
+
import json
|
| 5 |
+
import os
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
import torch
|
| 8 |
+
from ornith_g256.worker import OrnithG256Worker as Parent
|
| 9 |
+
from .draft_kernel import launch
|
| 10 |
+
|
| 11 |
+
R = Path(__file__).resolve().parents[1]
|
| 12 |
+
NAME = 'triton_red_fused_add_arange_bitwise_and_constant_pad_nd_fused_add_rms_norm_ge_mul_select_slice_unsqueeze_view_3'
|
| 13 |
+
KINDS = {
|
| 14 |
+
'3945f633918257a81244518749ca8c2613735dbedb656d49d811eba9e07a5f9a': 'middle',
|
| 15 |
+
'456e4167b1199b272395eae07c665ad6b501b39cb0018097a4abf87c72861bd7': 'final',
|
| 16 |
+
}
|
| 17 |
+
|
| 18 |
+
def install(output, cache_root):
|
| 19 |
+
out = Path(output)
|
| 20 |
+
state = dict(candidate_installed=True, kernels=[], calls={}, captured_rows={},
|
| 21 |
+
graph_replay_poison_check=False, sampler_changed=False)
|
| 22 |
+
def persist():
|
| 23 |
+
temp = out / 'kernel-provenance.tmp'
|
| 24 |
+
temp.write_text(json.dumps(state, indent=2))
|
| 25 |
+
temp.replace(out / 'kernel-provenance.json')
|
| 26 |
+
# Compile all configured graph shapes before graph capture, without RNG use.
|
| 27 |
+
for rows in [1, 8, 16, 32, 48, 64]:
|
| 28 |
+
residual = torch.zeros((rows, 2048), device='cuda', dtype=torch.float32)
|
| 29 |
+
base = torch.zeros((2, 2, 2048), device='cuda', dtype=torch.bfloat16)
|
| 30 |
+
coef = torch.zeros((rows, 512), device='cuda', dtype=torch.bfloat16)
|
| 31 |
+
hidden = torch.zeros_like(residual, dtype=torch.bfloat16)
|
| 32 |
+
mask = torch.full((1,), 15, device='cuda', dtype=torch.int32)
|
| 33 |
+
weight = torch.ones((2048,), device='cuda', dtype=torch.bfloat16)
|
| 34 |
+
out_residual = torch.empty_like(hidden)
|
| 35 |
+
out_norm = torch.empty_like(hidden)
|
| 36 |
+
for kind in ['middle', 'final']:
|
| 37 |
+
inputs = [residual, base, coef, hidden, mask, weight]
|
| 38 |
+
if kind == 'middle':
|
| 39 |
+
inputs.append(out_residual)
|
| 40 |
+
inputs.extend([out_norm, rows, 2048])
|
| 41 |
+
launch(kind, inputs, torch.cuda.current_stream().cuda_stream)
|
| 42 |
+
torch.cuda.synchronize()
|
| 43 |
+
graph = torch.cuda.CUDAGraph()
|
| 44 |
+
with torch.cuda.graph(graph):
|
| 45 |
+
launch(kind, inputs, torch.cuda.current_stream().cuda_stream)
|
| 46 |
+
out_norm.fill_(float('nan'))
|
| 47 |
+
out_residual.fill_(float('nan'))
|
| 48 |
+
graph.replay()
|
| 49 |
+
torch.cuda.synchronize()
|
| 50 |
+
assert torch.count_nonzero(out_norm) == 0
|
| 51 |
+
if kind == 'middle':
|
| 52 |
+
assert torch.count_nonzero(out_residual) == 0
|
| 53 |
+
del graph
|
| 54 |
+
state['graph_replay_poison_check'] = True
|
| 55 |
+
from torch._inductor.runtime.triton_heuristics import CachingAutotuner
|
| 56 |
+
original = CachingAutotuner.run
|
| 57 |
+
classified, seen = {}, set()
|
| 58 |
+
draft_root = str(Path(cache_root) / 'vllm/torch_compile_cache/torch_aot_compile')
|
| 59 |
+
|
| 60 |
+
def run(kernel, *inputs, stream, benchmark_run=False, **kwargs):
|
| 61 |
+
if kernel not in classified:
|
| 62 |
+
kind = None
|
| 63 |
+
if kernel.fn.__name__ == NAME:
|
| 64 |
+
relative = Path(kernel.filename).relative_to(draft_root)
|
| 65 |
+
assert len(relative.parts[0]) == 64 and 'inductor_cache' in relative.parts, kernel.filename
|
| 66 |
+
fn = next(n for n in ast.parse(kernel.fn.src).body if isinstance(n, ast.FunctionDef))
|
| 67 |
+
fn.decorator_list = []
|
| 68 |
+
digest = hashlib.sha256(ast.dump(fn, include_attributes=False).encode()).hexdigest()
|
| 69 |
+
assert digest in KINDS, digest
|
| 70 |
+
kind = KINDS[digest]
|
| 71 |
+
state['kernels'].append(dict(kind=kind, function=NAME, filename=kernel.filename, ast_sha256=digest))
|
| 72 |
+
classified[kernel] = kind
|
| 73 |
+
kind = classified[kernel]
|
| 74 |
+
if kind is None:
|
| 75 |
+
return original(kernel, *inputs, stream=stream, benchmark_run=benchmark_run, **kwargs)
|
| 76 |
+
assert not benchmark_run and not kwargs
|
| 77 |
+
launch(kind, inputs, stream)
|
| 78 |
+
rows = int(inputs[-2])
|
| 79 |
+
captured = torch.cuda.is_current_stream_capturing()
|
| 80 |
+
key = kind + ':' + str(rows)
|
| 81 |
+
state['calls'][key] = state['calls'].get(key, 0) + 1
|
| 82 |
+
if captured and rows not in state['captured_rows'].setdefault(kind, []):
|
| 83 |
+
state['captured_rows'][kind].append(rows)
|
| 84 |
+
if (kind, rows, captured) not in seen:
|
| 85 |
+
seen.add((kind, rows, captured))
|
| 86 |
+
persist()
|
| 87 |
+
return None
|
| 88 |
+
|
| 89 |
+
CachingAutotuner.run = run
|
| 90 |
+
persist()
|
bundle/plugin-site/ornith_g256/optimized/draft_kernel.py
ADDED
|
@@ -0,0 +1,50 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Single-pass convolution + residual RMSNorm; no scratch reread."""
|
| 2 |
+
import triton
|
| 3 |
+
import triton.language as tl
|
| 4 |
+
|
| 5 |
+
@triton.jit
|
| 6 |
+
def fused_norm_register(residual, base, coefficients, hidden, query_mask, weight,
|
| 7 |
+
out_residual, out_norm, rows: tl.constexpr,
|
| 8 |
+
RETURN_RESIDUAL: tl.constexpr, XBLOCK: tl.constexpr = 2):
|
| 9 |
+
row = tl.program_id(0) * XBLOCK + tl.arange(0, XBLOCK)[:, None]
|
| 10 |
+
col = tl.arange(0, 2048)[None, :]
|
| 11 |
+
valid = row < rows
|
| 12 |
+
mask = tl.load(query_mask)
|
| 13 |
+
b0 = tl.load(base + 4096 + col).to(tl.float32)
|
| 14 |
+
b1 = tl.load(base + 6144 + col).to(tl.float32)
|
| 15 |
+
c0 = tl.load(coefficients + 256 + 512 * row + col // 16, valid, 0).to(tl.float32)
|
| 16 |
+
c1 = tl.load(coefficients + 384 + 512 * row + col // 16, valid, 0).to(tl.float32)
|
| 17 |
+
x = tl.load(hidden + 2048 * row + col, valid, 0).to(tl.float32)
|
| 18 |
+
previous = tl.load(hidden + 2048 * (row - 1) + col, valid & (row >= 1), 0).to(tl.float32)
|
| 19 |
+
old = tl.load(residual + 2048 * row + col, valid, 0)
|
| 20 |
+
w = tl.load(weight + col).to(tl.float32)
|
| 21 |
+
coefficient0 = b0 + c0
|
| 22 |
+
coefficient1 = b1 + c1
|
| 23 |
+
term0 = coefficient0 * x
|
| 24 |
+
term1 = coefficient1 * previous
|
| 25 |
+
term1 = term1 * ((row & mask) >= 1).to(tl.float32)
|
| 26 |
+
convolved = term0 + term1
|
| 27 |
+
updated = convolved + old
|
| 28 |
+
square_sum = tl.sum(updated * updated, axis=1)[:, None]
|
| 29 |
+
inv = tl.rsqrt(square_sum / 2048.0 + 1e-6)
|
| 30 |
+
normalized = (updated * inv) * w
|
| 31 |
+
tl.store(residual + 2048 * row + col, updated, valid)
|
| 32 |
+
if RETURN_RESIDUAL:
|
| 33 |
+
tl.store(out_residual + 2048 * row + col, updated, valid)
|
| 34 |
+
tl.store(out_norm + 2048 * row + col, normalized, valid)
|
| 35 |
+
|
| 36 |
+
def launch(kind, inputs, stream):
|
| 37 |
+
import torch
|
| 38 |
+
assert stream == torch.cuda.current_stream().cuda_stream
|
| 39 |
+
if kind == 'middle':
|
| 40 |
+
residual, base, coef, hidden, mask, weight, out_residual, out_norm, rows, width = inputs
|
| 41 |
+
else:
|
| 42 |
+
residual, base, coef, hidden, mask, weight, out_norm, rows, width = inputs
|
| 43 |
+
out_residual = out_norm # ignored by the final-layer specialization
|
| 44 |
+
assert width == 2048
|
| 45 |
+
assert residual.dtype == torch.float32 and tuple(residual.shape) == (rows, 2048)
|
| 46 |
+
assert all(t.is_contiguous() for t in [residual, base, coef, hidden, mask, weight, out_residual, out_norm])
|
| 47 |
+
assert tuple(coef.shape) == (rows, 512) and tuple(base.shape) == (2, 2, 2048)
|
| 48 |
+
fused_norm_register[(triton.cdiv(rows, 2),)](residual, base, coef, hidden, mask, weight,
|
| 49 |
+
out_residual, out_norm, rows, kind == 'middle', XBLOCK=2, num_warps=8,
|
| 50 |
+
num_stages=1, enable_fp_fusion=False)
|
bundle/plugin-site/ornith_g256/optimized/grammar_parallel.py
ADDED
|
@@ -0,0 +1,122 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Parallelize independent vLLM grammar chains without changing their semantics.
|
| 2 |
+
|
| 3 |
+
Copyright 2026 Ciru. Licensed under Apache-2.0.
|
| 4 |
+
The original runtime implements every mask, reasoning transition and rollback.
|
| 5 |
+
Each worker receives a disjoint view of the original output buffer and exactly
|
| 6 |
+
one request, so the original speculative sequence remains ordered per matcher.
|
| 7 |
+
"""
|
| 8 |
+
from concurrent.futures import ThreadPoolExecutor
|
| 9 |
+
from copy import copy
|
| 10 |
+
from dataclasses import dataclass
|
| 11 |
+
import functools
|
| 12 |
+
import hashlib
|
| 13 |
+
import inspect
|
| 14 |
+
import os
|
| 15 |
+
from pathlib import Path
|
| 16 |
+
import time
|
| 17 |
+
|
| 18 |
+
BASELINE_SHA256 = '057f1ac57ec7d5f28b47a0ad85e03ae568951d408ceff345dbcac4c19b5d6ee9'
|
| 19 |
+
MIN_SERIAL_COST_NS = 750_000
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
@dataclass
|
| 23 |
+
class MaskState:
|
| 24 |
+
cost_per_request_ns: float = 0.0
|
| 25 |
+
pool: ThreadPoolExecutor | None = None
|
| 26 |
+
serial_batches: int = 0
|
| 27 |
+
parallel_batches: int = 0
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def _fill_one(original, manager, requests, req_id, scheduled, start, stop):
|
| 31 |
+
# A shallow facade shares immutable configuration/backend and owns its
|
| 32 |
+
# output-buffer view. All request-local state belongs to this one request.
|
| 33 |
+
facade = copy(manager)
|
| 34 |
+
facade._grammar_bitmask = manager._grammar_bitmask[start:stop]
|
| 35 |
+
begin = time.perf_counter_ns()
|
| 36 |
+
original(facade, requests, [req_id], scheduled)
|
| 37 |
+
return time.perf_counter_ns() - begin
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def _available_workers(count):
|
| 41 |
+
cpus = len(os.sched_getaffinity(0)) if hasattr(os, 'sched_getaffinity') else (os.cpu_count() or 1)
|
| 42 |
+
return max(1, min(count, cpus // 2, 8))
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def make_wrapper(original, grammar_type=None):
|
| 46 |
+
if grammar_type is None:
|
| 47 |
+
from vllm.v1.structured_output.backend_xgrammar import XgrammarGrammar
|
| 48 |
+
grammar_type = XgrammarGrammar
|
| 49 |
+
@functools.wraps(original)
|
| 50 |
+
def grammar_bitmask(self, requests, structured_output_request_ids, scheduled_spec_decode_tokens):
|
| 51 |
+
manager = self
|
| 52 |
+
request_ids = structured_output_request_ids
|
| 53 |
+
scheduled = scheduled_spec_decode_tokens
|
| 54 |
+
count = len(request_ids)
|
| 55 |
+
# Keep C0/C1, diffusion and the existing large-batch optimization on
|
| 56 |
+
# their original paths. This patch fills the small-batch serving gap.
|
| 57 |
+
if count < 2 or count > 128 or manager.vllm_config.model_config.is_diffusion:
|
| 58 |
+
return original(manager, requests, request_ids, scheduled)
|
| 59 |
+
|
| 60 |
+
state = getattr(manager, '_ciru_mask_state', None)
|
| 61 |
+
if state is None:
|
| 62 |
+
state = manager._ciru_mask_state = MaskState()
|
| 63 |
+
grammars = [requests[r].structured_output_request.grammar for r in request_ids]
|
| 64 |
+
if not all(isinstance(g, grammar_type) for g in grammars):
|
| 65 |
+
return original(manager, requests, request_ids, scheduled)
|
| 66 |
+
unique = len({id(g) for g in grammars}) == count
|
| 67 |
+
thinking = not manager.enable_in_reasoning and all(
|
| 68 |
+
requests[r].structured_output_request.reasoning_ended is False for r in request_ids
|
| 69 |
+
)
|
| 70 |
+
parallel = (
|
| 71 |
+
manager._grammar_bitmask is not None
|
| 72 |
+
and unique
|
| 73 |
+
and not thinking
|
| 74 |
+
and state.cost_per_request_ns * count >= MIN_SERIAL_COST_NS
|
| 75 |
+
and _available_workers(count) > 1
|
| 76 |
+
)
|
| 77 |
+
if not parallel:
|
| 78 |
+
begin = time.perf_counter_ns()
|
| 79 |
+
result = original(manager, requests, request_ids, scheduled)
|
| 80 |
+
state.cost_per_request_ns = (time.perf_counter_ns() - begin) / count
|
| 81 |
+
state.serial_batches += 1
|
| 82 |
+
return result
|
| 83 |
+
|
| 84 |
+
if state.pool is None:
|
| 85 |
+
state.pool = ThreadPoolExecutor(max_workers=_available_workers(manager.vllm_config.scheduler_config.max_num_seqs), thread_name_prefix='ciru-grammar')
|
| 86 |
+
futures = []
|
| 87 |
+
offset = 0
|
| 88 |
+
for req_id in request_ids:
|
| 89 |
+
rows = len(scheduled.get(req_id, ())) + 1
|
| 90 |
+
futures.append(state.pool.submit(_fill_one, original, manager, requests, req_id, scheduled, offset, offset + rows))
|
| 91 |
+
offset += rows
|
| 92 |
+
elapsed = []
|
| 93 |
+
first_error = None
|
| 94 |
+
# Drain every future before returning or raising: no worker can keep
|
| 95 |
+
# writing into a buffer the scheduler might reuse for its next step.
|
| 96 |
+
for future in futures:
|
| 97 |
+
try:
|
| 98 |
+
elapsed.append(future.result())
|
| 99 |
+
except BaseException as error:
|
| 100 |
+
if first_error is None:
|
| 101 |
+
first_error = error
|
| 102 |
+
state.parallel_batches += 1
|
| 103 |
+
if first_error is not None:
|
| 104 |
+
raise first_error
|
| 105 |
+
state.cost_per_request_ns = sum(elapsed) / count
|
| 106 |
+
return manager._grammar_bitmask[:offset].numpy()
|
| 107 |
+
return grammar_bitmask
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
def install():
|
| 111 |
+
"""Install once in the engine process, requiring the tested source identity."""
|
| 112 |
+
from vllm.v1.structured_output import StructuredOutputManager
|
| 113 |
+
if getattr(StructuredOutputManager, '_ciru_parallel_masks_installed', False):
|
| 114 |
+
return
|
| 115 |
+
source = Path(inspect.getsourcefile(StructuredOutputManager))
|
| 116 |
+
actual = hashlib.sha256(source.read_bytes()).hexdigest()
|
| 117 |
+
if actual != BASELINE_SHA256:
|
| 118 |
+
raise RuntimeError(f'Ciru parallel grammar requires its tested vLLM source; found {actual}')
|
| 119 |
+
StructuredOutputManager.grammar_bitmask = make_wrapper(StructuredOutputManager.grammar_bitmask)
|
| 120 |
+
StructuredOutputManager._ciru_parallel_masks_installed = True
|
| 121 |
+
from vllm.logger import init_logger
|
| 122 |
+
init_logger(__name__).info('Ciru parallel grammar masks R01 enabled; strict constraints and original per-request transitions retained')
|
bundle/plugin-site/ornith_g256/optimized/greedy.py
ADDED
|
@@ -0,0 +1,40 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Install the exact greedy walk beneath the unchanged public prefill wrapper."""
|
| 2 |
+
import ast, copy, hashlib, inspect, json, os
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
import torch
|
| 5 |
+
from .greedy_kernel import greedy_path
|
| 6 |
+
from ornith_g256.worker import OrnithG256Worker as Parent
|
| 7 |
+
|
| 8 |
+
EXPECTED_SOURCE='56b817daecb3740cad899a3bc1ebc481def465457d5e1a7a516e0535522e9e7f'
|
| 9 |
+
|
| 10 |
+
def install(out):
|
| 11 |
+
from ornith_g256 import prefill_draft
|
| 12 |
+
from vllm.v1.spec_decode.dflash import DFlashProposer
|
| 13 |
+
prefill_draft.install()
|
| 14 |
+
assert DFlashProposer._sample_draft_tokens is prefill_draft._sample_draft_tokens
|
| 15 |
+
original=prefill_draft._upstream_sample
|
| 16 |
+
source=Path(inspect.getsourcefile(original))
|
| 17 |
+
assert hashlib.sha256(source.read_bytes()).hexdigest()==EXPECTED_SOURCE
|
| 18 |
+
assert original.__module__=='vllm.v1.spec_decode.dflash'
|
| 19 |
+
tree=ast.parse(source.read_text())
|
| 20 |
+
cls=next(n for n in tree.body if isinstance(n,ast.ClassDef) and n.name=='DFlashProposer')
|
| 21 |
+
fn=copy.deepcopy(next(n for n in cls.body if isinstance(n,ast.FunctionDef) and n.name=='_sample_draft_tokens'))
|
| 22 |
+
fn.decorator_list=[]
|
| 23 |
+
index=next(i for i,n in enumerate(fn.body) if isinstance(n,ast.Assign)
|
| 24 |
+
and any(isinstance(t,ast.Name) and t.id=='rows' for t in n.targets))
|
| 25 |
+
fn.body.insert(index,ast.parse('if sampling_metadata.all_greedy:\n return qualified_path(scores, candidate_ids), None').body[0])
|
| 26 |
+
fn.name='_s13_sample'
|
| 27 |
+
out=Path(out)
|
| 28 |
+
globals_=dict(original.__globals__,qualified_path=greedy_path)
|
| 29 |
+
module=ast.fix_missing_locations(ast.Module(body=ast.parse('from __future__ import annotations').body+[fn],type_ignores=[]))
|
| 30 |
+
exec(compile(module,str(source)+':s13', 'exec'),globals_)
|
| 31 |
+
candidate=globals_['_s13_sample']
|
| 32 |
+
def dispatch(self,hidden,metadata):
|
| 33 |
+
if self.is_dflash2 and metadata.all_greedy:
|
| 34 |
+
return candidate(self,hidden,metadata)
|
| 35 |
+
return original(self,hidden,metadata)
|
| 36 |
+
prefill_draft._upstream_sample=dispatch
|
| 37 |
+
with (out/'sampler-installed.json').open('x') as f:
|
| 38 |
+
json.dump(dict(source_sha256=EXPECTED_SOURCE,prefill_wrapper_preserved=True,
|
| 39 |
+
non_greedy_original=True,delta='Greedy early return after unchanged score calculation'),f,indent=2)
|
| 40 |
+
|
bundle/plugin-site/ornith_g256/optimized/greedy_kernel.py
ADDED
|
@@ -0,0 +1,40 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Exact greedy DFlash2 path walk; candidate logits and edge scores unchanged."""
|
| 2 |
+
import torch
|
| 3 |
+
import triton
|
| 4 |
+
import triton.language as tl
|
| 5 |
+
|
| 6 |
+
|
| 7 |
+
@triton.jit
|
| 8 |
+
def _walk(S, I, O, S0: tl.constexpr, S1: tl.constexpr, S2: tl.constexpr,
|
| 9 |
+
S3: tl.constexpr, I0: tl.constexpr, I1: tl.constexpr, I2: tl.constexpr,
|
| 10 |
+
L: tl.constexpr, K: tl.constexpr):
|
| 11 |
+
row = tl.program_id(0)
|
| 12 |
+
col = tl.arange(0, K)
|
| 13 |
+
predecessor = tl.full((), 0, tl.int32)
|
| 14 |
+
for step in tl.static_range(L):
|
| 15 |
+
values = tl.load(S + row*S0 + step*S1 + predecessor*S2 + col*S3).to(tl.float32)
|
| 16 |
+
# Torch argmax selects the first NaN, or the first maximum otherwise.
|
| 17 |
+
nan_index = tl.min(tl.where(values != values, col, K), 0)
|
| 18 |
+
maximum = tl.max(tl.where(values != values, -float('inf'), values), 0)
|
| 19 |
+
max_index = tl.min(tl.where(values == maximum, col, K), 0)
|
| 20 |
+
selected = tl.where(nan_index < K, nan_index, max_index)
|
| 21 |
+
token = tl.load(I + row*I0 + step*I1 + selected*I2)
|
| 22 |
+
tl.store(O + row*L + step, token)
|
| 23 |
+
predecessor = selected
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def greedy_path(scores, candidate_ids):
|
| 27 |
+
b, length, k = candidate_ids.shape
|
| 28 |
+
if scores.shape != (b, length, k, k) or k != 16 or not 1 <= length <= 15:
|
| 29 |
+
raise ValueError('Unqualified DFlash2 score/ID shape')
|
| 30 |
+
if scores.dtype not in (torch.float32, torch.bfloat16, torch.float16):
|
| 31 |
+
raise ValueError('Unqualified score dtype')
|
| 32 |
+
if candidate_ids.dtype not in (torch.int32, torch.int64):
|
| 33 |
+
raise ValueError('Unqualified candidate-ID dtype')
|
| 34 |
+
if not scores.is_cuda or candidate_ids.device != scores.device:
|
| 35 |
+
raise ValueError('DFlash path requires matching GPU tensors')
|
| 36 |
+
output = torch.empty((b * length,), device=candidate_ids.device, dtype=candidate_ids.dtype)
|
| 37 |
+
if b:
|
| 38 |
+
_walk[(b,)](scores, candidate_ids, output, *scores.stride(),
|
| 39 |
+
*candidate_ids.stride(), length, k, num_warps=1)
|
| 40 |
+
return output
|
bundle/plugin-site/ornith_g256/optimized/launch.py
ADDED
|
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Published serving settings with the qualified wide prefill policy."""
|
| 2 |
+
from ornith_g256 import launch
|
| 3 |
+
|
| 4 |
+
original = launch.make_settings
|
| 5 |
+
|
| 6 |
+
|
| 7 |
+
def settings(args):
|
| 8 |
+
result = original(args)
|
| 9 |
+
result.update(worker_cls='ornith_g256.optimized.worker.Worker',
|
| 10 |
+
max_num_batched_tokens=8192, long_prefill_token_threshold=4480)
|
| 11 |
+
return result
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
launch.make_settings = settings
|
| 15 |
+
|
| 16 |
+
if __name__ == '__main__':
|
| 17 |
+
launch.main()
|
bundle/plugin-site/ornith_g256/optimized/mixed_qk_sum.py
ADDED
|
@@ -0,0 +1,66 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import triton
|
| 2 |
+
import triton.language as tl
|
| 3 |
+
|
| 4 |
+
@triton.jit
|
| 5 |
+
def mixed_qk_sum(in_ptr0, out_ptr0, out_ptr1, xnumel_0, xnumel_1, XBLOCK: tl.constexpr, R0_BLOCK: tl.constexpr, K_XBLOCK: tl.constexpr, K_RBLOCK: tl.constexpr):
|
| 6 |
+
pid = tl.program_id(0)
|
| 7 |
+
num_xblocks_0 = tl.cdiv(xnumel_0, XBLOCK)
|
| 8 |
+
num_xblocks_1 = num_xblocks_0 + tl.cdiv(xnumel_1, K_XBLOCK)
|
| 9 |
+
if pid < num_xblocks_0:
|
| 10 |
+
pid_offset = pid
|
| 11 |
+
r0_numel = 256
|
| 12 |
+
rnumel = r0_numel
|
| 13 |
+
RBLOCK: tl.constexpr = R0_BLOCK
|
| 14 |
+
xoffset = pid_offset * XBLOCK
|
| 15 |
+
xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
|
| 16 |
+
xmask = xindex < xnumel_0
|
| 17 |
+
r0_base = tl.arange(0, R0_BLOCK)[None, :]
|
| 18 |
+
rbase = r0_base
|
| 19 |
+
x0 = xindex % 16
|
| 20 |
+
x1 = xindex // 16
|
| 21 |
+
_tmp4 = tl.full([XBLOCK, R0_BLOCK], 0, tl.float32)
|
| 22 |
+
x3 = xindex
|
| 23 |
+
for r0_offset in tl.range(0, r0_numel, R0_BLOCK, num_stages=2):
|
| 24 |
+
r0_index = r0_offset + r0_base
|
| 25 |
+
r0_mask = r0_index < r0_numel
|
| 26 |
+
roffset = r0_offset
|
| 27 |
+
rindex = r0_index
|
| 28 |
+
r0_2 = r0_index
|
| 29 |
+
tmp0 = tl.load(in_ptr0 + (r0_2 + 512 * x0 + 9216 * x1), r0_mask & xmask, eviction_policy='evict_first', other=0.0).to(tl.float32)
|
| 30 |
+
tmp1 = tmp0.to(tl.float32)
|
| 31 |
+
tmp2 = tmp1 * tmp1
|
| 32 |
+
tmp3 = tl.broadcast_to(tmp2, [XBLOCK, R0_BLOCK])
|
| 33 |
+
tmp5 = _tmp4 + tmp3
|
| 34 |
+
_tmp4 = tl.where(r0_mask & xmask, tmp5, _tmp4)
|
| 35 |
+
tmp4 = tl.sum(_tmp4, 1)[:, None]
|
| 36 |
+
tl.store(out_ptr0 + x3, tmp4, xmask)
|
| 37 |
+
elif pid < num_xblocks_1:
|
| 38 |
+
pid_offset = pid - num_xblocks_0
|
| 39 |
+
r0_numel = 256
|
| 40 |
+
rnumel = r0_numel
|
| 41 |
+
RBLOCK: tl.constexpr = K_RBLOCK
|
| 42 |
+
xoffset = pid_offset * K_XBLOCK
|
| 43 |
+
xindex = xoffset + tl.arange(0, K_XBLOCK)[:, None]
|
| 44 |
+
xmask = xindex < xnumel_1
|
| 45 |
+
r0_base = tl.arange(0, K_RBLOCK)[None, :]
|
| 46 |
+
rbase = r0_base
|
| 47 |
+
x4 = xindex % 2
|
| 48 |
+
x5 = xindex // 2
|
| 49 |
+
_tmp10 = tl.full([K_XBLOCK, K_RBLOCK], 0, tl.float32)
|
| 50 |
+
x7 = xindex
|
| 51 |
+
for r0_offset in tl.range(0, r0_numel, K_RBLOCK, num_stages=2):
|
| 52 |
+
r0_index = r0_offset + r0_base
|
| 53 |
+
r0_mask = r0_index < r0_numel
|
| 54 |
+
roffset = r0_offset
|
| 55 |
+
rindex = r0_index
|
| 56 |
+
r0_6 = r0_index
|
| 57 |
+
tmp6 = tl.load(in_ptr0 + (8192 + r0_6 + 256 * x4 + 9216 * x5), r0_mask & xmask, eviction_policy='evict_first', other=0.0).to(tl.float32)
|
| 58 |
+
tmp7 = tmp6.to(tl.float32)
|
| 59 |
+
tmp8 = tmp7 * tmp7
|
| 60 |
+
tmp9 = tl.broadcast_to(tmp8, [K_XBLOCK, K_RBLOCK])
|
| 61 |
+
tmp11 = _tmp10 + tmp9
|
| 62 |
+
_tmp10 = tl.where(r0_mask & xmask, tmp11, _tmp10)
|
| 63 |
+
tmp10 = tl.sum(_tmp10, 1)[:, None]
|
| 64 |
+
tl.store(out_ptr1 + x7, tmp10, xmask)
|
| 65 |
+
else:
|
| 66 |
+
pass
|
bundle/plugin-site/ornith_g256/optimized/pipeline/SOURCE-PROOF.json
ADDED
|
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"chunk_scaled_dot_kkt.py": {
|
| 3 |
+
"parent_sha256": "7ae41df3755acc7e925a134d848f78bb5fc00b90069039d86403c5853323ea42",
|
| 4 |
+
"candidate_sha256": "71735d2b1c4a19eecbb6058708481817bb610426c48f96a5f646784afcf0c65c",
|
| 5 |
+
"reverse_exact_grid_program_id_edits_restores_identical_bytes": true,
|
| 6 |
+
"ast_equal_after_reverse": true
|
| 7 |
+
},
|
| 8 |
+
"solve_tril.py": {
|
| 9 |
+
"parent_sha256": "d1a4ff27623a938825052500afa7e2aa24fadd2fa4c9be6376c9128824f86f82",
|
| 10 |
+
"candidate_sha256": "bdd04b69c0f91562048d25834bd53d298dc1f1bd068b9494af88776b2bf4b1d1",
|
| 11 |
+
"reverse_exact_grid_program_id_edits_restores_identical_bytes": true,
|
| 12 |
+
"ast_equal_after_reverse": true
|
| 13 |
+
},
|
| 14 |
+
"wy_fast.py": {
|
| 15 |
+
"parent_sha256": "c9fc8259db0a8746dc04f1d799cb73333457118cf89a0bdce8dde8a80a28cd96",
|
| 16 |
+
"candidate_sha256": "4cc0be519bfcf3812d810a4bebb3dd45fb9b13a8d0ab62d0ad463a1a2e36f61f",
|
| 17 |
+
"reverse_exact_grid_program_id_edits_restores_identical_bytes": true,
|
| 18 |
+
"ast_equal_after_reverse": true
|
| 19 |
+
},
|
| 20 |
+
"chunk_o.py": {
|
| 21 |
+
"parent_sha256": "3b63b836a8cc561723c23db5a528a0b474d9cd78508d610f0b348c585a889a19",
|
| 22 |
+
"candidate_sha256": "a93f951048958265b4d1c3f9e67341329196e7a57525433934215fde4b16e61b",
|
| 23 |
+
"reverse_exact_grid_program_id_edits_restores_identical_bytes": true,
|
| 24 |
+
"ast_equal_after_reverse": true
|
| 25 |
+
},
|
| 26 |
+
"causal_conv1d.py": {
|
| 27 |
+
"parent_sha256": "230f3e554f5f1d609a37022578c84e42d5c71d46e684101278a7e21af4f7d418",
|
| 28 |
+
"candidate_sha256": "254c540f77b01694696a1ab535abd0e5a23d54e7211e202cf9beead04face16b",
|
| 29 |
+
"reverse_exact_grid_program_id_edits_restores_identical_bytes": true,
|
| 30 |
+
"ast_equal_after_reverse": true,
|
| 31 |
+
"preserves_HF_correctness_overlay": true,
|
| 32 |
+
"scope": "Existing guarded single-sequence non-APC scope only"
|
| 33 |
+
}
|
| 34 |
+
}
|
bundle/plugin-site/ornith_g256/optimized/pipeline/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
"""Retained Ornith runtime optimizations."""
|
bundle/plugin-site/ornith_g256/optimized/pipeline/adapter.py
ADDED
|
@@ -0,0 +1,172 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Isolated GDN spatial traversal adapter; leaves postconv/h/cumsum untouched.
|
| 2 |
+
|
| 3 |
+
Install with enabled=False during normal model warmup. Candidate dispatch reuses
|
| 4 |
+
only configurations already selected by the original Autotuner. A missing key
|
| 5 |
+
fails before compilation/benchmarking of alternative configurations.
|
| 6 |
+
"""
|
| 7 |
+
from __future__ import annotations
|
| 8 |
+
import copy
|
| 9 |
+
import hashlib
|
| 10 |
+
import importlib
|
| 11 |
+
import importlib.util
|
| 12 |
+
import inspect
|
| 13 |
+
from pathlib import Path
|
| 14 |
+
import sys
|
| 15 |
+
import types
|
| 16 |
+
import json
|
| 17 |
+
|
| 18 |
+
class ParentConfigurationRequired(Exception):
|
| 19 |
+
pass
|
| 20 |
+
|
| 21 |
+
ROOT = Path(__file__).resolve().parent
|
| 22 |
+
OPS = 'vllm.third_party.flash_linear_attention.ops'
|
| 23 |
+
SPECS = {
|
| 24 |
+
'chunk_scaled_dot_kkt_fwd': (OPS+'.chunk_scaled_dot_kkt','chunk_scaled_dot_kkt_fwd_kernel'),
|
| 25 |
+
'solve_tril': (OPS+'.solve_tril','merge_16x16_to_64x64_inverse_kernel'),
|
| 26 |
+
'recompute_w_u_fwd': (OPS+'.wy_fast','recompute_w_u_fwd_kernel'),
|
| 27 |
+
'chunk_fwd_o': (OPS+'.chunk_o','chunk_fwd_kernel_o'),
|
| 28 |
+
'causal_conv1d_fn': ('vllm.model_executor.layers.mamba.ops.causal_conv1d','_causal_conv1d_fwd_kernel'),
|
| 29 |
+
}
|
| 30 |
+
|
| 31 |
+
def exact_tuner_key(tuner, args, kwargs):
|
| 32 |
+
# Exact installed Autotuner.run key construction, including dtype order.
|
| 33 |
+
nargs = dict(zip(tuner.arg_names,args))
|
| 34 |
+
all_args = {**nargs,**kwargs}
|
| 35 |
+
selected = {k:v for k,v in all_args.items() if k in tuner.arg_names}
|
| 36 |
+
key = [selected[k] for k in tuner.keys if k in selected]
|
| 37 |
+
for value in selected.values():
|
| 38 |
+
if hasattr(value,'dtype'): key.append(str(value.dtype))
|
| 39 |
+
return tuple(key)
|
| 40 |
+
|
| 41 |
+
def safe_conv_scope(arguments):
|
| 42 |
+
# Only one live sequence, no APC intermediate cache writes and no x/state alias.
|
| 43 |
+
if any(arguments.get(n) is not None for n in ('block_idx_first_scheduled_token','block_idx_last_scheduled_token','initial_state_idx','num_computed_tokens')):
|
| 44 |
+
return False
|
| 45 |
+
x = arguments['x']; starts=arguments['query_start_loc']; state=arguments['conv_states']
|
| 46 |
+
if tuple(starts.shape)!=(2,) or x.ndim!=2 or x.shape[1]<=8:
|
| 47 |
+
return False
|
| 48 |
+
if arguments.get('has_initial_state') is not None and tuple(arguments['has_initial_state'].shape)!=(1,):
|
| 49 |
+
return False
|
| 50 |
+
if x.untyped_storage().data_ptr()==state.untyped_storage().data_ptr():
|
| 51 |
+
return False
|
| 52 |
+
return True
|
| 53 |
+
|
| 54 |
+
class SpatialPipeline:
|
| 55 |
+
def __init__(self, enabled):
|
| 56 |
+
from triton.runtime.autotuner import Autotuner
|
| 57 |
+
from triton.runtime.jit import JITFunction
|
| 58 |
+
self.enabled=enabled
|
| 59 |
+
self.patches=[]
|
| 60 |
+
self.original={}
|
| 61 |
+
self.candidate={}
|
| 62 |
+
self.tuners={}
|
| 63 |
+
self.calls={name:{'parent':0,'candidate':0,'scope_fallback':0} for name in SPECS}
|
| 64 |
+
self.activations=[]
|
| 65 |
+
self.configs=[]
|
| 66 |
+
self.seen=set()
|
| 67 |
+
proof=json.loads((ROOT/'SOURCE-PROOF.json').read_text())
|
| 68 |
+
for stage,(module_name,kernel_name) in SPECS.items():
|
| 69 |
+
parent=importlib.import_module(module_name)
|
| 70 |
+
filename=module_name.rsplit('.',1)[1]+'.py'
|
| 71 |
+
assert hashlib.sha256(Path(parent.__file__).read_bytes()).hexdigest()==proof[filename]['parent_sha256'],('Installed source changed',module_name)
|
| 72 |
+
path=ROOT/'candidate'/filename
|
| 73 |
+
assert hashlib.sha256(path.read_bytes()).hexdigest()==proof[filename]['candidate_sha256']
|
| 74 |
+
# Loading inside the real package preserves every unchanged relative import.
|
| 75 |
+
name=module_name.rsplit('.',1)[0]+'._gdn_spatial_'+filename[:-3]
|
| 76 |
+
spec=importlib.util.spec_from_file_location(name,path)
|
| 77 |
+
candidate=importlib.util.module_from_spec(spec)
|
| 78 |
+
sys.modules[name]=candidate
|
| 79 |
+
spec.loader.exec_module(candidate)
|
| 80 |
+
new_leaf=getattr(candidate,kernel_name)
|
| 81 |
+
while not isinstance(new_leaf,JITFunction): new_leaf=new_leaf.fn
|
| 82 |
+
old_kernel=getattr(parent,kernel_name)
|
| 83 |
+
def clone(node, *, stage=stage, new_leaf=new_leaf):
|
| 84 |
+
if isinstance(node,JITFunction):
|
| 85 |
+
assert node.arg_names==new_leaf.arg_names
|
| 86 |
+
return new_leaf
|
| 87 |
+
duplicate=copy.copy(node)
|
| 88 |
+
duplicate.fn=clone(node.fn)
|
| 89 |
+
if isinstance(node,Autotuner):
|
| 90 |
+
assert duplicate.cache is node.cache and duplicate.configs is node.configs
|
| 91 |
+
duplicate.base_fn=new_leaf.fn
|
| 92 |
+
self.tuners[stage]=(node,duplicate)
|
| 93 |
+
run_impl=type(node).run
|
| 94 |
+
def guarded_run(this,*args,**kwargs):
|
| 95 |
+
key=exact_tuner_key(this,args,kwargs)
|
| 96 |
+
if len(this.configs)>1 and key not in this.cache:
|
| 97 |
+
raise ParentConfigurationRequired(stage, key)
|
| 98 |
+
config=this.cache[key] if len(this.configs)>1 else this.configs[0]
|
| 99 |
+
record_key=(stage,repr(key),repr(config.all_kwargs()))
|
| 100 |
+
if record_key not in self.seen:
|
| 101 |
+
self.seen.add(record_key)
|
| 102 |
+
self.configs.append(dict(stage=stage,key=repr(key),parent_selected_config=config.all_kwargs(),same_config_object=True))
|
| 103 |
+
return run_impl(this,*args,**kwargs)
|
| 104 |
+
duplicate.run=types.MethodType(guarded_run,duplicate)
|
| 105 |
+
return duplicate
|
| 106 |
+
candidate.__dict__[kernel_name]=clone(old_kernel)
|
| 107 |
+
self.original[stage]=getattr(parent,stage)
|
| 108 |
+
self.candidate[stage]=getattr(candidate,stage)
|
| 109 |
+
from vllm.third_party.flash_linear_attention.ops import chunk
|
| 110 |
+
from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn as qwen
|
| 111 |
+
chunk_globals=chunk.chunk_gated_delta_rule_fwd.__globals__
|
| 112 |
+
core_globals=qwen.QwenGatedDeltaNetAttention._forward_core.__globals__
|
| 113 |
+
# Validate the whole hook set before the first global mutation.
|
| 114 |
+
for stage in SPECS:
|
| 115 |
+
namespace=core_globals if stage=='causal_conv1d_fn' else chunk_globals
|
| 116 |
+
if namespace.get(stage) is not self.original[stage]:
|
| 117 |
+
raise RuntimeError(f'Unexpected actual installed caller alias for {stage}')
|
| 118 |
+
for stage in SPECS:
|
| 119 |
+
namespace=core_globals if stage=='causal_conv1d_fn' else chunk_globals
|
| 120 |
+
original=self.original[stage]
|
| 121 |
+
signature=inspect.signature(original)
|
| 122 |
+
def dispatch(*args,_stage=stage,_signature=signature,**kwargs):
|
| 123 |
+
active=bool(self.enabled())
|
| 124 |
+
if active:
|
| 125 |
+
bound=_signature.bind(*args,**kwargs);bound.apply_defaults()
|
| 126 |
+
values=bound.arguments
|
| 127 |
+
if _stage=='causal_conv1d_fn':
|
| 128 |
+
safe=safe_conv_scope(values)
|
| 129 |
+
rows=int(values['x'].shape[1])
|
| 130 |
+
elif _stage=='solve_tril':
|
| 131 |
+
rows=int(values['A'].shape[1])
|
| 132 |
+
safe=values['A'].shape[-1]==64 and rows>8
|
| 133 |
+
else:
|
| 134 |
+
input_tensor=values['q'] if _stage=='chunk_fwd_o' else values['k']
|
| 135 |
+
rows=int(input_tensor.shape[1]);safe=rows>8
|
| 136 |
+
if not safe:
|
| 137 |
+
self.calls[_stage]['scope_fallback']+=1
|
| 138 |
+
active=False
|
| 139 |
+
else:
|
| 140 |
+
key=('activation',_stage,rows)
|
| 141 |
+
if key not in self.seen:
|
| 142 |
+
self.seen.add(key)
|
| 143 |
+
self.activations.append(dict(stage=_stage,rows=rows,scope='single-sequence non-APC' if _stage=='causal_conv1d_fn' else 'independent chunk/head tiles'))
|
| 144 |
+
arm='candidate' if active else 'parent'
|
| 145 |
+
self.calls[_stage][arm]+=1
|
| 146 |
+
if not active:
|
| 147 |
+
return self.original[_stage](*args, **kwargs)
|
| 148 |
+
try:
|
| 149 |
+
return self.candidate[_stage](*args, **kwargs)
|
| 150 |
+
except ParentConfigurationRequired:
|
| 151 |
+
self.calls[_stage]['candidate'] -= 1
|
| 152 |
+
self.calls[_stage]['parent'] += 1
|
| 153 |
+
self.calls[_stage]['tuner_cold'] = self.calls[_stage].get('tuner_cold', 0) + 1
|
| 154 |
+
return self.original[_stage](*args, **kwargs)
|
| 155 |
+
dispatch.__name__=stage
|
| 156 |
+
dispatch.__signature__=signature
|
| 157 |
+
self.patches.append((namespace,stage,original))
|
| 158 |
+
namespace[stage]=dispatch
|
| 159 |
+
|
| 160 |
+
def snapshot(self):
|
| 161 |
+
return dict(calls=copy.deepcopy(self.calls),activations=copy.deepcopy(self.activations),
|
| 162 |
+
parent_selected_configs=copy.deepcopy(self.configs),
|
| 163 |
+
shared_autotuner_cache={name:parent.cache is candidate.cache for name,(parent,candidate) in self.tuners.items()},
|
| 164 |
+
unchanged=['postconv','chunk_local_cumsum','chunk_gated_delta_rule_fwd_h','convolution update/decode','solve BT16/BT32'])
|
| 165 |
+
|
| 166 |
+
def restore(self):
|
| 167 |
+
for namespace,name,original in reversed(self.patches): namespace[name]=original
|
| 168 |
+
self.patches.clear()
|
| 169 |
+
|
| 170 |
+
def install(enabled):
|
| 171 |
+
"""Return a handle exposing snapshot()/restore(); enabled is a no-argument callable."""
|
| 172 |
+
return SpatialPipeline(enabled)
|
bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/causal_conv1d.py
ADDED
|
@@ -0,0 +1,1307 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
| 3 |
+
|
| 4 |
+
# Copyright (c) 2024, Tri Dao.
|
| 5 |
+
# Adapted from https://github.com/Dao-AILab/causal-conv1d/blob/main/causal_conv1d/causal_conv1d_interface.py
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
import numpy as np
|
| 9 |
+
import torch
|
| 10 |
+
|
| 11 |
+
from vllm.platforms import current_platform
|
| 12 |
+
from vllm.triton_utils import tl, triton
|
| 13 |
+
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"])
|
| 17 |
+
def _causal_conv1d_fwd_kernel( # continuous batching
|
| 18 |
+
# Pointers to matrices
|
| 19 |
+
x_ptr, # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences
|
| 20 |
+
w_ptr, # (dim, width)
|
| 21 |
+
bias_ptr,
|
| 22 |
+
initial_states_ptr, # conv_states_ptr
|
| 23 |
+
cache_indices_ptr, # (batch, n_blocks + padding) The second dimension contains
|
| 24 |
+
# the block indices relevant for each sequence
|
| 25 |
+
# plus potential 0-padding at the beginning and at the end
|
| 26 |
+
has_initial_states_ptr,
|
| 27 |
+
query_start_loc_ptr,
|
| 28 |
+
batch_ptr,
|
| 29 |
+
token_chunk_offset_ptr,
|
| 30 |
+
block_idx_first_scheduled_token, # (batch,)
|
| 31 |
+
block_idx_last_scheduled_token, # (batch,)
|
| 32 |
+
initial_state_idx, # (batch,)
|
| 33 |
+
num_computed_tokens, # (batch,)
|
| 34 |
+
num_accepted_tokens_ptr, # (batch,) or None
|
| 35 |
+
o_ptr, # (dim, seqlen) - actually pointing to x_ptr
|
| 36 |
+
# Matrix dimensions
|
| 37 |
+
dim: tl.constexpr,
|
| 38 |
+
num_cache_lines, # added to support vLLM larger cache lines
|
| 39 |
+
# Strides
|
| 40 |
+
stride_x_dim: tl.constexpr, # stride to get to next feature-value,
|
| 41 |
+
stride_x_token: tl.int64, # stride to get to next token (same feature-index, same sequence-index)
|
| 42 |
+
stride_w_dim: tl.constexpr, # stride to get to next dim-axis value
|
| 43 |
+
stride_w_width: tl.constexpr, # stride to get to next width-axis value
|
| 44 |
+
stride_istate_seq: tl.constexpr,
|
| 45 |
+
stride_istate_dim: tl.constexpr,
|
| 46 |
+
stride_istate_token: tl.constexpr,
|
| 47 |
+
stride_cache_indices: tl.constexpr,
|
| 48 |
+
stride_o_dim: tl.constexpr,
|
| 49 |
+
stride_o_token: tl.int64,
|
| 50 |
+
stride_block_m: tl.constexpr, # Stride block to align divided by BLOCK_M
|
| 51 |
+
# others
|
| 52 |
+
pad_slot_id: tl.constexpr,
|
| 53 |
+
null_block_id: tl.constexpr,
|
| 54 |
+
# Meta-parameters
|
| 55 |
+
HAS_BIAS: tl.constexpr,
|
| 56 |
+
KERNEL_WIDTH: tl.constexpr,
|
| 57 |
+
SILU_ACTIVATION: tl.constexpr,
|
| 58 |
+
IS_APC_ENABLED: tl.constexpr,
|
| 59 |
+
IS_SPEC_DECODING: tl.constexpr,
|
| 60 |
+
HAS_NULL_BLOCK: tl.constexpr,
|
| 61 |
+
NP2_STATELEN: tl.constexpr,
|
| 62 |
+
BLOCK_M: tl.constexpr,
|
| 63 |
+
BLOCK_N: tl.constexpr,
|
| 64 |
+
launch_pdl: tl.constexpr,
|
| 65 |
+
):
|
| 66 |
+
conv_states_ptr = initial_states_ptr
|
| 67 |
+
conv_state_indices_ptr = cache_indices_ptr
|
| 68 |
+
stride_conv_state_seq = stride_istate_seq
|
| 69 |
+
stride_conv_state_dim = stride_istate_dim
|
| 70 |
+
stride_conv_state_tok = stride_istate_token
|
| 71 |
+
state_len = (
|
| 72 |
+
KERNEL_WIDTH - 1
|
| 73 |
+
) # can be passed via argument if it's not the same as this value
|
| 74 |
+
|
| 75 |
+
if launch_pdl:
|
| 76 |
+
tl.extra.cuda.gdc_wait()
|
| 77 |
+
|
| 78 |
+
# one program handles one chunk in a single sequence
|
| 79 |
+
# rather than mixing sequences - to make updating initial_states across sequences efficiently
|
| 80 |
+
|
| 81 |
+
# single-sequence id
|
| 82 |
+
idx_seq = tl.load(batch_ptr + tl.program_id(1)).to(tl.int64)
|
| 83 |
+
|
| 84 |
+
if IS_SPEC_DECODING:
|
| 85 |
+
conv_state_token_offset = (
|
| 86 |
+
tl.load(num_accepted_tokens_ptr + idx_seq).to(tl.int64) - 1
|
| 87 |
+
)
|
| 88 |
+
else:
|
| 89 |
+
conv_state_token_offset = 0
|
| 90 |
+
chunk_offset = tl.load(token_chunk_offset_ptr + tl.program_id(1))
|
| 91 |
+
|
| 92 |
+
# BLOCK_N elements along the feature-dimension (channel)
|
| 93 |
+
idx_feats = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N)
|
| 94 |
+
|
| 95 |
+
if idx_seq == pad_slot_id:
|
| 96 |
+
if launch_pdl:
|
| 97 |
+
tl.extra.cuda.gdc_launch_dependents()
|
| 98 |
+
return
|
| 99 |
+
|
| 100 |
+
sequence_start_index = tl.load(query_start_loc_ptr + idx_seq)
|
| 101 |
+
sequence_end_index = tl.load(query_start_loc_ptr + idx_seq + 1)
|
| 102 |
+
# find the actual sequence length
|
| 103 |
+
seqlen = sequence_end_index - sequence_start_index
|
| 104 |
+
|
| 105 |
+
B_size: tl.constexpr = stride_block_m * BLOCK_M
|
| 106 |
+
|
| 107 |
+
if IS_APC_ENABLED:
|
| 108 |
+
# Handle the case if prefix caching is enabled.
|
| 109 |
+
# In particular, if prefix caching is enabled, the program write additional cache states to "cache_indices_ptr"
|
| 110 |
+
|
| 111 |
+
# Get the length of the completed sequence so far and compute the offset.
|
| 112 |
+
current_first_index = tl.load(block_idx_first_scheduled_token + idx_seq)
|
| 113 |
+
current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq)
|
| 114 |
+
sequence_completed_index = tl.load(num_computed_tokens + idx_seq)
|
| 115 |
+
|
| 116 |
+
# Compute the offset where the first stride_block_m-aligned first full block is
|
| 117 |
+
# Value in "token-space"
|
| 118 |
+
sequence_completed_offset_token = sequence_completed_index % B_size
|
| 119 |
+
seq_completed_offset = B_size - sequence_completed_offset_token
|
| 120 |
+
seq_end_offset = (seqlen - seq_completed_offset) % B_size
|
| 121 |
+
last_full_block_token_index = sequence_end_index - seq_end_offset
|
| 122 |
+
# If the sequence without the sequence_offset_index is stride_cache_chunk-aligned, then the last full chunk is the second-to-last one
|
| 123 |
+
if seq_end_offset == 0:
|
| 124 |
+
last_full_block_token_index = last_full_block_token_index - B_size
|
| 125 |
+
|
| 126 |
+
# Get the number of blocks to be filled for the current sequence
|
| 127 |
+
# If n_block_to_fill = 0, then only the state at the sequence end is stored
|
| 128 |
+
n_block_to_fill = current_last_index - current_first_index
|
| 129 |
+
|
| 130 |
+
# Get the index of the init block
|
| 131 |
+
conv_state_init_index = tl.load(initial_state_idx + idx_seq)
|
| 132 |
+
else:
|
| 133 |
+
n_block_to_fill = 0
|
| 134 |
+
current_last_index = 0
|
| 135 |
+
conv_state_init_index = 0
|
| 136 |
+
current_first_index = 0
|
| 137 |
+
last_full_block_token_index = 0
|
| 138 |
+
|
| 139 |
+
token_offset = BLOCK_M * chunk_offset
|
| 140 |
+
segment_len = min(BLOCK_M, seqlen - token_offset)
|
| 141 |
+
|
| 142 |
+
# base of the sequence
|
| 143 |
+
x_base = (
|
| 144 |
+
x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim
|
| 145 |
+
) # [BLOCK_N,]
|
| 146 |
+
|
| 147 |
+
# cache_idx
|
| 148 |
+
conv_states_input_coord = tl.load(
|
| 149 |
+
conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index
|
| 150 |
+
).to(tl.int64)
|
| 151 |
+
|
| 152 |
+
if HAS_NULL_BLOCK: # noqa
|
| 153 |
+
if conv_states_input_coord == null_block_id:
|
| 154 |
+
# not processing as this is a null block (padding)
|
| 155 |
+
if launch_pdl:
|
| 156 |
+
tl.extra.cuda.gdc_launch_dependents()
|
| 157 |
+
return
|
| 158 |
+
conv_states_base = (
|
| 159 |
+
conv_states_ptr
|
| 160 |
+
+ (conv_states_input_coord * stride_conv_state_seq)
|
| 161 |
+
+ (idx_feats * stride_conv_state_dim)
|
| 162 |
+
) # [BLOCK_N,]
|
| 163 |
+
|
| 164 |
+
w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,]
|
| 165 |
+
|
| 166 |
+
# Does 2 things:
|
| 167 |
+
# 1. READ prior-block init-state data - [done by every Triton programs]
|
| 168 |
+
# 2. update conv_state with new data [only by the Triton program handles chunk_offset=0]
|
| 169 |
+
if chunk_offset == 0:
|
| 170 |
+
# read from conv_states
|
| 171 |
+
load_init_state = tl.load(has_initial_states_ptr + idx_seq).to(tl.int1)
|
| 172 |
+
if load_init_state:
|
| 173 |
+
# load from conv_states
|
| 174 |
+
prior_tokens = (
|
| 175 |
+
conv_states_base
|
| 176 |
+
+ (state_len - 1 + conv_state_token_offset) * stride_conv_state_tok
|
| 177 |
+
)
|
| 178 |
+
mask_w = idx_feats < dim
|
| 179 |
+
if KERNEL_WIDTH == 2:
|
| 180 |
+
conv_states_ptrs = prior_tokens # [BLOCK_N]
|
| 181 |
+
col0 = tl.load(conv_states_ptrs, mask_w, 0.0)
|
| 182 |
+
if KERNEL_WIDTH == 3:
|
| 183 |
+
conv_states_ptrs = prior_tokens # [BLOCK_N]
|
| 184 |
+
col1 = tl.load(conv_states_ptrs, mask_w, 0.0)
|
| 185 |
+
conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N]
|
| 186 |
+
col0 = tl.load(conv_states_ptrs, mask_w, 0.0)
|
| 187 |
+
if KERNEL_WIDTH == 4:
|
| 188 |
+
conv_states_ptrs = prior_tokens # [BLOCK_N]
|
| 189 |
+
col2 = tl.load(conv_states_ptrs, mask_w, 0.0)
|
| 190 |
+
conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N]
|
| 191 |
+
col1 = tl.load(conv_states_ptrs, mask_w, 0.0)
|
| 192 |
+
conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N]
|
| 193 |
+
col0 = tl.load(conv_states_ptrs, mask_w, 0.0)
|
| 194 |
+
if KERNEL_WIDTH == 5:
|
| 195 |
+
conv_states_ptrs = prior_tokens # [BLOCK_N]
|
| 196 |
+
col3 = tl.load(conv_states_ptrs, mask_w, 0.0)
|
| 197 |
+
conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N]
|
| 198 |
+
col2 = tl.load(conv_states_ptrs, mask_w, 0.0)
|
| 199 |
+
conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N]
|
| 200 |
+
col1 = tl.load(conv_states_ptrs, mask_w, 0.0)
|
| 201 |
+
conv_states_ptrs = prior_tokens - 3 * stride_conv_state_tok # [BLOCK_N]
|
| 202 |
+
col0 = tl.load(conv_states_ptrs, mask_w, 0.0)
|
| 203 |
+
else:
|
| 204 |
+
# prior-tokens are zeros
|
| 205 |
+
if KERNEL_WIDTH >= 2: # STRATEGY1
|
| 206 |
+
# first chunk and does not have prior-token, so just set to 0
|
| 207 |
+
col0 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
|
| 208 |
+
if KERNEL_WIDTH >= 3: # STRATEGY1
|
| 209 |
+
col1 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
|
| 210 |
+
if KERNEL_WIDTH >= 4: # STRATEGY1
|
| 211 |
+
col2 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
|
| 212 |
+
if KERNEL_WIDTH >= 5: # STRATEGY1
|
| 213 |
+
col3 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty)
|
| 214 |
+
|
| 215 |
+
# STEP 2:
|
| 216 |
+
# here prepare data for updating conv_state
|
| 217 |
+
if (
|
| 218 |
+
state_len <= seqlen
|
| 219 |
+
): # SMALL_CACHE=True (only move part of 'x' into conv_state cache)
|
| 220 |
+
# just read from 'x'
|
| 221 |
+
# copy 'x' data to conv_state
|
| 222 |
+
# load only 'x' data (and set 0 before 'x' if seqlen < state_len)
|
| 223 |
+
idx_tokens_last = (seqlen - state_len) + tl.arange(
|
| 224 |
+
0, NP2_STATELEN
|
| 225 |
+
) # [BLOCK_M]
|
| 226 |
+
x_ptrs = (
|
| 227 |
+
x_ptr
|
| 228 |
+
+ ((sequence_start_index + idx_tokens_last) * stride_x_token)[:, None]
|
| 229 |
+
+ (idx_feats * stride_x_dim)[None, :]
|
| 230 |
+
) # [BLOCK_M,BLOCK_N,]
|
| 231 |
+
mask_x = (
|
| 232 |
+
(idx_tokens_last >= 0)[:, None]
|
| 233 |
+
& (idx_tokens_last < seqlen)[:, None]
|
| 234 |
+
& (idx_feats < dim)[None, :]
|
| 235 |
+
) # token-index # token-index # feature-index
|
| 236 |
+
loaded_x = tl.load(x_ptrs, mask_x, 0.0)
|
| 237 |
+
idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M]
|
| 238 |
+
|
| 239 |
+
# Compute the offset where the last block should be written in the conv_states
|
| 240 |
+
conv_states_output_coord = tl.load(
|
| 241 |
+
conv_state_indices_ptr
|
| 242 |
+
+ idx_seq * stride_cache_indices
|
| 243 |
+
+ current_last_index
|
| 244 |
+
).to(tl.int64)
|
| 245 |
+
|
| 246 |
+
conv_states_ptrs_target = (
|
| 247 |
+
conv_states_ptr
|
| 248 |
+
+ (conv_states_output_coord * stride_conv_state_seq) # Offset from seq
|
| 249 |
+
+ (idx_feats * stride_conv_state_dim)
|
| 250 |
+
)[None, :] + ( # [BLOCK_N,]
|
| 251 |
+
idx_tokens_conv * stride_conv_state_tok
|
| 252 |
+
)[:, None]
|
| 253 |
+
|
| 254 |
+
mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
|
| 255 |
+
tl.debug_barrier() # NOTE: use this due to bug in Triton compiler
|
| 256 |
+
tl.store(conv_states_ptrs_target, loaded_x, mask)
|
| 257 |
+
|
| 258 |
+
else:
|
| 259 |
+
if load_init_state:
|
| 260 |
+
# update conv_state by shifting left, i.e. take last few cols from conv_state + cols from 'x'
|
| 261 |
+
idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M]
|
| 262 |
+
|
| 263 |
+
conv_states_ptrs_source = (
|
| 264 |
+
conv_states_ptr
|
| 265 |
+
+ (conv_states_input_coord * stride_conv_state_seq)
|
| 266 |
+
+ (idx_feats * stride_conv_state_dim)[None, :]
|
| 267 |
+
+ (
|
| 268 |
+
(idx_tokens_conv + seqlen + conv_state_token_offset)
|
| 269 |
+
* stride_conv_state_tok
|
| 270 |
+
)[:, None]
|
| 271 |
+
) # [BLOCK_M, BLOCK_N]
|
| 272 |
+
mask = (
|
| 273 |
+
(conv_states_input_coord < num_cache_lines)
|
| 274 |
+
& ((idx_tokens_conv + seqlen) < state_len)[:, None]
|
| 275 |
+
& (idx_feats < dim)[None, :]
|
| 276 |
+
)
|
| 277 |
+
conv_state = tl.load(conv_states_ptrs_source, mask, other=0.0)
|
| 278 |
+
|
| 279 |
+
VAL = state_len - seqlen
|
| 280 |
+
|
| 281 |
+
x_ptrs = (
|
| 282 |
+
x_base[None, :]
|
| 283 |
+
+ ((idx_tokens_conv - VAL) * stride_x_token)[:, None]
|
| 284 |
+
) # [BLOCK_M, BLOCK_N]
|
| 285 |
+
|
| 286 |
+
mask_x = (
|
| 287 |
+
(idx_tokens_conv - VAL >= 0)[:, None]
|
| 288 |
+
& (idx_tokens_conv - VAL < seqlen)[:, None]
|
| 289 |
+
& (idx_feats < dim)[None, :]
|
| 290 |
+
) # token-index # token-index # feature-index
|
| 291 |
+
loaded_x = tl.load(x_ptrs, mask_x, 0.0)
|
| 292 |
+
|
| 293 |
+
tl.debug_barrier() # need this due to the bug in tl.where not enforcing this when data is the result of another tl.load
|
| 294 |
+
new_conv_state = tl.where(
|
| 295 |
+
mask, conv_state, loaded_x
|
| 296 |
+
) # BUG in 'tl.where' which requires a barrier before this
|
| 297 |
+
conv_states_ptrs_target = (
|
| 298 |
+
conv_states_base
|
| 299 |
+
+ (idx_tokens_conv * stride_conv_state_tok)[:, None]
|
| 300 |
+
) # [BLOCK_M, BLOCK_N]
|
| 301 |
+
mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[
|
| 302 |
+
None, :
|
| 303 |
+
]
|
| 304 |
+
tl.store(conv_states_ptrs_target, new_conv_state, mask)
|
| 305 |
+
else: # load_init_state == False
|
| 306 |
+
# update conv_state by shifting left, BUT
|
| 307 |
+
# set cols prior to 'x' as zeros + cols from 'x'
|
| 308 |
+
idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M]
|
| 309 |
+
|
| 310 |
+
VAL = state_len - seqlen
|
| 311 |
+
|
| 312 |
+
x_ptrs = (
|
| 313 |
+
x_base[None, :]
|
| 314 |
+
+ ((idx_tokens_conv - VAL) * stride_x_token)[:, None]
|
| 315 |
+
) # [BLOCK_M, BLOCK_N]
|
| 316 |
+
|
| 317 |
+
mask_x = (
|
| 318 |
+
(idx_tokens_conv - VAL >= 0)[:, None]
|
| 319 |
+
& (idx_tokens_conv - VAL < seqlen)[:, None]
|
| 320 |
+
& (idx_feats < dim)[None, :]
|
| 321 |
+
) # token-index # token-index # feature-index
|
| 322 |
+
new_conv_state = tl.load(x_ptrs, mask_x, 0.0)
|
| 323 |
+
|
| 324 |
+
conv_states_ptrs_target = (
|
| 325 |
+
conv_states_base
|
| 326 |
+
+ (idx_tokens_conv * stride_conv_state_tok)[:, None]
|
| 327 |
+
) # [BLOCK_M, BLOCK_N]
|
| 328 |
+
mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[
|
| 329 |
+
None, :
|
| 330 |
+
]
|
| 331 |
+
tl.store(conv_states_ptrs_target, new_conv_state, mask)
|
| 332 |
+
|
| 333 |
+
else: # chunk_offset > 0
|
| 334 |
+
# read prior-token data from `x`
|
| 335 |
+
load_init_state = True
|
| 336 |
+
prior_tokens = x_base + (token_offset - 1) * stride_x_token
|
| 337 |
+
mask_w = idx_feats < dim
|
| 338 |
+
if KERNEL_WIDTH == 2:
|
| 339 |
+
conv_states_ptrs = prior_tokens # [BLOCK_N]
|
| 340 |
+
col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
|
| 341 |
+
if KERNEL_WIDTH == 3:
|
| 342 |
+
conv_states_ptrs = prior_tokens # [BLOCK_N]
|
| 343 |
+
col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
|
| 344 |
+
conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N]
|
| 345 |
+
col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
|
| 346 |
+
if KERNEL_WIDTH == 4:
|
| 347 |
+
conv_states_ptrs = prior_tokens # [BLOCK_N]
|
| 348 |
+
col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
|
| 349 |
+
conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N]
|
| 350 |
+
col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
|
| 351 |
+
conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N]
|
| 352 |
+
col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
|
| 353 |
+
if KERNEL_WIDTH == 5:
|
| 354 |
+
# ruff: noqa: F841
|
| 355 |
+
conv_states_ptrs = prior_tokens # [BLOCK_N]
|
| 356 |
+
col3 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
|
| 357 |
+
conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N]
|
| 358 |
+
col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
|
| 359 |
+
conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N]
|
| 360 |
+
col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
|
| 361 |
+
conv_states_ptrs = prior_tokens - 3 * stride_x_token # [BLOCK_N]
|
| 362 |
+
col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca")
|
| 363 |
+
|
| 364 |
+
# Store intermediate states aligned with stride_block_m
|
| 365 |
+
# The additional states are cached starting from the last stride_block_m.
|
| 366 |
+
# For example:
|
| 367 |
+
# If n_block_to_fill = 0, then only the state at the sequence end is cached and the process below is not involved.
|
| 368 |
+
# If n_block_to_fill > 0, then the states at the sequence end and at the n_block_to_fill-last
|
| 369 |
+
# stride_block_m are cached.
|
| 370 |
+
# For example chunk_offset = n_block_to_fill stores the state at last_full_block
|
| 371 |
+
if (chunk_offset - 1) < n_block_to_fill:
|
| 372 |
+
# Store the states at the chunk boundaries from the start of the sequence
|
| 373 |
+
idx_tokens_last = (
|
| 374 |
+
last_full_block_token_index
|
| 375 |
+
- (n_block_to_fill - chunk_offset) * B_size
|
| 376 |
+
- state_len
|
| 377 |
+
) + tl.arange(0, NP2_STATELEN) # [BLOCK_M]
|
| 378 |
+
x_ptrs = (
|
| 379 |
+
x_ptr
|
| 380 |
+
+ (idx_tokens_last * stride_x_token)[:, None]
|
| 381 |
+
+ (idx_feats * stride_x_dim)[None, :]
|
| 382 |
+
) # [BLOCK_M,BLOCK_N,]
|
| 383 |
+
|
| 384 |
+
mask_x = (idx_tokens_last >= 0)[:, None] & (idx_feats < dim)[
|
| 385 |
+
None, :
|
| 386 |
+
] # token-index # token-index # feature-index
|
| 387 |
+
loaded_x = tl.load(x_ptrs, mask_x, 0.0)
|
| 388 |
+
idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M]
|
| 389 |
+
|
| 390 |
+
# cache_idx
|
| 391 |
+
conv_states_output_coord = tl.load(
|
| 392 |
+
conv_state_indices_ptr
|
| 393 |
+
+ idx_seq * stride_cache_indices
|
| 394 |
+
+ current_first_index
|
| 395 |
+
+ (chunk_offset - 1)
|
| 396 |
+
).to(tl.int64)
|
| 397 |
+
|
| 398 |
+
conv_states_ptrs_target = (
|
| 399 |
+
conv_states_ptr
|
| 400 |
+
+ (conv_states_output_coord * stride_conv_state_seq) # Offset from seq
|
| 401 |
+
+ (idx_feats * stride_conv_state_dim)
|
| 402 |
+
)[None, :] + ( # [BLOCK_N,]
|
| 403 |
+
idx_tokens_conv * stride_conv_state_tok
|
| 404 |
+
)[:, None]
|
| 405 |
+
|
| 406 |
+
mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :]
|
| 407 |
+
tl.debug_barrier() # NOTE: use this due to bug in Triton compiler
|
| 408 |
+
tl.store(conv_states_ptrs_target, loaded_x, mask)
|
| 409 |
+
|
| 410 |
+
if HAS_BIAS:
|
| 411 |
+
bias = bias_ptr + idx_feats
|
| 412 |
+
mask_bias = idx_feats < dim
|
| 413 |
+
acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to(
|
| 414 |
+
tl.float32
|
| 415 |
+
) # [BLOCK_N]
|
| 416 |
+
else:
|
| 417 |
+
acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32)
|
| 418 |
+
|
| 419 |
+
x_base_1d = x_base + token_offset * stride_x_token # starting of chunk
|
| 420 |
+
|
| 421 |
+
# PRE-LOAD WEIGHTS
|
| 422 |
+
mask_w = idx_feats < dim
|
| 423 |
+
if KERNEL_WIDTH >= 2:
|
| 424 |
+
w_ptrs = w_base + (0 * stride_w_width) # [BLOCK_N] tensor
|
| 425 |
+
w_col0 = tl.load(w_ptrs, mask_w, other=0.0)
|
| 426 |
+
w_ptrs = w_base + (1 * stride_w_width) # [BLOCK_N] tensor
|
| 427 |
+
w_col1 = tl.load(w_ptrs, mask_w, other=0.0)
|
| 428 |
+
if KERNEL_WIDTH >= 3:
|
| 429 |
+
w_ptrs = w_base + (2 * stride_w_width) # [BLOCK_N] tensor
|
| 430 |
+
w_col2 = tl.load(w_ptrs, mask_w, other=0.0)
|
| 431 |
+
if KERNEL_WIDTH >= 4:
|
| 432 |
+
w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor
|
| 433 |
+
w_col3 = tl.load(w_ptrs, mask_w, other=0.0)
|
| 434 |
+
mask_x_1d = idx_feats < dim
|
| 435 |
+
|
| 436 |
+
if launch_pdl:
|
| 437 |
+
tl.extra.cuda.gdc_launch_dependents()
|
| 438 |
+
|
| 439 |
+
for idx_token in range(segment_len):
|
| 440 |
+
acc = acc_preload
|
| 441 |
+
|
| 442 |
+
matrix_w = w_col0
|
| 443 |
+
matrix_x = col0
|
| 444 |
+
for j in tl.static_range(KERNEL_WIDTH):
|
| 445 |
+
if KERNEL_WIDTH == 2:
|
| 446 |
+
if j == 1: # KERNEL_WIDTH-1:
|
| 447 |
+
matrix_w = w_col1
|
| 448 |
+
x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N]
|
| 449 |
+
matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
|
| 450 |
+
elif KERNEL_WIDTH == 3:
|
| 451 |
+
if j == 1:
|
| 452 |
+
matrix_w = w_col1
|
| 453 |
+
matrix_x = col1
|
| 454 |
+
elif j == 2:
|
| 455 |
+
matrix_w = w_col2
|
| 456 |
+
x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N]
|
| 457 |
+
matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
|
| 458 |
+
elif KERNEL_WIDTH == 4:
|
| 459 |
+
if j == 1:
|
| 460 |
+
matrix_w = w_col1
|
| 461 |
+
matrix_x = col1
|
| 462 |
+
elif j == 2:
|
| 463 |
+
matrix_w = w_col2
|
| 464 |
+
matrix_x = col2
|
| 465 |
+
elif j == 3:
|
| 466 |
+
matrix_w = w_col3
|
| 467 |
+
x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N]
|
| 468 |
+
matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
|
| 469 |
+
|
| 470 |
+
acc += matrix_x.to(tl.float32) * matrix_w.to(tl.float32) # [BLOCK_N]
|
| 471 |
+
|
| 472 |
+
if KERNEL_WIDTH == 2:
|
| 473 |
+
col0 = matrix_x
|
| 474 |
+
elif KERNEL_WIDTH == 3:
|
| 475 |
+
col0 = col1
|
| 476 |
+
col1 = matrix_x
|
| 477 |
+
elif KERNEL_WIDTH == 4:
|
| 478 |
+
col0 = col1
|
| 479 |
+
col1 = col2
|
| 480 |
+
col2 = matrix_x
|
| 481 |
+
|
| 482 |
+
if SILU_ACTIVATION:
|
| 483 |
+
acc = acc / (1 + tl.exp(-acc))
|
| 484 |
+
mask_1d = (idx_token < segment_len) & (
|
| 485 |
+
idx_feats < dim
|
| 486 |
+
) # token-index # feature-index
|
| 487 |
+
o_ptrs = (
|
| 488 |
+
o_ptr
|
| 489 |
+
+ (sequence_start_index + token_offset + idx_token) * stride_o_token
|
| 490 |
+
+ (idx_feats * stride_o_dim)
|
| 491 |
+
)
|
| 492 |
+
|
| 493 |
+
tl.store(o_ptrs, acc, mask=mask_1d)
|
| 494 |
+
|
| 495 |
+
|
| 496 |
+
def causal_conv1d_fn(
|
| 497 |
+
x: torch.Tensor,
|
| 498 |
+
weight: torch.Tensor,
|
| 499 |
+
bias: torch.Tensor | None,
|
| 500 |
+
conv_states: torch.Tensor,
|
| 501 |
+
query_start_loc: torch.Tensor,
|
| 502 |
+
cache_indices: torch.Tensor | None = None,
|
| 503 |
+
has_initial_state: torch.Tensor | None = None,
|
| 504 |
+
activation: str | None = "silu",
|
| 505 |
+
pad_slot_id: int = PAD_SLOT_ID,
|
| 506 |
+
null_block_id: int = NULL_BLOCK_ID,
|
| 507 |
+
num_accepted_tokens: torch.Tensor | None = None,
|
| 508 |
+
block_idx_first_scheduled_token: torch.Tensor | None = None,
|
| 509 |
+
block_idx_last_scheduled_token: torch.Tensor | None = None,
|
| 510 |
+
initial_state_idx: torch.Tensor | None = None,
|
| 511 |
+
num_computed_tokens: torch.Tensor | None = None,
|
| 512 |
+
block_size_to_align=0,
|
| 513 |
+
metadata=None,
|
| 514 |
+
validate_data=False,
|
| 515 |
+
):
|
| 516 |
+
"""support varlen + continuous batching when x is 2D tensor
|
| 517 |
+
|
| 518 |
+
x: (dim,cu_seq_len)
|
| 519 |
+
cu_seq_len = total tokens of all seqs in that batch
|
| 520 |
+
sequences are concatenated from left to right for varlen
|
| 521 |
+
weight: (dim, width)
|
| 522 |
+
conv_states: (...,dim,width - 1) itype
|
| 523 |
+
updated inplace if cache_indices are not provided
|
| 524 |
+
[it use `cache_indices` to get the index to the cache of conv_state for that sequence
|
| 525 |
+
|
| 526 |
+
conv_state[cache_indices[i]] for seq-i - to be used as initial_state when has_initial_state[i] = True
|
| 527 |
+
and after that conv_state[cache_indices[i]] need to be shift-left and updated with values from 'x'
|
| 528 |
+
]
|
| 529 |
+
query_start_loc: (batch + 1) int32
|
| 530 |
+
The cumulative sequence lengths of the sequences in
|
| 531 |
+
the batch, used to index into sequence. prepended by 0.
|
| 532 |
+
if
|
| 533 |
+
x = [5, 1, 1, 1] <- continuous batching (batch=4)
|
| 534 |
+
then
|
| 535 |
+
query_start_loc = [0, 5, 6, 7, 8] <- the starting index of the next sequence; while the last value is
|
| 536 |
+
the ending index of the last sequence
|
| 537 |
+
[length(query_start_loc)-1 == batch]
|
| 538 |
+
for example: query_start_loc = torch.Tensor([0,10,16,17]),
|
| 539 |
+
x.shape=(dim,17)
|
| 540 |
+
cache_indices: (batch) int32
|
| 541 |
+
indicates the corresponding state index,
|
| 542 |
+
like so: conv_state = conv_states[cache_indices[batch_id]]
|
| 543 |
+
has_initial_state: (batch) bool
|
| 544 |
+
indicates whether should the kernel take the current state as initial
|
| 545 |
+
state for the calculations
|
| 546 |
+
[single boolean for each sequence in the batch: True or False]
|
| 547 |
+
bias: (dim,)
|
| 548 |
+
activation: either None or "silu" or "swish" or True
|
| 549 |
+
pad_slot_id: int
|
| 550 |
+
if cache_indices is passed, lets the kernel identify padded
|
| 551 |
+
entries that will not be processed,
|
| 552 |
+
for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id]
|
| 553 |
+
in this case, the kernel will not process entries at
|
| 554 |
+
indices 0 and 3
|
| 555 |
+
block_idx_first_scheduled_token: (batch,), dtype int32
|
| 556 |
+
The pointer into cache_indices, where the first cache block to be filled is located.
|
| 557 |
+
block_idx_last_scheduled_token: (batch,), dtype int32
|
| 558 |
+
The pointer into cache_indices, where the last cache block to be filled is located.
|
| 559 |
+
initial_state_idx: (batch,), dtype int32
|
| 560 |
+
The pointer into cache_indices, where the cache block containing the initial state is located.
|
| 561 |
+
num_computed_tokens: (batch,), dtype int32
|
| 562 |
+
The number of tokens already completed for each sequence
|
| 563 |
+
block_size_to_align: int
|
| 564 |
+
The block size to align the cached states to
|
| 565 |
+
out: same shape as `x`
|
| 566 |
+
"""
|
| 567 |
+
if isinstance(activation, bool) and activation:
|
| 568 |
+
activation = "silu"
|
| 569 |
+
|
| 570 |
+
args = None
|
| 571 |
+
# Store original dtype to cast back at the end
|
| 572 |
+
original_x_dtype = x.dtype
|
| 573 |
+
x = x.to(conv_states.dtype)
|
| 574 |
+
out = torch.empty_like(x)
|
| 575 |
+
if metadata is not None:
|
| 576 |
+
nums_dict = metadata.nums_dict
|
| 577 |
+
args = nums_dict
|
| 578 |
+
batch_ptr = metadata.batch_ptr
|
| 579 |
+
token_chunk_offset_ptr = metadata.token_chunk_offset_ptr
|
| 580 |
+
else:
|
| 581 |
+
seqlens = query_start_loc.diff().to("cpu")
|
| 582 |
+
args = seqlens
|
| 583 |
+
MAX_NUM_PROGRAMS = 1024
|
| 584 |
+
|
| 585 |
+
batch_ptr = torch.full(
|
| 586 |
+
(MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device
|
| 587 |
+
) # tracking which seq-idx the Triton program is handling
|
| 588 |
+
token_chunk_offset_ptr = torch.full(
|
| 589 |
+
(MAX_NUM_PROGRAMS,), PAD_SLOT_ID, dtype=torch.int32, device=x.device
|
| 590 |
+
) # tracking BLOCK_M-based index in the sequence the Triton program is handling
|
| 591 |
+
|
| 592 |
+
is_channel_last = (x.stride(0) == 1) & (x.stride(1) > 1)
|
| 593 |
+
dim, cu_seqlen = x.shape
|
| 594 |
+
_, width = weight.shape
|
| 595 |
+
state_len = width - 1
|
| 596 |
+
np2_statelen = triton.next_power_of_2(state_len)
|
| 597 |
+
|
| 598 |
+
padded_batch = query_start_loc.size(0) - 1
|
| 599 |
+
stride_x_dim = x.stride(0)
|
| 600 |
+
stride_x_token = x.stride(1)
|
| 601 |
+
stride_w_dim = weight.stride(0)
|
| 602 |
+
stride_w_width = weight.stride(1)
|
| 603 |
+
stride_istate_seq = 0
|
| 604 |
+
stride_istate_dim = 0
|
| 605 |
+
stride_istate_token = 0
|
| 606 |
+
num_cache_lines = 0
|
| 607 |
+
BLOCK_M = 8
|
| 608 |
+
if conv_states is not None:
|
| 609 |
+
# extensions to support vLLM:
|
| 610 |
+
# 1. conv_states is used to replaced initial_states
|
| 611 |
+
# 2. conv_states serve as a cache with num cache lines can be larger than batch size
|
| 612 |
+
# 3. mapping from sequence x[idx] to a cache line at index as specified via cache_indices[idx]
|
| 613 |
+
# 4. computation can be skipped if cache_indices[idx] == pad_slot_id
|
| 614 |
+
num_cache_lines = conv_states.size(0)
|
| 615 |
+
assert (
|
| 616 |
+
num_cache_lines == conv_states.shape[0]
|
| 617 |
+
and dim == conv_states.shape[1]
|
| 618 |
+
and width - 1 <= conv_states.shape[2]
|
| 619 |
+
)
|
| 620 |
+
stride_istate_seq = conv_states.stride(0)
|
| 621 |
+
stride_istate_dim = conv_states.stride(1)
|
| 622 |
+
stride_istate_token = conv_states.stride(2)
|
| 623 |
+
if out.dim() == 2:
|
| 624 |
+
stride_o_dim = out.stride(0)
|
| 625 |
+
stride_o_token = out.stride(1)
|
| 626 |
+
else:
|
| 627 |
+
stride_o_dim = out.stride(1)
|
| 628 |
+
stride_o_token = out.stride(2)
|
| 629 |
+
stride_cache_indices = cache_indices.stride(0) if cache_indices is not None else 0
|
| 630 |
+
|
| 631 |
+
if validate_data:
|
| 632 |
+
assert x.dim() == 2
|
| 633 |
+
assert query_start_loc is not None
|
| 634 |
+
assert query_start_loc.dim() == 1
|
| 635 |
+
assert x.stride(0) == 1 or x.stride(1) == 1
|
| 636 |
+
if bias is not None:
|
| 637 |
+
assert bias.dim() == 1
|
| 638 |
+
assert dim == bias.size(0)
|
| 639 |
+
if cache_indices is not None:
|
| 640 |
+
assert cache_indices.dim() == 1
|
| 641 |
+
assert padded_batch == cache_indices.size(0)
|
| 642 |
+
if has_initial_state is not None:
|
| 643 |
+
assert has_initial_state.size() == (padded_batch,)
|
| 644 |
+
assert conv_states is not None, (
|
| 645 |
+
"ERROR: `has_initial_state` is used, which needs also `conv_states`"
|
| 646 |
+
)
|
| 647 |
+
assert weight.stride(1) == 1
|
| 648 |
+
assert (dim, width) == weight.shape
|
| 649 |
+
assert is_channel_last, "Need to run in channel-last layout"
|
| 650 |
+
if block_size_to_align is not None and block_size_to_align > 0:
|
| 651 |
+
assert (block_size_to_align % BLOCK_M) == 0, (
|
| 652 |
+
"The mamba block size needs to be divisible by the BLOCK_M"
|
| 653 |
+
)
|
| 654 |
+
else:
|
| 655 |
+
block_size_to_align = BLOCK_M
|
| 656 |
+
|
| 657 |
+
if metadata is None:
|
| 658 |
+
|
| 659 |
+
def num_program(META, seqlens):
|
| 660 |
+
tot = 0
|
| 661 |
+
|
| 662 |
+
mlist = []
|
| 663 |
+
offsetlist = [] # type: ignore
|
| 664 |
+
|
| 665 |
+
nums = -(-seqlens // META["BLOCK_M"])
|
| 666 |
+
|
| 667 |
+
tot = nums.sum().item()
|
| 668 |
+
mlist = np.repeat(np.arange(len(nums)), nums)
|
| 669 |
+
for idx, num in enumerate(nums):
|
| 670 |
+
offsetlist.extend(
|
| 671 |
+
range(num)
|
| 672 |
+
) # chunk-idx if a sequence is split into multiple chunks
|
| 673 |
+
|
| 674 |
+
if META["batch_ptr"].nelement() < len(mlist):
|
| 675 |
+
newlen = len(mlist) + 1
|
| 676 |
+
META["batch_ptr"].resize_(newlen).fill_(PAD_SLOT_ID)
|
| 677 |
+
META["token_chunk_offset_ptr"].resize_(newlen).fill_(PAD_SLOT_ID)
|
| 678 |
+
|
| 679 |
+
if META["batch_ptr"].nelement() >= len(mlist):
|
| 680 |
+
META["batch_ptr"][0 : len(mlist)].copy_(
|
| 681 |
+
torch.from_numpy(np.array(mlist))
|
| 682 |
+
)
|
| 683 |
+
META["token_chunk_offset_ptr"][0 : len(mlist)].copy_(
|
| 684 |
+
torch.from_numpy(np.array(offsetlist))
|
| 685 |
+
)
|
| 686 |
+
|
| 687 |
+
META["batch_ptr"] = META["batch_ptr"].to(META["x_ptr"].device)
|
| 688 |
+
META["token_chunk_offset_ptr"] = META["token_chunk_offset_ptr"].to(
|
| 689 |
+
META["x_ptr"].device
|
| 690 |
+
)
|
| 691 |
+
return tot
|
| 692 |
+
else:
|
| 693 |
+
|
| 694 |
+
def num_program(META, nums_dict):
|
| 695 |
+
tot = nums_dict[META["BLOCK_M"]]["tot"]
|
| 696 |
+
|
| 697 |
+
mlist = nums_dict[META["BLOCK_M"]]["mlist"]
|
| 698 |
+
mlist_len = nums_dict[META["BLOCK_M"]]["mlist_len"]
|
| 699 |
+
|
| 700 |
+
offsetlist = nums_dict[META["BLOCK_M"]]["offsetlist"]
|
| 701 |
+
|
| 702 |
+
if nums_dict[META["BLOCK_M"]]["batch_ptr"] is not None:
|
| 703 |
+
META["batch_ptr"] = nums_dict[META["BLOCK_M"]]["batch_ptr"]
|
| 704 |
+
META["token_chunk_offset_ptr"] = nums_dict[META["BLOCK_M"]][
|
| 705 |
+
"token_chunk_offset_ptr"
|
| 706 |
+
]
|
| 707 |
+
else:
|
| 708 |
+
if META["batch_ptr"].nelement() < mlist_len:
|
| 709 |
+
newlen = mlist_len + 1
|
| 710 |
+
META["batch_ptr"].resize_(newlen).fill_(PAD_SLOT_ID)
|
| 711 |
+
META["token_chunk_offset_ptr"].resize_(newlen).fill_(PAD_SLOT_ID)
|
| 712 |
+
|
| 713 |
+
if META["batch_ptr"].nelement() >= mlist_len:
|
| 714 |
+
META["batch_ptr"][0:mlist_len].copy_(mlist)
|
| 715 |
+
META["token_chunk_offset_ptr"][0:mlist_len].copy_(offsetlist)
|
| 716 |
+
return tot
|
| 717 |
+
|
| 718 |
+
def grid(META):
|
| 719 |
+
return (
|
| 720 |
+
triton.cdiv(dim, META["BLOCK_N"]),
|
| 721 |
+
num_program(META, args),
|
| 722 |
+
)
|
| 723 |
+
|
| 724 |
+
if batch_ptr.device != x.device:
|
| 725 |
+
batch_ptr = batch_ptr.to(x.device)
|
| 726 |
+
token_chunk_offset_ptr = token_chunk_offset_ptr.to(x.device)
|
| 727 |
+
|
| 728 |
+
_causal_conv1d_fwd_kernel[grid](
|
| 729 |
+
# Pointers to matrices
|
| 730 |
+
x,
|
| 731 |
+
weight,
|
| 732 |
+
bias,
|
| 733 |
+
conv_states,
|
| 734 |
+
cache_indices,
|
| 735 |
+
has_initial_state,
|
| 736 |
+
query_start_loc,
|
| 737 |
+
batch_ptr,
|
| 738 |
+
token_chunk_offset_ptr,
|
| 739 |
+
block_idx_first_scheduled_token,
|
| 740 |
+
block_idx_last_scheduled_token,
|
| 741 |
+
initial_state_idx,
|
| 742 |
+
num_computed_tokens,
|
| 743 |
+
num_accepted_tokens,
|
| 744 |
+
out,
|
| 745 |
+
# Matrix dimensions
|
| 746 |
+
dim,
|
| 747 |
+
num_cache_lines,
|
| 748 |
+
# stride
|
| 749 |
+
stride_x_dim,
|
| 750 |
+
stride_x_token,
|
| 751 |
+
stride_w_dim,
|
| 752 |
+
stride_w_width,
|
| 753 |
+
stride_istate_seq,
|
| 754 |
+
stride_istate_dim,
|
| 755 |
+
stride_istate_token,
|
| 756 |
+
stride_cache_indices,
|
| 757 |
+
stride_o_dim,
|
| 758 |
+
stride_o_token,
|
| 759 |
+
block_size_to_align // BLOCK_M,
|
| 760 |
+
# others
|
| 761 |
+
pad_slot_id,
|
| 762 |
+
null_block_id,
|
| 763 |
+
# META
|
| 764 |
+
HAS_BIAS=bias is not None,
|
| 765 |
+
KERNEL_WIDTH=width,
|
| 766 |
+
SILU_ACTIVATION=activation in ["silu", "swish"],
|
| 767 |
+
IS_APC_ENABLED=block_idx_last_scheduled_token is not None,
|
| 768 |
+
IS_SPEC_DECODING=num_accepted_tokens is not None,
|
| 769 |
+
HAS_NULL_BLOCK=null_block_id is not None,
|
| 770 |
+
NP2_STATELEN=np2_statelen,
|
| 771 |
+
# launch_cooperative_grid=True
|
| 772 |
+
BLOCK_M=BLOCK_M,
|
| 773 |
+
BLOCK_N=256,
|
| 774 |
+
num_stages=2,
|
| 775 |
+
launch_pdl=current_platform.is_arch_support_pdl(),
|
| 776 |
+
)
|
| 777 |
+
return out.to(original_x_dtype)
|
| 778 |
+
|
| 779 |
+
|
| 780 |
+
@triton.jit(do_not_specialize_on_alignment=["num_cache_lines"])
|
| 781 |
+
def _causal_conv1d_update_kernel(
|
| 782 |
+
# Pointers to matrices
|
| 783 |
+
x_ptr, # (batch, dim, seqlen)
|
| 784 |
+
w_ptr, # (dim, width)
|
| 785 |
+
bias_ptr,
|
| 786 |
+
conv_state_ptr,
|
| 787 |
+
conv_state_indices_ptr,
|
| 788 |
+
num_accepted_tokens_ptr,
|
| 789 |
+
query_start_loc_ptr, # (batch + 1)
|
| 790 |
+
block_idx_last_scheduled_token, # (batch,)
|
| 791 |
+
initial_state_idx, # (batch,)
|
| 792 |
+
o_ptr, # (batch, dim, seqlen)
|
| 793 |
+
# Matrix dimensions
|
| 794 |
+
batch: int,
|
| 795 |
+
dim: tl.constexpr,
|
| 796 |
+
seqlen: tl.constexpr,
|
| 797 |
+
state_len: tl.constexpr,
|
| 798 |
+
num_cache_lines, # added to support vLLM larger cache lines
|
| 799 |
+
# Strides
|
| 800 |
+
stride_x_seq: tl.constexpr,
|
| 801 |
+
stride_x_dim: tl.constexpr,
|
| 802 |
+
stride_x_token: tl.int64,
|
| 803 |
+
stride_w_dim: tl.constexpr,
|
| 804 |
+
stride_w_width: tl.constexpr,
|
| 805 |
+
stride_conv_state_seq: tl.constexpr,
|
| 806 |
+
stride_conv_state_dim: tl.constexpr,
|
| 807 |
+
stride_conv_state_tok: tl.constexpr,
|
| 808 |
+
stride_state_indices: tl.constexpr,
|
| 809 |
+
stride_o_seq: tl.constexpr,
|
| 810 |
+
stride_o_dim: tl.constexpr,
|
| 811 |
+
stride_o_token: tl.int64,
|
| 812 |
+
# others
|
| 813 |
+
null_block_id: tl.constexpr,
|
| 814 |
+
# Meta-parameters
|
| 815 |
+
HAS_BIAS: tl.constexpr,
|
| 816 |
+
KERNEL_WIDTH: tl.constexpr,
|
| 817 |
+
SILU_ACTIVATION: tl.constexpr,
|
| 818 |
+
IS_VARLEN: tl.constexpr,
|
| 819 |
+
IS_APC_ENABLED: tl.constexpr,
|
| 820 |
+
IS_SPEC_DECODING: tl.constexpr,
|
| 821 |
+
NP2_STATELEN: tl.constexpr,
|
| 822 |
+
HAS_NULL_BLOCK: tl.constexpr,
|
| 823 |
+
BLOCK_N: tl.constexpr,
|
| 824 |
+
launch_pdl: tl.constexpr,
|
| 825 |
+
):
|
| 826 |
+
if launch_pdl:
|
| 827 |
+
tl.extra.cuda.gdc_wait()
|
| 828 |
+
|
| 829 |
+
# ruff: noqa: E501
|
| 830 |
+
idx_seq = tl.program_id(0)
|
| 831 |
+
if idx_seq >= batch:
|
| 832 |
+
if launch_pdl:
|
| 833 |
+
tl.extra.cuda.gdc_launch_dependents()
|
| 834 |
+
return
|
| 835 |
+
|
| 836 |
+
# [BLOCK_N,] elements along the feature-dimension (channel)
|
| 837 |
+
idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
|
| 838 |
+
|
| 839 |
+
if IS_APC_ENABLED:
|
| 840 |
+
# Get the state from the initial_state_idx
|
| 841 |
+
conv_state_init = tl.load(initial_state_idx + idx_seq)
|
| 842 |
+
current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq)
|
| 843 |
+
else:
|
| 844 |
+
conv_state_init = 0
|
| 845 |
+
current_last_index = 0
|
| 846 |
+
|
| 847 |
+
# cache_idx
|
| 848 |
+
conv_states_input_coord = tl.load(
|
| 849 |
+
conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init
|
| 850 |
+
).to(tl.int64)
|
| 851 |
+
|
| 852 |
+
if HAS_NULL_BLOCK: # noqa
|
| 853 |
+
if conv_states_input_coord == null_block_id:
|
| 854 |
+
# not processing as this is not the actual sequence
|
| 855 |
+
if launch_pdl:
|
| 856 |
+
tl.extra.cuda.gdc_launch_dependents()
|
| 857 |
+
return
|
| 858 |
+
|
| 859 |
+
if IS_VARLEN:
|
| 860 |
+
query_start_index = tl.load(query_start_loc_ptr + idx_seq).to(tl.int64)
|
| 861 |
+
query_end_index = tl.load(query_start_loc_ptr + (idx_seq + 1)).to(tl.int64)
|
| 862 |
+
# revise state_len and seqlen
|
| 863 |
+
state_len = state_len - (seqlen - (query_end_index - query_start_index))
|
| 864 |
+
seqlen = query_end_index - query_start_index
|
| 865 |
+
x_offset = query_start_index * stride_x_token
|
| 866 |
+
o_offset = query_start_index * stride_o_token
|
| 867 |
+
else:
|
| 868 |
+
query_start_index = idx_seq * seqlen
|
| 869 |
+
query_end_index = query_start_index + seqlen
|
| 870 |
+
x_offset = idx_seq * stride_x_seq
|
| 871 |
+
o_offset = idx_seq * stride_o_seq
|
| 872 |
+
|
| 873 |
+
if query_start_index == query_end_index:
|
| 874 |
+
if launch_pdl:
|
| 875 |
+
tl.extra.cuda.gdc_launch_dependents()
|
| 876 |
+
return
|
| 877 |
+
|
| 878 |
+
if IS_SPEC_DECODING:
|
| 879 |
+
# The rolling of conv state:
|
| 880 |
+
#
|
| 881 |
+
# Before forward, the conv_state is:
|
| 882 |
+
# [history1, history2, ..., historyM].
|
| 883 |
+
#
|
| 884 |
+
# After forward, the conv_state becomes:
|
| 885 |
+
# [history2, ..., historyM, draft1, draft2, ..., draftN].
|
| 886 |
+
#
|
| 887 |
+
# After acceptance, it becomes:
|
| 888 |
+
#
|
| 889 |
+
# - accept 1 tokens: [history2, ..., historyM, draft1]
|
| 890 |
+
# - accept 2 tokens: [history3, ..., historyM, draft1, draft2]
|
| 891 |
+
# - and so on.
|
| 892 |
+
conv_state_token_offset = (
|
| 893 |
+
tl.load(num_accepted_tokens_ptr + idx_seq).to(tl.int64) - 1
|
| 894 |
+
)
|
| 895 |
+
else:
|
| 896 |
+
conv_state_token_offset = 0
|
| 897 |
+
|
| 898 |
+
# STEP 1: READ init_state data
|
| 899 |
+
conv_states_base = (
|
| 900 |
+
conv_state_ptr
|
| 901 |
+
+ (conv_states_input_coord * stride_conv_state_seq)
|
| 902 |
+
+ (idx_feats * stride_conv_state_dim)
|
| 903 |
+
)
|
| 904 |
+
mask_w = idx_feats < dim
|
| 905 |
+
|
| 906 |
+
prior_tokens = conv_states_base + conv_state_token_offset * stride_conv_state_tok
|
| 907 |
+
if KERNEL_WIDTH >= 2:
|
| 908 |
+
conv_states_ptrs = prior_tokens # [BLOCK_N]
|
| 909 |
+
col0 = tl.load(conv_states_ptrs, mask_w, 0.0)
|
| 910 |
+
if KERNEL_WIDTH >= 3:
|
| 911 |
+
conv_states_ptrs = prior_tokens + 1 * stride_conv_state_tok # [BLOCK_N]
|
| 912 |
+
col1 = tl.load(conv_states_ptrs, mask_w, 0.0)
|
| 913 |
+
if KERNEL_WIDTH >= 4:
|
| 914 |
+
conv_states_ptrs = prior_tokens + 2 * stride_conv_state_tok # [BLOCK_N]
|
| 915 |
+
col2 = tl.load(conv_states_ptrs, mask_w, 0.0)
|
| 916 |
+
if KERNEL_WIDTH >= 5:
|
| 917 |
+
conv_states_ptrs = prior_tokens + 3 * stride_conv_state_tok # [BLOCK_N]
|
| 918 |
+
col3 = tl.load(conv_states_ptrs, mask_w, 0.0)
|
| 919 |
+
if KERNEL_WIDTH >= 6:
|
| 920 |
+
conv_states_ptrs = prior_tokens + 4 * stride_conv_state_tok # [BLOCK_N]
|
| 921 |
+
col4 = tl.load(conv_states_ptrs, mask_w, 0.0)
|
| 922 |
+
|
| 923 |
+
# STEP 2: assume state_len > seqlen
|
| 924 |
+
idx_tokens = tl.arange(0, NP2_STATELEN) # [BLOCK_M]
|
| 925 |
+
|
| 926 |
+
# With speculative decoding, the conv_state updates works in a sliding
|
| 927 |
+
# window manner, at each forward pass, the tokens are shift by 1, so we
|
| 928 |
+
# load since idx_tokens + 1.
|
| 929 |
+
conv_state_ptrs_source = (
|
| 930 |
+
conv_state_ptr
|
| 931 |
+
+ (conv_states_input_coord * stride_conv_state_seq)
|
| 932 |
+
+ conv_state_token_offset * stride_conv_state_tok
|
| 933 |
+
+ (idx_feats * stride_conv_state_dim)[None, :]
|
| 934 |
+
+ ((idx_tokens + (1 if IS_SPEC_DECODING else seqlen)) * stride_conv_state_tok)[
|
| 935 |
+
:, None
|
| 936 |
+
]
|
| 937 |
+
) # [BLOCK_M, BLOCK_N]
|
| 938 |
+
mask = (
|
| 939 |
+
(conv_states_input_coord < num_cache_lines)
|
| 940 |
+
& ((idx_tokens + seqlen) < state_len)[:, None]
|
| 941 |
+
& (idx_feats < dim)[None, :]
|
| 942 |
+
)
|
| 943 |
+
conv_state = tl.load(conv_state_ptrs_source, mask, other=0.0)
|
| 944 |
+
|
| 945 |
+
VAL = state_len - seqlen
|
| 946 |
+
x_base = x_ptr + x_offset + (idx_feats * stride_x_dim) # [BLOCK_N]
|
| 947 |
+
|
| 948 |
+
x_ptrs = (
|
| 949 |
+
x_base[None, :] + ((idx_tokens - VAL) * stride_x_token)[:, None]
|
| 950 |
+
) # [BLOCK_M, BLOCK_N]
|
| 951 |
+
|
| 952 |
+
mask_x = (
|
| 953 |
+
(idx_tokens - VAL >= 0)[:, None]
|
| 954 |
+
& (idx_tokens - VAL < seqlen)[:, None]
|
| 955 |
+
& (idx_feats < dim)[None, :]
|
| 956 |
+
) # token-index # token-index # feature-index
|
| 957 |
+
loaded_x = tl.load(x_ptrs, mask_x, 0.0)
|
| 958 |
+
tl.debug_barrier()
|
| 959 |
+
|
| 960 |
+
new_conv_state = tl.where(mask, conv_state, loaded_x)
|
| 961 |
+
|
| 962 |
+
# Get the state from the initial_state_idx
|
| 963 |
+
# cache_idx
|
| 964 |
+
conv_states_offset = tl.load(
|
| 965 |
+
conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index
|
| 966 |
+
).to(tl.int64)
|
| 967 |
+
conv_state_ptrs_target = (
|
| 968 |
+
conv_state_ptr
|
| 969 |
+
+ (conv_states_offset * stride_conv_state_seq) # Offset from seq
|
| 970 |
+
+ (idx_feats * stride_conv_state_dim)
|
| 971 |
+
)[None, :] + ( # [BLOCK_N,]
|
| 972 |
+
idx_tokens * stride_conv_state_tok
|
| 973 |
+
)[:, None]
|
| 974 |
+
mask = (idx_tokens < state_len)[:, None] & (idx_feats < dim)[None, :]
|
| 975 |
+
tl.store(conv_state_ptrs_target, new_conv_state, mask)
|
| 976 |
+
|
| 977 |
+
# STEP 3: init accumulator
|
| 978 |
+
if HAS_BIAS:
|
| 979 |
+
bias = bias_ptr + idx_feats
|
| 980 |
+
mask_bias = idx_feats < dim
|
| 981 |
+
acc_preload = tl.load(bias, mask=mask_bias, other=0.0).to(
|
| 982 |
+
tl.float32
|
| 983 |
+
) # [BLOCK_N]
|
| 984 |
+
else:
|
| 985 |
+
acc_preload = tl.zeros((BLOCK_N,), dtype=tl.float32)
|
| 986 |
+
|
| 987 |
+
# STEP 4:
|
| 988 |
+
# PRE-LOAD WEIGHTS
|
| 989 |
+
# first kernel column, configured for weights to handle BLOCK_N features in range
|
| 990 |
+
w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,]
|
| 991 |
+
mask_w = idx_feats < dim
|
| 992 |
+
if KERNEL_WIDTH >= 2:
|
| 993 |
+
w_ptrs = w_base + (0 * stride_w_width) # [BLOCK_N] tensor
|
| 994 |
+
w_col0 = tl.load(w_ptrs, mask_w, other=0.0)
|
| 995 |
+
w_ptrs = w_base + (1 * stride_w_width) # [BLOCK_N] tensor
|
| 996 |
+
w_col1 = tl.load(w_ptrs, mask_w, other=0.0)
|
| 997 |
+
if KERNEL_WIDTH >= 3:
|
| 998 |
+
w_ptrs = w_base + (2 * stride_w_width) # [BLOCK_N] tensor
|
| 999 |
+
w_col2 = tl.load(w_ptrs, mask_w, other=0.0)
|
| 1000 |
+
if KERNEL_WIDTH >= 4:
|
| 1001 |
+
w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor
|
| 1002 |
+
w_col3 = tl.load(w_ptrs, mask_w, other=0.0)
|
| 1003 |
+
if KERNEL_WIDTH >= 5:
|
| 1004 |
+
w_ptrs = w_base + (4 * stride_w_width) # [BLOCK_N] tensor
|
| 1005 |
+
w_col4 = tl.load(w_ptrs, mask_w, other=0.0)
|
| 1006 |
+
if KERNEL_WIDTH >= 6:
|
| 1007 |
+
w_ptrs = w_base + (5 * stride_w_width) # [BLOCK_N] tensor
|
| 1008 |
+
w_col5 = tl.load(w_ptrs, mask_w, other=0.0)
|
| 1009 |
+
|
| 1010 |
+
x_base_1d = x_base # starting of chunk [BLOCK_N]
|
| 1011 |
+
mask_x_1d = idx_feats < dim
|
| 1012 |
+
|
| 1013 |
+
# STEP 5: compute each token
|
| 1014 |
+
if launch_pdl:
|
| 1015 |
+
tl.extra.cuda.gdc_launch_dependents()
|
| 1016 |
+
|
| 1017 |
+
for idx_token in tl.range(seqlen):
|
| 1018 |
+
acc = acc_preload
|
| 1019 |
+
|
| 1020 |
+
matrix_w = w_col0
|
| 1021 |
+
matrix_x = col0
|
| 1022 |
+
for j in tl.static_range(KERNEL_WIDTH):
|
| 1023 |
+
if KERNEL_WIDTH == 2:
|
| 1024 |
+
if j == 1: # KERNEL_WIDTH-1:
|
| 1025 |
+
matrix_w = w_col1
|
| 1026 |
+
x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N]
|
| 1027 |
+
matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
|
| 1028 |
+
elif KERNEL_WIDTH == 3:
|
| 1029 |
+
if j == 1:
|
| 1030 |
+
matrix_w = w_col1
|
| 1031 |
+
matrix_x = col1
|
| 1032 |
+
elif j == 2:
|
| 1033 |
+
matrix_w = w_col2
|
| 1034 |
+
x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N]
|
| 1035 |
+
matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
|
| 1036 |
+
elif KERNEL_WIDTH == 4:
|
| 1037 |
+
if j == 1:
|
| 1038 |
+
matrix_w = w_col1
|
| 1039 |
+
matrix_x = col1
|
| 1040 |
+
elif j == 2:
|
| 1041 |
+
matrix_w = w_col2
|
| 1042 |
+
matrix_x = col2
|
| 1043 |
+
elif j == 3:
|
| 1044 |
+
matrix_w = w_col3
|
| 1045 |
+
x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N]
|
| 1046 |
+
matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
|
| 1047 |
+
elif KERNEL_WIDTH == 5:
|
| 1048 |
+
if j == 1:
|
| 1049 |
+
matrix_w = w_col1
|
| 1050 |
+
matrix_x = col1
|
| 1051 |
+
elif j == 2:
|
| 1052 |
+
matrix_w = w_col2
|
| 1053 |
+
matrix_x = col2
|
| 1054 |
+
elif j == 3:
|
| 1055 |
+
matrix_w = w_col3
|
| 1056 |
+
matrix_x = col3
|
| 1057 |
+
elif j == 4:
|
| 1058 |
+
matrix_w = w_col4
|
| 1059 |
+
x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N]
|
| 1060 |
+
matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
|
| 1061 |
+
elif KERNEL_WIDTH == 6:
|
| 1062 |
+
if j == 1:
|
| 1063 |
+
matrix_w = w_col1
|
| 1064 |
+
matrix_x = col1
|
| 1065 |
+
elif j == 2:
|
| 1066 |
+
matrix_w = w_col2
|
| 1067 |
+
matrix_x = col2
|
| 1068 |
+
elif j == 3:
|
| 1069 |
+
matrix_w = w_col3
|
| 1070 |
+
matrix_x = col3
|
| 1071 |
+
elif j == 4:
|
| 1072 |
+
matrix_w = w_col4
|
| 1073 |
+
matrix_x = col4
|
| 1074 |
+
elif j == 5:
|
| 1075 |
+
matrix_w = w_col5
|
| 1076 |
+
x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N]
|
| 1077 |
+
matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
|
| 1078 |
+
|
| 1079 |
+
acc += matrix_x.to(tl.float32) * matrix_w.to(tl.float32) # [BLOCK_N]
|
| 1080 |
+
|
| 1081 |
+
if KERNEL_WIDTH == 2:
|
| 1082 |
+
col0 = matrix_x
|
| 1083 |
+
elif KERNEL_WIDTH == 3:
|
| 1084 |
+
col0 = col1
|
| 1085 |
+
col1 = matrix_x
|
| 1086 |
+
elif KERNEL_WIDTH == 4:
|
| 1087 |
+
col0 = col1
|
| 1088 |
+
col1 = col2
|
| 1089 |
+
col2 = matrix_x
|
| 1090 |
+
elif KERNEL_WIDTH == 5:
|
| 1091 |
+
col0 = col1
|
| 1092 |
+
col1 = col2
|
| 1093 |
+
col2 = col3
|
| 1094 |
+
col3 = matrix_x
|
| 1095 |
+
elif KERNEL_WIDTH == 6:
|
| 1096 |
+
col0 = col1
|
| 1097 |
+
col1 = col2
|
| 1098 |
+
col2 = col3
|
| 1099 |
+
col3 = col4
|
| 1100 |
+
col4 = matrix_x
|
| 1101 |
+
|
| 1102 |
+
if SILU_ACTIVATION:
|
| 1103 |
+
acc = acc / (1 + tl.exp(-acc))
|
| 1104 |
+
mask_1d = (idx_token < seqlen) & (
|
| 1105 |
+
idx_feats < dim
|
| 1106 |
+
) # token-index # feature-index
|
| 1107 |
+
o_ptrs = (
|
| 1108 |
+
o_ptr + o_offset + idx_token * stride_o_token + (idx_feats * stride_o_dim)
|
| 1109 |
+
)
|
| 1110 |
+
|
| 1111 |
+
tl.store(o_ptrs, acc, mask=mask_1d)
|
| 1112 |
+
|
| 1113 |
+
|
| 1114 |
+
def causal_conv1d_update(
|
| 1115 |
+
x: torch.Tensor,
|
| 1116 |
+
conv_state: torch.Tensor,
|
| 1117 |
+
weight: torch.Tensor,
|
| 1118 |
+
bias: torch.Tensor | None = None,
|
| 1119 |
+
activation: bool | str | None = None,
|
| 1120 |
+
conv_state_indices: torch.Tensor | None = None,
|
| 1121 |
+
num_accepted_tokens: torch.Tensor | None = None,
|
| 1122 |
+
query_start_loc: torch.Tensor | None = None,
|
| 1123 |
+
max_query_len: int = -1,
|
| 1124 |
+
null_block_id: int = NULL_BLOCK_ID,
|
| 1125 |
+
block_idx_last_scheduled_token: torch.Tensor | None = None,
|
| 1126 |
+
initial_state_idx: torch.Tensor | None = None,
|
| 1127 |
+
validate_data=False,
|
| 1128 |
+
out: torch.Tensor | None = None,
|
| 1129 |
+
):
|
| 1130 |
+
"""
|
| 1131 |
+
x: Input tensor which can take the following shapes:
|
| 1132 |
+
|
| 1133 |
+
- `[batch, dim]` - single token prediction
|
| 1134 |
+
- `[batch, dim, seqlen]` - single or multiple tokens prediction
|
| 1135 |
+
- `[num_tokens, dim]` - continuous batching, where num_tokens is
|
| 1136 |
+
the total tokens of all sequences in that batch
|
| 1137 |
+
|
| 1138 |
+
conv_state: (..., dim, state_len), where state_len >= width - 1
|
| 1139 |
+
weight: (dim, width)
|
| 1140 |
+
bias: (dim,)
|
| 1141 |
+
conv_state_indices: (batch,), dtype int32
|
| 1142 |
+
If not None, the conv_state is a larger tensor along the batch dim,
|
| 1143 |
+
and we are selecting the batch coords specified by conv_state_indices.
|
| 1144 |
+
Useful for a continuous batching scenario.
|
| 1145 |
+
block_idx_last_scheduled_token: (batch,), dtype int32
|
| 1146 |
+
The pointer into conv_state_indices, where the last cache block to be filled is located.
|
| 1147 |
+
initial_state_idx: (batch,), dtype int32
|
| 1148 |
+
The pointer into conv_state_indices, where the cache block containing the initial state is located.
|
| 1149 |
+
num_accepted_tokens: (batch,), dtype int32
|
| 1150 |
+
If not None, it indicates the number of accepted tokens for each
|
| 1151 |
+
sequence in the batch.
|
| 1152 |
+
This is used in speculative decoding, where the conv_state is updated
|
| 1153 |
+
in a sliding window manner.
|
| 1154 |
+
query_start_loc: (batch + 1,) int32
|
| 1155 |
+
If not None, the inputs is given in a varlen fashion and this indicates
|
| 1156 |
+
the starting index of each sequence in the batch.
|
| 1157 |
+
max_query_len: int
|
| 1158 |
+
If query_start_loc is not None, this indicates the maximum query
|
| 1159 |
+
length in the batch.
|
| 1160 |
+
null_block_id: int
|
| 1161 |
+
Block ID used to identify padded entries in
|
| 1162 |
+
conv_state_indices. Block 0 is the null block.
|
| 1163 |
+
for example: conv_state_indices = [null_block_id, 1, 20, null_block_id]
|
| 1164 |
+
in this case, the kernel will not process entries at
|
| 1165 |
+
indices 0 and 3
|
| 1166 |
+
out: optional output tensor with the same shape as `x`. When omitted,
|
| 1167 |
+
the input is overwritten.
|
| 1168 |
+
"""
|
| 1169 |
+
if validate_data:
|
| 1170 |
+
assert null_block_id is not None
|
| 1171 |
+
assert x.stride(1) == 1
|
| 1172 |
+
if isinstance(activation, bool):
|
| 1173 |
+
activation = "silu" if activation is True else None
|
| 1174 |
+
elif activation is not None:
|
| 1175 |
+
assert activation in ["silu", "swish"]
|
| 1176 |
+
|
| 1177 |
+
original_x_dtype = x.dtype
|
| 1178 |
+
x = x.to(conv_state.dtype)
|
| 1179 |
+
if out is None:
|
| 1180 |
+
out = x
|
| 1181 |
+
else:
|
| 1182 |
+
if out.shape != x.shape:
|
| 1183 |
+
raise ValueError(
|
| 1184 |
+
f"`out` shape {tuple(out.shape)} must match `x` shape {tuple(x.shape)}."
|
| 1185 |
+
)
|
| 1186 |
+
if out.dtype != original_x_dtype or out.device != x.device:
|
| 1187 |
+
raise ValueError(
|
| 1188 |
+
"`out` must have the same dtype and device as the input `x`."
|
| 1189 |
+
)
|
| 1190 |
+
unsqueeze = query_start_loc is None and x.dim() == 2
|
| 1191 |
+
if unsqueeze:
|
| 1192 |
+
# make it (batch, dim, seqlen) with seqlen == 1
|
| 1193 |
+
x = x.unsqueeze(-1)
|
| 1194 |
+
out = out.unsqueeze(-1)
|
| 1195 |
+
if query_start_loc is None:
|
| 1196 |
+
batch, dim, seqlen = x.shape
|
| 1197 |
+
else:
|
| 1198 |
+
assert conv_state_indices is not None
|
| 1199 |
+
batch = conv_state_indices.size(0)
|
| 1200 |
+
dim = x.size(1)
|
| 1201 |
+
seqlen = max_query_len
|
| 1202 |
+
_, width = weight.shape
|
| 1203 |
+
# conv_state: (..., dim, state_len), where state_len >= width - 1
|
| 1204 |
+
num_cache_lines, _, state_len = conv_state.size()
|
| 1205 |
+
|
| 1206 |
+
if validate_data:
|
| 1207 |
+
assert dim == weight.size(0)
|
| 1208 |
+
assert state_len >= width - 1
|
| 1209 |
+
# when above happens, we don't shift-left to keep any records in conv_state
|
| 1210 |
+
assert dim == conv_state.size(1)
|
| 1211 |
+
if conv_state_indices is None:
|
| 1212 |
+
assert conv_state.size(0) >= batch
|
| 1213 |
+
else:
|
| 1214 |
+
assert batch == conv_state_indices.shape[0], (
|
| 1215 |
+
f"ERROR: conv_state_indices should have shape ({batch},*) but got {conv_state_indices.shape}"
|
| 1216 |
+
)
|
| 1217 |
+
|
| 1218 |
+
assert num_cache_lines >= batch
|
| 1219 |
+
assert weight.stride(1) == 1 # Need this
|
| 1220 |
+
|
| 1221 |
+
stride_w_dim, stride_w_width = weight.stride()
|
| 1222 |
+
|
| 1223 |
+
if query_start_loc is None:
|
| 1224 |
+
# X (batch, dim, seqlen)
|
| 1225 |
+
stride_x_seq, stride_x_dim, stride_x_token = x.stride()
|
| 1226 |
+
stride_o_seq, stride_o_dim, stride_o_token = out.stride()
|
| 1227 |
+
else:
|
| 1228 |
+
# X (dim, cu_seqlen)
|
| 1229 |
+
stride_x_token, stride_x_dim = x.stride()
|
| 1230 |
+
stride_x_seq = 0
|
| 1231 |
+
stride_o_token, stride_o_dim = out.stride()
|
| 1232 |
+
stride_o_seq = 0
|
| 1233 |
+
|
| 1234 |
+
stride_istate_seq, stride_istate_dim, stride_istate_token = conv_state.stride()
|
| 1235 |
+
stride_state_indices = (
|
| 1236 |
+
conv_state_indices.stride(0) if conv_state_indices is not None else 0
|
| 1237 |
+
)
|
| 1238 |
+
if num_accepted_tokens is not None:
|
| 1239 |
+
state_len = width - 1 + (seqlen - 1) # effective state_len needed
|
| 1240 |
+
else:
|
| 1241 |
+
state_len = width - 1
|
| 1242 |
+
np2_statelen = triton.next_power_of_2(state_len)
|
| 1243 |
+
|
| 1244 |
+
def grid(META):
|
| 1245 |
+
return (
|
| 1246 |
+
batch,
|
| 1247 |
+
triton.cdiv(dim, META["BLOCK_N"]),
|
| 1248 |
+
)
|
| 1249 |
+
|
| 1250 |
+
_causal_conv1d_update_kernel[grid](
|
| 1251 |
+
# Pointers to matrices
|
| 1252 |
+
x,
|
| 1253 |
+
weight,
|
| 1254 |
+
bias,
|
| 1255 |
+
conv_state,
|
| 1256 |
+
conv_state_indices,
|
| 1257 |
+
num_accepted_tokens,
|
| 1258 |
+
query_start_loc,
|
| 1259 |
+
block_idx_last_scheduled_token,
|
| 1260 |
+
initial_state_idx,
|
| 1261 |
+
out,
|
| 1262 |
+
# Matrix dimensions
|
| 1263 |
+
batch,
|
| 1264 |
+
dim,
|
| 1265 |
+
seqlen,
|
| 1266 |
+
state_len,
|
| 1267 |
+
num_cache_lines,
|
| 1268 |
+
# stride
|
| 1269 |
+
stride_x_seq,
|
| 1270 |
+
stride_x_dim,
|
| 1271 |
+
stride_x_token,
|
| 1272 |
+
stride_w_dim,
|
| 1273 |
+
stride_w_width,
|
| 1274 |
+
stride_istate_seq,
|
| 1275 |
+
stride_istate_dim,
|
| 1276 |
+
stride_istate_token,
|
| 1277 |
+
stride_state_indices,
|
| 1278 |
+
stride_o_seq,
|
| 1279 |
+
stride_o_dim,
|
| 1280 |
+
stride_o_token,
|
| 1281 |
+
# others
|
| 1282 |
+
null_block_id,
|
| 1283 |
+
# META
|
| 1284 |
+
HAS_BIAS=bias is not None,
|
| 1285 |
+
KERNEL_WIDTH=width,
|
| 1286 |
+
SILU_ACTIVATION=activation in ["silu", "swish"],
|
| 1287 |
+
IS_VARLEN=query_start_loc is not None,
|
| 1288 |
+
IS_APC_ENABLED=block_idx_last_scheduled_token is not None,
|
| 1289 |
+
IS_SPEC_DECODING=num_accepted_tokens is not None,
|
| 1290 |
+
NP2_STATELEN=np2_statelen,
|
| 1291 |
+
HAS_NULL_BLOCK=null_block_id is not None,
|
| 1292 |
+
BLOCK_N=256,
|
| 1293 |
+
launch_pdl=current_platform.is_arch_support_pdl(),
|
| 1294 |
+
)
|
| 1295 |
+
if unsqueeze:
|
| 1296 |
+
out = out.squeeze(-1)
|
| 1297 |
+
return out.to(original_x_dtype)
|
| 1298 |
+
|
| 1299 |
+
|
| 1300 |
+
if current_platform.is_cpu():
|
| 1301 |
+
from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
|
| 1302 |
+
causal_conv1d_fn_cpu,
|
| 1303 |
+
causal_conv1d_update_cpu,
|
| 1304 |
+
)
|
| 1305 |
+
|
| 1306 |
+
causal_conv1d_fn = causal_conv1d_fn_cpu # type: ignore
|
| 1307 |
+
causal_conv1d_update = causal_conv1d_update_cpu # type: ignore
|
bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/chunk_o.py
ADDED
|
@@ -0,0 +1,190 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
| 3 |
+
# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
|
| 4 |
+
#
|
| 5 |
+
# This file contains code copied from the flash-linear-attention project.
|
| 6 |
+
# The original source code was licensed under the MIT license and included
|
| 7 |
+
# the following copyright notice:
|
| 8 |
+
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
|
| 9 |
+
|
| 10 |
+
# ruff: noqa: E501
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
import torch
|
| 14 |
+
|
| 15 |
+
from vllm.triton_utils import tl, triton
|
| 16 |
+
|
| 17 |
+
from .index import prepare_chunk_indices
|
| 18 |
+
from .op import exp
|
| 19 |
+
from .utils import FLA_CHUNK_SIZE, check_shared_mem, is_nvidia_hopper
|
| 20 |
+
|
| 21 |
+
BKV_LIST = [64, 128] if check_shared_mem() else [32, 64]
|
| 22 |
+
NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8]
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
@triton.heuristics(
|
| 26 |
+
{
|
| 27 |
+
"USE_G": lambda args: args["g"] is not None,
|
| 28 |
+
"IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
|
| 29 |
+
}
|
| 30 |
+
)
|
| 31 |
+
@triton.autotune(
|
| 32 |
+
configs=[
|
| 33 |
+
triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages)
|
| 34 |
+
for BK in BKV_LIST
|
| 35 |
+
for BV in BKV_LIST
|
| 36 |
+
for num_warps in NUM_WARPS
|
| 37 |
+
for num_stages in [2, 3, 4]
|
| 38 |
+
],
|
| 39 |
+
key=["H", "K", "V", "BT"],
|
| 40 |
+
)
|
| 41 |
+
@triton.jit(do_not_specialize=["T"])
|
| 42 |
+
def chunk_fwd_kernel_o(
|
| 43 |
+
q,
|
| 44 |
+
k,
|
| 45 |
+
v,
|
| 46 |
+
h,
|
| 47 |
+
g,
|
| 48 |
+
o,
|
| 49 |
+
cu_seqlens,
|
| 50 |
+
chunk_indices,
|
| 51 |
+
scale,
|
| 52 |
+
T,
|
| 53 |
+
H: tl.constexpr,
|
| 54 |
+
Hg: tl.constexpr,
|
| 55 |
+
K: tl.constexpr,
|
| 56 |
+
V: tl.constexpr,
|
| 57 |
+
BT: tl.constexpr,
|
| 58 |
+
BK: tl.constexpr,
|
| 59 |
+
BV: tl.constexpr,
|
| 60 |
+
USE_G: tl.constexpr,
|
| 61 |
+
IS_VARLEN: tl.constexpr,
|
| 62 |
+
):
|
| 63 |
+
i_v, i_t, i_bh = tl.program_id(0), tl.program_id(2), tl.program_id(1)
|
| 64 |
+
i_b, i_h = i_bh // H, i_bh % H
|
| 65 |
+
|
| 66 |
+
if IS_VARLEN:
|
| 67 |
+
i_tg = i_t
|
| 68 |
+
i_n, i_t = (
|
| 69 |
+
tl.load(chunk_indices + i_t * 2).to(tl.int32),
|
| 70 |
+
tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
|
| 71 |
+
)
|
| 72 |
+
bos, eos = (
|
| 73 |
+
tl.load(cu_seqlens + i_n).to(tl.int32),
|
| 74 |
+
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
|
| 75 |
+
)
|
| 76 |
+
T = eos - bos
|
| 77 |
+
NT = tl.cdiv(T, BT)
|
| 78 |
+
else:
|
| 79 |
+
NT = tl.cdiv(T, BT)
|
| 80 |
+
i_tg = i_b * NT + i_t
|
| 81 |
+
bos, eos = i_b * T, i_b * T + T
|
| 82 |
+
|
| 83 |
+
# offset calculation
|
| 84 |
+
q += (bos * Hg + i_h // (H // Hg)) * K
|
| 85 |
+
k += (bos * Hg + i_h // (H // Hg)) * K
|
| 86 |
+
v += (bos * H + i_h) * V
|
| 87 |
+
o += (bos * H + i_h) * V
|
| 88 |
+
h += (i_tg * H + i_h).to(tl.int64) * V * K
|
| 89 |
+
|
| 90 |
+
b_o = tl.zeros([BT, BV], dtype=tl.float32)
|
| 91 |
+
b_A = tl.zeros([BT, BT], dtype=tl.float32)
|
| 92 |
+
|
| 93 |
+
for i_k in range(tl.cdiv(K, BK)):
|
| 94 |
+
p_q = tl.make_block_ptr(
|
| 95 |
+
q, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)
|
| 96 |
+
)
|
| 97 |
+
p_k = tl.make_block_ptr(
|
| 98 |
+
k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)
|
| 99 |
+
)
|
| 100 |
+
p_h = tl.make_block_ptr(
|
| 101 |
+
h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)
|
| 102 |
+
)
|
| 103 |
+
# [BT, BK]
|
| 104 |
+
b_q = tl.load(p_q, boundary_check=(0, 1))
|
| 105 |
+
# [BK, BT]
|
| 106 |
+
b_k = tl.load(p_k, boundary_check=(0, 1))
|
| 107 |
+
# [BV, BK]
|
| 108 |
+
b_h = tl.load(p_h, boundary_check=(0, 1))
|
| 109 |
+
|
| 110 |
+
# [BT, BK] @ [BK, BV] -> [BT, BV]
|
| 111 |
+
b_o += tl.dot(b_q, tl.trans(b_h))
|
| 112 |
+
# [BT, BK] @ [BK, BT] -> [BT, BT]
|
| 113 |
+
b_A += tl.dot(b_q, b_k)
|
| 114 |
+
|
| 115 |
+
if USE_G:
|
| 116 |
+
g += bos * H + i_h
|
| 117 |
+
p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,))
|
| 118 |
+
b_g = tl.load(p_g, boundary_check=(0,))
|
| 119 |
+
b_o = b_o * exp(b_g)[:, None]
|
| 120 |
+
b_A = b_A * exp(b_g[:, None] - b_g[None, :])
|
| 121 |
+
|
| 122 |
+
o_t = i_t * BT + tl.arange(0, BT)
|
| 123 |
+
m_t = o_t < T
|
| 124 |
+
m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t)
|
| 125 |
+
b_A = tl.where(m_A, b_A, 0)
|
| 126 |
+
|
| 127 |
+
p_v = tl.make_block_ptr(
|
| 128 |
+
v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)
|
| 129 |
+
)
|
| 130 |
+
p_o = tl.make_block_ptr(
|
| 131 |
+
o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)
|
| 132 |
+
)
|
| 133 |
+
b_v = tl.load(p_v, boundary_check=(0, 1))
|
| 134 |
+
|
| 135 |
+
# to fix mma -> mma layout conversion
|
| 136 |
+
# already solved by triton v3.2 or higher
|
| 137 |
+
b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale
|
| 138 |
+
tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
def chunk_fwd_o(
|
| 142 |
+
q: torch.Tensor,
|
| 143 |
+
k: torch.Tensor,
|
| 144 |
+
v: torch.Tensor,
|
| 145 |
+
h: torch.Tensor,
|
| 146 |
+
g: torch.Tensor | None = None, # cumsum of log decay
|
| 147 |
+
scale: float | None = None,
|
| 148 |
+
cu_seqlens: torch.Tensor | None = None,
|
| 149 |
+
chunk_indices: torch.Tensor | None = None,
|
| 150 |
+
chunk_size: int = FLA_CHUNK_SIZE,
|
| 151 |
+
core_attn_out: torch.Tensor | None = None,
|
| 152 |
+
) -> torch.Tensor:
|
| 153 |
+
B, T, Hg, K, V = *q.shape, v.shape[-1]
|
| 154 |
+
H = v.shape[-2]
|
| 155 |
+
BT = chunk_size
|
| 156 |
+
if chunk_indices is None and cu_seqlens is not None:
|
| 157 |
+
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
|
| 158 |
+
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
|
| 159 |
+
if scale is None:
|
| 160 |
+
scale = k.shape[-1] ** -0.5
|
| 161 |
+
|
| 162 |
+
if core_attn_out is not None:
|
| 163 |
+
assert core_attn_out.numel() >= v.numel(), (
|
| 164 |
+
f"core_attn_out too small: {core_attn_out.numel()} < {v.numel()}"
|
| 165 |
+
)
|
| 166 |
+
o = core_attn_out[: v.numel()].view(*v.shape)
|
| 167 |
+
else:
|
| 168 |
+
o = torch.empty_like(v)
|
| 169 |
+
|
| 170 |
+
def grid(meta):
|
| 171 |
+
return (triton.cdiv(V, meta["BV"]), B * H, NT)
|
| 172 |
+
|
| 173 |
+
chunk_fwd_kernel_o[grid](
|
| 174 |
+
q,
|
| 175 |
+
k,
|
| 176 |
+
v,
|
| 177 |
+
h,
|
| 178 |
+
g,
|
| 179 |
+
o,
|
| 180 |
+
cu_seqlens,
|
| 181 |
+
chunk_indices,
|
| 182 |
+
scale,
|
| 183 |
+
T=T,
|
| 184 |
+
H=H,
|
| 185 |
+
Hg=Hg,
|
| 186 |
+
K=K,
|
| 187 |
+
V=V,
|
| 188 |
+
BT=BT,
|
| 189 |
+
)
|
| 190 |
+
return o
|
bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/chunk_scaled_dot_kkt.py
ADDED
|
@@ -0,0 +1,176 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
| 3 |
+
# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
|
| 4 |
+
#
|
| 5 |
+
# This file contains code copied from the flash-linear-attention project.
|
| 6 |
+
# The original source code was licensed under the MIT license and included
|
| 7 |
+
# the following copyright notice:
|
| 8 |
+
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
|
| 9 |
+
# ruff: noqa: E501
|
| 10 |
+
|
| 11 |
+
import torch
|
| 12 |
+
|
| 13 |
+
from vllm.platforms import current_platform
|
| 14 |
+
from vllm.triton_utils import tl, triton
|
| 15 |
+
|
| 16 |
+
from .index import prepare_chunk_indices
|
| 17 |
+
from .op import exp
|
| 18 |
+
from .utils import FLA_CHUNK_SIZE
|
| 19 |
+
|
| 20 |
+
# On RDNA (gfx11xx/gfx12xx) WMMA only
|
| 21 |
+
# accepts 16-bit/int inputs, so a widened (e.g. fp32) tl.dot is lowered to a
|
| 22 |
+
# software matmul (~190x amdgcn-stage blowup). There we cast both operands down
|
| 23 |
+
# to k's native storage dtype (bf16/fp16) so fast WMMA is used instead.
|
| 24 |
+
_CAST_DOT_TO_K_DTYPE = False
|
| 25 |
+
if current_platform.is_rocm():
|
| 26 |
+
from vllm.platforms.rocm import on_gfx1x
|
| 27 |
+
|
| 28 |
+
_CAST_DOT_TO_K_DTYPE = on_gfx1x()
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
@triton.heuristics(
|
| 32 |
+
{
|
| 33 |
+
"USE_G": lambda args: args["g"] is not None,
|
| 34 |
+
"IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
|
| 35 |
+
}
|
| 36 |
+
)
|
| 37 |
+
@triton.autotune(
|
| 38 |
+
configs=[
|
| 39 |
+
triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages)
|
| 40 |
+
for BK in [32, 64, 128]
|
| 41 |
+
for num_warps in [2, 4, 8]
|
| 42 |
+
for num_stages in [2, 3, 4]
|
| 43 |
+
],
|
| 44 |
+
key=["H", "K", "BT", "IS_VARLEN"],
|
| 45 |
+
)
|
| 46 |
+
@triton.jit(do_not_specialize=["T"])
|
| 47 |
+
def chunk_scaled_dot_kkt_fwd_kernel(
|
| 48 |
+
k,
|
| 49 |
+
beta,
|
| 50 |
+
g,
|
| 51 |
+
A,
|
| 52 |
+
cu_seqlens,
|
| 53 |
+
chunk_indices,
|
| 54 |
+
T,
|
| 55 |
+
H: tl.constexpr,
|
| 56 |
+
Hg: tl.constexpr,
|
| 57 |
+
K: tl.constexpr,
|
| 58 |
+
BT: tl.constexpr,
|
| 59 |
+
BK: tl.constexpr,
|
| 60 |
+
IS_VARLEN: tl.constexpr,
|
| 61 |
+
USE_G: tl.constexpr,
|
| 62 |
+
CAST_DOT_TO_K_DTYPE: tl.constexpr,
|
| 63 |
+
):
|
| 64 |
+
i_t, i_bh = tl.program_id(1), tl.program_id(0)
|
| 65 |
+
i_b, i_h = i_bh // H, i_bh % H
|
| 66 |
+
if IS_VARLEN:
|
| 67 |
+
i_n, i_t = (
|
| 68 |
+
tl.load(chunk_indices + i_t * 2).to(tl.int32),
|
| 69 |
+
tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
|
| 70 |
+
)
|
| 71 |
+
bos, eos = (
|
| 72 |
+
tl.load(cu_seqlens + i_n).to(tl.int32),
|
| 73 |
+
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
|
| 74 |
+
)
|
| 75 |
+
T = eos - bos
|
| 76 |
+
else:
|
| 77 |
+
bos, eos = i_b * T, i_b * T + T
|
| 78 |
+
o_t = i_t * BT + tl.arange(0, BT)
|
| 79 |
+
m_t = o_t < T
|
| 80 |
+
|
| 81 |
+
p_beta = tl.make_block_ptr(
|
| 82 |
+
beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)
|
| 83 |
+
)
|
| 84 |
+
b_beta = tl.load(p_beta, boundary_check=(0,))
|
| 85 |
+
|
| 86 |
+
b_A = tl.zeros([BT, BT], dtype=tl.float32)
|
| 87 |
+
for i_k in range(tl.cdiv(K, BK)):
|
| 88 |
+
p_k = tl.make_block_ptr(
|
| 89 |
+
k + (bos * Hg + i_h // (H // Hg)) * K,
|
| 90 |
+
(T, K),
|
| 91 |
+
(Hg * K, 1),
|
| 92 |
+
(i_t * BT, i_k * BK),
|
| 93 |
+
(BT, BK),
|
| 94 |
+
(1, 0),
|
| 95 |
+
)
|
| 96 |
+
b_k = tl.load(p_k, boundary_check=(0, 1))
|
| 97 |
+
b_kb = b_k * b_beta[:, None]
|
| 98 |
+
if CAST_DOT_TO_K_DTYPE:
|
| 99 |
+
# RDNA: force operands to k's native dtype so WMMA is used.
|
| 100 |
+
b_A += tl.dot(b_kb.to(b_k.dtype), tl.trans(b_k))
|
| 101 |
+
else:
|
| 102 |
+
# Keep the promoted precision of the beta-scaled operand (WGMMA/MFMA).
|
| 103 |
+
b_A += tl.dot(b_kb, tl.trans(b_k).to(b_kb.dtype))
|
| 104 |
+
|
| 105 |
+
if USE_G:
|
| 106 |
+
p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
|
| 107 |
+
b_g = tl.load(p_g, boundary_check=(0,))
|
| 108 |
+
b_g_diff = b_g[:, None] - b_g[None, :]
|
| 109 |
+
b_A = b_A * exp(b_g_diff)
|
| 110 |
+
|
| 111 |
+
m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t)
|
| 112 |
+
b_A = tl.where(m_A, b_A, 0)
|
| 113 |
+
p_A = tl.make_block_ptr(
|
| 114 |
+
A + (bos * H + i_h) * BT, (T, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0)
|
| 115 |
+
)
|
| 116 |
+
tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1))
|
| 117 |
+
|
| 118 |
+
|
| 119 |
+
def chunk_scaled_dot_kkt_fwd(
|
| 120 |
+
k: torch.Tensor,
|
| 121 |
+
g: torch.Tensor | None = None,
|
| 122 |
+
beta: torch.Tensor | None = None,
|
| 123 |
+
cu_seqlens: torch.Tensor | None = None,
|
| 124 |
+
chunk_indices: torch.Tensor | None = None,
|
| 125 |
+
chunk_size: int = FLA_CHUNK_SIZE,
|
| 126 |
+
output_dtype: torch.dtype = torch.float32,
|
| 127 |
+
) -> torch.Tensor:
|
| 128 |
+
r"""
|
| 129 |
+
Compute beta * K * K^T.
|
| 130 |
+
|
| 131 |
+
Args:
|
| 132 |
+
k (torch.Tensor):
|
| 133 |
+
The key tensor of shape `[B, T, H, K]`.
|
| 134 |
+
beta (torch.Tensor):
|
| 135 |
+
The beta tensor of shape `[B, T, H]`.
|
| 136 |
+
g (torch.Tensor):
|
| 137 |
+
The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`.
|
| 138 |
+
cu_seqlens (torch.Tensor):
|
| 139 |
+
The cumulative sequence lengths of the input tensor.
|
| 140 |
+
Default: None
|
| 141 |
+
chunk_indices (torch.Tensor):
|
| 142 |
+
Pre-computed chunk indices. If None and cu_seqlens is provided,
|
| 143 |
+
computed internally. Default: None
|
| 144 |
+
chunk_size (int):
|
| 145 |
+
The chunk size. Default: 64.
|
| 146 |
+
output_dtype (torch.dtype):
|
| 147 |
+
The dtype of the output tensor. Default: `torch.float32`
|
| 148 |
+
|
| 149 |
+
Returns:
|
| 150 |
+
beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size.
|
| 151 |
+
"""
|
| 152 |
+
# This kernel is slightly different from fla to support Q/K with different head numbers.
|
| 153 |
+
# In fla, Q/K always have the same head number, so Hg is always equal to H.
|
| 154 |
+
B, T, Hg, K = k.shape
|
| 155 |
+
H = beta.shape[-1]
|
| 156 |
+
BT = chunk_size
|
| 157 |
+
if chunk_indices is None and cu_seqlens is not None:
|
| 158 |
+
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
|
| 159 |
+
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
|
| 160 |
+
|
| 161 |
+
A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype)
|
| 162 |
+
chunk_scaled_dot_kkt_fwd_kernel[(B * H, NT)](
|
| 163 |
+
k=k,
|
| 164 |
+
g=g,
|
| 165 |
+
beta=beta,
|
| 166 |
+
A=A,
|
| 167 |
+
cu_seqlens=cu_seqlens,
|
| 168 |
+
chunk_indices=chunk_indices,
|
| 169 |
+
T=T,
|
| 170 |
+
H=H,
|
| 171 |
+
Hg=Hg,
|
| 172 |
+
K=K,
|
| 173 |
+
BT=BT,
|
| 174 |
+
CAST_DOT_TO_K_DTYPE=_CAST_DOT_TO_K_DTYPE,
|
| 175 |
+
)
|
| 176 |
+
return A
|
bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/solve_tril.py
ADDED
|
@@ -0,0 +1,558 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
| 3 |
+
# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
|
| 4 |
+
#
|
| 5 |
+
# This file contains code copied from the flash-linear-attention project.
|
| 6 |
+
# The original source code was licensed under the MIT license and included
|
| 7 |
+
# the following copyright notice:
|
| 8 |
+
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
|
| 9 |
+
# ruff: noqa: E501
|
| 10 |
+
|
| 11 |
+
import os
|
| 12 |
+
|
| 13 |
+
import torch
|
| 14 |
+
|
| 15 |
+
from vllm.triton_utils import tl, triton
|
| 16 |
+
|
| 17 |
+
from .index import prepare_chunk_indices
|
| 18 |
+
from .op import make_tensor_descriptor
|
| 19 |
+
from .utils import input_guard, is_amd, is_tma_supported
|
| 20 |
+
|
| 21 |
+
FLA_TRIL_PRECISION = os.environ.get("FLA_TRIL_PRECISION", "ieee")
|
| 22 |
+
ALLOWED_TRIL_PRECISIONS = ["ieee", "tf32"] if is_amd else ["ieee", "tf32", "tf32x3"]
|
| 23 |
+
assert FLA_TRIL_PRECISION in ALLOWED_TRIL_PRECISIONS, (
|
| 24 |
+
f"FLA_TRIL_PRECISION must be one of {ALLOWED_TRIL_PRECISIONS}, but got {FLA_TRIL_PRECISION}"
|
| 25 |
+
)
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
|
| 29 |
+
@triton.autotune(
|
| 30 |
+
configs=[
|
| 31 |
+
triton.Config({}, num_warps=num_warps, num_stages=num_stages)
|
| 32 |
+
for num_warps in [1, 2, 4, 8]
|
| 33 |
+
for num_stages in [2, 3, 4, 5]
|
| 34 |
+
],
|
| 35 |
+
key=["BT"],
|
| 36 |
+
)
|
| 37 |
+
@triton.jit(do_not_specialize=["T"])
|
| 38 |
+
def solve_tril_16x16_kernel(
|
| 39 |
+
A,
|
| 40 |
+
Ai,
|
| 41 |
+
cu_seqlens,
|
| 42 |
+
chunk_indices,
|
| 43 |
+
T,
|
| 44 |
+
H: tl.constexpr,
|
| 45 |
+
BT: tl.constexpr,
|
| 46 |
+
USE_TMA: tl.constexpr,
|
| 47 |
+
IS_VARLEN: tl.constexpr,
|
| 48 |
+
DOT_PRECISION: tl.constexpr,
|
| 49 |
+
):
|
| 50 |
+
i_t, i_bh = tl.program_id(0), tl.program_id(1)
|
| 51 |
+
i_b, i_h = i_bh // H, i_bh % H
|
| 52 |
+
if IS_VARLEN:
|
| 53 |
+
i_n, i_t = (
|
| 54 |
+
tl.load(chunk_indices + i_t * 2).to(tl.int32),
|
| 55 |
+
tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
|
| 56 |
+
)
|
| 57 |
+
bos, eos = (
|
| 58 |
+
tl.load(cu_seqlens + i_n).to(tl.int32),
|
| 59 |
+
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
|
| 60 |
+
)
|
| 61 |
+
T = eos - bos
|
| 62 |
+
else:
|
| 63 |
+
bos, eos = i_b * T, i_b * T + T
|
| 64 |
+
o_i = tl.arange(0, 16)
|
| 65 |
+
m_A = o_i[:, None] > o_i[None, :]
|
| 66 |
+
m_I = o_i[:, None] == o_i[None, :]
|
| 67 |
+
|
| 68 |
+
A = A + (bos * H + i_h) * BT
|
| 69 |
+
Ai = Ai + (bos * H + i_h) * 16
|
| 70 |
+
|
| 71 |
+
offset = (i_t * 16) % BT
|
| 72 |
+
if not USE_TMA:
|
| 73 |
+
p_A = tl.make_block_ptr(
|
| 74 |
+
A, (T, BT), (H * BT, 1), (i_t * 16, offset), (16, 16), (1, 0)
|
| 75 |
+
)
|
| 76 |
+
# [16, 16]
|
| 77 |
+
b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32)
|
| 78 |
+
else:
|
| 79 |
+
desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16])
|
| 80 |
+
desc_o = make_tensor_descriptor(Ai, [T, 16], [H * 16, 1], [16, 16])
|
| 81 |
+
b_A = desc.load([i_t * 16, offset]).to(tl.float32)
|
| 82 |
+
b_A = -tl.where(m_A, b_A, 0)
|
| 83 |
+
|
| 84 |
+
for i in range(2, min(16, T - i_t * 16)):
|
| 85 |
+
# [16]
|
| 86 |
+
b_a = -tl.load(A + (i_t * 16 + i) * H * BT + o_i + offset)
|
| 87 |
+
b_a = b_a + tl.sum(b_a[:, None] * b_A, 0)
|
| 88 |
+
b_A = tl.where((o_i == i)[:, None], b_a, b_A)
|
| 89 |
+
b_A += m_I
|
| 90 |
+
if not USE_TMA:
|
| 91 |
+
p_Ai = tl.make_block_ptr(
|
| 92 |
+
Ai, (T, 16), (H * 16, 1), (i_t * 16, 0), (16, 16), (1, 0)
|
| 93 |
+
)
|
| 94 |
+
tl.store(
|
| 95 |
+
p_Ai,
|
| 96 |
+
b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"),
|
| 97 |
+
boundary_check=(0, 1),
|
| 98 |
+
)
|
| 99 |
+
else:
|
| 100 |
+
desc_o.store([i_t * 16, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne"))
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
|
| 104 |
+
@triton.autotune(
|
| 105 |
+
configs=[
|
| 106 |
+
triton.Config({}, num_warps=num_warps, num_stages=num_stages)
|
| 107 |
+
for num_warps in [1, 2, 4, 8]
|
| 108 |
+
for num_stages in [2, 3, 4, 5]
|
| 109 |
+
],
|
| 110 |
+
key=["H", "BT", "IS_VARLEN"],
|
| 111 |
+
)
|
| 112 |
+
@triton.jit(do_not_specialize=["T"])
|
| 113 |
+
def merge_16x16_to_32x32_inverse_kernel(
|
| 114 |
+
A,
|
| 115 |
+
Ai,
|
| 116 |
+
cu_seqlens,
|
| 117 |
+
chunk_indices,
|
| 118 |
+
T,
|
| 119 |
+
H: tl.constexpr,
|
| 120 |
+
BT: tl.constexpr,
|
| 121 |
+
USE_TMA: tl.constexpr,
|
| 122 |
+
IS_VARLEN: tl.constexpr,
|
| 123 |
+
DOT_PRECISION: tl.constexpr,
|
| 124 |
+
):
|
| 125 |
+
i_t, i_bh = tl.program_id(0), tl.program_id(1)
|
| 126 |
+
i_b, i_h = i_bh // H, i_bh % H
|
| 127 |
+
if IS_VARLEN:
|
| 128 |
+
i_n, i_t = (
|
| 129 |
+
tl.load(chunk_indices + i_t * 2).to(tl.int32),
|
| 130 |
+
tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
|
| 131 |
+
)
|
| 132 |
+
bos, eos = (
|
| 133 |
+
tl.load(cu_seqlens + i_n).to(tl.int32),
|
| 134 |
+
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
|
| 135 |
+
)
|
| 136 |
+
T = eos - bos
|
| 137 |
+
else:
|
| 138 |
+
bos, eos = i_b * T, i_b * T + T
|
| 139 |
+
|
| 140 |
+
o_i = tl.arange(0, 16)
|
| 141 |
+
m_A = o_i[:, None] > o_i[None, :]
|
| 142 |
+
m_I = o_i[:, None] == o_i[None, :]
|
| 143 |
+
A += (bos * H + i_h) * BT
|
| 144 |
+
Ai += (bos * H + i_h) * BT
|
| 145 |
+
|
| 146 |
+
if not USE_TMA:
|
| 147 |
+
p_A_11 = tl.make_block_ptr(
|
| 148 |
+
A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)
|
| 149 |
+
)
|
| 150 |
+
p_A_22 = tl.make_block_ptr(
|
| 151 |
+
A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)
|
| 152 |
+
)
|
| 153 |
+
b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32)
|
| 154 |
+
b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32)
|
| 155 |
+
else:
|
| 156 |
+
desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16])
|
| 157 |
+
desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16])
|
| 158 |
+
b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32)
|
| 159 |
+
b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32)
|
| 160 |
+
|
| 161 |
+
# [16, 16]
|
| 162 |
+
b_Ai_11 = -tl.where(m_A, b_Ai_11, 0)
|
| 163 |
+
b_Ai_22 = -tl.where(m_A, b_Ai_22, 0)
|
| 164 |
+
|
| 165 |
+
for i in range(2, min(16, T - i_t * BT)):
|
| 166 |
+
b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i)
|
| 167 |
+
b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0)
|
| 168 |
+
b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11)
|
| 169 |
+
for i in range(16 + 2, min(32, T - i_t * BT)):
|
| 170 |
+
b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16)
|
| 171 |
+
b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0)
|
| 172 |
+
b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22)
|
| 173 |
+
|
| 174 |
+
b_Ai_11 += m_I
|
| 175 |
+
b_Ai_22 += m_I
|
| 176 |
+
|
| 177 |
+
if not USE_TMA:
|
| 178 |
+
p_A_21 = tl.make_block_ptr(
|
| 179 |
+
A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)
|
| 180 |
+
)
|
| 181 |
+
b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32)
|
| 182 |
+
else:
|
| 183 |
+
b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32)
|
| 184 |
+
|
| 185 |
+
b_Ai_21 = -tl.dot(
|
| 186 |
+
tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION),
|
| 187 |
+
b_Ai_11,
|
| 188 |
+
input_precision=DOT_PRECISION,
|
| 189 |
+
)
|
| 190 |
+
|
| 191 |
+
if not USE_TMA:
|
| 192 |
+
p_Ai_11 = tl.make_block_ptr(
|
| 193 |
+
Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)
|
| 194 |
+
)
|
| 195 |
+
p_Ai_21 = tl.make_block_ptr(
|
| 196 |
+
Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)
|
| 197 |
+
)
|
| 198 |
+
p_Ai_22 = tl.make_block_ptr(
|
| 199 |
+
Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)
|
| 200 |
+
)
|
| 201 |
+
tl.store(
|
| 202 |
+
p_Ai_11,
|
| 203 |
+
b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"),
|
| 204 |
+
boundary_check=(0, 1),
|
| 205 |
+
)
|
| 206 |
+
tl.store(
|
| 207 |
+
p_Ai_22,
|
| 208 |
+
b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"),
|
| 209 |
+
boundary_check=(0, 1),
|
| 210 |
+
)
|
| 211 |
+
tl.store(
|
| 212 |
+
p_Ai_21,
|
| 213 |
+
b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"),
|
| 214 |
+
boundary_check=(0, 1),
|
| 215 |
+
)
|
| 216 |
+
else:
|
| 217 |
+
desc_o.store(
|
| 218 |
+
[i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne")
|
| 219 |
+
)
|
| 220 |
+
desc_o.store(
|
| 221 |
+
[i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne")
|
| 222 |
+
)
|
| 223 |
+
desc_o.store(
|
| 224 |
+
[i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne")
|
| 225 |
+
)
|
| 226 |
+
|
| 227 |
+
|
| 228 |
+
@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
|
| 229 |
+
@triton.autotune(
|
| 230 |
+
configs=[
|
| 231 |
+
triton.Config({}, num_warps=num_warps, num_stages=num_stages)
|
| 232 |
+
for num_warps in [2, 4, 8]
|
| 233 |
+
for num_stages in [2, 3, 4, 5]
|
| 234 |
+
],
|
| 235 |
+
key=["H", "BT", "IS_VARLEN"],
|
| 236 |
+
)
|
| 237 |
+
@triton.jit(do_not_specialize=["T"])
|
| 238 |
+
def merge_16x16_to_64x64_inverse_kernel(
|
| 239 |
+
A,
|
| 240 |
+
Ai,
|
| 241 |
+
cu_seqlens,
|
| 242 |
+
chunk_indices,
|
| 243 |
+
T,
|
| 244 |
+
H: tl.constexpr,
|
| 245 |
+
BT: tl.constexpr,
|
| 246 |
+
USE_TMA: tl.constexpr,
|
| 247 |
+
IS_VARLEN: tl.constexpr,
|
| 248 |
+
DOT_PRECISION: tl.constexpr,
|
| 249 |
+
):
|
| 250 |
+
i_t, i_bh = tl.program_id(1), tl.program_id(0)
|
| 251 |
+
i_b, i_h = i_bh // H, i_bh % H
|
| 252 |
+
if IS_VARLEN:
|
| 253 |
+
i_n, i_t = (
|
| 254 |
+
tl.load(chunk_indices + i_t * 2).to(tl.int32),
|
| 255 |
+
tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
|
| 256 |
+
)
|
| 257 |
+
bos, eos = (
|
| 258 |
+
tl.load(cu_seqlens + i_n).to(tl.int32),
|
| 259 |
+
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
|
| 260 |
+
)
|
| 261 |
+
T = eos - bos
|
| 262 |
+
else:
|
| 263 |
+
bos, eos = i_b * T, i_b * T + T
|
| 264 |
+
|
| 265 |
+
o_i = tl.arange(0, 16)
|
| 266 |
+
m_A = o_i[:, None] > o_i[None, :]
|
| 267 |
+
m_I = o_i[:, None] == o_i[None, :]
|
| 268 |
+
A += (bos * H + i_h) * BT
|
| 269 |
+
Ai += (bos * H + i_h) * BT
|
| 270 |
+
|
| 271 |
+
if not USE_TMA:
|
| 272 |
+
p_A_11 = tl.make_block_ptr(
|
| 273 |
+
A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)
|
| 274 |
+
)
|
| 275 |
+
p_A_22 = tl.make_block_ptr(
|
| 276 |
+
A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)
|
| 277 |
+
)
|
| 278 |
+
p_A_33 = tl.make_block_ptr(
|
| 279 |
+
A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)
|
| 280 |
+
)
|
| 281 |
+
p_A_44 = tl.make_block_ptr(
|
| 282 |
+
A, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)
|
| 283 |
+
)
|
| 284 |
+
b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32)
|
| 285 |
+
b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32)
|
| 286 |
+
b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32)
|
| 287 |
+
b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32)
|
| 288 |
+
else:
|
| 289 |
+
desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16])
|
| 290 |
+
desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16])
|
| 291 |
+
b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32)
|
| 292 |
+
b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32)
|
| 293 |
+
b_Ai_33 = desc.load([i_t * BT + 32, 32]).to(tl.float32)
|
| 294 |
+
b_Ai_44 = desc.load([i_t * BT + 48, 48]).to(tl.float32)
|
| 295 |
+
|
| 296 |
+
# [16, 16]
|
| 297 |
+
b_Ai_11 = -tl.where(m_A, b_Ai_11, 0)
|
| 298 |
+
b_Ai_22 = -tl.where(m_A, b_Ai_22, 0)
|
| 299 |
+
b_Ai_33 = -tl.where(m_A, b_Ai_33, 0)
|
| 300 |
+
b_Ai_44 = -tl.where(m_A, b_Ai_44, 0)
|
| 301 |
+
|
| 302 |
+
for i in range(2, min(16, T - i_t * BT)):
|
| 303 |
+
b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i)
|
| 304 |
+
b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0)
|
| 305 |
+
b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11)
|
| 306 |
+
for i in range(16 + 2, min(32, T - i_t * BT)):
|
| 307 |
+
b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16)
|
| 308 |
+
b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0)
|
| 309 |
+
b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22)
|
| 310 |
+
for i in range(32 + 2, min(48, T - i_t * BT)):
|
| 311 |
+
b_a_33 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 32)
|
| 312 |
+
b_a_33 += tl.sum(b_a_33[:, None] * b_Ai_33, 0)
|
| 313 |
+
b_Ai_33 = tl.where((o_i == i - 32)[:, None], b_a_33, b_Ai_33)
|
| 314 |
+
for i in range(48 + 2, min(64, T - i_t * BT)):
|
| 315 |
+
b_a_44 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 48)
|
| 316 |
+
b_a_44 += tl.sum(b_a_44[:, None] * b_Ai_44, 0)
|
| 317 |
+
b_Ai_44 = tl.where((o_i == i - 48)[:, None], b_a_44, b_Ai_44)
|
| 318 |
+
b_Ai_11 += m_I
|
| 319 |
+
b_Ai_22 += m_I
|
| 320 |
+
b_Ai_33 += m_I
|
| 321 |
+
b_Ai_44 += m_I
|
| 322 |
+
|
| 323 |
+
if not USE_TMA:
|
| 324 |
+
p_A_21 = tl.make_block_ptr(
|
| 325 |
+
A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)
|
| 326 |
+
)
|
| 327 |
+
p_A_31 = tl.make_block_ptr(
|
| 328 |
+
A, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)
|
| 329 |
+
)
|
| 330 |
+
p_A_32 = tl.make_block_ptr(
|
| 331 |
+
A, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)
|
| 332 |
+
)
|
| 333 |
+
p_A_41 = tl.make_block_ptr(
|
| 334 |
+
A, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)
|
| 335 |
+
)
|
| 336 |
+
p_A_42 = tl.make_block_ptr(
|
| 337 |
+
A, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)
|
| 338 |
+
)
|
| 339 |
+
p_A_43 = tl.make_block_ptr(
|
| 340 |
+
A, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)
|
| 341 |
+
)
|
| 342 |
+
b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32)
|
| 343 |
+
b_A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32)
|
| 344 |
+
b_A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32)
|
| 345 |
+
b_A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32)
|
| 346 |
+
b_A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32)
|
| 347 |
+
b_A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32)
|
| 348 |
+
else:
|
| 349 |
+
b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32)
|
| 350 |
+
b_A_31 = desc.load([i_t * BT + 32, 0]).to(tl.float32)
|
| 351 |
+
b_A_32 = desc.load([i_t * BT + 32, 16]).to(tl.float32)
|
| 352 |
+
b_A_41 = desc.load([i_t * BT + 48, 0]).to(tl.float32)
|
| 353 |
+
b_A_42 = desc.load([i_t * BT + 48, 16]).to(tl.float32)
|
| 354 |
+
b_A_43 = desc.load([i_t * BT + 48, 32]).to(tl.float32)
|
| 355 |
+
|
| 356 |
+
b_Ai_21 = -tl.dot(
|
| 357 |
+
tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION),
|
| 358 |
+
b_Ai_11,
|
| 359 |
+
input_precision=DOT_PRECISION,
|
| 360 |
+
)
|
| 361 |
+
b_Ai_32 = -tl.dot(
|
| 362 |
+
tl.dot(b_Ai_33, b_A_32, input_precision=DOT_PRECISION),
|
| 363 |
+
b_Ai_22,
|
| 364 |
+
input_precision=DOT_PRECISION,
|
| 365 |
+
)
|
| 366 |
+
b_Ai_43 = -tl.dot(
|
| 367 |
+
tl.dot(b_Ai_44, b_A_43, input_precision=DOT_PRECISION),
|
| 368 |
+
b_Ai_33,
|
| 369 |
+
input_precision=DOT_PRECISION,
|
| 370 |
+
)
|
| 371 |
+
|
| 372 |
+
b_Ai_31 = -tl.dot(
|
| 373 |
+
b_Ai_33,
|
| 374 |
+
tl.dot(b_A_31, b_Ai_11, input_precision=DOT_PRECISION)
|
| 375 |
+
+ tl.dot(b_A_32, b_Ai_21, input_precision=DOT_PRECISION),
|
| 376 |
+
input_precision=DOT_PRECISION,
|
| 377 |
+
)
|
| 378 |
+
b_Ai_42 = -tl.dot(
|
| 379 |
+
b_Ai_44,
|
| 380 |
+
tl.dot(b_A_42, b_Ai_22, input_precision=DOT_PRECISION)
|
| 381 |
+
+ tl.dot(b_A_43, b_Ai_32, input_precision=DOT_PRECISION),
|
| 382 |
+
input_precision=DOT_PRECISION,
|
| 383 |
+
)
|
| 384 |
+
b_Ai_41 = -tl.dot(
|
| 385 |
+
b_Ai_44,
|
| 386 |
+
tl.dot(b_A_41, b_Ai_11, input_precision=DOT_PRECISION)
|
| 387 |
+
+ tl.dot(b_A_42, b_Ai_21, input_precision=DOT_PRECISION)
|
| 388 |
+
+ tl.dot(b_A_43, b_Ai_31, input_precision=DOT_PRECISION),
|
| 389 |
+
input_precision=DOT_PRECISION,
|
| 390 |
+
)
|
| 391 |
+
|
| 392 |
+
if not USE_TMA:
|
| 393 |
+
p_Ai_11 = tl.make_block_ptr(
|
| 394 |
+
Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)
|
| 395 |
+
)
|
| 396 |
+
p_Ai_22 = tl.make_block_ptr(
|
| 397 |
+
Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)
|
| 398 |
+
)
|
| 399 |
+
p_Ai_33 = tl.make_block_ptr(
|
| 400 |
+
Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)
|
| 401 |
+
)
|
| 402 |
+
p_Ai_44 = tl.make_block_ptr(
|
| 403 |
+
Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)
|
| 404 |
+
)
|
| 405 |
+
p_Ai_21 = tl.make_block_ptr(
|
| 406 |
+
Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)
|
| 407 |
+
)
|
| 408 |
+
p_Ai_31 = tl.make_block_ptr(
|
| 409 |
+
Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)
|
| 410 |
+
)
|
| 411 |
+
p_Ai_32 = tl.make_block_ptr(
|
| 412 |
+
Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)
|
| 413 |
+
)
|
| 414 |
+
p_Ai_41 = tl.make_block_ptr(
|
| 415 |
+
Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)
|
| 416 |
+
)
|
| 417 |
+
p_Ai_42 = tl.make_block_ptr(
|
| 418 |
+
Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)
|
| 419 |
+
)
|
| 420 |
+
p_Ai_43 = tl.make_block_ptr(
|
| 421 |
+
Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)
|
| 422 |
+
)
|
| 423 |
+
tl.store(
|
| 424 |
+
p_Ai_11,
|
| 425 |
+
b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"),
|
| 426 |
+
boundary_check=(0, 1),
|
| 427 |
+
)
|
| 428 |
+
tl.store(
|
| 429 |
+
p_Ai_22,
|
| 430 |
+
b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"),
|
| 431 |
+
boundary_check=(0, 1),
|
| 432 |
+
)
|
| 433 |
+
tl.store(
|
| 434 |
+
p_Ai_33,
|
| 435 |
+
b_Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"),
|
| 436 |
+
boundary_check=(0, 1),
|
| 437 |
+
)
|
| 438 |
+
tl.store(
|
| 439 |
+
p_Ai_44,
|
| 440 |
+
b_Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"),
|
| 441 |
+
boundary_check=(0, 1),
|
| 442 |
+
)
|
| 443 |
+
tl.store(
|
| 444 |
+
p_Ai_21,
|
| 445 |
+
b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"),
|
| 446 |
+
boundary_check=(0, 1),
|
| 447 |
+
)
|
| 448 |
+
tl.store(
|
| 449 |
+
p_Ai_31,
|
| 450 |
+
b_Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"),
|
| 451 |
+
boundary_check=(0, 1),
|
| 452 |
+
)
|
| 453 |
+
tl.store(
|
| 454 |
+
p_Ai_32,
|
| 455 |
+
b_Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"),
|
| 456 |
+
boundary_check=(0, 1),
|
| 457 |
+
)
|
| 458 |
+
tl.store(
|
| 459 |
+
p_Ai_41,
|
| 460 |
+
b_Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"),
|
| 461 |
+
boundary_check=(0, 1),
|
| 462 |
+
)
|
| 463 |
+
tl.store(
|
| 464 |
+
p_Ai_42,
|
| 465 |
+
b_Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"),
|
| 466 |
+
boundary_check=(0, 1),
|
| 467 |
+
)
|
| 468 |
+
tl.store(
|
| 469 |
+
p_Ai_43,
|
| 470 |
+
b_Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"),
|
| 471 |
+
boundary_check=(0, 1),
|
| 472 |
+
)
|
| 473 |
+
else:
|
| 474 |
+
desc_o.store(
|
| 475 |
+
[i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne")
|
| 476 |
+
)
|
| 477 |
+
desc_o.store(
|
| 478 |
+
[i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne")
|
| 479 |
+
)
|
| 480 |
+
desc_o.store(
|
| 481 |
+
[i_t * BT + 32, 32], b_Ai_33.to(desc_o.dtype, fp_downcast_rounding="rtne")
|
| 482 |
+
)
|
| 483 |
+
desc_o.store(
|
| 484 |
+
[i_t * BT + 48, 48], b_Ai_44.to(desc_o.dtype, fp_downcast_rounding="rtne")
|
| 485 |
+
)
|
| 486 |
+
desc_o.store(
|
| 487 |
+
[i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne")
|
| 488 |
+
)
|
| 489 |
+
desc_o.store(
|
| 490 |
+
[i_t * BT + 32, 0], b_Ai_31.to(desc_o.dtype, fp_downcast_rounding="rtne")
|
| 491 |
+
)
|
| 492 |
+
desc_o.store(
|
| 493 |
+
[i_t * BT + 32, 16], b_Ai_32.to(desc_o.dtype, fp_downcast_rounding="rtne")
|
| 494 |
+
)
|
| 495 |
+
desc_o.store(
|
| 496 |
+
[i_t * BT + 48, 0], b_Ai_41.to(desc_o.dtype, fp_downcast_rounding="rtne")
|
| 497 |
+
)
|
| 498 |
+
desc_o.store(
|
| 499 |
+
[i_t * BT + 48, 16], b_Ai_42.to(desc_o.dtype, fp_downcast_rounding="rtne")
|
| 500 |
+
)
|
| 501 |
+
desc_o.store(
|
| 502 |
+
[i_t * BT + 48, 32], b_Ai_43.to(desc_o.dtype, fp_downcast_rounding="rtne")
|
| 503 |
+
)
|
| 504 |
+
|
| 505 |
+
|
| 506 |
+
@input_guard
|
| 507 |
+
def solve_tril(
|
| 508 |
+
A: torch.Tensor,
|
| 509 |
+
cu_seqlens: torch.Tensor | None = None,
|
| 510 |
+
chunk_indices: torch.Tensor | None = None,
|
| 511 |
+
output_dtype: torch.dtype = torch.float,
|
| 512 |
+
) -> torch.Tensor:
|
| 513 |
+
"""
|
| 514 |
+
Compute the inverse of the matrix I + A
|
| 515 |
+
A should be strictly lower triangular, i.e., A.triu() == 0.
|
| 516 |
+
|
| 517 |
+
Args:
|
| 518 |
+
A (torch.Tensor):
|
| 519 |
+
[B, T, H, BT], where BT should only be 16, 32, or 64.
|
| 520 |
+
cu_seqlens (torch.Tensor):
|
| 521 |
+
The cumulative sequence lengths of the input tensor. Default: `None`.
|
| 522 |
+
chunk_indices (torch.Tensor):
|
| 523 |
+
Pre-computed chunk indices. Default: `None`.
|
| 524 |
+
output_dtype (torch.dtype):
|
| 525 |
+
The dtype of the output tensor. Default: `torch.float`.
|
| 526 |
+
If `None`, the output dtype will be the same as the input dtype.
|
| 527 |
+
|
| 528 |
+
Returns:
|
| 529 |
+
(I + A)^-1 with the same shape as A
|
| 530 |
+
"""
|
| 531 |
+
assert A.shape[-1] in [16, 32, 64]
|
| 532 |
+
output_dtype = A.dtype if output_dtype is None else output_dtype
|
| 533 |
+
|
| 534 |
+
B, T, H, BT = A.shape
|
| 535 |
+
if chunk_indices is None and cu_seqlens is not None:
|
| 536 |
+
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
|
| 537 |
+
NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)
|
| 538 |
+
|
| 539 |
+
Ai = torch.zeros_like(A, dtype=output_dtype)
|
| 540 |
+
if BT == 16:
|
| 541 |
+
merge_fn = solve_tril_16x16_kernel
|
| 542 |
+
elif BT == 32:
|
| 543 |
+
merge_fn = merge_16x16_to_32x32_inverse_kernel
|
| 544 |
+
elif BT == 64:
|
| 545 |
+
merge_fn = merge_16x16_to_64x64_inverse_kernel
|
| 546 |
+
|
| 547 |
+
merge_fn[(B * H, NT) if BT == 64 else (NT, B * H)](
|
| 548 |
+
A=A,
|
| 549 |
+
Ai=Ai,
|
| 550 |
+
cu_seqlens=cu_seqlens,
|
| 551 |
+
chunk_indices=chunk_indices,
|
| 552 |
+
T=T,
|
| 553 |
+
H=H,
|
| 554 |
+
BT=BT,
|
| 555 |
+
USE_TMA=is_tma_supported,
|
| 556 |
+
DOT_PRECISION=FLA_TRIL_PRECISION,
|
| 557 |
+
)
|
| 558 |
+
return Ai
|
bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/wy_fast.py
ADDED
|
@@ -0,0 +1,158 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 2 |
+
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
| 3 |
+
# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
|
| 4 |
+
#
|
| 5 |
+
# This file contains code copied from the flash-linear-attention project.
|
| 6 |
+
# The original source code was licensed under the MIT license and included
|
| 7 |
+
# the following copyright notice:
|
| 8 |
+
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
|
| 9 |
+
|
| 10 |
+
# ruff: noqa: E501
|
| 11 |
+
|
| 12 |
+
import torch
|
| 13 |
+
|
| 14 |
+
from vllm.triton_utils import tl, triton
|
| 15 |
+
|
| 16 |
+
from .index import prepare_chunk_indices
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
|
| 20 |
+
@triton.autotune(
|
| 21 |
+
configs=[
|
| 22 |
+
triton.Config({}, num_warps=num_warps, num_stages=num_stages)
|
| 23 |
+
for num_warps in [2, 4, 8]
|
| 24 |
+
for num_stages in [2, 3, 4]
|
| 25 |
+
],
|
| 26 |
+
key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"],
|
| 27 |
+
)
|
| 28 |
+
@triton.jit(do_not_specialize=["T"])
|
| 29 |
+
def recompute_w_u_fwd_kernel(
|
| 30 |
+
k,
|
| 31 |
+
v,
|
| 32 |
+
beta,
|
| 33 |
+
w,
|
| 34 |
+
u,
|
| 35 |
+
A,
|
| 36 |
+
g,
|
| 37 |
+
cu_seqlens,
|
| 38 |
+
chunk_indices,
|
| 39 |
+
T,
|
| 40 |
+
H: tl.constexpr,
|
| 41 |
+
Hg: tl.constexpr,
|
| 42 |
+
K: tl.constexpr,
|
| 43 |
+
V: tl.constexpr,
|
| 44 |
+
BT: tl.constexpr,
|
| 45 |
+
BK: tl.constexpr,
|
| 46 |
+
BV: tl.constexpr,
|
| 47 |
+
IS_VARLEN: tl.constexpr,
|
| 48 |
+
):
|
| 49 |
+
i_t, i_bh = tl.program_id(1), tl.program_id(0)
|
| 50 |
+
i_b, i_h = i_bh // H, i_bh % H
|
| 51 |
+
if IS_VARLEN:
|
| 52 |
+
i_n, i_t = (
|
| 53 |
+
tl.load(chunk_indices + i_t * 2).to(tl.int32),
|
| 54 |
+
tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
|
| 55 |
+
)
|
| 56 |
+
bos, eos = (
|
| 57 |
+
tl.load(cu_seqlens + i_n).to(tl.int32),
|
| 58 |
+
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
|
| 59 |
+
)
|
| 60 |
+
T = eos - bos
|
| 61 |
+
else:
|
| 62 |
+
bos, eos = i_b * T, i_b * T + T
|
| 63 |
+
p_beta = tl.make_block_ptr(
|
| 64 |
+
beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)
|
| 65 |
+
)
|
| 66 |
+
p_g = tl.make_block_ptr(g + (bos * H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,))
|
| 67 |
+
p_A = tl.make_block_ptr(
|
| 68 |
+
A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
|
| 69 |
+
)
|
| 70 |
+
b_beta = tl.load(p_beta, boundary_check=(0,))
|
| 71 |
+
b_A = tl.load(p_A, boundary_check=(0, 1))
|
| 72 |
+
b_g = tl.exp(tl.load(p_g, boundary_check=(0,)))
|
| 73 |
+
|
| 74 |
+
for i_v in range(tl.cdiv(V, BV)):
|
| 75 |
+
p_v = tl.make_block_ptr(
|
| 76 |
+
v + (bos * H + i_h) * V,
|
| 77 |
+
(T, V),
|
| 78 |
+
(H * V, 1),
|
| 79 |
+
(i_t * BT, i_v * BV),
|
| 80 |
+
(BT, BV),
|
| 81 |
+
(1, 0),
|
| 82 |
+
)
|
| 83 |
+
p_u = tl.make_block_ptr(
|
| 84 |
+
u + (bos * H + i_h) * V,
|
| 85 |
+
(T, V),
|
| 86 |
+
(H * V, 1),
|
| 87 |
+
(i_t * BT, i_v * BV),
|
| 88 |
+
(BT, BV),
|
| 89 |
+
(1, 0),
|
| 90 |
+
)
|
| 91 |
+
b_v = tl.load(p_v, boundary_check=(0, 1))
|
| 92 |
+
b_vb = (b_v * b_beta[:, None]).to(b_v.dtype)
|
| 93 |
+
b_u = tl.dot(b_A, b_vb, allow_tf32=False)
|
| 94 |
+
tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1))
|
| 95 |
+
|
| 96 |
+
for i_k in range(tl.cdiv(K, BK)):
|
| 97 |
+
p_k = tl.make_block_ptr(
|
| 98 |
+
k + (bos * Hg + i_h // (H // Hg)) * K,
|
| 99 |
+
(T, K),
|
| 100 |
+
(Hg * K, 1),
|
| 101 |
+
(i_t * BT, i_k * BK),
|
| 102 |
+
(BT, BK),
|
| 103 |
+
(1, 0),
|
| 104 |
+
)
|
| 105 |
+
p_w = tl.make_block_ptr(
|
| 106 |
+
w + (bos * H + i_h) * K,
|
| 107 |
+
(T, K),
|
| 108 |
+
(H * K, 1),
|
| 109 |
+
(i_t * BT, i_k * BK),
|
| 110 |
+
(BT, BK),
|
| 111 |
+
(1, 0),
|
| 112 |
+
)
|
| 113 |
+
b_k = tl.load(p_k, boundary_check=(0, 1))
|
| 114 |
+
b_kb = (b_k * b_beta[:, None] * b_g[:, None]).to(b_k.dtype)
|
| 115 |
+
b_w = tl.dot(b_A, b_kb)
|
| 116 |
+
tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))
|
| 117 |
+
|
| 118 |
+
|
| 119 |
+
def recompute_w_u_fwd(
|
| 120 |
+
k: torch.Tensor,
|
| 121 |
+
v: torch.Tensor,
|
| 122 |
+
beta: torch.Tensor,
|
| 123 |
+
g_cumsum: torch.Tensor,
|
| 124 |
+
A: torch.Tensor,
|
| 125 |
+
cu_seqlens: torch.Tensor | None,
|
| 126 |
+
chunk_indices: torch.Tensor | None = None,
|
| 127 |
+
) -> tuple[torch.Tensor, torch.Tensor]:
|
| 128 |
+
B, T, Hg, K, V = *k.shape, v.shape[-1]
|
| 129 |
+
H = v.shape[-2]
|
| 130 |
+
BT = A.shape[-1]
|
| 131 |
+
|
| 132 |
+
if chunk_indices is None and cu_seqlens is not None:
|
| 133 |
+
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
|
| 134 |
+
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
|
| 135 |
+
BK = 64
|
| 136 |
+
BV = 64
|
| 137 |
+
u = torch.empty_like(v)
|
| 138 |
+
w = k.new_empty(B, T, H, K)
|
| 139 |
+
recompute_w_u_fwd_kernel[(B * H, NT)](
|
| 140 |
+
k=k,
|
| 141 |
+
v=v,
|
| 142 |
+
beta=beta,
|
| 143 |
+
w=w,
|
| 144 |
+
u=u,
|
| 145 |
+
A=A,
|
| 146 |
+
g=g_cumsum,
|
| 147 |
+
cu_seqlens=cu_seqlens,
|
| 148 |
+
chunk_indices=chunk_indices,
|
| 149 |
+
T=T,
|
| 150 |
+
H=H,
|
| 151 |
+
Hg=Hg,
|
| 152 |
+
K=K,
|
| 153 |
+
V=V,
|
| 154 |
+
BT=BT,
|
| 155 |
+
BK=BK,
|
| 156 |
+
BV=BV,
|
| 157 |
+
)
|
| 158 |
+
return w, u
|