jcbtc commited on
Commit
d824ec9
·
verified ·
1 Parent(s): f6f9d19

Release Ciru v4 (4.0.0): integrated runtime and patch notes

Browse files

Publish 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
Files changed (50) hide show
  1. INSTALL.md +28 -9
  2. README.md +72 -15
  3. RELEASE.json +137 -295
  4. RUNTIME-FIXES.md +37 -1
  5. SOURCE-PROVENANCE.json +159 -0
  6. bundle/LICENSES/XGRAMMAR-LICENSE +201 -0
  7. bundle/native/libornith_attention_iu4.so +2 -2
  8. bundle/native/libornith_dense_g256.so +0 -0
  9. bundle/native/libornith_dense_g256_n32.so +0 -0
  10. bundle/native/libornith_persistent_iu4.so +2 -2
  11. bundle/native/libornith_routed_storage_n32.so +2 -2
  12. bundle/packaging/serve.sh +1 -1
  13. bundle/plugin-site/TOOL-INTEGRATION.json +205 -0
  14. bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/direct_url.json +0 -1
  15. bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/uv_build.json +0 -1
  16. bundle/plugin-site/ciru_ornith_g256-1.0.2.dist-info/uv_cache.json +0 -1
  17. bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/INSTALLER +0 -0
  18. bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/METADATA +1 -1
  19. bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/RECORD +42 -17
  20. bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/REQUESTED +0 -0
  21. bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/WHEEL +0 -0
  22. bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/entry_points.txt +0 -0
  23. 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
  24. bundle/plugin-site/{ciru_ornith_g256-1.0.2.dist-info → ciru_ornith_g256-4.0.0.dist-info}/top_level.txt +0 -0
  25. bundle/plugin-site/ornith_g256/__init__.py +6 -0
  26. bundle/plugin-site/ornith_g256/_vllm_correctness/chat_completion_serving_r03.py +1319 -0
  27. bundle/plugin-site/ornith_g256/_vllm_correctness/manifest.json +7 -1
  28. bundle/plugin-site/ornith_g256/_xgrammar_native/build-receipt.json +29 -0
  29. bundle/plugin-site/ornith_g256/_xgrammar_native/libxgrammar-candidate.so +3 -0
  30. bundle/plugin-site/ornith_g256/_xgrammar_native/selection.json +4 -0
  31. bundle/plugin-site/ornith_g256/attention_compact.py +5 -5
  32. bundle/plugin-site/ornith_g256/attention_iu4.py +20 -9
  33. bundle/plugin-site/ornith_g256/attention_iu4_persistent.py +2 -2
  34. bundle/plugin-site/ornith_g256/optimized/__init__.py +1 -0
  35. bundle/plugin-site/ornith_g256/optimized/compile_scope.py +30 -0
  36. bundle/plugin-site/ornith_g256/optimized/draft.py +90 -0
  37. bundle/plugin-site/ornith_g256/optimized/draft_kernel.py +50 -0
  38. bundle/plugin-site/ornith_g256/optimized/grammar_parallel.py +122 -0
  39. bundle/plugin-site/ornith_g256/optimized/greedy.py +40 -0
  40. bundle/plugin-site/ornith_g256/optimized/greedy_kernel.py +40 -0
  41. bundle/plugin-site/ornith_g256/optimized/launch.py +17 -0
  42. bundle/plugin-site/ornith_g256/optimized/mixed_qk_sum.py +66 -0
  43. bundle/plugin-site/ornith_g256/optimized/pipeline/SOURCE-PROOF.json +34 -0
  44. bundle/plugin-site/ornith_g256/optimized/pipeline/__init__.py +1 -0
  45. bundle/plugin-site/ornith_g256/optimized/pipeline/adapter.py +172 -0
  46. bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/causal_conv1d.py +1307 -0
  47. bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/chunk_o.py +190 -0
  48. bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/chunk_scaled_dot_kkt.py +176 -0
  49. bundle/plugin-site/ornith_g256/optimized/pipeline/candidate/solve_tril.py +558 -0
  50. 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 runtime 1.0.2
4
 
