Publish Nyx INT4+FP8 production checkpoint

.gitattributes

@@ -1,35 +1,11 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
 *.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.safetensor filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text
 *.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2026 Alibaba Cloud
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

NOTICE

@@ -0,0 +1,13 @@
+Nyx
+Copyright 2026 devanshbatham
+
+Derived from Qwen/Qwen3.5-35B-A3B, revision
+59d61f3ce65a6d9863b86d2e96597125219dc754.
+
+Qwen3.5 was created by the Qwen team at Alibaba Cloud and is distributed under
+the Apache License 2.0. Nyx adds task-specific decision adaptation, packaging,
+runtime integration, and INT4+FP8 quantization. The original Qwen license is
+retained in LICENSE.
+
+Qwen and Alibaba Cloud do not endorse and are not affiliated with this release.
+

README.md

@@ -0,0 +1,124 @@
+---
+license: apache-2.0
+library_name: sglang
+pipeline_tag: text-generation
+base_model:
+  - Qwen/Qwen3.5-35B-A3B
+tags:
+  - nyx
+  - qwen3.5
+  - mixture-of-experts
+  - int4
+  - fp8
+  - rocm
+  - sglang
+---
+
+# Nyx
+
+Nyx is a task-specialized Qwen3.5-35B-A3B decision model published by
+[`devanshbatham`](https://huggingface.co/devanshbatham). This repository contains the exact
+accepted production checkpoint in its native SGLang presharded format.
+
+![Nyx and Jev benchmark comparison](assets/nyx-vs-jev-benchmarks.png)
+
+## Model
+
+| Property | Value |
+|---|---|
+| Base model | [`Qwen/Qwen3.5-35B-A3B`](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) |
+| Base revision | `59d61f3ce65a6d9863b86d2e96597125219dc754` |
+| Architecture | Qwen3.5-35B-A3B text MoE |
+| Quantization | Per-channel INT4 routed-expert storage; FP8 expert compute; dynamic FP8 dense layers |
+| Runtime precision | BF16 activations and KV cache; FP32 SSM state |
+| Measured model memory | 19.43 GB |
+| Qualified hardware | AMD MI325X (`gfx942`) |
+| Publisher | `devanshbatham` |
+
+The original architecture, tokenizer, and base weights are from Qwen. Nyx contains subsequent
+task-specific decision adaptation and quantization. See [NOTICE](NOTICE) and [provenance.json](provenance.json).
+
+## Compatibility
+
+This is **not** a standard Transformers `safetensors` checkpoint and cannot be loaded with
+`AutoModel.from_pretrained`. The weight files use SGLang's pinned presharded loader plus the
+included runtime patches. Only the pinned SGLang container on AMD MI325X has been qualified.
+Other GPUs, SGLang versions, quantization backends, and conversions are unverified.
+
+## Run
+
+Requirements: Linux, Docker, ROCm-visible `/dev/kfd` and `/dev/dri`, and an AMD MI325X.
+
+```bash
+hf download devanshbatham/nyx --local-dir nyx
+cd nyx
+python3 verify.py
+bash serve.sh
+```
+
+The server binds to `127.0.0.1:30000` by default. Override the port with `NYX_PORT`. Keep the
+model backend private and place an authenticated API gateway and TLS termination in front of it
+for networked production use.
+
+```bash
+curl http://127.0.0.1:30000/v1/chat/completions \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "model": "/model",
+    "messages": [{"role": "user", "content": "Classify as positive or negative: I loved it."}],
+    "temperature": 0,
+    "max_tokens": 8
+  }'
+```
+
+The production System One gateway and typed clients are maintained separately at
+[`devanshbatham/nyx`](https://github.com/devanshbatham/nyx).
+
+## Evaluation
+
+The frozen public comparison used 2,277 identical cases, concurrency 8, and counted every failure
+as incorrect. All three providers completed every case with zero errors and zero retries.
+
+| Benchmark | Nyx INT4+FP8 | Qwen base BF16 | Jev 1.13 |
+|---|---:|---:|---:|
+| AG News | 84.75% | 85.75% | 86.50% |
+| TREC coarse | 92.00% | 94.00% | 93.00% |
+| CoLA | 79.00% | 79.75% | 77.75% |
+| RTE | 93.14% | 93.86% | 91.34% |
+| Emotion | 43.33% | 44.67% | 47.00% |
+| IMDb | 96.50% | 96.25% | 97.00% |
+| **Dataset-macro accuracy** | **81.45%** | **82.38%** | **82.10%** |
+| **Pooled accuracy** | **82.96%** | **83.93%** | **83.62%** |
+
+Nyx versus Jev exact paired McNemar p=`0.23154`. Base versus Nyx p=`0.03896` before correction
+for three comparisons and does not cross the Bonferroni-adjusted `0.0167` threshold. Public
+pretraining contamination is unknown for every model. These six benchmarks do not establish
+universal superiority. See [the evaluation record](benchmarks/random-six-report.md).
+
+On the separately authored 2,000-case H1 development workload, Nyx INT4+FP8 scored 85.65% versus
+86.65% for Nyx FP8. The compact model used 19.43 GB versus 34.58 GB model memory and was accepted
+with the measured 1.00 percentage-point tradeoff. H1 is development evidence, not an independent
+external test.
+
+## Training disclosure
+
+Nyx was post-trained for structured decision/classification tasks using task-specific examples,
+including teacher probability targets. The publisher has affirmed the right to distribute the
+resulting weights. Raw training records and private evaluation data are not included.
+
+## Limitations
+
+- Optimized for structured classification and decision workloads; general assistant quality was not established.
+- Qualified only with the included SGLang/ROCm runtime on one AMD MI325X.
+- The public evaluation is finite and may overlap model pretraining data.
+- Outputs can be incorrect or biased. Use independent validation and human review for consequential decisions.
+- No claim of bit-exact parity with BF16, Jev, or the original Qwen model is made.
+- The included launcher exposes only a loopback model backend; production authentication is the operator's responsibility.
+
+## License and attribution
+
+Apache License 2.0. Nyx is derived from
+[`Qwen/Qwen3.5-35B-A3B`](https://huggingface.co/Qwen/Qwen3.5-35B-A3B), also distributed under
+Apache-2.0. Qwen and Alibaba Cloud are credited as the original model authors. This project is not
+endorsed by or affiliated with the Qwen authors.
+

benchmarks/random-six-report.md

@@ -0,0 +1,43 @@
+# Six public classification benchmarks: Nyx vs base vs Jev
+
+Frozen paired cohort: **2,277 cases**; SHA-256 `046a38c24747fd724003aca7e5f1d1b6e148d599092222361c0717ecff84ecad`; concurrency 8; every failure counts wrong.
+
+| Dataset | Cases | Nyx INT4+FP8 | Qwen base BF16 | Jev 1.13 |
+|---|---:|---:|---:|---:|
+| ag_news | 400 | 84.75% (339/400) | 85.75% (343/400) | 86.50% (346/400) |
+| trec_coarse | 500 | 92.00% (460/500) | 94.00% (470/500) | 93.00% (465/500) |
+| cola | 400 | 79.00% (316/400) | 79.75% (319/400) | 77.75% (311/400) |
+| rte | 277 | 93.14% (258/277) | 93.86% (260/277) | 91.34% (253/277) |
+| emotion | 300 | 43.33% (130/300) | 44.67% (134/300) | 47.00% (141/300) |
+| imdb | 400 | 96.50% (386/400) | 96.25% (385/400) | 97.00% (388/400) |
+
+| Aggregate | Nyx INT4+FP8 | Qwen base BF16 | Jev 1.13 |
+|---|---:|---:|---:|
+| Dataset-macro accuracy | 81.45% | 82.38% | 82.10% |
+| Dataset-macro F1 | 81.06% | 82.02% | 81.85% |
+| Pooled micro accuracy | 82.96% | 83.93% | 83.62% |
+
+## Paired comparison
+
+Exact two-sided McNemar tests use the identical case-level predictions. Base beat Nyx on 63 discordant cases versus 41 in the other direction (p=0.03896, nominal and not corrected for three comparisons). Jev beat Nyx 76 to 61 (p=0.23154), while base beat Jev 66 to 59 (p=0.59168). Thus only base-versus-Nyx crosses an uncorrected 0.05 threshold; none crosses a Bonferroni-adjusted 0.0167 threshold.
+
+| Operational timing | Nyx INT4+FP8 | Qwen base BF16 | Jev 1.13 |
+|---|---:|---:|---:|
+| Wall time | 20.21 s | 23.82 s | 53.67 s |
+| Throughput | 112.67 req/s | 95.58 req/s | 42.43 req/s |
+| Median request | 66.25 ms | 78.23 ms | 182.33 ms |
+| p95 request | 89.22 ms | 101.82 ms | 245.12 ms |
+
+## Method and caveats
+
+- Dataset choices and examples were frozen before provider scoring; request bodies and order are byte-equivalent across providers.
+- AG News, CoLA, Emotion, and IMDb are fixed class-stratified samples; TREC test and RTE validation are complete.
+- Accuracy is reported per dataset. Dataset-macro accuracy weights each benchmark equally; pooled accuracy weights individual cases.
+- Public-benchmark pretraining contamination is unknown for all three models. This finite suite does not prove universal classification quality.
+- Timing is operational only: Nyx and base are local on one MI325X, while Jev is remote HTTPS on unmatched infrastructure.
+
+## Artifacts
+
+- Manifest: `artifacts/random-six-headtohead-v2/manifest.json`
+- Comparison JSON: `artifacts/random-six-headtohead-v2/comparison.json`
+- Raw provider results: `artifacts/random-six-headtohead-v2/result-nyx-int4fp8.json`, `artifacts/random-six-headtohead-v2/result-qwen-base-bf16.json`, `artifacts/random-six-headtohead-v2/result-jev-1.13.json`

benchmarks/summary.json

@@ -0,0 +1,32 @@
+{
+  "format_version": 1,
+  "cases_sha256": "046a38c24747fd724003aca7e5f1d1b6e148d599092222361c0717ecff84ecad",
+  "cases": 2277,
+  "concurrency": 8,
+  "errors": {
+    "nyx-int4fp8": 0,
+    "qwen-base-bf16": 0,
+    "jev-1.13": 0
+  },
+  "dataset_macro_accuracy": {
+    "nyx-int4fp8": 0.8145402125952668,
+    "qwen-base-bf16": 0.8237991375852386,
+    "jev-1.13": 0.8209762334536702
+  },
+  "pooled_accuracy": {
+    "nyx-int4fp8": 0.8296003513394817,
+    "qwen-base-bf16": 0.839262187088274,
+    "jev-1.13": 0.8361879666227492
+  },
+  "paired_mcnemar_two_sided_p": {
+    "nyx-int4fp8_vs_qwen-base-bf16": 0.03895793340131187,
+    "nyx-int4fp8_vs_jev-1.13": 0.23153716588439505,
+    "qwen-base-bf16_vs_jev-1.13": 0.591683927094163
+  },
+  "limitations": [
+    "Public benchmark contamination is unknown.",
+    "Fixed samples are not complete official-suite scores except TREC test and RTE validation.",
+    "Jev timing used unmatched remote infrastructure."
+  ]
+}
+

build_manifest.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python3
+"""Build the deterministic release manifest for this repository."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from pathlib import Path
+
+
+ROOT = Path(__file__).resolve().parent
+EXCLUDE = {"release-manifest.json", "build_manifest.py", "verify.py"}
+
+
+def digest(path: Path) -> str:
+    checksum = hashlib.sha256()
+    with path.open("rb") as stream:
+        for block in iter(lambda: stream.read(8 * 1024 * 1024), b""):
+            checksum.update(block)
+    return checksum.hexdigest()
+
+
+files = {}
+for path in sorted(ROOT.rglob("*")):
+    if not path.is_file() or ".git" in path.parts:
+        continue
+    relative = path.relative_to(ROOT).as_posix()
+    if relative in EXCLUDE:
+        continue
+    files[relative] = {"bytes": path.stat().st_size, "sha256": digest(path)}
+
+payload = {
+    "format_version": 1,
+    "repository": "devanshbatham/nyx",
+    "artifact": "nyx-35b-a3b-int4fp8",
+    "files": files,
+    "total_bytes": sum(item["bytes"] for item in files.values()),
+}
+(ROOT / "release-manifest.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
+print(json.dumps({"files": len(files), "total_bytes": payload["total_bytes"]}))
+

chat_template.jinja

@@ -0,0 +1,154 @@
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count, is_system_content=false) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain images.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain videos.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Video ' ~ video_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- else %}
+                {{- raise_exception('Unexpected item type in content.') }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif content is none or content is undefined %}
+        {{- '' }}
+    {%- else %}
+        {{- raise_exception('Unexpected content type.') }}
+    {%- endif %}
+{%- endmacro %}
+{%- if not messages %}
+    {{- raise_exception('No messages provided.') }}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- else %}
+        {{- '<think>\n' }}
+    {%- endif %}
+{%- endif %}

config.json

@@ -0,0 +1,96 @@
+{
+  "architectures": [
+    "Qwen3_5MoeForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attn_output_gate": true,
+  "bos_token_id": null,
+  "dtype": "bfloat16",
+  "eos_token_id": 248044,
+  "full_attention_interval": 4,
+  "head_dim": 256,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "layer_types": [
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention",
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention"
+  ],
+  "linear_conv_kernel_dim": 4,
+  "linear_key_head_dim": 128,
+  "linear_num_key_heads": 16,
+  "linear_num_value_heads": 32,
+  "linear_value_head_dim": 128,
+  "mamba_ssm_dtype": "float32",
+  "max_position_embeddings": 262144,
+  "mlp_only_layers": [],
+  "model_type": "qwen3_5_moe_text",
+  "moe_intermediate_size": 512,
+  "mtp_num_hidden_layers": 1,
+  "mtp_use_dedicated_embeddings": false,
+  "num_attention_heads": 16,
+  "num_experts": 256,
+  "num_experts_per_tok": 8,
+  "num_hidden_layers": 40,
+  "num_key_value_heads": 2,
+  "output_router_logits": false,
+  "pad_token_id": null,
+  "partial_rotary_factor": 0.25,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "mrope_interleaved": true,
+    "mrope_section": [
+      11,
+      11,
+      10
+    ],
+    "partial_rotary_factor": 0.25,
+    "rope_theta": 10000000,
+    "rope_type": "default"
+  },
+  "router_aux_loss_coef": 0.001,
+  "shared_expert_intermediate_size": 512,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.17.0",
+  "use_cache": true,
+  "vocab_size": 248320
+}

generation_config.json

@@ -0,0 +1,13 @@
+{
+  "bos_token_id": 248044,
+  "do_sample": true,
+  "eos_token_id": [
+    248046,
+    248044
+  ],
+  "pad_token_id": 248044,
+  "temperature": 1.0,
+  "top_k": 20,
+  "top_p": 0.95,
+  "transformers_version": "5.17.0"
+}

provenance.json

@@ -0,0 +1,30 @@
+{
+  "format_version": 1,
+  "model": "devanshbatham/nyx",
+  "publisher": "devanshbatham",
+  "base_model": "Qwen/Qwen3.5-35B-A3B",
+  "base_revision": "59d61f3ce65a6d9863b86d2e96597125219dc754",
+  "license": "Apache-2.0",
+  "quantization": {
+    "routed_expert_storage": "per-channel INT4",
+    "expert_compute": "FP8",
+    "dense_layers": "dynamic FP8",
+    "activations": "BF16",
+    "kv_cache": "BF16",
+    "ssm_state": "FP32"
+  },
+  "qualified_runtime": {
+    "hardware": "AMD MI325X gfx942",
+    "sglang_image": "lmsysorg/sglang@sha256:a0ffcdd013af6d79c9b7c4348c42b14a79a48f6c43dd3ca0235ec3802da45f75",
+    "tensor_parallel_size": 1
+  },
+  "accepted_evidence": {
+    "model_memory_gb": 19.43,
+    "package_bytes": 19657745446,
+    "h1_correct": 1713,
+    "h1_total": 2000,
+    "h1_result_sha256": "b10697d6a03fb19c56e7193eb85b932c92aee2a27d4b6a088f684051dff1c01e",
+    "source_checksum_manifest_sha256": "829aa54dd4f440c10d85c6253a41891a69ba1717ace2d193cbfab3c1a7546c6e"
+  }
+}
+

release-manifest.json

@@ -0,0 +1,116 @@
+{
+  "artifact": "nyx-35b-a3b-int4fp8",
+  "files": {
+    ".gitattributes": {
+      "bytes": 448,
+      "sha256": "f3b77d221a265fd892b0993f88aa613596c2f58adf13dcee9b1d3a14156a14f8"
+    },
+    "LICENSE": {
+      "bytes": 11544,
+      "sha256": "bbedc3fda3305820b977265f01b8619d87570a6739de3a5582c3464840f1e57a"
+    },
+    "NOTICE": {
+      "bytes": 459,
+      "sha256": "08640c173e604546300dd76b47c842858ffcc236942b45c6cfdd894c851d8a77"
+    },
+    "README.md": {
+      "bytes": 4989,
+      "sha256": "f8965349f25854cf6818f780b0a28151150f0b78eab15064a435e695defa0316"
+    },
+    "assets/nyx-vs-jev-benchmarks.png": {
+      "bytes": 1757593,
+      "sha256": "679a775f1f891713884549dcf6c2eca849f5b4c40064bc1ffb70a507e1add9c8"
+    },
+    "benchmarks/random-six-report.md": {
+      "bytes": 2731,
+      "sha256": "1dd4c4a7cad8a4158a4e5d9f2ddd0fa01d1f8c3df629585a1591229c8b4d9a98"
+    },
+    "benchmarks/summary.json": {
+      "bytes": 955,
+      "sha256": "e3a56d1bc57b8d1837ef9abbd5852d2d45a83cbc345b2a638f1e8e5018599a3f"
+    },
+    "chat_template.jinja": {
+      "bytes": 7756,
+      "sha256": "a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715"
+    },
+    "config.json": {
+      "bytes": 2336,
+      "sha256": "27e3575d049c5a0187b783aecff914135be455439134368aba8cdcc6d678ebb7"
+    },
+    "generation_config.json": {
+      "bytes": 214,
+      "sha256": "df6f86c3fdce573ecdb55cec35f502ddda70e70abe78288a78164aef120232e7"
+    },
+    "provenance.json": {
+      "bytes": 980,
+      "sha256": "85773821fad372dd98cf10b02c4103e78406f0b45d13f2875772470a9791f1d6"
+    },
+    "runtime/entrypoint.sh": {
+      "bytes": 276,
+      "sha256": "742824765cb2f5361a476674e16a5d347abb9ae905e2798a2183121ea1ee4b07"
+    },
+    "runtime/label-token-ids.json": {
+      "bytes": 761,
+      "sha256": "8cc5610b02d5ba14358ebef59cd16f567767632fb6ea03de1565b18f1c003a40"
+    },
+    "runtime/patches/qwen35-text-quark-generated-scales.patch": {
+      "bytes": 1680,
+      "sha256": "920a8cabd8d2f6565e4f0c3296b1b0e154e48c227e4fd12f019033792cbe567f"
+    },
+    "runtime/patches/sglang-batch-result-processor.py": {
+      "bytes": 57803,
+      "sha256": "f7f58ef7f377e3df111c87d1803b6e0008bdf76726d80cbf1012275839e77da7"
+    },
+    "runtime/patches/sglang-logits-processor.py": {
+      "bytes": 54217,
+      "sha256": "aed498368a4402a84a4ef71a5962f1ff78cad265d4becbe5b3b677f1cb8ad9e2"
+    },
+    "runtime/patches/sglang-qwen35-text.py": {
+      "bytes": 12308,
+      "sha256": "52cdbda541ff612353adf5d5b04a6668954040b48cfa5209f59abbd662a023e1"
+    },
+    "runtime/patches/sglang-qwen35.py": {
+      "bytes": 106934,
+      "sha256": "f6975a9b3303002c25418b0d5b32c6bd40502ba6568ab7d5c847c6983705c01e"
+    },
+    "serve.sh": {
+      "bytes": 2065,
+      "sha256": "5f308855054cd6dd86fde051cc9771e02e1e160dc0ddbad1f5eed8e7b81a4b65"
+    },
+    "tokenizer.json": {
+      "bytes": 19989325,
+      "sha256": "06b9509352d2af50381ab2247e083b80d32d5c0aba91c272ca9ff729b6a0e523"
+    },
+    "tokenizer_config.json": {
+      "bytes": 1123,
+      "sha256": "bee8eba30f0eb4af73c0fe2cd06d0f89b657d7819941c438157ec42f7c80ea87"
+    },
+    "weights/TP-1-sig-e793a9432f30cd4e/READY": {
+      "bytes": 70,
+      "sha256": "b20cd5cb19a0b44c0189bb79b85ee0300f0453a31bfef7b39ac825b409cbb578"
+    },
+    "weights/TP-1-sig-e793a9432f30cd4e/checksum.json": {
+      "bytes": 359616,
+      "sha256": "829aa54dd4f440c10d85c6253a41891a69ba1717ace2d193cbfab3c1a7546c6e"
+    },
+    "weights/TP-1-sig-e793a9432f30cd4e/model-00000-common.safetensor": {
+      "bytes": 5268047400,
+      "sha256": "16e3906e9a161a6337a4a0f6d06e352d09360b37a9d4c4ea341ea978192ee289"
+    },
+    "weights/TP-1-sig-e793a9432f30cd4e/model-00001-common.safetensor": {
+      "bytes": 5120198752,
+      "sha256": "7ad7affbcc8912528ff2859568b51ee9427cc17119312ea8750cf118544ea5d3"
+    },
+    "weights/TP-1-sig-e793a9432f30cd4e/model-00002-common.safetensor": {
+      "bytes": 5254941984,
+      "sha256": "fa4a08a1e0e0de30543d343a629c6710e59575cd6f98f5beae789f34a59b79b7"
+    },
+    "weights/TP-1-sig-e793a9432f30cd4e/model-00003-common.safetensor": {
+      "bytes": 4014197624,
+      "sha256": "dfb557a70a97b91208837dac145569c844971062970f41466738e61e77d53eb3"
+    }
+  },
+  "format_version": 1,
+  "repository": "devanshbatham/nyx",
+  "total_bytes": 19679761943
+}

runtime/entrypoint.sh

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+target=/sgl-workspace/sglang/python/sglang/srt/models/qwen3_5_text.py
+cp /model/runtime/patches/sglang-qwen35-text.py "$target"
+patch --batch "$target" < /model/runtime/patches/qwen35-text-quark-generated-scales.patch
+exec python3 "$@"
+

runtime/label-token-ids.json

@@ -0,0 +1 @@
+[32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 5840, 1803, 1646, 1745, 13276, 8018, 1825, 28946, 15015, 29595, 11568, 939, 1354, 1058, 18183, 2456, 88898, 905, 1846, 802, 33869, 7839, 14006, 2860, 2926, 22828, 6844, 9798, 4738, 9265, 11261, 19278, 36513, 93801, 8335, 14544, 85266, 9110, 28000, 15137, 4525, 25261, 12717, 7116, 17078, 14497, 57339, 74909, 52072, 19305, 4887, 12607, 3580, 6281, 2036, 9362, 8533, 2080, 10911, 2925, 3040, 9690, 27731, 8023, 6901, 8702, 6211, 1123, 16307, 18990, 64045, 63037, 33380, 6151, 3392, 5449, 3967, 1113, 5095, 51923, 49600, 17099, 51483, 17756, 16037, 8135, 30237, 5683, 9992, 7444, 5751, 10284, 20887, 59884, 52396, 16103, 67547, 18535, 8006, 7263, 1425, 6878]

runtime/patches/qwen35-text-quark-generated-scales.patch

@@ -0,0 +1,32 @@
+--- qwen3_5_text.py
++++ qwen3_5_text.py
+@@ -242,17 +242,24 @@
+             loaded_params.add("lm_head.weight")
+ 
+         missing = set(params_dict) - loaded_params
+-        if (self.quant_config is not None
+-                and self.quant_config.get_name() == "fp8"
++        quant_name = (
++            self.quant_config.get_name() if self.quant_config is not None else None
++        )
++        if (quant_name in {"fp8", "quark_int4fp8_moe"}
+                 and not self.quant_config.is_checkpoint_fp8_serialized):
+-            # Online BF16->FP8 creates these scales itself. They are deliberately
+-            # absent from the source BF16 checkpoint; all model weights still
+-            # have to load. Do not apply this exception to serialized FP8.
++            # Online quantizers create their scales from source weights. They are
++            # deliberately absent from the BF16 checkpoint; all source model
++            # weights still have to load. Never apply this to serialized FP8.
+             generated_suffixes = (
+                 ".mlp.experts.w13_weight_scale", ".mlp.experts.w2_weight_scale",
+                 ".mlp.experts.w13_weight_scale1", ".mlp.experts.w2_weight_scale1",
+                 ".attn.k_scale", ".attn.v_scale",
+             )
++            if quant_name == "quark_int4fp8_moe":
++                generated_suffixes += (
++                    ".mlp.experts.w13_fp8_scale", ".mlp.experts.w2_fp8_scale",
++                    ".mlp.experts.w13_int4_scale", ".mlp.experts.w2_int4_scale",
++                )
+             generated = {name for name in missing if name.endswith(generated_suffixes)}
+             missing -= generated
+         if missing:

runtime/patches/sglang-batch-result-processor.py

@@ -0,0 +1,1378 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
+
+import torch
+
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.environ import envs
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.managers.schedule_batch import (
+    FINISH_ABORT,
+    FINISH_MATCHED_TOKEN,
+    Req,
+    ScheduleBatch,
+    mamba_lazy_spec_in_window,
+)
+from sglang.srt.mem_cache.common import (
+    maybe_cache_unfinished_req,
+    release_kv_cache,
+)
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    get_required_capture_hidden_mode,
+    get_server_return_hidden_states_mode,
+)
+from sglang.srt.runtime_context import (
+    get_disagg,
+    get_memory,
+    get_observability,
+    mamba_extra_buffer_lazy_enabled,
+    mamba_track_grid,
+    max_speculative_num_draft_tokens,
+)
+from sglang.srt.sampling.sampling_observer import CommittedTokens
+from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
+from sglang.srt.state_capturer.indexer_topk import get_global_indexer_capturer
+from sglang.srt.state_capturer.routed_experts import get_global_experts_capturer
+
+if TYPE_CHECKING:
+    from sglang.srt.beam_search.coordinator import BeamCoordinator
+    from sglang.srt.configs.model_config import ModelConfig
+    from sglang.srt.disaggregation.decode_kvcache_offload_manager import (
+        DecodeKVCacheOffloadManager,
+    )
+    from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator
+    from sglang.srt.managers.scheduler_components.logprob_result_processor import (
+        SchedulerLogprobResultProcessor,
+    )
+    from sglang.srt.managers.scheduler_components.metrics_reporter import (
+        SchedulerMetricsReporter,
+    )
+    from sglang.srt.managers.scheduler_components.output_streamer import (
+        SchedulerOutputStreamer,
+    )
+    from sglang.srt.managers.tp_worker import BaseTpWorker
+    from sglang.srt.managers.utils import (
+        EmbeddingBatchResult,
+        GenerationBatchResult,
+    )
+    from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+    from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+    from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
+    from sglang.srt.observability.metrics_collector import SchedulerMetricsCollector
+    from sglang.srt.sampling.sampling_observer import HostAuxiliaryOutput
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(kw_only=True, slots=True, frozen=True)
+class SchedulerBatchResultProcessor:
+    is_generation: bool
+    disaggregation_mode: DisaggregationMode
+    enable_overlap: bool
+    enable_overlap_mlx: bool
+    model_config: ModelConfig
+    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
+    tree_cache: BasePrefixCache
+    hisparse_coordinator: Optional[HiSparseCoordinator]
+    req_to_token_pool: ReqToTokenPool
+    decode_offload_manager: Optional[DecodeKVCacheOffloadManager]
+    metrics_collector: SchedulerMetricsCollector
+    metrics_reporter: SchedulerMetricsReporter
+    draft_worker: BaseTpWorker
+    model_worker: BaseTpWorker
+    logprob_result_processor: SchedulerLogprobResultProcessor
+    output_streamer: SchedulerOutputStreamer
+    beam_coordinator: BeamCoordinator
+    abort_request: Callable
+
+    def process_batch_result_prebuilt(self, batch: ScheduleBatch):
+        assert self.disaggregation_mode == DisaggregationMode.DECODE
+        use_free_group = get_disagg().disaggregation_decode_enable_radix_cache
+        if use_free_group:
+            self.token_to_kv_pool_allocator.free_group_begin()
+        for req in batch.reqs:
+            req.time_stats.set_decode_prebuilt_finish_time()
+            req.update_finish_state()
+            if req.finished():
+                req.time_stats.set_quick_finish_time()
+                if get_memory().enable_hisparse:
+                    self.hisparse_coordinator.request_finished(req)
+                release_kv_cache(req, self.tree_cache)
+
+        # Note: Logprobs should be handled on the prefill engine.
+        self.output_streamer.stream_output(batch.reqs, batch.return_logprob)
+        if use_free_group:
+            self.token_to_kv_pool_allocator.free_group_end()
+
+    def _maybe_collect_routed_experts(self, req: Req):
+        """Collect routed experts for a finished request.
+
+        Returns immediately if `return_routed_experts` was not set on the
+        request, so non-opted-in reqs don't pay the host-gather cost.
+
+        Honors the caller's absolute start so the response covers
+        `[start_len, seqlen - 1)`. The default start_len is 0, which returns
+        the full sequence.
+
+        Logs a soft warning if the resulting tensor's row count differs from
+        the expected `seqlen - 1 - start_len`, to catch silent regressions.
+        """
+        if not req.return_routed_experts:
+            return
+        capturer = get_global_experts_capturer()
+        if capturer is None:
+            return
+        start_len = req.routed_experts_start_len
+        seqlen = len(req.origin_input_ids) + len(req.output_ids_through_stop)
+        req.routed_experts = capturer.get_topk(
+            req_pool_idx=req.kv.req_pool_idx,
+            seqlen=seqlen,
+            req_to_token_pool=self.req_to_token_pool,
+            start_len=start_len,
+        )
+
+        expected_rows = max(0, seqlen - 1 - start_len)
+        if (
+            req.routed_experts is not None
+            and req.routed_experts.shape[0] != expected_rows
+        ):
+            logger.warning(
+                "routed_experts row-count mismatch for req %s: got %d, expected %d "
+                "(seqlen=%d, raw_seqlen=%d, cached_tokens=%d, start_len=%s). "
+                "This indicates a silent bug.",
+                req.rid,
+                req.routed_experts.shape[0],
+                expected_rows,
+                seqlen,
+                req.seqlen,
+                req.cached_tokens,
+                req.routed_experts_start_len,
+            )
+
+    def _maybe_collect_indexer_topk(self, req: Req):
+        capturer = get_global_indexer_capturer()
+        if capturer is None:
+            return
+        seqlen = len(req.origin_input_ids) + len(req.output_ids_through_stop)
+        req.indexer_topk = capturer.get_topk(
+            req_pool_idx=req.kv.req_pool_idx,
+            seqlen=seqlen,
+            req_to_token_pool=self.req_to_token_pool,
+        )
+
+    def _maybe_collect_customized_info(
+        self,
+        i: int,
+        req: Req,
+        logits_output: LogitsProcessorOutput,
+    ):
+        if logits_output is not None and logits_output.customized_info is not None:
+            if req.customized_info is None:
+                req.customized_info = {}
+            for k, v in logits_output.customized_info.items():
+                if k not in req.customized_info:
+                    req.customized_info[k] = []
+                # Copy the element so it doesn't retain the entire batch
+                # tensor/array via a view reference.
+                elem = v[i]
+                if isinstance(elem, torch.Tensor):
+                    elem = elem.clone()
+                elif hasattr(elem, "copy") and callable(elem.copy):
+                    elem = elem.copy()
+                req.customized_info[k].append(elem)
+
+    @staticmethod
+    def _visible_output_len(req: Req) -> int:
+        return req.finished_len if req.finished_len is not None else len(req.output_ids)
+
+    @classmethod
+    def snapshot_auxiliary_output_starts(
+        cls,
+        batch: ScheduleBatch,
+        result: GenerationBatchResult,
+    ) -> Optional[List[int]]:
+        if result.auxiliary_host_output is None:
+            return None
+        return [cls._visible_output_len(req) for req in batch.reqs]
+
+    @classmethod
+    def _build_auxiliary_commits(
+        cls,
+        batch: ScheduleBatch,
+        output_starts: List[int],
+    ) -> List[Optional[CommittedTokens]]:
+        commits: List[Optional[CommittedTokens]] = []
+        for req, output_start in zip(batch.reqs, output_starts, strict=True):
+            output_end = cls._visible_output_len(req)
+            if output_end < output_start:
+                raise RuntimeError("committed output length moved backwards")
+            if output_end == output_start:
+                commits.append(None)
+                continue
+            commits.append(
+                CommittedTokens(
+                    output_index=output_start,
+                    token_ids=tuple(req.output_ids[output_start:output_end]),
+                )
+            )
+        return commits
+
+    @classmethod
+    def consume_auxiliary_output(
+        cls,
+        batch: ScheduleBatch,
+        output: HostAuxiliaryOutput,
+        output_starts: List[int],
+    ) -> None:
+        output.consume(batch, cls._build_auxiliary_commits(batch, output_starts))
+
+    def process_batch_result_prefill(
+        self,
+        batch: ScheduleBatch,
+        result: Union[GenerationBatchResult, EmbeddingBatchResult],
+    ):
+        skip_stream_req = None
+        self.token_to_kv_pool_allocator.free_group_begin()
+
+        if self.is_generation:
+            if result.copy_done is not None:
+                result.copy_done.synchronize()
+            auxiliary_output_starts = self.snapshot_auxiliary_output_starts(
+                batch, result
+            )
+            auxiliary_output = result.auxiliary_host_output
+            if result.routed_experts_output is not None:
+                result.routed_experts_output.finalize()
+                result.routed_experts_output = None
+            if result.indexer_topk_output is not None:
+                result.indexer_topk_output.finalize()
+                result.indexer_topk_output = None
+
+            (
+                logits_output,
+                next_token_ids,
+                extend_input_len_per_req,
+                extend_logprob_start_len_per_req,
+            ) = (
+                result.logits_output,
+                result.next_token_ids,
+                result.extend_input_len_per_req,
+                result.extend_logprob_start_len_per_req,
+            )
+
+            # Move next_token_ids and logprobs to cpu
+            next_token_ids = next_token_ids.tolist()
+            self.move_logprobs_to_cpu(batch=batch, logits_output=logits_output)
+
+            self._validate_pp_skip_output_comm(batch, result)
+
+            hidden_state_offset = 0
+            prefill_hidden_capture_mode = self._get_prefill_hidden_capture_mode(
+                batch,
+            )
+
+            # Check finish conditions
+            logprob_pt = 0
+
+            for i, (req, next_token_id) in enumerate(zip(batch.reqs, next_token_ids)):
+                if (
+                    batch.return_hidden_states
+                    and logits_output.hidden_states is not None
+                ):
+                    assert extend_input_len_per_req is not None
+                    hidden_state_offset = self._append_prefill_hidden_states(
+                        req=req,
+                        logits_output=logits_output,
+                        hidden_state_offset=hidden_state_offset,
+                        capture_hidden_mode=prefill_hidden_capture_mode,
+                        extend_input_len=extend_input_len_per_req[i],
+                        store=(
+                            not req.finished()
+                            and not req.is_retracted
+                            and req.inflight_middle_chunks <= 0
+                        ),
+                    )
+
+                if (
+                    req.finished() and req.inflight_middle_chunks <= 0
+                ) or req.is_retracted:
+                    # Decode req in a mixed batch, or a retracted req. Keep an
+                    # aborted middle chunk in the chunked branch long enough to
+                    # drain its accounting without streaming it.
+                    continue
+
+                if req.inflight_middle_chunks <= 0:
+                    req.time_stats.set_prefill_finished_time()
+
+                    if req.beam_group is not None:
+                        # The relay point already replaced the sampled-token
+                        # append; the group owns all finish semantics.
+                        self.beam_coordinator.commit_prefill(
+                            req, up_to_tick=batch.forward_iter
+                        )
+                    else:
+                        # req output_ids are set here
+                        req.output_ids.append(next_token_id)
+
+                        self._maybe_update_reasoning_tokens(req, next_token_id)
+
+                        req.update_finish_state()
+                    # A mixed spec tail committed its pending bonus token; advance
+                    # so the next spec prepare_for_decode reserves from the right base.
+                    if (
+                        not req.finished()
+                        and batch.decoding_reqs
+                        and req in batch.decoding_reqs
+                        and not batch.spec_algorithm.is_none()
+                    ):
+                        req.kv.kv_committed_len += 1
+                    if req.finished():
+                        self._maybe_collect_routed_experts(req)
+                        self._maybe_collect_indexer_topk(req)
+                        release_kv_cache(req, self.tree_cache)
+                        req.time_stats.set_completion_time()
+                    elif not batch.decoding_reqs or req not in batch.decoding_reqs:
+                        maybe_cache_unfinished_req(req, self.tree_cache)
+                        if get_memory().enable_hisparse:
+                            self.hisparse_coordinator.admit_request_into_staging(req)
+
+                    self._maybe_collect_customized_info(i, req, logits_output)
+
+                    if batch.return_logprob:
+                        logprob_pt = self._apply_prefill_logprobs(
+                            req=req,
+                            i=i,
+                            logits_output=logits_output,
+                            extend_input_len_per_req=extend_input_len_per_req,
+                            extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
+                            next_token_ids=next_token_ids,
+                            logprob_pt=logprob_pt,
+                        )
+
+                    if req.return_sampling_mask:
+                        self.add_sampling_mask_return_values(i, req, logits_output)
+
+                    if req.grammar is not None:
+                        self._apply_prefill_grammar(
+                            req=req,
+                            next_token_id=next_token_id,
+                            already_advanced=result.grammar_advanced,
+                        )
+
+                else:
+                    # being chunked reqs' prefill is not finished
+                    req.inflight_middle_chunks -= 1
+                    # There is only at most one request being currently chunked.
+                    # Because this request does not finish prefill,
+                    # we don't want to stream the request currently being chunked.
+                    skip_stream_req = req
+
+                    # Incrementally update input logprobs.
+                    if batch.return_logprob:
+                        logprob_pt = self._apply_chunked_prefill_logprobs(
+                            req=req,
+                            i=i,
+                            logits_output=logits_output,
+                            extend_input_len_per_req=extend_input_len_per_req,
+                            extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
+                            logprob_pt=logprob_pt,
+                        )
+
+                    req.time_stats.set_last_chunked_prefill_finish_time()
+
+            if auxiliary_output is not None:
+                self.consume_auxiliary_output(
+                    batch,
+                    auxiliary_output,
+                    auxiliary_output_starts,
+                )
+
+        else:  # embedding or reward model
+            if result.copy_done is not None:
+                result.copy_done.synchronize()
+
+            embeddings = self._convert_embeddings(result=result)
+            phs = result.pooled_hidden_states
+
+            if phs is not None:
+                if isinstance(phs, list):
+                    phs = [t.cpu().detach() for t in phs]
+                else:
+                    phs = phs.cpu().detach()
+
+            # Check finish conditions
+            for i, req in enumerate(batch.reqs):
+                if req.is_retracted:
+                    continue
+
+                req.embedding = embeddings[i]
+                if req.return_pooled_hidden_states and phs is not None:
+                    req.pooled_hidden_state = phs[i]
+                if req.inflight_middle_chunks <= 0:
+                    req.time_stats.set_prefill_finished_time()
+                    # Dummy output token for embedding models
+                    req.output_ids.append(0)
+                    req.update_finish_state()
+
+                    if req.finished():
+                        release_kv_cache(req, self.tree_cache)
+                        req.time_stats.set_completion_time()
+                    else:
+                        maybe_cache_unfinished_req(req, self.tree_cache)
+                else:
+                    # being chunked reqs' prefill is not finished
+                    req.inflight_middle_chunks -= 1
+                    req.time_stats.set_last_chunked_prefill_finish_time()
+
+        self.token_to_kv_pool_allocator.free_group_end()
+        self.output_streamer.stream_output(
+            batch.reqs, batch.return_logprob, skip_stream_req
+        )
+
+        can_run_cuda_graph = result.can_run_cuda_graph
+        # None on decode->extend converted batches; they are decode work and
+        # have no prefill stats to report.
+        if batch.prefill_stats is not None:
+            self.metrics_reporter.report_prefill_stats(
+                batch=batch,
+                prefill_stats=batch.prefill_stats,
+                can_run_cuda_graph=can_run_cuda_graph,
+                dp_cooperation_info=batch.dp_cooperation_info,
+            )
+
+    def _convert_embeddings(self, *, result: EmbeddingBatchResult) -> list:
+        is_sparse = envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set()
+
+        embeddings = result.embeddings
+
+        if is_sparse:
+            batch_ids, token_ids = embeddings.indices()
+            values = embeddings.values()
+
+            embeddings = [{} for _ in range(embeddings.size(0))]
+            for i in range(batch_ids.shape[0]):
+                embeddings[batch_ids[i].item()][token_ids[i].item()] = values[i].item()
+        else:
+            if isinstance(embeddings, torch.Tensor):
+                embeddings = embeddings.tolist()
+            else:
+                embeddings = [tensor.tolist() for tensor in embeddings]
+        return embeddings
+
+    def move_logprobs_to_cpu(
+        self,
+        *,
+        batch: ScheduleBatch,
+        logits_output: LogitsProcessorOutput,
+    ) -> None:
+        if batch.return_logprob:
+            if logits_output.next_token_logprobs is not None:
+                logits_output.next_token_logprobs = (
+                    logits_output.next_token_logprobs.tolist()
+                )
+            if logits_output.input_token_logprobs is not None:
+                logits_output.input_token_logprobs = tuple(
+                    logits_output.input_token_logprobs.tolist()
+                )
+            if logits_output.next_token_top_logprobs_val:
+                logits_output.next_token_top_logprobs_val = [
+                    v.tolist() for v in logits_output.next_token_top_logprobs_val
+                ]
+                logits_output.next_token_top_logprobs_idx = [
+                    x.tolist() for x in logits_output.next_token_top_logprobs_idx
+                ]
+            if logits_output.next_token_token_ids_logprobs_val:
+                logits_output.next_token_token_ids_logprobs_val = [
+                    v.tolist() if isinstance(v, torch.Tensor) else v
+                    for v in logits_output.next_token_token_ids_logprobs_val
+                ]
+
+    def _apply_prefill_logprobs(
+        self,
+        *,
+        req: Req,
+        i: int,
+        logits_output: LogitsProcessorOutput,
+        extend_input_len_per_req: Optional[List[int]],
+        extend_logprob_start_len_per_req: Optional[List[int]],
+        next_token_ids: List[int],
+        logprob_pt: int,
+    ) -> int:
+        assert extend_logprob_start_len_per_req is not None
+        assert extend_input_len_per_req is not None
+        extend_logprob_start_len = extend_logprob_start_len_per_req[i]
+        extend_input_len = extend_input_len_per_req[i]
+
+        num_input_logprobs = self.logprob_result_processor.calculate_num_input_logprobs(
+            req,
+            extend_input_len,
+            extend_logprob_start_len,
+        )
+
+        if req.return_logprob:
+            self.logprob_result_processor.add_logprob_return_values(
+                i,
+                req,
+                logprob_pt,
+                next_token_ids,
+                num_input_logprobs,
+                logits_output,
+            )
+        logprob_pt += num_input_logprobs
+        return logprob_pt
+
+    @staticmethod
+    def _validate_pp_skip_output_comm(
+        batch: ScheduleBatch,
+        result: Union[GenerationBatchResult, EmbeddingBatchResult],
+    ):
+        """Validate PP skip output comm correctness.
+
+        - When skip=True: all reqs must be middle chunks (inflight_middle_chunks > 0)
+          so placeholder zeros are never consumed via req.output_ids.append().
+        - When skip=False: at least one req should consume next_token_ids
+          (inflight_middle_chunks <= 0), otherwise warn.
+        """
+        if not envs.SGLANG_PP_SKIP_PURE_CHUNKED_OUTPUT_COMM.get():
+            return
+
+        if not getattr(result, "skipped_output_comm", False):
+            if batch.forward_mode.is_extend() and not batch.forward_mode.is_prebuilt():
+                has_consumed_output = any(
+                    req.inflight_middle_chunks <= 0
+                    for req in batch.reqs
+                    if not req.finished() and not req.is_retracted
+                )
+                if not has_consumed_output and len(batch.reqs) > 0:
+                    chunks = list([r.inflight_middle_chunks for r in batch.reqs])
+                    logger.warning(
+                        f"PP non-skip output comm: no req consumed next_token_ids. "
+                        f"contains_last_prefill_chunk={batch.contains_last_prefill_chunk}, "
+                        f"num_reqs={len(batch.reqs)}, all inflight_middle_chunks={chunks}"
+                    )
+            return
+
+        for req in batch.reqs:
+            if not req.finished() and not req.is_retracted:
+                assert req.inflight_middle_chunks > 0, (
+                    f"PP skip output comm invariant violated: req {req.rid} "
+                    f"has inflight_middle_chunks={req.inflight_middle_chunks} "
+                    f"but output was skipped (contains_last_prefill_chunk="
+                    f"{batch.contains_last_prefill_chunk}). "
+                    f"Placeholder zeros would be appended to output_ids."
+                )
+
+    def _append_prefill_hidden_states(
+        self,
+        *,
+        req: Req,
+        logits_output: LogitsProcessorOutput,
+        hidden_state_offset: int,
+        capture_hidden_mode: CaptureHiddenMode,
+        extend_input_len: int,
+        store: bool = True,
+    ) -> int:
+        if capture_hidden_mode.is_full():
+            start = hidden_state_offset
+            hidden_state_offset += extend_input_len
+            if not store or not req.return_hidden_states:
+                return hidden_state_offset
+
+            req_hidden_states = logits_output.hidden_states[start:hidden_state_offset]
+            if req.return_hidden_states is True:
+                req.hidden_states.append(req_hidden_states.cpu().clone().tolist())
+            elif req.return_hidden_states == "last":
+                req.hidden_states.append(req_hidden_states[-1].cpu().tolist())
+        elif capture_hidden_mode.is_last():
+            index = hidden_state_offset
+            hidden_state_offset += 1
+            if store and req.return_hidden_states:
+                req.hidden_states.append(
+                    logits_output.hidden_states[index].cpu().tolist()
+                )
+        else:
+            raise ValueError(
+                f"Unexpected hidden states capture mode: {capture_hidden_mode}"
+            )
+        return hidden_state_offset
+
+    @staticmethod
+    def _append_decode_hidden_states(
+        *,
+        req: Req,
+        hidden_states: torch.Tensor,
+        start: int,
+        accept_len: int,
+    ) -> None:
+        if accept_len <= 0:
+            return
+
+        if req.return_hidden_states == "last":
+            valid_accept_len = accept_len
+            if req.finished_len is not None:
+                step_start = len(req.output_ids) - accept_len
+                valid_accept_len = max(
+                    0,
+                    min(accept_len, req.finished_len - step_start),
+                )
+            if valid_accept_len > 0:
+                req.hidden_states[:] = [
+                    hidden_states[start + valid_accept_len - 1].cpu().tolist()
+                ]
+        else:
+            req.hidden_states.extend(
+                hidden_states[start : start + accept_len].cpu().tolist()
+            )
+
+    @staticmethod
+    def _get_prefill_hidden_capture_mode(batch: ScheduleBatch) -> CaptureHiddenMode:
+        return get_required_capture_hidden_mode(
+            max(
+                batch.return_hidden_states_mode,
+                get_server_return_hidden_states_mode(),
+            ),
+            batch.spec_info,
+        )
+
+    def _apply_prefill_grammar(
+        self, *, req: Req, next_token_id: int, already_advanced: bool = False
+    ) -> None:
+        # The grammar barrier may have already advanced the FSM over this prefilled
+        # token (spec overlap path); only advance if not, but always sync
+        # grammar.finished.
+        if not already_advanced:
+            # FIXME: this try-except block is for handling unexpected xgrammar issue.
+            try:
+                req.grammar.accept_token(next_token_id)
+            except ValueError as e:
+                # Grammar accept_token can raise ValueError if the token is not in the
+                # grammar. This can happen if the grammar is not set correctly or the
+                # token is invalid.
+                logger.error(
+                    f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
+                )
+                req.to_finish = FINISH_ABORT()
+        req.grammar.finished = req.finished()
+
+    def _apply_chunked_prefill_logprobs(
+        self,
+        *,
+        req: Req,
+        i: int,
+        logits_output: LogitsProcessorOutput,
+        extend_input_len_per_req: Optional[List[int]],
+        extend_logprob_start_len_per_req: Optional[List[int]],
+        logprob_pt: int,
+    ) -> int:
+        extend_logprob_start_len = extend_logprob_start_len_per_req[i]
+        extend_input_len = extend_input_len_per_req[i]
+        if extend_logprob_start_len < extend_input_len:
+            # Update input logprobs.
+            num_input_logprobs = (
+                self.logprob_result_processor.calculate_num_input_logprobs(
+                    req,
+                    extend_input_len,
+                    extend_logprob_start_len,
+                )
+            )
+            if req.return_logprob:
+                self.logprob_result_processor.add_input_logprob_return_values(
+                    i,
+                    req,
+                    logits_output,
+                    logprob_pt,
+                    num_input_logprobs,
+                    last_prefill_chunk=False,
+                )
+            logprob_pt += num_input_logprobs
+        return logprob_pt
+
+    def _resolve_spec_v2_tokens(
+        self,
+        result: GenerationBatchResult,
+        batch: ScheduleBatch,
+    ) -> List[List[int]]:
+        """Resolve the padded next token ids for spec-v2 (overlap and non-overlap)."""
+        assert result.next_token_ids.is_cpu
+        assert result.accept_lens.is_cpu
+
+        next_token_ids = result.next_token_ids.tolist()
+        accept_lens = result.accept_lens.tolist()
+        result.num_correct_drafts = sum(accept_lens) - len(batch.reqs)
+        result.num_correct_drafts_per_req_cpu = [x - 1 for x in accept_lens]
+
+        block_accept_lens = (
+            result.block_accept_lens.tolist()
+            if result.block_accept_lens is not None
+            else None
+        )
+        result.num_block_accept_tokens = (
+            sum(block_accept_lens) if block_accept_lens else 0
+        )
+        cap_lens = result.cap_lens.tolist() if result.cap_lens is not None else None
+        result.num_cap_tokens = sum(cap_lens) if cap_lens else 0
+
+        # Feed the adaptive controller now that accept_lens is on CPU,
+        # instead of doing a synchronous GPU→CPU copy in the worker hot path.
+        # BaseSpecWorker provides a no-op default for non-adaptive workers.
+        self.model_worker.on_verify_complete_cpu(
+            result.num_correct_drafts_per_req_cpu, batch_size=len(batch.reqs)
+        )
+
+        # Advance the grammar FSM over this batch's committed tokens (idempotent):
+        # the EAGLE overlap path already did this inside verify() via the grammar
+        # barrier; otherwise advance now. advance_grammar_fsm self-gates on per-req
+        # grammar (the queued batch.copy() does not carry has_grammar) and consumes
+        # result.grammar_retained_tokens below instead of re-advancing.
+        self.advance_grammar_fsm(result, batch)
+
+        predict_tokens = []
+        # In adaptive spec-v2, the worker state may already have switched when this
+        # delayed result is processed. Use the draft token count recorded on result.
+        stride = result.speculative_num_draft_tokens
+        assert stride is not None, "spec-v2 result missing speculative_num_draft_tokens"
+
+        for i, req in enumerate(batch.reqs):
+            accept_tokens = next_token_ids[i * stride : i * stride + accept_lens[i]]
+
+            if req.is_retracted or req.finished():
+                # Nothing to settle: no worker pre-claims the bonus, so
+                # kv_committed_len already holds the committed prefix.
+                pass
+            else:
+                if req.grammar is not None:
+                    # FSM already advanced + truncated by advance_grammar_fsm; reuse
+                    # the retained (grammar-legal) run instead of advancing again.
+                    accept_tokens = result.grammar_retained_tokens[i]
+
+                # Commit the full accepted run (drafts + bonus).
+                num_accept_tokens = len(accept_tokens)
+                req.kv.kv_committed_len += num_accept_tokens
+                req.spec_verify_ct += 1
+
+                num_correct_drafts = result.num_correct_drafts_per_req_cpu[i]
+                req.spec_num_correct_drafts += num_correct_drafts
+                req.update_spec_correct_drafts_histogram(num_correct_drafts)
+
+                if block_accept_lens is not None:
+                    req.spec_num_block_accept_tokens += block_accept_lens[i]
+                if cap_lens is not None:
+                    req.spec_num_cap_tokens += cap_lens[i]
+                    req.update_spec_cap_lens_histogram(cap_lens[i])
+
+            predict_tokens.append(accept_tokens)
+
+        return predict_tokens
+
+    def _accept_grammar_tokens(
+        self, req: Req, tokens: Union[int, List[int]]
+    ) -> List[int]:
+        """Advance the grammar over the accepted token(s), stopping at the token
+        that terminates it.
+
+        ``tokens`` is a single sampled token (normal decode) or the whole
+        verified run (spec decode). Returns the retained prefix; for spec the
+        suffix past grammar completion is dropped so it is never committed to KV
+        nor emitted. Advances the grammar FSM only -- ``grammar.finished`` is
+        synced by the caller once the finish state is updated.
+        """
+        if isinstance(tokens, int):
+            tokens = [tokens]
+        retained = []
+        try:
+            for token_id in tokens:
+                req.grammar.accept_token(token_id)
+                retained.append(token_id)
+                if req.grammar.is_terminated():
+                    break
+        except ValueError as e:
+            # accept_token raises ValueError if the token is not in the grammar
+            # (misconfigured grammar or invalid token); abort the request.
+            logger.error(
+                f"Grammar accept_token failed for req {req.rid} with token "
+                f"{tokens}: {e}"
+            )
+            req.to_finish = FINISH_ABORT()
+        return retained
+
+    def advance_grammar_fsm(
+        self, result: GenerationBatchResult, batch: ScheduleBatch
+    ) -> None:
+        """Advance each req's grammar FSM over the tokens THIS batch committed, and
+        (for decode) memoize the grammar-truncated run on ``result``.
+
+        This is the single place the spec-v2 FSM advances. It is idempotent
+        (``result.grammar_advanced``) so it runs either eagerly — inside ``verify(N)``
+        via the scheduler's grammar barrier, so the advance overlaps the target-verify
+        forward — or lazily from the result processors on the non-overlap / non-EAGLE
+        paths. It handles both decode (the accepted spec run) and extend (the single
+        prefilled token) results, so the barrier can resolve whatever the previous
+        batch was — e.g. the extend->decode boundary.
+        """
+        if result.grammar_advanced or not batch.has_grammar:
+            return
+        is_decode = batch.forward_mode.is_decode()
+        if not (is_decode or batch.forward_mode.is_extend()):
+            return
+        if result.copy_done is not None:
+            result.copy_done.synchronize()
+        next_token_ids = result.next_token_ids.tolist()
+
+        if not is_decode:
+            # Extend: advance over the single token each completed-prefill req emitted
+            # (mirrors process_batch_result_prefill's per-req token indexing).
+            for i, req in enumerate(batch.reqs):
+                if (
+                    req.grammar is None
+                    or req.is_retracted
+                    or req.finished()
+                    or req.inflight_middle_chunks > 0
+                ):
+                    continue
+                self._accept_grammar_tokens(req, next_token_ids[i])
+            result.grammar_advanced = True
+            return
+
+        # Decode: only the spec-v2 path reaches here (the grammar barrier for
+        # spec-overlap workers and _resolve_spec_v2_tokens). Non-spec grammar decode
+        # advances its FSM in process_batch_result_decode and has no accept_lens, so
+        # bail out defensively.
+        if result.accept_lens is None:
+            return
+        accept_lens = result.accept_lens.tolist()
+        stride = result.speculative_num_draft_tokens
+        assert stride is not None, "spec-v2 result missing speculative_num_draft_tokens"
+        retained = [None] * len(batch.reqs)
+        for i, req in enumerate(batch.reqs):
+            if req.grammar is None or req.is_retracted or req.finished():
+                continue
+            accept_tokens = next_token_ids[i * stride : i * stride + accept_lens[i]]
+            # Stop accepting once the grammar terminates so the over-drafted suffix
+            # is never committed to KV nor emitted; this advances the FSM.
+            retained[i] = self._accept_grammar_tokens(req, accept_tokens)
+        result.grammar_retained_tokens = retained
+        result.grammar_advanced = True
+
+    def process_batch_result_idle(
+        self,
+        batch: ScheduleBatch,
+        result: GenerationBatchResult,
+    ):
+        if result.copy_done is not None:
+            result.copy_done.synchronize()
+
+        self.output_streamer._stream_output_generation(
+            batch.reqs, batch.return_logprob, is_idle_batch=True
+        )
+
+    def process_batch_result_decode(
+        self,
+        batch: ScheduleBatch,
+        result: GenerationBatchResult,
+    ):
+        if result.copy_done is not None:
+            result.copy_done.synchronize()
+        auxiliary_output_starts = self.snapshot_auxiliary_output_starts(batch, result)
+        auxiliary_output = result.auxiliary_host_output
+        if result.routed_experts_output is not None:
+            result.routed_experts_output.finalize()
+            result.routed_experts_output = None
+        if result.indexer_topk_output is not None:
+            result.indexer_topk_output.finalize()
+            result.indexer_topk_output = None
+
+        logits_output, next_token_ids, can_run_cuda_graph = (
+            result.logits_output,
+            result.next_token_ids,
+            result.can_run_cuda_graph,
+        )
+
+        next_token_ids, next_token_logprobs = self._normalize_decode_outputs(
+            batch=batch,
+            result=result,
+            logits_output=logits_output,
+            next_token_ids=next_token_ids,
+        )
+
+        self.metrics_reporter.num_generated_tokens += len(batch.reqs)
+        if not batch.spec_algorithm.is_none():
+            self.metrics_reporter.update_spec_metrics(
+                batch.batch_size(),
+                result.num_correct_drafts,
+                num_block_accept_tokens=result.num_block_accept_tokens,
+                num_cap_tokens=result.num_cap_tokens,
+            )
+        if get_observability().enable_metrics:
+            self.metrics_collector.increment_decode_cuda_graph_pass(
+                value=can_run_cuda_graph
+            )
+
+        self.token_to_kv_pool_allocator.free_group_begin()
+
+        # Folds the relay point's selection into the DAG and sets the finish
+        # states the loop below observes. Beam + spec is rejected at admission.
+        newly_finished_beam_groups = set()
+        if batch.spec_algorithm.is_none() and logits_output is not None:
+            newly_finished_beam_groups = self.beam_coordinator.commit_decode(batch)
+
+        for i, req in enumerate(batch.reqs):
+            req: Req
+
+            if req.beam_group is not None:
+                # Under overlap a finished row reappears for one overshoot tick;
+                # gate on the committing tick so this runs exactly once.
+                if req.finished() and (
+                    id(req.beam_group) not in newly_finished_beam_groups
+                ):
+                    continue
+                req.time_stats.set_last_decode_finish_time()
+                self._handle_finish_state_updated_req(
+                    req, batch, result, i, logits_output
+                )
+                continue
+
+            if (self.enable_overlap or self.enable_overlap_mlx) and (
+                req.finished() or req.is_retracted
+            ):
+                # NOTE: This (req.finished() or req.is_retracted) should only happen when overlap scheduling is enabled.
+                # And all the over-allocated tokens will be freed in `release_kv_cache`.
+                continue
+
+            # next_token_id is a per-req list: 1 token for non-spec, the verified
+            # run for spec (already grammar-truncated in _resolve_spec_v2_tokens).
+            next_token_id = next_token_ids[i]
+            is_spec = not batch.spec_algorithm.is_none()
+
+            req.output_ids.extend(next_token_id)
+            new_accept_len = len(next_token_id)
+
+            self._maybe_update_reasoning_tokens(req, next_token_id)
+            req.time_stats.set_last_decode_finish_time()
+            req.update_finish_state(new_accept_len)
+
+            self._handle_finish_state_updated_req(req, batch, result, i, logits_output)
+
+            if req.return_logprob:
+                self._apply_decode_logprobs(
+                    req=req,
+                    i=i,
+                    batch=batch,
+                    next_token_id=next_token_id,
+                    next_token_logprobs=next_token_logprobs,
+                    logits_output=logits_output,
+                )
+
+            if req.return_sampling_mask:
+                # return_sampling_mask + speculative decoding is rejected at
+                # request entry, so this remains one support mask per token.
+                self.add_sampling_mask_return_values(i, req, logits_output)
+
+            if req.return_hidden_states and logits_output.hidden_states is not None:
+                # hidden_states is [bs * stride, hidden_dim], one row per emitted
+                # token; stride = speculative_num_draft_tokens for spec, 1 for non-spec.
+                stride = result.speculative_num_draft_tokens or 1
+                accept_len = len(next_token_id)
+                start = i * stride
+                self._append_decode_hidden_states(
+                    req=req,
+                    hidden_states=logits_output.hidden_states,
+                    start=start,
+                    accept_len=accept_len,
+                )
+
+            if req.grammar is not None:
+                if not is_spec:
+                    # Normal decode advances the grammar for its single token
+                    # here; spec already advanced it in _resolve_spec_v2_tokens.
+                    self._accept_grammar_tokens(req, next_token_id)
+                req.grammar.finished = req.finished()
+
+        if auxiliary_output is not None:
+            self.consume_auxiliary_output(
+                batch,
+                auxiliary_output,
+                auxiliary_output_starts,
+            )
+
+        self.output_streamer.stream_output(batch.reqs, batch.return_logprob)
+        self.token_to_kv_pool_allocator.free_group_end()
+
+        self.metrics_reporter.forward_ct_decode = (
+            self.metrics_reporter.forward_ct_decode + 1
+        ) % (1 << 30)
+        self.metrics_reporter.report_decode_stats(
+            can_run_cuda_graph,
+            running_batch=batch,
+            num_correct_drafts=result.num_correct_drafts,
+        )
+
+    def _normalize_decode_outputs(
+        self,
+        *,
+        batch: ScheduleBatch,
+        result: GenerationBatchResult,
+        logits_output: LogitsProcessorOutput,
+        next_token_ids: Union[torch.Tensor, List[int]],
+    ) -> Tuple[Union[List[int], List[List[int]]], Optional[List[float]]]:
+        next_token_logprobs = None
+        # Normalize to a uniform per-req list of accepted tokens (List[List[int]]):
+        # spec unpacks the padded verify output; non-spec wraps its single token.
+        if not batch.spec_algorithm.is_none():
+            next_token_ids = self._resolve_spec_v2_tokens(result, batch)
+        else:
+            # CUDA workers return a device tensor, MLX a host list[int]; both -> list.
+            ids = (
+                next_token_ids.tolist()
+                if torch.is_tensor(next_token_ids)
+                else next_token_ids
+            )
+            next_token_ids = [[t] for t in ids]
+
+        if batch.return_logprob:
+            next_token_logprobs = logits_output.next_token_logprobs.tolist()
+            if logits_output.next_token_top_logprobs_val:
+                logits_output.next_token_top_logprobs_val = [
+                    v.tolist() for v in logits_output.next_token_top_logprobs_val
+                ]
+                logits_output.next_token_top_logprobs_idx = [
+                    x.tolist() for x in logits_output.next_token_top_logprobs_idx
+                ]
+
+            if logits_output.next_token_token_ids_logprobs_val:
+                logits_output.next_token_token_ids_logprobs_val = [
+                    v.tolist() if isinstance(v, torch.Tensor) else v
+                    for v in logits_output.next_token_token_ids_logprobs_val
+                ]
+        return next_token_ids, next_token_logprobs
+
+    def _apply_decode_logprobs(
+        self,
+        *,
+        req: Req,
+        i: int,
+        batch: ScheduleBatch,
+        next_token_id: Union[int, List[int]],
+        next_token_logprobs: list,
+        logits_output: LogitsProcessorOutput,
+    ) -> None:
+        # accepted_ids is already a per-req list; non-spec logprobs are flat, so
+        # the scalar logprob still needs wrapping.
+        if not batch.spec_algorithm.is_none():
+            accepted_logprobs = next_token_logprobs[i]
+            accepted_ids = next_token_id
+            max_accept = len(accepted_logprobs)
+        else:
+            accepted_logprobs = [next_token_logprobs[i]]
+            accepted_ids = next_token_id
+            max_accept = 1
+
+        for j, tok_id in enumerate(accepted_ids):
+            req.logprob.output_token_logprobs_val.append(accepted_logprobs[j])
+            req.logprob.output_token_logprobs_idx.append(tok_id)
+            if req.logprob.top_logprobs_num > 0:
+                flat_idx = i * max_accept + j
+                req.logprob.output_top_logprobs_val.append(
+                    logits_output.next_token_top_logprobs_val[flat_idx]
+                )
+                req.logprob.output_top_logprobs_idx.append(
+                    logits_output.next_token_top_logprobs_idx[flat_idx]
+                )
+            if req.logprob.token_ids_logprob is not None:
+                flat_idx = i * max_accept + j
+                req.logprob.output_token_ids_logprobs_val.append(
+                    logits_output.next_token_token_ids_logprobs_val[flat_idx]
+                )
+                req.logprob.output_token_ids_logprobs_idx.append(
+                    logits_output.next_token_token_ids_logprobs_idx[flat_idx]
+                )
+
+    def add_sampling_mask_return_values(
+        self,
+        i: int,
+        req: Req,
+        output: LogitsProcessorOutput,
+    ) -> None:
+        """Attach sparse sampling support metadata to the return values."""
+        mask = output.next_token_sampling_mask_idx
+        logprobs = output.next_token_sampling_logprobs
+        req.output_token_sampling_mask.append(None if mask is None else mask[i])
+        req.output_token_sampling_logprobs.append(
+            None if logprobs is None else logprobs[i]
+        )
+
+    def _handle_finish_state_updated_req(
+        self,
+        req: Req,
+        batch: ScheduleBatch,
+        result: GenerationBatchResult,
+        i: int,
+        logits_output: LogitsProcessorOutput,
+    ):
+        lazy = mamba_extra_buffer_lazy_enabled()
+        known_mamba_boundary = None
+        completed_mamba_boundary = None
+        lookahead = 0
+        if batch.mamba_track_mask_cpu is not None:
+            completed_mamba_boundary = bool(batch.mamba_track_mask_cpu[i])
+            lookahead = req.decode_batch_idx - batch.mamba_decode_batch_idx_cpu[i]
+            assert lookahead in (0, 1), (
+                f"mamba result lookahead={lookahead} for req {req.rid}; "
+                "overlap advanced more than one decode batch"
+            )
+            if lookahead == 0:
+                known_mamba_boundary = bool(batch.mamba_track_mask_cpu[i])
+            else:
+                known_mamba_boundary = bool(batch.mamba_track_mask_next_cpu[i])
+
+            if completed_mamba_boundary and not lazy:
+                req.kv.mamba_last_track_idx = batch.mamba_track_buffer_indices[i]
+                req.kv.mamba_last_track_seqlen = req.kv.kv_committed_len - lookahead
+            elif (
+                req.finished()
+                and lazy
+                and lookahead == 1
+                and known_mamba_boundary
+                and req.kv.mamba_next_track_idx == req.kv.mamba_last_track_idx
+            ):
+                req.mamba_lazy_is_insert = False
+
+        # Called here (after update_finish_state) so req.finished() is valid
+        # for mamba_lazy_post_decode_at_boundary inside.
+        should_update = completed_mamba_boundary if lazy else known_mamba_boundary
+        if should_update is None or should_update:
+            self._mamba_prefix_cache_update(
+                req,
+                batch,
+                result,
+                i,
+                known_boundary=not lazy and known_mamba_boundary is True,
+            )
+
+        if (
+            get_disagg().disaggregation_decode_enable_offload_kvcache
+            and not req.finished()
+        ):
+            self.decode_offload_manager.offload_kv_cache(req)
+
+        if req.finished():
+            # isinstance narrowing: create_worker may also return plain
+            # TpModelWorker-based drafts, which carry no spec-worker hooks.
+            if isinstance(self.draft_worker, BaseSpecWorker):
+                self.draft_worker.note_request_finished(
+                    rid=req.rid,
+                    natural_stop=isinstance(req.finished_reason, FINISH_MATCHED_TOKEN),
+                )
+
+            # delete feature to save memory
+            if req.multimodal_inputs is not None and req.session is None:
+                req.multimodal_inputs.release_features()
+            self._maybe_collect_routed_experts(req)
+            self._maybe_collect_indexer_topk(req)
+
+            if get_disagg().disaggregation_decode_enable_offload_kvcache:
+                # Asynchronously offload KV cache; release_kv_cache will be called after Device->Host transfer completes
+                if not self.decode_offload_manager.offload_kv_cache(req):
+                    self.decode_offload_manager.finalize_release_on_finish(req)
+            else:
+                if get_memory().enable_hisparse:
+                    self.hisparse_coordinator.request_finished(req)
+                prepare_release = getattr(
+                    self.model_worker, "prepare_for_kv_cache_release", None
+                )
+                if callable(prepare_release):
+                    prepare_release(req)
+                is_insert = (
+                    req.mamba_lazy_is_insert
+                    if mamba_extra_buffer_lazy_enabled()
+                    else True
+                )
+                release_kv_cache(req, self.tree_cache, is_insert=is_insert)
+
+            req.time_stats.set_completion_time()
+
+        self._maybe_collect_customized_info(i, req, logits_output)
+
+    def _maybe_update_reasoning_tokens(
+        self,
+        req: Req,
+        next_token_id: Union[int, List[int]],
+    ):
+        think_end_ids = self.model_config.think_end_ids
+        if req.require_reasoning and think_end_ids:
+            req.update_reasoning_tokens(next_token_id, think_end_ids)
+
+    def _mamba_prefix_cache_update(
+        self,
+        req: Req,
+        batch: ScheduleBatch,
+        result: GenerationBatchResult,
+        i: int,
+        known_boundary: bool = False,
+    ) -> None:
+        """Update mamba track state at ping-pong boundaries.
+
+        Non-lazy: swap the ping-pong index so the next forward writes to
+        the alternate slot.
+        Lazy: keep the same index (prealloc handles the swap) and run
+        post-decode cleanup to free the temporary second slot.
+        """
+        if req.kv.mamba_ping_pong_track_buffer is None:
+            return
+
+        lazy = mamba_extra_buffer_lazy_enabled()
+        if known_boundary:
+            self._mamba_assert_committed_len_lookahead(req)
+            track_seqlen = req.kv.kv_committed_len
+            assert track_seqlen % mamba_track_grid(self.tree_cache.page_size) == 0
+            at_boundary = True
+        else:
+            at_boundary, track_seqlen = self._mamba_check_track_boundary(
+                req, batch, result, i
+            )
+
+        if lazy and not batch.spec_algorithm.is_none():
+            # For spec, at_boundary means this step actually crossed an interval.
+            self._mamba_lazy_spec_update(req, batch, i, at_boundary, track_seqlen)
+            return
+
+        if not at_boundary:
+            return
+
+        track_idx = req.kv.mamba_next_track_idx
+        if not known_boundary and batch.mamba_track_buffer_indices is not None:
+            track_idx = batch.mamba_track_buffer_indices[i]
+        if not known_boundary:
+            req.kv.mamba_last_track_seqlen = track_seqlen
+        if lazy:
+            self.mamba_lazy_post_decode_at_boundary(req, batch, track_idx)
+        else:
+            if not known_boundary:
+                req.kv.mamba_last_track_idx = track_idx
+            req.kv.mamba_next_track_idx = (
+                batch.req_to_token_pool.get_mamba_ping_pong_other_idx(track_idx)
+            )
+
+    def _mamba_lazy_spec_update(
+        self, req: Req, batch: ScheduleBatch, i: int, crossed: bool, track_seqlen: int
+    ) -> None:
+        """Lazy + spec post-processing.
+
+        Running req with a confirmed crossing: the commit scattered the
+        crossing state into the planned position; promote it to keep unless
+        the plan already pointed at the keep slot (in-place fallback, or
+        promoted by an earlier overlapped confirmation).
+        Finishing req: donate the keep slot only if no scatter wrote it
+        (this step) or may still write it (the in-flight next step).
+        """
+        positions = batch.mamba_lazy_spec_track_positions_cpu
+        planned_pos = (
+            positions[i]
+            if positions is not None and i < len(positions)
+            else None  # filtered/merged snapshot without a plan: be conservative
+        )
+
+        if req.finished():
+            # Skip the donation if a scatter wrote or may still write the keep slot.
+            keep_written_by_this_step = (
+                crossed and planned_pos == req.kv.mamba_next_track_idx
+            )
+            other_idx = 1 - req.kv.mamba_next_track_idx
+            # Recompute the in-flight verify's plan (kv_committed_len is
+            # frozen since its prepare, so the recompute is exact).
+            keep_may_be_written_in_flight = req.kv.mamba_ping_pong_track_buffer[
+                other_idx
+            ].item() == -1 and mamba_lazy_spec_in_window(
+                req,
+                mamba_track_grid(self.tree_cache.page_size),
+                max_speculative_num_draft_tokens(),
+            )
+            if (
+                planned_pos is None
+                or keep_written_by_this_step
+                or keep_may_be_written_in_flight
+            ):
+                req.mamba_lazy_is_insert = False
+            return
+
+        if not crossed or planned_pos is None:
+            return
+        if planned_pos != req.kv.mamba_next_track_idx:
+            # Promote pending -> keep: free the old checkpoint, repoint.
+            pool = batch.req_to_token_pool
+            keep_idx = req.kv.mamba_next_track_idx
+            keep_val = req.kv.mamba_ping_pong_track_buffer[keep_idx]
+            pool.mamba_allocator.free(keep_val.unsqueeze(0))
+            pool.set_mamba_ping_pong_slot(req, keep_idx, -1)
+            req.kv.mamba_next_track_idx = planned_pos
+        # else: in-place fallback, or promoted by an earlier confirmation —
+        # keep holds the track_seqlen state either way.
+        req.kv.mamba_last_track_idx = planned_pos
+        req.kv.mamba_last_track_seqlen = track_seqlen
+
+    @staticmethod
+    def _mamba_assert_committed_len_lookahead(req: Req) -> None:
+        """Alarm if overlap advances beyond the scheduler's modeled window."""
+        assert req.output_ids, (
+            "mamba track boundary reached before a decode token was appended "
+            f"(req {req.rid}); output_ids is empty"
+        )
+        token_seq_len = len(req.origin_input_ids) + len(req.output_ids) - 1
+        assert (req.kv.kv_committed_len - token_seq_len) in (0, 1), (
+            f"mamba track boundary: kv_committed_len={req.kv.kv_committed_len} "
+            f"leads seq_len={token_seq_len} by more than one (req {req.rid}); "
+            "overlap lookahead wider than assumed"
+        )
+
+    def _mamba_check_track_boundary(self, req, batch, result, i):
+        """Check if this decode step crosses a mamba track interval boundary.
+
+        Returns (at_boundary, track_seqlen).  The boundary condition
+        matches what the forward's tracking mask used:
+        ``prepare_for_decode`` increments both ``seq_lens_cpu`` and
+        ``kv_committed_len`` by 1, then checks
+        ``seq_lens_cpu % interval == 0``.  Subtracting the overlap
+        lookahead from ``kv_committed_len`` reproduces that check.
+
+        For spec decode, the boundary is detected by comparing the
+        accepted seq_len range against interval boundaries.
+        """
+        interval = mamba_track_grid(self.tree_cache.page_size)
+
+        if batch.spec_algorithm.is_none():
+            lookahead = req.decode_batch_idx - batch.mamba_decode_batch_idx_cpu[i]
+            committed_len = req.kv.kv_committed_len - lookahead
+            if committed_len % interval == 0:
+                return True, committed_len
+        elif result.num_correct_drafts_per_req_cpu is not None:
+            cur = req.seqlen - 1
+            prev = cur - result.num_correct_drafts_per_req_cpu[i] - 1
+            if cur // interval != prev // interval:
+                return True, cur // interval * interval
+
+        return False, 0
+
+    def mamba_lazy_post_decode_at_boundary(
+        self, req: Req, batch: ScheduleBatch, track_idx: int
+    ):
+        """Commit a completed lazy-mode boundary and free its old slot."""
+        req.kv.mamba_last_track_idx = track_idx
+        req.kv.mamba_next_track_idx = track_idx
+        other_idx = 1 - track_idx
+        other_val = req.kv.mamba_ping_pong_track_buffer[other_idx].item()
+        if other_val != -1:
+            pool = batch.req_to_token_pool
+            pool.mamba_allocator.free(
+                req.kv.mamba_ping_pong_track_buffer[other_idx].unsqueeze(0)
+            )
+            pool.set_mamba_ping_pong_slot(req, other_idx, -1)

runtime/patches/sglang-logits-processor.py

@@ -0,0 +1,1272 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Logits processing."""
+
+import dataclasses
+import logging
+from contextlib import contextmanager
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from sglang.kernels.ops.activation.softcap import (
+    softcap_inplace_logits as fused_softcap,
+)
+from sglang.srt.beam_search.logits_capture import BeamLogitsCapture
+from sglang.srt.distributed import get_tp_group
+from sglang.srt.distributed.device_communicators import triton_symm_mem_ag
+from sglang.srt.environ import envs
+from sglang.srt.layers import layernorm_sp
+from sglang.srt.layers.aux_hidden_states import (
+    AuxHiddenStates,
+    pack_aux_hidden_states,
+)
+from sglang.srt.layers.dp_attention import (
+    DpPaddingMode,
+    attn_tp_all_gather,
+    attn_tp_all_gather_into_tensor,
+    dp_gather_replicate,
+    dp_scatter,
+    get_dp_device,
+    get_dp_dtype,
+    get_dp_hidden_size,
+)
+from sglang.srt.layers.logprob_processor import (
+    InputLogprobProcessor,
+    LogprobStage,
+    get_token_ids_logprobs_raw,
+    get_top_logprobs_raw,
+)
+from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+)
+from sglang.srt.runtime_context import get_exec, get_parallel
+from sglang.srt.sampling.sampling_observer import DeviceAuxiliaryOutput
+from sglang.srt.utils.common import (
+    is_cpu,
+    is_npu,
+    is_pin_memory_available,
+    use_intel_amx_backend,
+)
+
+logger = logging.getLogger(__name__)
+
+_is_npu = is_npu()
+_is_cpu = is_cpu()
+
+_UNQUANTIZED_LM_HEAD_METHODS = {
+    "UnquantizedEmbeddingMethod",
+    "UnquantizedLinearMethod",
+    "PackWeightMethod",
+}
+
+# None outside a FlashInfer autotune pass; inside one, whether that pass runs the
+# LM head. Not-None means the forward's output is discarded -- attention backends
+# read that via get_in_autotune_dummy_run() to skip a cross-node exchange.
+# Skipping the LM head skips its [batch * dp_size, vocab] all-gather, which OOMs
+# under DP attention with a tight mem_fraction_static.
+_autotune_run_lm_head: Optional[bool] = None
+
+
+def _trace_e2e_logits(stage: str, **fields) -> None:
+    if not envs.SGLANG_TRACE_LOGITS_E2E.get():
+        return
+    try:
+        parallel = get_parallel()
+        rank = f"dp={parallel.attn_dp_rank} " f"tp={parallel.tp_rank}"
+    except Exception:
+        rank = "rank=unknown"
+    details = " ".join(f"{key}={value}" for key, value in fields.items())
+    print(f"SGLANG_TRACE_LOGITS_E2E {rank} stage={stage} {details}", flush=True)
+
+
+def _has_lm_head_runtime_attrs(lm_head, attr_names: Tuple[str, ...]) -> bool:
+    return all(hasattr(lm_head, attr_name) for attr_name in attr_names)
+
+
+def should_apply_lm_head_quant_method(lm_head, quant_method) -> bool:
+    if (
+        quant_method is None
+        or not hasattr(lm_head, "weight")
+        or not callable(getattr(quant_method, "apply", None))
+    ):
+        return False
+
+    method_name = type(quant_method).__name__
+    if method_name in _UNQUANTIZED_LM_HEAD_METHODS:
+        return False
+
+    # A shared target lm_head can retain the draft's stale ModelOpt method; use it
+    # only when the runtime tensor layout matches that method.
+    if method_name == "ModelOptFp4LinearMethod":
+        if lm_head.weight.dtype == torch.int32 and _has_lm_head_runtime_attrs(
+            lm_head,
+            (
+                "weight_scale",
+                "weight_global_scale",
+                "workspace",
+                "input_size_per_partition",
+                "output_size_per_partition",
+            ),
+        ):
+            return True
+        return lm_head.weight.dtype == torch.uint8 and _has_lm_head_runtime_attrs(
+            lm_head,
+            (
+                "weight_scale_interleaved",
+                "alpha",
+                "input_scale_inv",
+                "input_size_per_partition",
+                "output_size_per_partition",
+            ),
+        )
+    if method_name == "ModelOptNvFp4A16LinearMethod":
+        return lm_head.weight.dtype == torch.int32 and _has_lm_head_runtime_attrs(
+            lm_head,
+            (
+                "weight_scale",
+                "weight_global_scale",
+                "workspace",
+                "input_size_per_partition",
+                "output_size_per_partition",
+            ),
+        )
+    if method_name == "ModelOptFp8LinearMethod":
+        return (
+            lm_head.weight.dtype == torch.float8_e4m3fn
+            and _has_lm_head_runtime_attrs(lm_head, ("weight_scale", "input_scale"))
+        )
+
+    return True
+
+
+# FlashInfer autotune skips the unprofiled LM-head all-gather; its
+# [batch * dp_size, vocab] output can OOM under tight DP-attention memory.
+_in_autotune_dummy_run = False
+
+
+def get_in_autotune_dummy_run() -> bool:
+    return _autotune_run_lm_head is not None
+
+
+@contextmanager
+def autotune_dummy_run_mode(*, run_lm_head: bool):
+    global _autotune_run_lm_head
+    _autotune_run_lm_head = run_lm_head
+    try:
+        yield
+    finally:
+        _autotune_run_lm_head = None
+
+
+@dataclasses.dataclass
+class LogitsProcessorOutput:
+    ## Part 1: This part will be assigned in python/sglang/srt/layers/logits_processor.py::LogitsProcessor
+    # The logits of the next tokens.       shape: [#seq, vocab_size]
+    # Can be None for certain prefill-only requests (e.g., multi-item scoring) that don't need next token generation
+    next_token_logits: Optional[torch.Tensor]
+    # Used by speculative decoding (EAGLE)
+    # The last hidden layers
+    hidden_states: Optional[torch.Tensor] = None
+
+    ## Part 2: This part will be assigned in python/sglang/srt/layers/sampler.py::Sampler
+    # he log probs of output tokens, if SGLANG_RETURN_ORIGINAL_LOGPROB = True, will get the log probs before applying temperature. If False, will get the log probs before applying temperature.
+    next_token_logprobs: Optional[torch.Tensor] = None
+    # The logprobs and ids of the top-k tokens in output positions. shape: [#seq, k]
+    next_token_top_logprobs_val: Optional[List] = None
+    next_token_top_logprobs_idx: Optional[List] = None
+    # The logprobs and ids of the requested token ids in output positions. shape: [#seq, n] (n is the number of requested token ids)
+    # Can contain either lists or GPU tensors (for delayed copy optimization in prefill-only requests)
+    next_token_token_ids_logprobs_val: Optional[
+        List[Union[List[float], torch.Tensor]]
+    ] = None
+    next_token_token_ids_logprobs_idx: Optional[List] = None
+    # Sparse top-k/top-p/min-p support ids and selected-token logprob after
+    # truncation/renormalization. Only populated when requested.
+    next_token_sampling_mask_idx: Optional[List[Optional[List[int]]]] = None
+    next_token_sampling_logprobs: Optional[List[Optional[float]]] = None
+
+    ## Part 3: Prefill-only. This part will be assigned in python/sglang/srt/layers/logits_processor.py::LogitsProcessor
+    # The logprobs of input tokens.        shape: [#token]
+    input_token_logprobs: Optional[torch.Tensor] = None
+    # The logprobs and ids of the top-k tokens in input positions.  shape: [#seq, #token, k]
+    input_top_logprobs_val: Optional[List] = None
+    input_top_logprobs_idx: Optional[List] = None
+    # The logprobs and ids of the requested token ids in input positions. shape: [#seq, n] (n is the number of requested token ids)
+    # Can contain either lists or GPU tensors (for delayed GPU-to-CPU transfer optimization)
+    input_token_ids_logprobs_val: Optional[List[Union[List[float], torch.Tensor]]] = (
+        None
+    )
+    input_token_ids_logprobs_idx: Optional[List] = None
+
+    ## Part 4: Diffusion LLM only.
+    full_logits: Optional[torch.Tensor] = None
+
+    # Beam search only: raw pre-sample logits for the scheduler-side joint
+    # selection; see beam_search.logits_capture.
+    beam: Optional[BeamLogitsCapture] = None
+
+    ## Part 5: Customized Info
+    customized_info: Optional[Dict[str, List[Any]]] = None
+
+    ## Part 6: Temporary variables
+    # FIXME: These fields are not logits-related but are passed through here as a
+    # workaround since ForwardBatch is local to forward_batch_generation().
+    # They should be moved to GenerationBatchResult to keep this class clean.
+    mm_input_embeds: Optional[torch.Tensor] = None
+
+    # Scheduler-local output copied alongside the ordinary generation result.
+    auxiliary_device_output: Optional[DeviceAuxiliaryOutput] = None
+
+
+@dataclasses.dataclass
+class LogitsMetadata:
+    forward_mode: ForwardMode
+    capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.NULL
+    next_token_logits_buffer: Optional[torch.Tensor] = None
+
+    extend_return_logprob: bool = False
+    extend_return_top_logprob: bool = False
+    extend_token_ids_logprob: bool = False
+    extend_seq_lens: Optional[torch.Tensor] = None
+    extend_seq_lens_cpu: Optional[List[int]] = None
+    extend_logprob_start_lens_cpu: Optional[List[int]] = None
+    extend_logprob_pruned_lens_cpu: Optional[List[int]] = None
+    top_logprobs_nums: Optional[List[int]] = None
+    extend_input_logprob_token_ids_gpu: Optional[torch.Tensor] = None
+    token_ids_logprobs: Optional[List[List[int]]] = None
+
+    # logits and logprobs post processing
+    temperature: torch.Tensor = None
+    top_p: torch.Tensor = None
+
+    # DP attention metadata. Not needed when DP attention is not used.
+    # Number of tokens in the request.
+    global_num_tokens_gpu: Optional[torch.Tensor] = None
+    # The start position of local hidden states.
+    dp_local_start_pos: Optional[torch.Tensor] = None
+    dp_local_num_tokens: Optional[torch.Tensor] = None
+    global_dp_buffer_len: Optional[int] = None
+    # Number of tokens to sample per DP rank
+    global_num_tokens_for_logprob_cpu: Optional[torch.Tensor] = None
+    global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] = None
+    # The gather mode for DP attention
+    dp_padding_mode: Optional[DpPaddingMode] = None
+
+    # Whether this batch is prefill-only (no token generation needed)
+    is_prefill_only: bool = False
+
+    mm_input_embeds: Optional[torch.Tensor] = None
+
+    # DRAFT_EXTEND_V2: when set, lm_head runs only on these rows (see
+    # EagleDraftExtendInput.select_index).
+    draft_extend_select_index: Optional[torch.Tensor] = None
+
+    @classmethod
+    def from_forward_batch(cls, forward_batch: ForwardBatch):
+        if (
+            forward_batch.forward_mode.is_extend()
+            and forward_batch.return_logprob
+            and not forward_batch.forward_mode.is_target_verify()
+        ):
+            extend_return_top_logprob = any(
+                x > 0 for x in forward_batch.top_logprobs_nums
+            )
+            extend_token_ids_logprob = any(
+                x is not None for x in forward_batch.token_ids_logprobs
+            )
+            extend_return_logprob = False
+            extend_logprob_pruned_lens_cpu = []
+            for extend_len, start_len in zip(
+                forward_batch.extend_seq_lens_cpu,
+                forward_batch.extend_logprob_start_lens_cpu,
+            ):
+                if extend_len - start_len > 0:
+                    extend_return_logprob = True
+                extend_logprob_pruned_lens_cpu.append(extend_len - start_len)
+        else:
+            extend_return_logprob = extend_return_top_logprob = (
+                extend_token_ids_logprob
+            ) = extend_logprob_pruned_lens_cpu = False
+
+        if forward_batch.forward_mode.is_draft_extend_v2():
+            draft_extend_select_index = forward_batch.spec_info.select_index
+        else:
+            draft_extend_select_index = None
+
+        return cls(
+            forward_mode=forward_batch.forward_mode,
+            capture_hidden_mode=forward_batch.capture_hidden_mode,
+            next_token_logits_buffer=forward_batch.next_token_logits_buffer,
+            extend_return_logprob=extend_return_logprob,
+            extend_return_top_logprob=extend_return_top_logprob,
+            extend_token_ids_logprob=extend_token_ids_logprob,
+            extend_seq_lens=forward_batch.extend_seq_lens,
+            extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
+            extend_logprob_start_lens_cpu=forward_batch.extend_logprob_start_lens_cpu,
+            extend_logprob_pruned_lens_cpu=extend_logprob_pruned_lens_cpu,
+            top_logprobs_nums=forward_batch.top_logprobs_nums,
+            token_ids_logprobs=forward_batch.token_ids_logprobs,
+            extend_input_logprob_token_ids_gpu=forward_batch.extend_input_logprob_token_ids_gpu,
+            is_prefill_only=forward_batch.is_prefill_only,
+            global_num_tokens_gpu=forward_batch.global_num_tokens_gpu,
+            dp_local_start_pos=forward_batch.dp_local_start_pos,
+            dp_local_num_tokens=forward_batch.dp_local_num_tokens,
+            global_dp_buffer_len=forward_batch.global_dp_buffer_len,
+            global_num_tokens_for_logprob_cpu=forward_batch.global_num_tokens_for_logprob_cpu,
+            global_num_tokens_for_logprob_gpu=forward_batch.global_num_tokens_for_logprob_gpu,
+            dp_padding_mode=DpPaddingMode.SUM_LEN,
+            mm_input_embeds=forward_batch.mm_input_embeds,
+            draft_extend_select_index=draft_extend_select_index,
+        )
+
+    def compute_dp_attention_metadata(self):
+        cumtokens = torch.cumsum(self.global_num_tokens_for_logprob_gpu, dim=0)
+        dp_rank = get_parallel().attn_dp_rank
+        if dp_rank == 0:
+            dp_local_start_pos = torch.zeros_like(
+                self.global_num_tokens_for_logprob_gpu[0]
+            )
+        else:
+            dp_local_start_pos = cumtokens[dp_rank - 1]
+
+        self.dp_local_start_pos = dp_local_start_pos
+        self.dp_local_num_tokens = self.global_num_tokens_for_logprob_gpu[dp_rank]
+
+        hidden_size = get_dp_hidden_size()
+        dtype = get_dp_dtype()
+        device = get_dp_device()
+
+        if self.global_num_tokens_for_logprob_cpu is not None:
+            # create a smaller buffer to reduce peak memory usage
+            self.global_dp_buffer_len = sum(self.global_num_tokens_for_logprob_cpu)
+        else:
+            self.global_dp_buffer_len = self.global_dp_buffer_len
+
+        self.gathered_buffer = torch.empty(
+            (
+                self.global_dp_buffer_len,
+                hidden_size,
+            ),
+            dtype=dtype,
+            device=device,
+        )
+
+
+class LogitsProcessor(nn.Module):
+    def __init__(
+        self,
+        config,
+        skip_all_gather: bool = False,
+        logit_scale: Optional[float] = None,
+        return_full_logits: bool = False,
+    ):
+        super().__init__()
+        self.config = config
+        # Decision-only optional projection. Public API renormalizes selected labels.
+        import os, json
+        self.jev_selected_ids = None
+        if os.getenv("JEV_SELECTED_HEAD", "0") == "1":
+            if get_parallel().tp_size != 1:
+                raise ValueError("Selected decision head requires TP1")
+            with open("/jev-label-token-ids.json") as label_file:
+                self.jev_selected_ids = json.load(label_file)
+        self.jev_head_cache = None
+        self.vocab_size = config.vocab_size
+        self.logit_scale = logit_scale
+        self.use_attn_tp_group = get_parallel().enable_dp_lm_head
+        self.use_tp_lm_head_all_to_all = get_parallel().enable_tp_lm_head_all_to_all
+        self.use_fp32_lm_head = get_exec().features.enable_fp32_lm_head
+        if self.use_attn_tp_group:
+            self.attn_tp_size = get_parallel().attn_tp_size
+            self.do_tensor_parallel_all_gather = (
+                not skip_all_gather and self.attn_tp_size > 1
+            )
+            self.do_tensor_parallel_all_gather_dp_attn = False
+        else:
+            self.do_tensor_parallel_all_gather = (
+                not skip_all_gather and get_parallel().tp_size > 1
+            )
+            self.do_tensor_parallel_all_gather_dp_attn = (
+                self.do_tensor_parallel_all_gather and get_parallel().attn_dp_size != 1
+            )
+        self.final_logit_softcapping = getattr(
+            self.config, "final_logit_softcapping", None
+        )
+        if (
+            self.final_logit_softcapping is not None
+            and self.final_logit_softcapping < 0
+        ):
+            self.final_logit_softcapping = None
+
+        self.return_full_logits = return_full_logits
+        self.enable_mis = get_exec().features.enable_mis
+        self.rl_on_policy_target = get_exec().deterministic.rl_on_policy_target
+
+        self._logits_gatherer = triton_symm_mem_ag.MultimemAllGatherer(
+            max_tokens=triton_symm_mem_ag.recommended_max_tokens(
+                include_prefill=False, floor=128
+            ),
+            enabled=self.do_tensor_parallel_all_gather and not self.use_attn_tp_group,
+            skip_entry_sync=True,
+        )
+
+        self.input_logprob_processor = InputLogprobProcessor()
+
+    def forward(
+        self,
+        input_ids,
+        hidden_states,
+        lm_head: VocabParallelEmbedding,
+        logits_metadata: Union[LogitsMetadata, ForwardBatch],
+        aux_hidden_states: Optional[AuxHiddenStates] = None,
+        hidden_states_before_norm: Optional[torch.Tensor] = None,
+    ) -> LogitsProcessorOutput:
+        # Extract MIS indices before ForwardBatch → LogitsMetadata conversion
+        multi_item_delimiter_indices = None
+        if isinstance(logits_metadata, ForwardBatch):
+            multi_item_delimiter_indices = logits_metadata.multi_item_delimiter_indices
+            logits_metadata = LogitsMetadata.from_forward_batch(logits_metadata)
+
+        # Autotune dummy run discards this output. `is False` not `not`: None
+        # means no autotune pass, which must not skip. Placed before the MIS /
+        # DLLM / common dispatch so all three LM-head paths are skipped.
+        if _autotune_run_lm_head is False:
+            return LogitsProcessorOutput(next_token_logits=None)
+
+        # Under LayerNorm SP the decoder loop leaves these sequence-sharded; undo
+        # that before the LM head, which must not participate.
+        hidden_states, hidden_states_before_norm = layernorm_sp.maybe_exit_gather(
+            hidden_states=hidden_states,
+            hidden_states_before_norm=hidden_states_before_norm,
+            input_ids=input_ids,
+            forward_mode=logits_metadata.forward_mode,
+        )
+
+        # Multi-item scoring only for prefill-only requests with pre-computed indices.
+        if multi_item_delimiter_indices is not None and logits_metadata.is_prefill_only:
+            return self.compute_logprobs_for_multi_item_scoring(
+                input_ids,
+                hidden_states,
+                lm_head,
+                logits_metadata,
+                multi_item_delimiter_indices,
+            )
+
+        # Diffusion LLM only.
+        if logits_metadata.forward_mode.is_dllm_extend():
+            return self._get_dllm_logits(hidden_states, lm_head, logits_metadata)
+
+        # Get the last hidden states and last logits for the next token prediction
+        (
+            pruned_states,
+            pruned_states_before_norm,
+            aux_pruned_states,
+            sample_indices,
+            input_logprob_indices,
+            token_to_seq_idx,
+        ) = self._get_pruned_states(
+            hidden_states,
+            hidden_states_before_norm,
+            aux_hidden_states,
+            logits_metadata,
+        )
+
+        hidden_states_to_store = self._get_hidden_states_to_store(
+            hidden_states,
+            hidden_states_before_norm,
+            aux_hidden_states,
+            pruned_states,
+            pruned_states_before_norm,
+            aux_pruned_states,
+            sample_indices,
+            logits_metadata,
+        )
+        del hidden_states
+
+        if not logits_metadata.extend_return_logprob:
+            # Compute logits for both input and sampled tokens.
+            logits = self._get_logits(pruned_states, lm_head, logits_metadata)
+            sampled_logits = (
+                logits[sample_indices] if sample_indices is not None else logits
+            )
+
+            # Decode mode or extend mode without return_logprob.
+            return LogitsProcessorOutput(
+                next_token_logits=sampled_logits,
+                hidden_states=hidden_states_to_store,
+                mm_input_embeds=logits_metadata.mm_input_embeds,
+            )
+
+        logprobs_result, sampled_logits = self.input_logprob_processor.forward(
+            pruned_states=pruned_states,
+            sample_indices=sample_indices,
+            input_logprob_indices=input_logprob_indices,
+            token_to_seq_idx=token_to_seq_idx,
+            lm_head=lm_head,
+            get_logits_fn=self._get_logits,
+            logits_metadata=logits_metadata,
+            skip_chunking_for_dp_attn=self.do_tensor_parallel_all_gather_dp_attn,
+        )
+
+        logits_output = LogitsProcessorOutput(
+            next_token_logits=sampled_logits,
+            hidden_states=hidden_states_to_store,
+            mm_input_embeds=logits_metadata.mm_input_embeds,
+        )
+        logprobs_result.write_input_to(logits_output)
+        return logits_output
+
+    def _get_pruned_states(
+        self,
+        hidden_states: torch.Tensor,
+        hidden_states_before_norm: Optional[torch.Tensor],
+        aux_hidden_states: Optional[AuxHiddenStates],
+        logits_metadata: LogitsMetadata,
+    ):
+        pruned_states_before_norm: Optional[torch.Tensor] = None
+        aux_pruned_states = None
+        token_to_seq_idx = []
+
+        if (
+            logits_metadata.forward_mode.is_decode_or_idle()
+            or logits_metadata.forward_mode.is_target_verify()
+            or logits_metadata.forward_mode.is_draft_extend_v2()
+        ):
+            if logits_metadata.draft_extend_select_index is not None:
+                # Only next_token_logits narrows to [bs, vocab]; the
+                # FULL-capture hidden stays unpruned.
+                pruned_states = hidden_states[logits_metadata.draft_extend_select_index]
+            else:
+                pruned_states = hidden_states
+            pruned_states_before_norm = hidden_states_before_norm
+            if aux_hidden_states is not None:
+                aux_pruned_states = (
+                    aux_hidden_states
+                    if isinstance(aux_hidden_states, torch.Tensor)
+                    else [hidden for hidden in aux_hidden_states]
+                )
+            sample_indices = None
+            input_logprob_indices = None
+
+        elif (
+            logits_metadata.forward_mode.is_extend()
+            and not logits_metadata.extend_return_logprob
+        ):
+            # Prefill without input logprobs.
+            last_index = torch.cumsum(logits_metadata.extend_seq_lens, dim=0) - 1
+            pruned_states = hidden_states[last_index]
+            if hidden_states_before_norm is not None:
+                pruned_states_before_norm = hidden_states_before_norm[last_index]
+            if aux_hidden_states is not None:
+                aux_pruned_states = (
+                    aux_hidden_states[last_index]
+                    if isinstance(aux_hidden_states, torch.Tensor)
+                    else [hidden[last_index] for hidden in aux_hidden_states]
+                )
+            sample_indices = None
+            input_logprob_indices = None
+        else:
+            # Prefill with input logprobs.
+            # Find 4 different indices.
+            # 1. pruned_states: hidden states that we want logprobs from.
+            # 2. sample_indices: Indices that have sampled tokens.
+            # 3. input_logprob_indices: Indices that have input logprob tokens.
+            # 4. token_to_seq_idx: map each token to its sequence index
+            #
+            # Example
+            # -------
+            # Suppose a batch (flattened by sequence):
+            # [t00, t01, t02, t03, t10, t11, t12, t13, t14, t20, t21, t22, t23, t24, t25]
+            # extend_seq_lens_cpu           = [4, 5, 6]
+            # extend_logprob_start_lens_cpu = [0, 5, 3]
+            #
+            # Then, the indices are:
+            # pruned_states         -> [t00, t01, t02, t03, t14, t23, t24, t25]
+            # sample_indices        -> [3, 4, 7]
+            # input_logprob_indices -> [0, 1, 2, 3, 5, 6, 7]
+            # token_to_seq_idx      -> [0, 0, 0, 0, 1, 2, 2, 2]
+            #
+            # If chunk is enabled and chunk_size = 3, the chunks will be computed in a chunked manner:
+            # [t00, t01, t02], [t03, t14, t23], [t24, t25]
+
+            sample_index_pt = -1
+            sample_indices = []
+            input_logprob_indices_pt = 0
+            input_logprob_indices = []
+            pt, pruned_states_list, pruned_states_before_norm_list = 0, [], []
+            is_packed_aux_hidden_states = isinstance(aux_hidden_states, torch.Tensor)
+            aux_pruned_states_lists = None
+            if aux_hidden_states is not None:
+                aux_pruned_states_lists = (
+                    []
+                    if is_packed_aux_hidden_states
+                    else [[] for _ in aux_hidden_states]
+                )
+
+            for idx, (extend_logprob_start_len, extend_len) in enumerate(
+                zip(
+                    logits_metadata.extend_logprob_start_lens_cpu,
+                    logits_metadata.extend_seq_lens_cpu,
+                )
+            ):
+                # It can happen in chunked prefill. We still need to sample 1 token,
+                # But we don't want to include it in input logprob.
+                if extend_len == extend_logprob_start_len:
+                    start_len = extend_logprob_start_len - 1
+                else:
+                    start_len = extend_logprob_start_len
+
+                # We always need at least 1 token to sample because that's required
+                # by a caller.
+                assert extend_len > start_len
+                pruned_states_list.append(
+                    hidden_states[pt + start_len : pt + extend_len]
+                )
+                if hidden_states_before_norm is not None:
+                    pruned_states_before_norm_list.append(
+                        hidden_states_before_norm[pt + start_len : pt + extend_len]
+                    )
+                if aux_pruned_states_lists is not None:
+                    if is_packed_aux_hidden_states:
+                        aux_pruned_states_lists.append(
+                            aux_hidden_states[pt + start_len : pt + extend_len]
+                        )
+                    else:
+                        for j, hidden in enumerate(aux_hidden_states):
+                            aux_pruned_states_lists[j].append(
+                                hidden[pt + start_len : pt + extend_len]
+                            )
+                # Map each token to its sequence index, for chunked computation
+                # of input logprobs
+                token_to_seq_idx.extend([idx] * (extend_len - start_len))
+                pt += extend_len
+                sample_index_pt += extend_len - start_len
+                sample_indices.append(sample_index_pt)
+                input_logprob_indices.extend(
+                    [
+                        input_logprob_indices_pt + i
+                        for i in range(extend_len - extend_logprob_start_len)
+                    ]
+                )
+                input_logprob_indices_pt += extend_len - start_len
+
+            pruned_states = torch.cat(pruned_states_list)
+            if hidden_states_before_norm is not None:
+                pruned_states_before_norm = torch.cat(pruned_states_before_norm_list)
+            if aux_pruned_states_lists is not None:
+                aux_pruned_states = (
+                    torch.cat(aux_pruned_states_lists)
+                    if is_packed_aux_hidden_states
+                    else [torch.cat(lst) for lst in aux_pruned_states_lists]
+                )
+
+            # Build the index tensors via pinned host memory + non-blocking H2D
+            # so the small copy doesn't drain the stream.
+            sample_indices = torch.tensor(
+                sample_indices,
+                dtype=torch.int64,
+                pin_memory=is_pin_memory_available(),
+            ).to(pruned_states.device, non_blocking=True)
+            input_logprob_indices = torch.tensor(
+                input_logprob_indices,
+                dtype=torch.int64,
+                pin_memory=is_pin_memory_available(),
+            ).to(pruned_states.device, non_blocking=True)
+
+        return (
+            pruned_states,
+            pruned_states_before_norm,
+            aux_pruned_states,
+            sample_indices,
+            input_logprob_indices,
+            token_to_seq_idx,
+        )
+
+    def _get_hidden_states_to_store(
+        self,
+        hidden_states: torch.Tensor,
+        hidden_states_before_norm: Optional[torch.Tensor],
+        aux_hidden_states: Optional[AuxHiddenStates],
+        pruned_states: torch.Tensor,
+        pruned_states_before_norm: Optional[torch.Tensor],
+        aux_pruned_states: Optional[AuxHiddenStates],
+        sample_indices: Optional[torch.Tensor],
+        logits_metadata: LogitsMetadata,
+    ) -> Optional[torch.Tensor]:
+        hidden_states_to_store: Optional[torch.Tensor] = None
+        hidden_states_to_store_before_norm: Optional[torch.Tensor] = None
+        if logits_metadata.capture_hidden_mode.need_capture():
+            if logits_metadata.capture_hidden_mode.is_full():
+                if aux_hidden_states is not None:
+                    hidden_states_to_store = pack_aux_hidden_states(aux_hidden_states)
+                else:
+                    hidden_states_to_store = hidden_states
+                hidden_states_to_store_before_norm = hidden_states_before_norm
+            elif logits_metadata.capture_hidden_mode.is_last():
+                # Get the last token hidden states. If sample_indices is None,
+                # pruned states only contain the last tokens already.
+                if aux_hidden_states is not None:
+                    assert aux_pruned_states is not None
+                    aux_pruned_states = pack_aux_hidden_states(aux_pruned_states)
+                    hidden_states_to_store = (
+                        aux_pruned_states[sample_indices]
+                        if sample_indices is not None
+                        else aux_pruned_states
+                    )
+                else:
+                    hidden_states_to_store = (
+                        pruned_states[sample_indices]
+                        if sample_indices is not None
+                        else pruned_states
+                    )
+                    if hidden_states_before_norm is not None:
+                        hidden_states_to_store_before_norm = (
+                            pruned_states_before_norm[sample_indices]
+                            if sample_indices is not None
+                            else pruned_states_before_norm
+                        )
+            else:
+                assert False, "Should never reach"
+
+        if hidden_states_to_store_before_norm is not None:
+            # NOTE: when hidden_states_before_norm is provided, we always
+            # prefer to return it.
+            hidden_states_to_store = hidden_states_to_store_before_norm
+
+        return hidden_states_to_store
+
+    def _get_logits(
+        self,
+        hidden_states: torch.Tensor,
+        lm_head: VocabParallelEmbedding,
+        logits_metadata: LogitsMetadata,
+        embedding_bias: Optional[torch.Tensor] = None,
+        use_logits_buffer: bool = True,
+    ) -> torch.Tensor:
+        """Get logits from hidden_states.
+
+        If sampled_logits_only is True, it means hidden_states only contain the
+        last position (e.g., extend without input logprobs). The caller should
+        guarantee the given hidden_states follow this constraint.
+        """
+        _trace_e2e_logits(
+            "get_logits_enter",
+            hidden_shape=tuple(hidden_states.shape),
+            dp_gather=self.do_tensor_parallel_all_gather_dp_attn,
+            tp_gather=self.do_tensor_parallel_all_gather,
+        )
+        hidden_states, local_hidden_states = self._gather_dp_attn_hidden_states(
+            hidden_states, logits_metadata
+        )
+        _trace_e2e_logits(
+            "dp_hidden_gather_returned",
+            global_shape=tuple(hidden_states.shape),
+            local_shape=tuple(local_hidden_states.shape),
+        )
+
+        if envs.SGLANG_TRACE_LOGITS_E2E_SYNC.get():
+            _trace_e2e_logits("pre_lm_head_sync_enter")
+            torch.cuda.synchronize()
+            _trace_e2e_logits("pre_lm_head_sync_returned")
+
+        _trace_e2e_logits("lm_head_enter", hidden_shape=tuple(hidden_states.shape))
+        logits = self._compute_lm_head(hidden_states, lm_head, embedding_bias)
+        _trace_e2e_logits("lm_head_returned", logits_shape=tuple(logits.shape))
+        if envs.SGLANG_TRACE_LOGITS_E2E_SYNC.get():
+            _trace_e2e_logits("post_lm_head_sync_enter")
+            torch.cuda.synchronize()
+            _trace_e2e_logits("post_lm_head_sync_returned")
+
+        if self.logit_scale is not None:
+            logits.mul_(self.logit_scale)
+
+        used_tp_lm_head_all_to_all = False
+        if self.do_tensor_parallel_all_gather:
+            _trace_e2e_logits(
+                "tp_logits_gather_enter", logits_shape=tuple(logits.shape)
+            )
+            if self.use_attn_tp_group:
+                logits = self._gather_attn_tp_logits(logits)
+            elif self._can_use_tp_lm_head_all_to_all(
+                logits, local_hidden_states, lm_head, logits_metadata
+            ):
+                logits = self._tp_lm_head_all_to_all(logits)
+                used_tp_lm_head_all_to_all = True
+            else:
+                logits = self._logits_gatherer(logits)
+            _trace_e2e_logits(
+                "tp_logits_gather_returned", logits_shape=tuple(logits.shape)
+            )
+
+        if not used_tp_lm_head_all_to_all:
+            _trace_e2e_logits(
+                "dp_logits_scatter_enter", logits_shape=tuple(logits.shape)
+            )
+            logits = self._scatter_dp_attn_logits(
+                logits, local_hidden_states, logits_metadata
+            )
+            _trace_e2e_logits(
+                "dp_logits_scatter_returned", logits_shape=tuple(logits.shape)
+            )
+
+        logits = self._copy_logits_to_buffer(
+            logits, logits_metadata, use_buffer=use_logits_buffer
+        )
+
+        if self.final_logit_softcapping:
+            if not (_is_npu or _is_cpu):
+                fused_softcap(logits, self.final_logit_softcapping)
+            else:
+                logits = self.final_logit_softcapping * torch.tanh(
+                    logits / self.final_logit_softcapping
+                )
+
+        return logits
+
+    def _compute_lm_head(
+        self,
+        hidden_states: torch.Tensor,
+        lm_head: VocabParallelEmbedding,
+        embedding_bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if self.jev_selected_ids is not None:
+            if not hasattr(lm_head, "weight") or lm_head.weight.dtype not in (torch.bfloat16, torch.float16, torch.float32) or hasattr(lm_head, "set_lora"):
+                raise ValueError("Selected head requires a frozen unquantized output projection")
+            if embedding_bias is not None or self.final_logit_softcapping:
+                raise ValueError("Selected head does not support output bias or softcapping")
+            if self.jev_head_cache is None:
+                ids = torch.tensor(self.jev_selected_ids, device=lm_head.weight.device, dtype=torch.long)
+                weight = lm_head.weight.index_select(0, ids).contiguous()
+                self.jev_head_cache = (ids, weight)
+            ids, weight = self.jev_head_cache
+            if self.use_fp32_lm_head:
+                selected = torch.matmul(hidden_states.float(), weight.float().T)
+            else:
+                selected = torch.matmul(hidden_states.to(weight.dtype), weight.T)
+            logits = torch.full((hidden_states.shape[0], self.vocab_size), -float("inf"), device=hidden_states.device, dtype=selected.dtype)
+            return logits.index_copy_(1, ids, selected)
+        quant_method = getattr(lm_head, "quant_method", None)
+        if hasattr(lm_head, "set_lora") and hasattr(lm_head, "apply_lora"):
+            # This is a LoRA-wrapped module, use its forward method
+            logits = lm_head(hidden_states)
+        elif should_apply_lm_head_quant_method(lm_head, quant_method):
+            logits = quant_method.apply(lm_head, hidden_states, embedding_bias)
+        elif hasattr(lm_head, "weight"):
+            # Normal linear layer
+            if self.use_fp32_lm_head:
+                # Avoid materializing FP32 copies for same-dtype CUDA FP16/BF16
+                # inputs. Retain explicit FP32 casts for unsupported devices or
+                # dtype combinations.
+                use_mm_out_dtype = (
+                    hidden_states.is_cuda
+                    and hidden_states.dtype == lm_head.weight.dtype
+                    and hidden_states.dtype in (torch.float16, torch.bfloat16)
+                )
+                if use_mm_out_dtype:
+                    logits = torch.mm(
+                        hidden_states,
+                        lm_head.weight.T,
+                        out_dtype=torch.float32,
+                    )
+                else:
+                    logits = torch.matmul(
+                        hidden_states.to(torch.float32),
+                        lm_head.weight.to(torch.float32).T,
+                    )
+            elif use_intel_amx_backend(lm_head):
+                logits = torch.ops.sgl_kernel.weight_packed_linear(
+                    hidden_states.to(lm_head.weight.dtype),
+                    lm_head.weight,
+                    None,  # bias
+                    True,  # is_vnni
+                )
+            elif self.rl_on_policy_target is not None:
+                # Due to tie-weight, we may not be able to change lm_head's weight dtype
+                logits = torch.matmul(
+                    hidden_states.bfloat16(), lm_head.weight.T.bfloat16()
+                )
+            else:
+                logits = torch.matmul(
+                    hidden_states.to(lm_head.weight.dtype), lm_head.weight.T
+                )
+        else:
+            # GGUF models
+            # TODO: use weight_packed_linear for GGUF models
+            if self.use_fp32_lm_head:
+                with torch.cuda.amp.autocast(enabled=False):
+                    logits = lm_head.quant_method.apply(
+                        lm_head, hidden_states.to(torch.float32), embedding_bias
+                    )
+            else:
+                logits = lm_head.quant_method.apply(
+                    lm_head, hidden_states, embedding_bias
+                )
+        return logits
+
+    def _gather_dp_attn_hidden_states(
+        self, hidden_states: torch.Tensor, logits_metadata: LogitsMetadata
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.do_tensor_parallel_all_gather_dp_attn:
+            _trace_e2e_logits(
+                "dp_metadata_enter",
+                local_shape=tuple(hidden_states.shape),
+                global_counts_cpu=logits_metadata.global_num_tokens_for_logprob_cpu,
+            )
+            logits_metadata.compute_dp_attention_metadata()
+            _trace_e2e_logits(
+                "dp_metadata_returned",
+                buffer_shape=tuple(logits_metadata.gathered_buffer.shape),
+                local_start=logits_metadata.dp_local_start_pos,
+                local_tokens=logits_metadata.dp_local_num_tokens,
+            )
+            local_hidden_states = hidden_states
+            hidden_states = logits_metadata.gathered_buffer
+            _trace_e2e_logits(
+                "dp_hidden_gather_enter",
+                global_shape=tuple(hidden_states.shape),
+                local_shape=tuple(local_hidden_states.shape),
+            )
+            dp_gather_replicate(hidden_states, local_hidden_states, logits_metadata)
+            _trace_e2e_logits("dp_hidden_gather_collective_returned")
+            return hidden_states, local_hidden_states
+        return hidden_states, hidden_states
+
+    def _gather_attn_tp_logits(self, logits: torch.Tensor) -> torch.Tensor:
+        if self.vocab_size % self.attn_tp_size == 0:
+            global_logits = torch.empty(
+                (
+                    self.attn_tp_size,
+                    logits.shape[0],
+                    self.vocab_size // self.attn_tp_size,
+                ),
+                device=logits.device,
+                dtype=logits.dtype,
+            )
+            attn_tp_all_gather_into_tensor(global_logits, logits)
+            global_logits = global_logits.permute(1, 0, 2).reshape(
+                logits.shape[0], self.vocab_size
+            )
+        else:
+            global_logits = torch.empty(
+                (self.vocab_size, logits.shape[0]),
+                device=logits.device,
+                dtype=logits.dtype,
+            )
+            global_logits = global_logits.T
+            attn_tp_all_gather(
+                list(global_logits.tensor_split(self.attn_tp_size, dim=-1)),
+                logits,
+            )
+        return global_logits
+
+    def _can_use_tp_lm_head_all_to_all(
+        self,
+        logits: torch.Tensor,
+        local_hidden_states: torch.Tensor,
+        lm_head: VocabParallelEmbedding,
+        logits_metadata: LogitsMetadata,
+    ) -> bool:
+        if not self.use_tp_lm_head_all_to_all:
+            return False
+
+        tp_size = get_parallel().tp_size
+        base_lm_head = getattr(lm_head, "base_layer", lm_head)
+        if getattr(base_lm_head, "tp_size", None) != tp_size:
+            # Tied embeddings may be replicated across DP ranks (tp_size=1),
+            # even though the logits processor runs in a larger global TP
+            # group. Such logits are full-vocabulary rather than TP shards and
+            # therefore do not satisfy the all-to-all layout contract.
+            return False
+
+        # Every participant must make the same collective choice. Decode CUDA
+        # graphs omit CPU counts and fill every GPU count with the same padded
+        # bucket size. Eager batches carry the same global CPU count list on
+        # every rank, so they are also safe when all entries are equal.
+        global_counts_cpu = logits_metadata.global_num_tokens_for_logprob_cpu
+        is_equal_padded_graph_layout = global_counts_cpu is None and (
+            logits_metadata.global_num_tokens_for_logprob_gpu is not None
+        )
+        is_equal_eager_layout = (
+            global_counts_cpu is not None
+            and len(global_counts_cpu) == tp_size
+            and len(global_counts_cpu) > 0
+            and all(count == global_counts_cpu[0] for count in global_counts_cpu)
+        )
+        if not (is_equal_padded_graph_layout or is_equal_eager_layout):
+            return False
+
+        local_rows = local_hidden_states.shape[0]
+        return local_rows > 0 and logits.shape[0] == local_rows * tp_size
+
+    def _tp_lm_head_all_to_all(self, logits: torch.Tensor) -> torch.Tensor:
+        """Exchange only the row block owned by each destination DP rank."""
+        logits = logits.contiguous()
+        all_to_all_output = torch.empty_like(logits)
+        get_tp_group().all_to_all_single(all_to_all_output.view(-1), logits.view(-1))
+        return _reassemble_tp_lm_head_all_to_all_output(
+            all_to_all_output, get_parallel().tp_size
+        )
+
+    def _scatter_dp_attn_logits(
+        self,
+        logits: torch.Tensor,
+        local_hidden_states: torch.Tensor,
+        logits_metadata: LogitsMetadata,
+    ) -> torch.Tensor:
+        if self.do_tensor_parallel_all_gather_dp_attn:
+            global_logits = logits
+            logits = torch.empty(
+                (local_hidden_states.shape[0], global_logits.shape[1]),
+                device=global_logits.device,
+                dtype=global_logits.dtype,
+            )
+            dp_scatter(logits, global_logits, logits_metadata)
+        return logits
+
+    def _copy_logits_to_buffer(
+        self,
+        logits: torch.Tensor,
+        logits_metadata: LogitsMetadata,
+        use_buffer: bool = True,
+    ) -> torch.Tensor:
+        logits_buffer = logits_metadata.next_token_logits_buffer if use_buffer else None
+        if logits.shape[-1] > self.vocab_size:
+            logits = logits[:, : self.vocab_size]
+        logits_width = logits.shape[-1]
+        # The shared logits buffer is keyed by vocab width and rows; skip it
+        # when this batch has a different logits shape than the graph buffer.
+        if logits_buffer is not None and tuple(logits_buffer.shape) == tuple(
+            logits.shape
+        ):
+            assert logits_buffer.dtype == torch.float
+            logits_buffer.copy_(logits)
+            logits = logits_buffer
+        else:
+            logits = logits.float()
+        return logits
+
+    def _get_dllm_logits(
+        self,
+        hidden_states: torch.Tensor,
+        lm_head: VocabParallelEmbedding,
+        logits_metadata: LogitsMetadata,
+    ) -> LogitsProcessorOutput:
+        assert self.return_full_logits
+        full_logits = self._get_logits(hidden_states, lm_head, logits_metadata)
+        return LogitsProcessorOutput(
+            full_logits=full_logits,
+            next_token_logits=None,
+        )
+
+    def compute_logprobs_for_multi_item_scoring(
+        self,
+        input_ids,
+        hidden_states,
+        lm_head: VocabParallelEmbedding,
+        logits_metadata: Union[LogitsMetadata, ForwardBatch],
+        multi_item_delimiter_indices: List[torch.Tensor],
+    ):
+        """
+        Compute logprobs for multi-item scoring using pre-computed delimiter indices.
+
+        Sequence format: Query<delimiter>Item1<delimiter>Item2<delimiter>...
+        Scoring positions: Extracts logprobs at positions before each <delimiter>
+
+        Args:
+            input_ids: Input token IDs. Shape: [total_sequence_length].
+            hidden_states: Hidden states from the model. Shape: [sequence_length, hidden_dim].
+            lm_head: Language model head for computing logits.
+            logits_metadata: Metadata containing batch info and logprob specs.
+            multi_item_delimiter_indices: Pre-computed delimiter positions per request (CPU tensors).
+        """
+        # Compute positions just before each delimiter.
+        # Build offset-adjusted indices on CPU, then do a single CPU→GPU transfer.
+        device = input_ids.device
+        all_tensors = []
+        if logits_metadata.extend_seq_lens_cpu is not None:
+            offset = 0
+            for req_seq_len, indices_tensor in zip(
+                logits_metadata.extend_seq_lens_cpu, multi_item_delimiter_indices
+            ):
+                if len(indices_tensor) > 0:
+                    # Note: if the first delimiter is at position 0 (empty query),
+                    # indices - 1 wraps to -1. This is harmless — the first
+                    # delimiter entry is always discarded by
+                    # _process_multi_item_scoring_results.
+                    all_tensors.append(indices_tensor + (offset - 1))
+                offset += req_seq_len
+        else:
+            all_tensors.append(multi_item_delimiter_indices[0] - 1)
+        multi_item_indices = torch.cat(all_tensors).to(device, non_blocking=True)
+
+        # Extract hidden states at delimiter positions for multi-item scoring
+        sliced_hidden = hidden_states[multi_item_indices]
+
+        sliced_logits = self._get_logits(sliced_hidden, lm_head, logits_metadata)
+        sliced_logprobs = torch.nn.functional.log_softmax(sliced_logits, dim=-1)
+
+        # Initialize return values
+        input_token_ids_logprobs_val = []
+        input_token_ids_logprobs_idx = []
+        input_top_logprobs_val = None
+        input_top_logprobs_idx = None
+
+        # Recalculate extend_logprob_pruned_lens_cpu to match delimiter counts per request
+        if (
+            logits_metadata.token_ids_logprobs
+            or logits_metadata.extend_return_top_logprob
+        ):
+            logits_metadata.extend_logprob_pruned_lens_cpu = [
+                len(t) for t in multi_item_delimiter_indices
+            ]
+
+        # Get the logprobs of specified token ids
+        if logits_metadata.extend_token_ids_logprob:
+            (
+                input_token_ids_logprobs_val,
+                input_token_ids_logprobs_idx,
+            ) = get_token_ids_logprobs_raw(
+                sliced_logprobs,
+                logits_metadata.token_ids_logprobs,
+                stage=LogprobStage.PREFILL,
+                extend_logprob_pruned_lens_cpu=logits_metadata.extend_logprob_pruned_lens_cpu,
+                no_copy_to_cpu=True,
+            )
+
+        # Get the logprob of top-k tokens
+        if logits_metadata.extend_return_top_logprob:
+            (
+                input_top_logprobs_val,
+                input_top_logprobs_idx,
+            ) = get_top_logprobs_raw(
+                sliced_logprobs,
+                logits_metadata.top_logprobs_nums,
+                stage=LogprobStage.PREFILL,
+                extend_logprob_pruned_lens_cpu=logits_metadata.extend_logprob_pruned_lens_cpu,
+            )
+
+        # MIS scores come from input_token_ids_logprobs_val (label-token logprobs),
+        # not from per-position input_token_logprobs. However, the shared logprob
+        # pipeline (add_input_logprob_return_values) asserts input_token_logprobs is
+        # non-None, converts it to a tuple, slices it, and validates its length —
+        # all before score_request() ever sees the result. We can't set it to None
+        # without changing those shared asserts, so we fill with zeros to satisfy
+        # the pipeline. score_request() ignores this field entirely.
+        input_token_logprobs = torch.zeros(multi_item_indices.shape[0], device=device)
+
+        return LogitsProcessorOutput(
+            next_token_logits=None,
+            input_token_logprobs=input_token_logprobs,
+            input_top_logprobs_val=input_top_logprobs_val,
+            input_top_logprobs_idx=input_top_logprobs_idx,
+            input_token_ids_logprobs_val=input_token_ids_logprobs_val,
+            input_token_ids_logprobs_idx=input_token_ids_logprobs_idx,
+            mm_input_embeds=logits_metadata.mm_input_embeds,
+        )
+
+
+def _reassemble_tp_lm_head_all_to_all_output(
+    all_to_all_output: torch.Tensor, tp_size: int
+) -> torch.Tensor:
+    """Convert source-major all-to-all output to row-major full-vocab logits.
+
+    Each source TP rank contributes ``[local_rows, vocab_shard]`` for this
+    destination DP rank. ``all_to_all_single`` concatenates those contributions
+    along dim 0, while the sampler expects the vocab shards concatenated along
+    dim 1.
+    """
+    assert all_to_all_output.shape[0] % tp_size == 0
+    local_rows = all_to_all_output.shape[0] // tp_size
+    vocab_shard = all_to_all_output.shape[1]
+    return (
+        all_to_all_output.view(tp_size, local_rows, vocab_shard)
+        .permute(1, 0, 2)
+        .reshape(local_rows, tp_size * vocab_shard)
+    )
+
+
+def _has_lm_head_runtime_attrs(lm_head, attr_names: Tuple[str, ...]) -> bool:
+    return all(hasattr(lm_head, attr_name) for attr_name in attr_names)
+
+
+def should_apply_lm_head_quant_method(lm_head, quant_method) -> bool:
+    if (
+        quant_method is None
+        or not hasattr(lm_head, "weight")
+        or not callable(getattr(quant_method, "apply", None))
+    ):
+        return False
+
+    method_name = type(quant_method).__name__
+    if method_name in _UNQUANTIZED_LM_HEAD_METHODS:
+        return False
+
+    # Some draft models share an unquantized target lm_head tensor while still
+    # carrying the draft model's stale ModelOpt quant_method. Only use the
+    # ModelOpt lm_head kernel when the runtime quantization state matches it.
+    if method_name == "ModelOptFp4LinearMethod":
+        if quant_method.quant_mode == "w4a16":
+            return lm_head.weight.dtype == torch.uint8 and _has_lm_head_runtime_attrs(
+                lm_head,
+                (
+                    "weight_scale_interleaved",
+                    "alpha",
+                    "input_size_per_partition",
+                    "output_size_per_partition",
+                ),
+            )
+        if lm_head.weight.dtype == torch.int32 and _has_lm_head_runtime_attrs(
+            lm_head,
+            (
+                "weight_scale",
+                "weight_global_scale",
+                "workspace",
+                "input_size_per_partition",
+                "output_size_per_partition",
+            ),
+        ):
+            return True
+        return lm_head.weight.dtype == torch.uint8 and _has_lm_head_runtime_attrs(
+            lm_head,
+            (
+                "weight_scale_interleaved",
+                "alpha",
+                "input_scale_inv",
+                "input_size_per_partition",
+                "output_size_per_partition",
+            ),
+        )
+    if method_name == "ModelOptNvFp4A16LinearMethod":
+        return lm_head.weight.dtype == torch.int32 and _has_lm_head_runtime_attrs(
+            lm_head,
+            (
+                "weight_scale",
+                "weight_global_scale",
+                "workspace",
+                "input_size_per_partition",
+                "output_size_per_partition",
+            ),
+        )
+    if method_name == "ModelOptFp8LinearMethod":
+        return (
+            lm_head.weight.dtype == torch.float8_e4m3fn
+            and _has_lm_head_runtime_attrs(lm_head, ("weight_scale", "input_scale"))
+        )
+
+    return True

runtime/patches/sglang-qwen35-text.py

@@ -0,0 +1,298 @@
+# Copyright 2025 Qwen Team
+# Copyright 2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import logging
+from typing import Iterable, Optional, Set, Tuple, Union
+
+import torch
+from torch import nn
+
+from sglang.srt.distributed import get_pp_group
+from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
+from sglang.srt.layers.logits_processor import LogitsProcessor
+from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.layers.utils import PPMissingLayer
+from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.models import qwen3_5
+from sglang.srt.models.qwen2_moe import Qwen2MoeSparseMoeBlock
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.utils import LazyValue, add_prefix
+
+logger = logging.getLogger(__name__)
+
+_MODEL_PREFIX = "model."
+
+
+class Qwen3_5ForCausalLM(nn.Module):
+    body_cls = qwen3_5.Qwen3_5ForCausalLM
+
+    packed_modules_mapping = qwen3_5.Qwen3_5ForCausalLM.packed_modules_mapping
+    supported_lora_modules = qwen3_5.Qwen3_5ForCausalLM.supported_lora_modules
+
+    @classmethod
+    def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
+        # The body decides; it is handed this config and quantization verbatim.
+        return cls.body_cls.shared_experts_fusion_disable_reason(
+            hf_config, quant_config
+        )
+
+    def __init__(
+        self,
+        config,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.quant_config = quant_config
+        self.pp_group = get_pp_group()
+
+        if quant_config is not None and hasattr(quant_config, "packed_modules_mapping"):
+            quant_config.packed_modules_mapping = self.packed_modules_mapping
+
+        self.model = self.body_cls(
+            config=config,
+            quant_config=quant_config,
+            prefix=add_prefix("model", prefix),
+        )
+
+        if self.pp_group.is_last_rank:
+            if self.pp_group.world_size == 1 and config.tie_word_embeddings:
+                self.lm_head = self.model.embed_tokens
+            else:
+                self.lm_head = ParallelLMHead(
+                    config.vocab_size,
+                    config.hidden_size,
+                    quant_config=quant_config,
+                    org_num_embeddings=config.vocab_size,
+                    prefix=add_prefix("lm_head", prefix),
+                    use_attn_tp_group=get_parallel().enable_dp_lm_head,
+                )
+        else:
+            self.lm_head = PPMissingLayer()
+
+        self.logits_processor = LogitsProcessor(config)
+
+        # Text-only checkpoints retain mrope_section, but identical position rows
+        # make its rotary embedding equivalent to 1-D RoPE.
+        rope_config = getattr(config, "rope_parameters", None) or getattr(
+            config, "rope_scaling", None
+        )
+        self.is_mrope_enabled = bool(rope_config) and "mrope_section" in rope_config
+
+        self.capture_aux_hidden_states = False
+
+    @property
+    def start_layer(self) -> int:
+        return self.model.start_layer
+
+    @property
+    def end_layer(self) -> int:
+        return self.model.end_layer
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.embed_tokens
+
+    def prepare_before_cuda_graph_capture(self, model_runner) -> None:
+        """Forward model-owned warmup to the text backbone."""
+        self.model.prepare_before_cuda_graph_capture(model_runner)
+
+    def set_dflash_layers_to_capture(self, layer_ids: list[int]) -> None:
+        if self.pp_group.world_size > 1:
+            raise NotImplementedError("DFLASH/DSPARK aux hidden capture requires PP=1.")
+        num_layers = len(self.model.layers)
+        if sorted(set(layer_ids)) != list(layer_ids) or not all(
+            0 <= layer_id < num_layers - 1 for layer_id in layer_ids
+        ):
+            raise ValueError(
+                "target_layer_ids must be unique, strictly increasing, and in "
+                f"[0, {num_layers - 1}); got {layer_ids}"
+            )
+        self.capture_aux_hidden_states = True
+        self.model.set_dflash_layers_to_capture(
+            [layer_id + 1 for layer_id in layer_ids]
+        )
+
+    def get_embed_and_head(self):
+        # PP splits embedding and lm_head across first/last stages; the draft keeps
+        # its own copy of whichever half its stage cannot receive.
+        embed = (
+            None
+            if isinstance(self.model.embed_tokens, PPMissingLayer)
+            else self.model.embed_tokens.weight
+        )
+        head = None if isinstance(self.lm_head, PPMissingLayer) else self.lm_head.weight
+        return embed, head
+
+    def set_embed_and_head(self, embed, head):
+        del self.model.embed_tokens.weight
+        del self.lm_head.weight
+        self.model.embed_tokens.weight = embed
+        self.lm_head.weight = head
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+    def set_dflash_layers_to_capture(self, layers_to_capture: list[int]):
+        if not self.pp_group.is_last_rank:
+            return
+        if layers_to_capture is None:
+            raise ValueError(
+                "DFLASH requires explicit layer ids for aux hidden capture."
+            )
+        self.capture_aux_hidden_states = True
+        self.model.set_dflash_layers_to_capture(layers_to_capture)
+
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        input_embeds: Optional[torch.Tensor] = None,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
+        **kwargs,
+    ) -> Union[torch.Tensor, PPProxyTensors]:
+        if self.is_mrope_enabled:
+            positions = forward_batch.mrope_positions
+
+        hidden_states = self.model(
+            input_ids,
+            positions,
+            forward_batch,
+            input_embeds,
+            pp_proxy_tensors=pp_proxy_tensors,
+        )
+
+        if not self.pp_group.is_last_rank:
+            return hidden_states
+
+        aux_hidden_states = None
+        if isinstance(hidden_states, tuple):
+            hidden_states, aux_hidden_states = hidden_states
+
+        return self.logits_processor(
+            input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states
+        )
+
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
+        params_dict = dict(self.named_parameters())
+        loaded_params: Set[str] = set()
+
+        # Keep prefix stripping lazy: materializing mmap-backed checkpoint tensors
+        # across model processes can exceed host memory before GPU copies finish.
+        def portable_weights():
+            # HF native export packs experts and shared gate/up matrices. Expand
+            # lazily into the loader's per-expert format; no tensor values change.
+            for name, tensor in weights:
+                name = name.replace("model.language_model.", "model.")
+                if ".mlp.experts.gate_up_proj" in name:
+                    stem = name.split(".mlp.experts.gate_up_proj")[0]
+                    for expert in range(tensor.shape[0]):
+                        gate, up = tensor[expert].chunk(2, dim=0)
+                        yield f"{stem}.mlp.experts.{expert}.gate_proj.weight", gate
+                        yield f"{stem}.mlp.experts.{expert}.up_proj.weight", up
+                elif ".mlp.experts.down_proj" in name:
+                    stem = name.split(".mlp.experts.down_proj")[0]
+                    for expert in range(tensor.shape[0]):
+                        yield f"{stem}.mlp.experts.{expert}.down_proj.weight", tensor[expert]
+                elif ".mlp.shared_expert.gate_up_proj.weight" in name:
+                    gate, up = tensor.chunk(2, dim=0)
+                    yield name.replace("gate_up_proj", "gate_proj"), gate
+                    yield name.replace("gate_up_proj", "up_proj"), up
+                else:
+                    yield name, tensor
+
+        def body_weights():
+            fused_shared = max((getattr(layer.mlp, "num_fused_shared_experts", 0) for layer in self.model.layers), default=0)
+            for name, loaded_weight in portable_weights():
+                if fused_shared and ".mlp.shared_expert." in name:
+                    name = name.replace(".mlp.shared_expert.", f".mlp.experts.{self.config.num_experts}.")
+                if name.startswith(_MODEL_PREFIX):
+                    yield name[len(_MODEL_PREFIX) :], loaded_weight
+                elif name == "lm_head.weight":
+                    if self.config.tie_word_embeddings:
+                        continue
+                    if "lm_head.weight" not in params_dict:
+                        continue
+                    param = params_dict["lm_head.weight"]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight)
+                    loaded_params.add("lm_head.weight")
+
+        body_loaded = self.model.load_weights(body_weights())
+        loaded_params.update(f"{_MODEL_PREFIX}{n}" for n in body_loaded)
+
+        if self.config.tie_word_embeddings and self.pp_group.is_last_rank:
+            loaded_params.add("lm_head.weight")
+
+        missing = set(params_dict) - loaded_params
+        if (self.quant_config is not None
+                and self.quant_config.get_name() == "fp8"
+                and not self.quant_config.is_checkpoint_fp8_serialized):
+            # Online BF16->FP8 creates these scales itself. They are deliberately
+            # absent from the source BF16 checkpoint; all model weights still
+            # have to load. Do not apply this exception to serialized FP8.
+            generated_suffixes = (
+                ".mlp.experts.w13_weight_scale", ".mlp.experts.w2_weight_scale",
+                ".mlp.experts.w13_weight_scale1", ".mlp.experts.w2_weight_scale1",
+                ".attn.k_scale", ".attn.v_scale",
+            )
+            generated = {name for name in missing if name.endswith(generated_suffixes)}
+            missing -= generated
+        if missing:
+            raise RuntimeError(f"Incomplete text checkpoint load: {sorted(missing)[:20]}")
+        return loaded_params
+
+
+class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM):
+    body_cls = qwen3_5.Qwen3_5MoeForCausalLM
+
+    packed_modules_mapping = qwen3_5.Qwen3_5MoeForCausalLM.packed_modules_mapping
+    supported_lora_modules = qwen3_5.Qwen3_5MoeForCausalLM.supported_lora_modules
+
+    def __init__(
+        self,
+        config,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__(config=config, quant_config=quant_config, prefix=prefix)
+
+        self._routed_experts_weights_of_layer = LazyValue(
+            lambda: {
+                layer_id: layer.mlp.get_moe_weights()
+                for layer_id, layer in enumerate(self.model.layers)
+                if isinstance(layer.mlp, Qwen2MoeSparseMoeBlock)
+            }
+        )
+
+    @property
+    def routed_experts_weights_of_layer(self):
+        return self._routed_experts_weights_of_layer.value
+
+    @classmethod
+    def get_model_config_for_expert_location(cls, config):
+        return ModelConfigForExpertLocation(
+            num_layers=config.num_hidden_layers,
+            num_logical_experts=config.num_experts,
+            num_groups=None,
+        )
+
+
+EntryClass = [Qwen3_5MoeForCausalLM, Qwen3_5ForCausalLM]

serve.sh

@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+model_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+image='lmsysorg/sglang@sha256:a0ffcdd013af6d79c9b7c4348c42b14a79a48f6c43dd3ca0235ec3802da45f75'
+container_name=${NYX_CONTAINER_NAME:-nyx}
+port=${NYX_PORT:-30000}
+cache_dir=${NYX_CACHE_DIR:-/tmp/nyx-sglang-cache}
+memory_fraction=${NYX_MEM_FRACTION:-0.75}
+
+mkdir -p "$cache_dir/aiter-jit" "$cache_dir/sglang"
+
+exec docker run --rm --name "$container_name" \
+  --device /dev/kfd --device /dev/dri --group-add video \
+  --ipc=host -p "127.0.0.1:${port}:30000" \
+  -e JEV_SELECTED_HEAD=0 \
+  -e SGLANG_USE_AITER_UNIFIED_ATTN=0 \
+  -e SGLANG_MAMBA_SSM_DTYPE=float32 \
+  -e AITER_JIT_DIR=/aiter-jit \
+  -v "$model_dir:/model:ro" \
+  -v "$cache_dir/aiter-jit:/aiter-jit" \
+  -v "$cache_dir/sglang:/root/.cache" \
+  -v "$model_dir/runtime/label-token-ids.json:/jev-label-token-ids.json:ro" \
+  -v "$model_dir/runtime/patches/sglang-logits-processor.py:/sgl-workspace/sglang/python/sglang/srt/layers/logits_processor.py:ro" \
+  -v "$model_dir/runtime/patches/sglang-qwen35.py:/sgl-workspace/sglang/python/sglang/srt/models/qwen3_5.py:ro" \
+  -v "$model_dir/runtime/patches/sglang-batch-result-processor.py:/sgl-workspace/sglang/python/sglang/srt/managers/scheduler_components/batch_result_processor.py:ro" \
+  --entrypoint /model/runtime/entrypoint.sh "$image" \
+  -m sglang.launch_server \
+  --model-path /model \
+  --load-format presharded \
+  --model-loader-extra-config '{"presharded_path":"/model/weights","max_file_bytes":5368709120,"verify_on_load":true}' \
+  --host 0.0.0.0 --port 30000 \
+  --dtype bfloat16 \
+  --quantization quark_int4fp8_moe \
+  --fp8-gemm-backend aiter \
+  --moe-runner-backend aiter \
+  --tp-size 1 --mem-fraction-static "$memory_fraction" \
+  --context-length 262144 --max-running-requests 64 \
+  --chunked-prefill-size 4096 \
+  --mamba-radix-cache-strategy extra_buffer \
+  --page-size 64 --attention-backend aiter \
+  --cuda-graph-backend-prefill breakable \
+  --cuda-graph-max-bs-decode 64 \
+  --cuda-graph-bs-prefill 128 256 512 1024 2048 4096
+

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:06b9509352d2af50381ab2247e083b80d32d5c0aba91c272ca9ff729b6a0e523
+size 19989325

tokenizer_config.json

@@ -0,0 +1,32 @@
+{
+  "add_prefix_space": false,
+  "audio_bos_token": "<|audio_start|>",
+  "audio_eos_token": "<|audio_end|>",
+  "audio_token": "<|audio_pad|>",
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "image_token": "<|image_pad|>",
+  "is_local": true,
+  "local_files_only": true,
+  "model_max_length": 262144,
+  "model_specific_special_tokens": {
+    "audio_bos_token": "<|audio_start|>",
+    "audio_eos_token": "<|audio_end|>",
+    "audio_token": "<|audio_pad|>",
+    "image_token": "<|image_pad|>",
+    "video_token": "<|video_pad|>",
+    "vision_bos_token": "<|vision_start|>",
+    "vision_eos_token": "<|vision_end|>"
+  },
+  "pad_token": "<|endoftext|>",
+  "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null,
+  "video_token": "<|video_pad|>",
+  "vision_bos_token": "<|vision_start|>",
+  "vision_eos_token": "<|vision_end|>"
+}

verify.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python3
+"""Fail-closed integrity check for the published Nyx package."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from pathlib import Path
+
+
+ROOT = Path(__file__).resolve().parent
+
+
+def digest(path: Path) -> str:
+    checksum = hashlib.sha256()
+    with path.open("rb") as stream:
+        for block in iter(lambda: stream.read(8 * 1024 * 1024), b""):
+            checksum.update(block)
+    return checksum.hexdigest()
+
+
+manifest = json.loads((ROOT / "release-manifest.json").read_text())
+if manifest.get("repository") != "devanshbatham/nyx":
+    raise SystemExit("Unexpected repository identity")
+
+for name, metadata in manifest["files"].items():
+    path = ROOT / name
+    if not path.is_file() or path.is_symlink():
+        raise SystemExit(f"Missing or unsafe release file: {name}")
+    if path.stat().st_size != metadata["bytes"]:
+        raise SystemExit(f"Size mismatch: {name}")
+    if digest(path) != metadata["sha256"]:
+        raise SystemExit(f"SHA-256 mismatch: {name}")
+
+print(json.dumps({
+    "verified": True,
+    "files": len(manifest["files"]),
+    "total_bytes": manifest["total_bytes"],
+    "repository": manifest["repository"],
+}))
+

weights/TP-1-sig-e793a9432f30cd4e/READY

@@ -0,0 +1 @@
+{"plan_version": 1, "world_size": 1, "created_at": 1789971566.5357497}

weights/TP-1-sig-e793a9432f30cd4e/model-00000-common.safetensor

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:16e3906e9a161a6337a4a0f6d06e352d09360b37a9d4c4ea341ea978192ee289
+size 5268047400

weights/TP-1-sig-e793a9432f30cd4e/model-00001-common.safetensor

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7ad7affbcc8912528ff2859568b51ee9427cc17119312ea8750cf118544ea5d3
+size 5120198752

weights/TP-1-sig-e793a9432f30cd4e/model-00002-common.safetensor

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa4a08a1e0e0de30543d343a629c6710e59575cd6f98f5beae789f34a59b79b7
+size 5254941984

weights/TP-1-sig-e793a9432f30cd4e/model-00003-common.safetensor

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfb557a70a97b91208837dac145569c844971062970f41466738e61e77d53eb3
+size 4014197624