5
- Download the current `bundle/plugin-site/`, `bundle/serve.sh` and
6
- `bundle/serve-vision.sh` contents and restart the model process. Remove older `ciru_ornith_g256-*.dist-info/` directories after the
7
- new `ciru_ornith_g256-1.0.2.dist-info/` directory is present, leaving only the
8
- current metadata. The pinned runtime wheels, native libraries and IU4 weights
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 is text-only. For optional image input, 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#optional-vision--image-input) and [vision validation limits](VISION.md).
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
- The [Ciru source repository](https://github.com/ciru-ai/ornith-ciru-halo-agent) contains the model plugin, all eight native kernel sources, and the corresponding build script. 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.
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 runtime 1.0.2 — September 15, 2026.** Automatic tool calls now use
27
- native schema constraints to prevent missing required arguments and malformed
28
- tool output. IU4 weights and inference kernels are unchanged.
29
- Image input is enabled using the upstream BF16 vision weights. See [patch notes](RUNTIME-FIXES.md).
 
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
- **Measured on AMD Ryzen AI Max+ 395 / Radeon 8060S (gfx1151):**
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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/)** · **[Source and build instructions](https://github.com/ciru-ai/ornith-ciru-halo-agent)** · **[Credits](CREDITS.md)**
50
 
51
- ## September 13 runtime fixes
52
 
53
- **Existing installations should update the serving plugin and restart.** 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).
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 is enabled by default in runtime 1.0.2. The explicit image launcher is also available:
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
- ## Coding speed and concurrency
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
- ## Fast prompt ingestion
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
- ## Returning agents and large histories
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
- ## Quality and actual agent work
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
- "model": "Ornith1.5 Ciru Halo Agent",
3
- "source_repository": "https://github.com/ciru-ai/ornith-ciru-halo-agent",
4
- "source_commit": "65e32f00c8f13e5ec1f8b941639dc37fdc07424c",
5
- "benchmarks": "https://llm.ciru.ai/research/ornith-strix/",
6
- "target_revision": "10fbf86fed7ecee4a061f8b499a618f46001cac1",
7
- "draft_revision": "9b4852c05fd00b672b7434b1bb105bc03c8682b0",
8
- "files_bytes": {
9
- "INSTALL.md": 5977,
10
- "README.md": 18618,
11
- "CREDITS.md": 6275,
12
- ".gitattributes": 220,
13
- "NOTICE": 1088,
14
- "BUILD.md": 8422,
15
- "NATIVE-BUILD-CHECK.json": 1700,
16
- "bundle/paths.env": 268,
17
- "bundle/serve.sh": 1070,
18
- "runtime/INSTALL-ORNITH-RUNTIME.sh": 794,
19
- "runtime/INSTALL-RUNTIME.sh": 4428,
20
- "runtime/README.md": 1167,
21
- "runtime/UPSTREAM-RUNTIME.md": 14159,
22
- "runtime/nixos-module.nix": 349,
23
- "runtime/requirements-runtime.lock": 2129,
24
- "runtime/runtime-env.sh": 4973,
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
- "native_sha256": {
199
- "libornith_attention_iu4.so": "2905806824bece62ce9d9140608859df869f3e3d3663c67bf79ad4ea495cd61f",
200
- "libornith_dense_g256.so": "855838687084170807c5ba4556e9771f9b1b26f3556f48c0a16dbd8bbee41c9c",
201
- "libornith_dense_g256_n32.so": "24dbf9b1eaff06ff6530e4f37292347bcf073b3c07f7e16ea6ab111711162a40",
202
- "libornith_head_i8_tile.so": "2c03de5ed015a052b407436c7e42b50fa33ef76cb46dccab9b226546ca1d1e09",
203
- "libornith_persistent_iu4.so": "40e5cf35157c3e716de8067f7eef1312a3b9e787dd52a6f2e232ccd2d07fa68c",
204
- "libornith_routed_direct.so": "5b3ca6826547e0ff1c417e733229e104721cec403ccd8f88f7fcdb1726d6f4ed",
205
- "libornith_routed_n32.so": "48415e71f6c85e3aefcecf99d72607bb05900244b4aaa9a7ba2c58378d92f5fb",
206
- "libornith_routed_storage_n32.so": "3ea288cd0e37a4c373038744c3cc4cd4eac09675ac137e953d5ae7a6cd43dd3d"
207
- },
208
- "checks": {
209
- "clean_runtime_install": "passed on Ciru",
210
- "release_launcher_dry_run": "passed using new runtime",
211
- "native_source_build": "8/8 compiled; distributed measured libraries unchanged",
212
- "new_gpu_inference_run": false
213
- },
214
- "image_sha256": "ddbb4819065a652a6feb20ce760c8292b546800cd2e48d01c560809ad5b58765",
215
- "optional_vision": {
216
- "launcher": "bundle/serve-vision.sh",
217
- "encoder_projector": "bundle/models/target/protected-00.safetensors",
218
- "format": "native BF16 safetensors; no GGUF mmproj",
219
- "vision_tensors": 333,
220
- "images_per_request": 1,
221
- "image_max_pixels": 1048576,
222
- "encoder_attention_backend": "TRITON_ATTN",
223
- "validation": "Published weight header and launcher settings checked; earlier image-profile results documented separately in VISION.md",
224
- "historical_profile": true
225
- },
226
- "draft_depth_override": {
227
- "flag": "--draft-tokens",
228
- "choices": [
229
- 0,
230
- 7,
231
- 15
232
- ],
233
- "scope": "C1 at all supported context lengths; C2-8 remains 7",
234
- "validation": "CPU-only settings and policy tests; no model loaded"
235
- },
236
- "runtime_fixes": {
237
- "date": "2026-09-13",
238
- "notes": "RUNTIME-FIXES.md",
239
- "weights_changed": false,
240
- "runtime_wheels_changed": false,
241
- "native_libraries_changed": false,
242
- "changed_files_sha256": {
243
- "bundle/plugin-site/ornith_g256/__init__.py": "a15cbb0679a9cc4766a77c57594e941e9a911381f42c815cfdecb189c8b6635c",
244
- "bundle/plugin-site/ornith_g256/graph_phase_guard.py": "6cabaf7d4e3be04bbac7a642544d38948de2e07ea7ca38f5927f91bf3ad39163",
245
- "bundle/plugin-site/ornith_g256/launch.py": "8c17d8a738ec8b4424a1e451883cd15bd0a480525532de4c05c4b3e2ac8f22bb",
246
- "bundle/plugin-site/ornith_g256/parser_contract.py": "2bc5885394786931b2c00693f580ee99da220c6adffc00651fc8227ca9bb1d0d",
247
- "bundle/plugin-site/ornith_g256/strict_qwen/__init__.py": "ee8a5055249f1413fc22f46421ab77b5485aa39b199c68eb5b56c75dcc40922d",
248
- "bundle/plugin-site/ornith_g256/strict_qwen/qwen3.py": "33c653447a144b72257e508ccd92610db6d30578e28f98c69c99281773dead1c",
249
- "bundle/plugin-site/ornith_g256/strict_qwen/qwen3_contract.py": "74a0b03e3532883fa4d70494fbaecae7760e293cbb6b4e588640884f3093052a",
250
- "bundle/plugin-site/ciru_ornith_g256-0.0.2a0.dist-info/RECORD": "6605e7606bb5c9f323d2984473ad643ba747616e78c3ebea6938fa872d531f9b",
251
- "README.md": "83af97ae9c3ab0ea1185df77aa3436d0792ca730e89268fbaef5aea8b5f6b19a",
252
- "RUNTIME-FIXES.md": "084dfd99a09d31a7b4ee30d17087a9cc1d341ab1fced5da644fc9e4c0d6144f9"
253
- },
254
- "checks": {
255
- "shipped_runtime_guard_and_parser_import": "passed without GPU initialization",
256
- "shipped_runtime_launcher_argument_inspection": "passed using existing released weight paths",
257
- "prior_ornith_gpu_requests": "21/21 Hermes on Strix Halo; includes HE0-9 10/10",
258
- "final_parser_saved_response_replays": "6/6 ASGI and 3/3 Pi",
259
- "new_public_package_gpu_run": false
260
- },
261
- "native_parameter_whitespace": "Original Qwen XML wrapping-newline convention retained; exact boundary whitespace not guaranteed."
262
- },
263
- "runtime_version": "1.0.2",
264
- "release_date": "2026-09-15",
265
- "accuracy_patch": {
266
- "version": "1.0.1",
267
- "details": "ACCURACY-PATCH-1.0.1.json",
268
- "weights_changed": false,
269
- "native_libraries_changed": false,
270
- "fixes": [
271
- "FP32 products in BF16 causal convolution",
272
- "Accepted GDN state and convolution history recovery",
273
- "Persistent accepted-state source/count buffers for FULL graphs"
274
- ],
275
- "single_request_throughput_tradeoff": "Approximately 7% lower in the Ornith screen, accepted by the owner for this correctness release."
276
- },
277
- "runtime_accuracy_source_commit": "2c50cb2aa31abd0537e9c84b7ec43307e3ea40c7",
278
- "runtime_tool_schema_source_commit": "65e32f00c8f13e5ec1f8b941639dc37fdc07424c",
279
- "tool_schema_patch": {
280
- "version": "1.0.2",
281
- "details": "TOOL-SCHEMA-PATCH-1.0.2.json",
282
- "default_auto_constraints": true,
283
- "prompt_visible_tools_unchanged": true,
284
- "weights_changed": false,
285
- "native_libraries_changed": false,
286
- "runtime_wheels_changed": false
287
- },
288
- "vision": {
289
- "enabled_by_default": true,
290
- "upstream_supported": true,
291
- "bf16_tensors": 333,
292
- "image_limit": 1,
293
- "image_max_pixels": 1048576,
294
- "encoder_backend": "TRITON_ATTN",
295
- "text_only_flag": "--text-only",
296
- "live_image_checks_passed": 2,
297
- "details": "VISION.md"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
- ## Update an existing installation
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:2905806824bece62ce9d9140608859df869f3e3d3663c67bf79ad4ea495cd61f
3
- size 186640
 
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:40e5cf35157c3e716de8067f7eef1312a3b9e787dd52a6f2e232ccd2d07fa68c
3
- size 103792
 
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:3ea288cd0e37a4c373038744c3cc4cd4eac09675ac137e953d5ae7a6cd43dd3d
3
- size 158536
 
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: 1.0.2
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
- ciru_ornith_g256-1.0.2.dist-info/INSTALLER,sha256=5hhM4Q4mYTT9z6QB6PGpUAW81PGNFrYrdXMj4oM_6ak,2
2
- ciru_ornith_g256-1.0.2.dist-info/METADATA,sha256=hoIAew6fYBXVq6WqaJFx3YgykAk-iDq3d2OU9LaQhRc,254
3
- ciru_ornith_g256-1.0.2.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
4
- ciru_ornith_g256-1.0.2.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
5
- ciru_ornith_g256-1.0.2.dist-info/direct_url.json,sha256=MwFMTtCGcu6fn-TruM1lo2XsI7q0gT0JZbfW7sb3OEA,105
6
- ciru_ornith_g256-1.0.2.dist-info/entry_points.txt,sha256=ESe1wopdkiGfr5DNWZcSemlgPr3uvRre1OyHkl88QqU,121
7
- ciru_ornith_g256-1.0.2.dist-info/licenses/LICENSE-APACHE-2.0,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
8
- ciru_ornith_g256-1.0.2.dist-info/top_level.txt,sha256=EdCMLXnn8tDNBUR3JhkMvOi8hEw6IuhYFYEcfdhekxU,12
9
- ciru_ornith_g256-1.0.2.dist-info/uv_build.json,sha256=RBNvo1WzZ4oRRq0W9-hknpT7T8If536DEMBg9hyq_4o,2
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=Dd-kifT56tE1AG4pqmdedSKvpc7hBZbCEV4H8TGGR4w,1607
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=ZrrGSCCgtUdW2amQLV-AZmdaorCz6w7HP_G0O0NAeWM,10412
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=NCNbEQTq6WxVmGs1mFvimK7F7GbE-oFCQJRewD7ax_U,2575
24
- ornith_g256/attention_iu4_persistent.py,sha256=yivkrUYx6ie2L88HXng732woKS9ZD0nuh0jpkwmmzQ4,12418
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=En2VXcVk520ONHYD_aQmSaisvXJHfIMfg5Z5g--n63o,13579
62
  ornith_g256/worker_base.py,sha256=oA7xKVFy5lXIs8fXkM_QMWEUqUkuSD3t_eoDbTt7-Gg,4601
63
- ciru_ornith_g256-1.0.2.dist-info/RECORD,,
 
 
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.1",
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 <= 2048
111
  or torch.cuda.is_current_stream_capturing()):
112
- raise ValueError('Configure target compact prefill before capture, within C8/256K/2048 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,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 == 1120
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
- 'Q1120 and aligned K>=4096, existing BF16 fallback elsewhere')
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
- specs = [((16, 1120, 32), torch.int32), ((16, 1120), torch.float16),
 
 
 
 
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((1120, 16, 256), dtype=torch.bfloat16, device=device)
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
- # C1 max_query_len1120 and exact CPU metadata max_seq_len are required.
29
- # The unchanged consumer uses its original contiguous fixed1120 output.
30
  # No allocation, tensor copy to CPU, synchronization or cache mutation here.
31
- if not 1120 <= max_keys <= self.max_keys or max_keys % 32:
 
 
32
  raise ValueError("IU4 prefill requires aligned keys within reserved capacity")
 
 
 
 
33
  stream = torch.cuda.current_stream().cuda_stream
34
- rc = self.library.iu4_prepare(query.data_ptr(), key.data_ptr(), value.data_ptr(),
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 = self.library.iu4_attention(*self.pointers, self.output.data_ptr(), max_keys, stream)
39
  if rc:
40
  raise RuntimeError(f'IU4 attention launch failed: HIP {rc}')
41
- output[:1120].copy_(self.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(2048, 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,7 +122,7 @@ class State:
122
  bank = self.banks.get(cache.data_ptr())
123
  if bank is None:
124
  return
125
- if slots.numel() > 2048 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())
 
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