Add files using upload-large-folder tool

LICENSE

@@ -0,0 +1,61 @@
+Source: https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-agreement/ (retrieved 2026-10-20)
+
+NVIDIA Open Model Agreement
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+Release Date: April 2, 2026
+
+NVIDIA Works released under this License are intended to be used permissively and enable the further development of AI technologies. Subject to the terms of this License, NVIDIA confirms that:
+
+- Works are commercially usable.
+
+- You are free to create and distribute Derivative Works.
+
+- NVIDIA does not claim ownership to any outputs generated using the Works or Derivative Works.
+
+By using, reproducing, modifying, distributing, performing or displaying any portion or element of the Works or Derivative Works, or otherwise accepting the terms of this License, you agree to be bound by this License.
+
+1. Definitions.
+
+“License” shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 10 of this document.
+
+“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.
+
+“Work” shall mean the work of authorship, including machine learning model, software, checkpoints, learnt weights, algorithms, parameters, configuration files and documentation, 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.
+
+2. Grant of License. Subject to the terms and conditions of this License, NVIDIA hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable 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.
+
+If You institute patent or copyright litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or an output from the Work constitutes direct or contributory patent or copyright infringement, then any licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed.
+
+3. 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 a copy of this License; and
+
+b. 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
+
+c. 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 following attribution notice within a “Notice” text file with such copies and the following statement: “Licensed by NVIDIA Corporation under the NVIDIA Open Model Agreement.”
+
+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.
+
+4. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of NVIDIA, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file.
+
+5. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, NVIDIA provides the Work 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.
+
+6. 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 NVIDIA 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, a Derivative Work or an output from the Work or Derivative 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 NVIDIA has been advised of the possibility of such damages.
+
+7. 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 NVIDIA, and only if You agree to indemnify, defend, and hold NVIDIA harmless for any liability incurred by, or claims asserted against, NVIDIA by reason of your accepting any such warranty or additional liability. You will indemnify and hold harmless NVIDIA from and against any claim by any third party arising out of or related to your use or distribution of the Works, Derivative Works thereof, or output from the Works or Derivative Works.
+
+8. Feedback. NVIDIA appreciates your feedback, and You agree that NVIDIA may use it without restriction or compensation to You.
+
+9. Governing Law. This Agreement will be governed in all respects by the laws of the United States and the laws of the State of Delaware, without regard to conflict of laws principles or the United Nations Convention on Contracts for the International Sale of Goods. The state and federal courts residing in Santa Clara County, California will have exclusive jurisdiction over any dispute or claim arising out of or related to this Agreement, and the parties irrevocably consent to personal jurisdiction and venue in those courts; except that, either party may apply for injunctive remedies or an equivalent type of urgent legal relief in any jurisdiction.
+
+10. Trade and Compliance. You agree to comply with all applicable export, import, trade and economic sanctions laws and regulations, as amended, including without limitation U.S. Export Administration Regulations and Office of Foreign Assets Control regulations. These laws include restrictions on destinations, end-users and end-use.
+
+(v. March 9, 2026)
+
+End of terms and conditions

LICENSE.Qwen

@@ -0,0 +1,16 @@
+Qwen Community License 1.0
+
+Copyright (c) 2026 Qwen
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software, including the model weights, parameters, configuration files, inference code and associated documentation files (collectively, the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, sell, deploy, host, fine-tune, and create derivative works from (collectively, "Use" or "Using") copies of the Software; and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+1. The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. If the Software (or any derivative works thereof) is Used for any of the licensee's commercial products or services that have more than 100,000,000 monthly active users or US$ 20,000,000 (or equivalent in other currencies) monthly revenue, respective model name must be prominently displayed on the user interface of such product or service; and,
+
+2. If the licensee or any of its affiliates conducts a Model as a Service or AI Work Assistant business, the licensee shall obtain a separate license from Qwen before Using the Software or its derivative works for any commercial purpose. The foregoing requirement shall not apply to the licensee's internal Use of the Software, provided that such Use does not make the Software, its outputs, or its underlying model capabilities available to any third party.
+
+"Model as a Service" means giving a third party access to language model inference or fine-tuning (e.g., via API or a hosted endpoint) in a manner that allows such third parties to exercise meaningful control over the inputs, parameters, or training data. This does not include the mere relaying of requests to models hosted by other third parties. 
+“AI Work Assistant” means an independent AI-powered product primarily designed for AI-assisted coding or office productivity (e.g., Qoder and QwenWork). It does not include: (a) a single-purpose AI tool (such as an AI translation tool); (b) an AI assistant primarily designed for a domain other than coding or office productivity (such as Taobao AI Shopping Assistant or AMap AI Chat); or (c) an AI assistant that is a feature of a product whose primary purpose is not AI-assisted coding or office productivity.
+
+THE SOFTWARE AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL QWEN, ITS AFFILIATES OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. THE USE OF THE SOFTWARE MUST COMPLY WITH APPLICABLE LAWS AND REGULATIONS, AND MUST NOT INFRINGE THE INTELLECTUAL PROPERTY RIGHTS OF ANY THIRD PARTY.
+
+For any questions regarding this license, please contact model-business@notice.qwencloud.com.

NOTICE

@@ -0,0 +1,13 @@
+Qwen3.8-Flash-Next FP4-Local for one DGX Spark (diffbot)
+
+Model weights: Licensed by NVIDIA Corporation under the NVIDIA Open Model Agreement.
+The files config.json, generation_config.json, the tokenizer and processor files, chat_template.jinja and
+tensors/ are NVIDIA's nvidia/Qwen3.8-Flash-Next-FP4-Local release, redistributed unmodified; NVIDIA's model
+card is kept as NVIDIA_MODEL_CARD.md. A copy of the NVIDIA Open Model Agreement is in LICENSE.
+
+The underlying base model, Qwen/Qwen3.8-Flash-Next, is Copyright (c) 2026 Qwen and governed by the Qwen
+Community License 1.0 (LICENSE.Qwen).
+
+Additions by Diffbot (recipe/): serving scripts, benchmarks and the KL client are MIT (recipe/LICENSE);
+the files in recipe/patches/ are modified or new files for vLLM and FlashInfer's b12x, Apache-2.0
+(recipe/THIRD_PARTY_NOTICES.md). recipe/serve/draft_vocab_64k.json is Diffbot's token ranking.

NVIDIA_MODEL_CARD.md

@@ -0,0 +1,414 @@
+---
+pipeline_tag: image-text-to-text
+base_model:
+- Qwen/Qwen3.8-Flash-Next
+license: other
+license_name: nvidia-open-model-agreement
+library_name: Model Optimizer
+tags:
+- nvidia
+- ModelOpt
+- Qwen3.8
+- quantized
+- FP4
+- fp4
+---
+
+
+# Model Overview
+
+**IMPORTANT**: the NVFP4 scale factors in this checkpoint are in a new lossless compression format, but upstream vLLM support for this format has not yet landed. **You must use one of the containers provided below.**
+
+## Description:
+The NVIDIA Qwen3.8-Flash-Next-FP4-Local model is the quantized version of Alibaba's Qwen3.8-Flash-Next model, which is an auto-regressive language model that uses an optimized transformer architecture. Qwen3.8-Flash-Next is a causal language model with a vision encoder, hybrid attention (Gated DeltaNet and Qwen Sparse Attention), Mixture-of-Experts, gated residual streams, and n-gram embeddings. For more information, please check [here](https://huggingface.co/Qwen/Qwen3.8-Flash-Next). The NVIDIA Qwen3.8-Flash-Next-FP4-Local model is quantized with [Model Optimizer](https://github.com/NVIDIA/Model-Optimizer) using [quantization-aware distillation (QAD)](https://developer.nvidia.com/blog/how-quantization-aware-training-enables-low-precision-accuracy-recovery/).
+
+The checkpoint occupies approximately **96 GiB on disk** and is best suited for local AI hardware. It has been validated with vLLM on the following hardware: <br>
+* **NVIDIA RTX PRO 6000 Blackwell** — single GPU, with PLE (n-gram embedding table) offload. 850K tokens of FP8 KV Cache <br>
+* **NVIDIA DGX Spark** — Can run with or without PLE offload. 1M tokens of FP8 KV Cache with PLE offload <br>
+
+This model is ready for commercial or non-commercial use.  <br>
+
+## Third-Party Community Consideration
+This model is not owned or developed by NVIDIA. This model has been developed and built to a third-party's requirements for this application and use case; see link to Non-NVIDIA [(Qwen3.8-Flash-Next) Model Card](https://huggingface.co/Qwen/Qwen3.8-Flash-Next) from Qwen.
+
+### License/Terms of Use:
+**GOVERNING DOWNLOAD TERMS:** Use of this model is governed by [NVIDIA Open Model Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-agreement/).
+
+**ADDITIONAL INFORMATION** : The Qwen3.8-Flash-Next base model is governed by the [Qwen Community License 1.0](https://huggingface.co/Qwen/Qwen3.8-Flash-Next/blob/main/LICENSE). **Qwen3.8-Flash-Next**.
+
+### Deployment Geography:
+Global <br>
+
+### Use Case: <br>
+Developers looking to take off-the-shelf, pre-quantized models for deployment in AI Agent systems, chatbots, RAG systems, and other AI-powered applications. <br>
+
+### Release Date:  <br>
+Hugging Face 10/05/2026 via https://huggingface.co/nvidia/Qwen3.8-Flash-Next-FP4-Local <br>
+
+## References
+NVIDIA Model Optimizer: https://github.com/NVIDIA/Model-Optimizer
+
+## Model Architecture:
+**Architecture Type:** Transformers  <br>
+**Network Architecture:** Hybrid Attention (Gated DeltaNet and Qwen Sparse Attention) with Mixture-of-Experts (`Qwen4ExpForConditionalGeneration`) <br>
+**Number of Model Parameters:** 125B in total and 6B activated, plus 51B n-gram embedding and 4B MTP <br>
+**This model was developed based on [Qwen3.8-Flash-Next](https://huggingface.co/Qwen/Qwen3.8-Flash-Next)** <br>
+
+## Input:
+**Input Type(s):** Text, Image, Video <br>
+**Input Format(s):** String, Red, Green, Blue (RGB), Video (MP4/WebM) <br>
+**Input Parameters:** One-Dimensional (1D), Two-Dimensional (2D), Three-Dimensional (3D) <br>
+**Other Properties Related to Input:** Context length up to 262K natively and extensible up to 1,000,000 tokens. <br>
+
+## Output:
+**Output Type(s):** Text <br>
+**Output Format:** String <br>
+**Output Parameters:** One-Dimensional (1D): Sequences <br>
+**Other Properties Related to Output:** None <br>
+
+Our AI models are designed and/or optimized to run on NVIDIA GPU-accelerated systems. By leveraging NVIDIA's hardware (e.g. GPU cores) and software frameworks (e.g., CUDA libraries), the model achieves faster training and inference times compared to CPU-only solutions. <br>
+
+## Software Integration:
+**Supported Runtime Engine(s):** <br>
+* **vLLM** <br>
+
+**Supported Hardware Microarchitecture Compatibility:** <br>
+* NVIDIA Blackwell <br>
+
+**Preferred Operating System(s):** <br>
+* Linux <br>
+
+The integration of foundation and fine-tuned models into AI systems requires additional testing using use-case-specific data to ensure safe and effective deployment. Following the V-model methodology, iterative testing and validation at both unit and system levels are essential to mitigate risks, meet technical and functional requirements, and ensure compliance with safety and ethical standards before deployment.
+
+## Model Version(s):
+The model version is FP4 QAD 1.0 and is quantized with nvidia-modelopt **v0.47.0** <br>
+
+## Training and Evaluation Datasets:
+
+## Calibration Dataset:
+**Link:** Subset of [Nemotron-Post-Training-Dataset-v3](https://huggingface.co/collections/nvidia/nemotron-post-training-v3), [Nemotron-Post-Training-Dataset-v2](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2) and subset of [Nemotron-Pretraining-Dataset](https://huggingface.co/collections/nvidia/nemotron-pre-training-datasets) <br>
+**Data Collection Method by dataset:** Automated. <br>
+**Labeling Method by dataset:** Automated. <br>
+**Properties:**  NVIDIA's Nemotron-Post-Training-v3 is an open collection of approximately 50 curated Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL) datasets designed for the Nemotron 3 model family (Nano, Super, and Ultra). NVIDIA's  Nemotron-Post-Training-Dataset-v2 is a post-training dataset curated by NVIDIA containing multi-turn conversations across diverse topics. The Nemotron-Pre-Training-Datasets collection on Hugging Face provides large-scale data corpora used to train the NVIDIA Nemotron model family <br>
+
+## Training Dataset:
+**Link**: Subset of [Nemotron-Post-Training-Dataset-v3](https://huggingface.co/collections/nvidia/nemotron-post-training-v3), [Nemotron-Post-Training-Dataset-v2](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2) and subset of [Nemotron-Pretraining-Dataset](https://huggingface.co/collections/nvidia/nemotron-pre-training-datasets); only prompts from the datasets were used for data synthesis. <br>
+**Data Modality:** Text <br>
+**Text Training Data Size:** 500 Million Tokens, repeated for 3 epochs <br>
+**Data Collection Method by dataset:** Hybrid: Automated, manually-collected, Synthetic <br>
+**Labeling Method by dataset:** Hybrid: Automated, manually-labelled, Synthetic <br>
+**Properties:** The Qwen3.8-Flash-Next-FP4-Local model was trained exclusively on a post-training corpus includes synthetic data for reasoning, code, science, tool use, instruction following, structured outputs, and multilingual tasks. Dataset disclosures are published in the `nvidia/nemotron-post-training-v2` and subset of `nvidia/nemotron-post-training-v3` and subset of `nvidia/nemotron-pre-training-datasets` Hugging Face collections. <br>
+
+## Inference:
+**Acceleration Engine:** vLLM
+**Test Hardware:** NVIDIA DGX Spark and NVIDIA RTX PRO 6000 Blackwell
+
+## Evaluation
+
+**Datasets:** GPQA Diamond, ToolEvalBench <br>
+**Data Collection Method by dataset:** Hybrid: Automated, Manually-Collected <br>
+**Labeling Method by dataset:** Hybrid: Manually-Labeled, Automated <br>
+**Properties:** GPQA Diamond contains 198 graduate-level multiple-choice questions written by domain experts in biology, physics, and chemistry. ToolEvalBench is a tool-calling quality benchmark that runs deterministic multi-turn agent scenarios against mock tools, scoring tool selection, parameter precision, multi-step chains, error recovery, and structured output. <br>
+
+The accuracy benchmark results are presented in the table below:
+
+| Precision | GPQA Diamond | ToolEvalBench |
+|:---|:---:|:---:|
+| BF16 | 92.2% | 89.4% |
+| **FP4-Local** | **91.5%** | **88.7%** |
+
+> Baseline: [Qwen3.8-Flash-Next](https://huggingface.co/Qwen/Qwen3.8-Flash-Next).
+
+## Quantization-Aware Distillation
+The BF16 teacher generates responses and supplies token-level probability and hidden-state targets. The trainable NVFP4 and MXFP8 weights are quantized and reconstructed for each forward pass. Gradients update the underlying weights so their low-precision representations better match the teacher.
+
+The objective combines next-token probability matching using total variation, hidden-state matching, and separate losses for end-of-response and end-of-thinking behavior. These boundary losses target the teacher's decisions to continue, finish reasoning, and stop responding.
+
+Training uses the Quatrain distillation trainer: 2,500 trunk updates followed by 1,500 joint-refinement updates, which also train the n-gram embedding tables and simulate MXFP8 shared-expert weights. Attention remains frozen, using its MXFP8 weights during joint refinement.
+
+### What's quantized
+
+| Component | Stored Precision |
+|:---|:---|
+| Decoder routed experts | QAD NVFP4; E2M1 values, E4M3 block scales, FP32 global/input scales |
+| Decoder attention | MXFP8 projection selection from the pinned source |
+| MTP routed experts | NVFP4, BF16 activations; block scales stored directly |
+| Vision attention linears | MXFP8 E4M3 values; UE8M0 scales per 32 weights; 54 matrices |
+| Vision MLP and merger linears | NVFP4 E2M1 values; E4M3 scales per 16 weights and FP32 global scale; 56 matrices; BF16 activations |
+| Other vision tensors | BF16; 223 tensors, including normalization, biases, position embeddings, and convolutions |
+| Normalization parameters | 100 FP32 tensors (97 hyperconnection and 3 PLE norms); 203 BF16 tensors; original bytes preserved |
+
+## Inference:
+**Acceleration Engine:** **vLLM** <br>
+**Test Hardware:** **NVIDIA DGX Spark and NVIDIA RTX PRO 6000 Blackwell** <br>
+
+## Usage
+
+This checkpoint requires community vLLM containers (provided in cooperation with the "Local Inference Lab" community organization) until all features are available in upstream vLLM.
+
+### NVIDIA DGX Spark
+
+vLLM container: [`eugr/spark-vllm-b12x:latest`](https://hub.docker.com/r/eugr/spark-vllm-b12x/tags?name=latest)
+
+Start the container. Set `HF_TOKEN` in the host environment first.
+
+```sh
+mkdir -p $HOME/.cache/huggingface $HOME/.cache/vllm $HOME/.cache/flashinfer $HOME/.triton $HOME/.cache/b12x
+
+docker run --rm -it \
+  --gpus all \
+  --ipc host \
+  --network host \
+  --shm-size=16g \
+  --ulimit memlock=-1 \
+  --ulimit stack=67108864 \
+  --security-opt seccomp=unconfined \
+  --security-opt apparmor=unconfined \
+  -e HF_TOKEN=$HF_TOKEN \
+  -v ~/.cache/huggingface:/root/.cache/huggingface \
+  -v $HOME/.cache/vllm:/root/.cache/vllm \
+  -v $HOME/.cache/flashinfer:/root/.cache/flashinfer \
+  -v $HOME/.triton:/root/.triton \
+  -v $HOME/.cache/b12x:/root/.cache/b12x \
+  --entrypoint /bin/bash \
+  eugr/spark-vllm-b12x:latest
+```
+
+Then launch vLLM inside the container:
+
+```sh
+export VLLM_PLE_TABLE_MEMORY="disk"
+export B12X_DISK_BACKEND="io_uring"
+export CUTE_DSL_ARCH="sm_121a"
+export SAFETENSORS_FAST_GPU="1"
+export VLLM_WORKER_MULTIPROC_METHOD="spawn"
+export VLLM_SSM_CONV_STATE_LAYOUT="DS"
+export VLLM_USE_AOT_COMPILE="1"
+export VLLM_USE_MEGA_AOT_ARTIFACT="1"
+export VLLM_USE_V2_MODEL_RUNNER="1"
+export B12X_POLICY_MODE="auto"
+export VLLM_MXFP8_LM_HEAD="1"
+export VLLM_B12X_MOE_FP4_LAYER_MAX_INPUT_SCALE="w13"
+export VLLM_B12X_MOE_FP4_FORCE_A16=1
+export B12X_W4A16_A4_PREFILL=1
+
+vllm serve nvidia/Qwen3.8-Flash-Next-FP4-Local \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --tensor-parallel-size 1 \
+    --pipeline-parallel-size 1 \
+    --mamba-cache-mode align \
+    --enable-prefix-caching \
+    --enable-chunked-prefill \
+    --dtype bfloat16 \
+    --kv-cache-dtype fp8 \
+    --quantization modelopt_mixed \
+    --block-size 16 \
+    --load-format b12x \
+    --max-model-len 262144 \
+    --max-num-seqs 8 \
+    --max-num-batched-tokens 4096 \
+    --speculative-config '{"method":"mtp","num_speculative_tokens":4}' \
+    --gdn-decode-kernel b12x \
+    --linear-backend b12x \
+    --moe-backend b12x \
+    --no-enable-flashinfer-autotune \
+    --mm-encoder-tp-mode data \
+    --reasoning-parser qwen3 \
+    --tool-call-parser qwen3_xml \
+    --enable-auto-tool-choice \
+    --compilation-config '{"pass_config":{"fuse_act_quant":true}}' \
+    --gpu-memory-utilization 0.8
+```
+
+### NVIDIA RTX PRO 6000 Blackwell
+
+vLLM container: [`ghcr.io/local-inference-lab/vllm:release-10152026`](https://github.com/local-inference-lab/blackwell-llm-docker/pkgs/container/vllm)
+
+Simplified run command:
+
+```sh
+  docker run -d --name qwen38-tp1 --init --gpus '"device=0"' \
+    --restart on-failure --network host --ipc host --shm-size 32g --stop-timeout 30 \
+    --ulimit memlock=-1 --ulimit stack=67108864:67108864 \
+    -v model-cache:/root/.cache/huggingface -v qwen38-runtime:/cache \
+    -e PROFILE=qwen38-flash-next -e PORT=8000 \
+    ghcr.io/local-inference-lab/vllm:release-10152026
+```
+
+Two step verbose commands:
+
+Start the container. Set `HF_TOKEN` in the host environment first.
+
+```sh
+docker run --rm -it \
+  --gpus all \
+  --ipc host \
+  --network host \
+  --shm-size=16g \
+  --ulimit memlock=-1 \
+  --ulimit stack=67108864 \
+  -e HF_TOKEN=$HF_TOKEN \
+  -v ~/.cache/huggingface:/root/.cache/huggingface \
+  --entrypoint /bin/bash \
+  ghcr.io/local-inference-lab/vllm:release-10152026
+```
+
+Then launch vLLM inside the container:
+
+```sh
+export B12X_DENSE_SPLITK_TURBO=1
+export B12X_DYNAMIC_DETERMINISTIC_OUTPUT=0
+export B12X_DYNAMIC_DIRECT_EXPERT_SCALES=1
+export B12X_DYNAMIC_SKIP_SPLIT_BARRIER_RESET=1
+export B12X_DYNAMIC_SPLIT_COMPUTE_MAC=224
+export B12X_DYNAMIC_SPLIT_FAST_PREPARE=1
+export B12X_DYNAMIC_SPLIT_LOW_SMEM=1
+export B12X_DYNAMIC_SPLIT_ROUTE_COMPUTE=1
+export B12X_DYNAMIC_WORK_SOURCE=persistent_grid
+export B12X_W4A16_A4_PREFILL=1
+export B12X_DISK_BACKEND=io_uring
+export B12X_MHC_PDL=1
+export B12X_PCIE_ONESHOT_BLOCK_LIMIT=4
+export B12X_PCIE_ONESHOT_PDL=1
+export B12X_PCIE_ONESHOT_THREADS=512
+export CUDA_BINARY_LOADER_THREAD_COUNT=8
+export CUDA_DEVICE_ORDER=PCI_BUS_ID
+export CUDA_MODULE_LOADING=LAZY
+export CUTE_DSL_ARCH=sm_120a
+export INSTANTTENSOR_BACKEND=BUFFERED
+export NCCL_BUFFSIZE=2097152
+export NCCL_CUMEM_ENABLE=0
+export NCCL_IB_DISABLE=1
+export NCCL_MAX_NCHANNELS=16
+export NCCL_MIN_NCHANNELS=16
+export NCCL_NET_PLUGIN=spcx
+export NCCL_P2P_LEVEL=SYS
+export NCCL_PROTO=LL,LL128,Simple
+export OMP_NUM_THREADS=2
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+export SAFETENSORS_FAST_GPU=1
+export TORCHINDUCTOR_LOOP_ORDERING_AFTER_FUSION=0
+export TORCH_CUDA_ARCH_LIST=12.0a
+export VLLM_B12X_DENSE_ACTIVATION_MODE=auto
+export VLLM_B12X_MOE_FP4_FORCE_A16=1
+export VLLM_CAUSAL_CONV1D_UPDATE_HOIST=1
+export VLLM_COMPUTE_NANS_IN_LOGITS=0
+export VLLM_DISABLED_KERNELS=MarlinFP8ScaledMMLinearKernel
+export VLLM_ENABLE_PCIE_ALLREDUCE=1
+export VLLM_GDN_SPEC_DECODE_METADATA_FASTPATH=1
+export VLLM_LM_HEAD_A16=1
+export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0
+export VLLM_MTP_NVFP4_LM_HEAD=1
+export VLLM_MXFP8_LM_HEAD=0
+export VLLM_PCIE_ALLREDUCE_BACKEND=b12x
+export VLLM_PLE_CPU_OFFLOAD=1
+export VLLM_QWEN3_8_FLASH_NEXT_HC_TP=0
+export VLLM_QWEN3_8_FLASH_NEXT_MTP_COMPACT=1
+export VLLM_QWEN3_8_FLASH_NEXT_OVERLAP=1
+export VLLM_SSM_CONV_STATE_LAYOUT=DS
+export VLLM_USE_FASTOKENS=1
+export VLLM_USE_V2_MODEL_RUNNER=1
+export VLLM_WORKER_MULTIPROC_METHOD=spawn
+export TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=1
+export TORCH_NCCL_USE_COMM_NONBLOCKING=0
+
+vllm serve nvidia/Qwen3.8-Flash-Next-FP4-Local \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --async-scheduling \
+    --block-size 64 \
+    --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE","custom_ops":["all"]}' \
+    --decode-context-parallel-size 1 \
+    --dtype bfloat16 \
+    --enable-auto-tool-choice \
+    --enable-chunked-prefill \
+    --no-enable-flashinfer-autotune \
+    --enable-prefix-caching \
+    --enable-prompt-tokens-details \
+    --gdn-decode-kernel b12x \
+    --gpu-memory-utilization 0.975 \
+    --kv-cache-dtype fp8 \
+    --no-language-model-only \
+    --linear-backend b12x \
+    --load-format instanttensor \
+    --mamba-cache-mode align \
+    --mamba-ssm-cache-dtype auto \
+    --max-cudagraph-capture-size 64 \
+    --max-model-len 262144 \
+    --max-num-batched-tokens 6019 \
+    --max-num-seqs 16 \
+    --mm-encoder-tp-mode data \
+    --mm-processor-cache-gb 0.0 \
+    --moe-backend b12x \
+    --pipeline-parallel-size 1 \
+    --quantization modelopt_mixed \
+    --reasoning-parser qwen3 \
+    --speculative-config '{"method":"mtp","draft_sample_method":"probabilistic","rejection_sample_method":"standard","moe_backend":"b12x","num_speculative_tokens":3}' \
+    --tensor-parallel-size 1 \
+    --tool-call-parser qwen3_xml
+```
+
+## Ethical Considerations
+
+NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications.  When downloaded or used in accordance with our terms of service, developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse.
+
+For more detailed information on ethical considerations for this model, please see the Model Card++ Explainability, Bias, Safety & Security, and Privacy Subcards below.
+
+Please report model quality, risk, security vulnerabilities or NVIDIA AI Concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).
+
+SUBCARDS:
+## Explainability
+
+
+| Field                                    | Response                                                                                                                       |
+|:----------------------------------------:|:-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------:|
+| **Intended Task/Domain:**                | Text generation, reasoning, summarization, and question answering.                                                                                                                                                                                                                                                                                                                                                                                       |
+| **Model Type:**                          | Text and Image-to-text transformer                                                                                                                                                                                                                                                                                                                                                                                                                       |
+| **Intended Users:**                      | This model is intended for developers, researchers, and customers building/utilizing LLMs, while balancing accuracy and efficiency.                                                                                                                                                                                                                                                                                                                     |
+| **Output:**                              | Text String(s)                                                                                                                                                                                                                                                                                                                                                                                                                                           |
+| **Describe how the model works:**        | Generates text by predicting the next word or token based on the context provided in the input sequence using multiple self-attention layers.                                                                                                                                                                                                                                                                                                            |
+| **Name the adversely impacted groups this has been tested to deliver comparable outcomes regardless of:**        | Not Applicable    |
+| **Technical Limitations & Mitigation:**  | The model was trained on data that contains toxic language and societal biases originally crawled from the internet. Therefore, the model may amplify those biases and return toxic responses especially when prompted with toxic prompts. Therefore, before deploying any applications of this model, developers should perform safety testing and tuning tailored to their specific applications of the model.                                       |
+| **Verified to have met prescribed NVIDIA quality standards:** | Yes                                                                                                                                                                                                                                                                                                                                                                                                                                |
+| **Performance Metrics:**                 | Accuracy, Throughput, and user-side throughput                                                                                                                                                                                                                                                                                                                                                                                                           |
+| **Potential Known Risk**                 | The model may generate answers that may be inaccurate, omit key information, or include irrelevant or redundant text producing socially unacceptable or undesirable text, even if the prompt itself does not include anything explicitly offensive.                                                                                                                                                                                                       |
+| **Licensing:**                           | **GOVERNING DOWNLOAD TERMS:** Use of this model is governed by [NVIDIA Open Model Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-agreement/).
+
+**ADDITIONAL INFORMATION** : The Qwen3.8-Flash-Next base model is governed by the [Qwen Community License 1.0](https://huggingface.co/Qwen/Qwen3.8-Flash-Next/blob/main/LICENSE).                                                                                                                            |
+
+
+## Bias
+
+| Field                                                                                                           | Response |
+|:---------------------------------------------------------------------------------------------------------------:|:--------:|
+| Participation considerations from adversely impacted groups [protected classes](https://www.senate.ca.gov/protected-classes) in model design and testing:       | None     |
+| Measures taken to mitigate against unwanted bias                                                                | None     |
+
+
+## Safety & Security
+
+| Field                              | Response                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                         |
+|:-----------------------------------:|:--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------:|
+| Model Application Field(s):         | Chat, Instruction Following, Chatbot Development, Code Generation, Reasoning                                                                                                                                                                                                                                                                                                                                                                                                                                                                     |
+| Describe the life critical impact (if present) | Not Applicable                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            |
+| Use Case Restrictions:              | Abide by the **GOVERNING DOWNLOAD TERMS:** Use of this model is governed by [NVIDIA Open Model Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-agreement/).
+
+**ADDITIONAL INFORMATION** : The Qwen3.8-Flash-Next base model is governed by the [Qwen Community License 1.0](https://huggingface.co/Qwen/Qwen3.8-Flash-Next/blob/main/LICENSE).                                                                                                                                                                                                                   |
+| Model and Dataset Restrictions:     | The Principle of least privilege (PoLP) is applied limiting access for dataset generation. Restrictions enforce dataset access during training, and dataset license constraints adhered to. Model checkpoints are made available on Hugging Face, and may become available on cloud providers' model catalog.                                                                                                                                                                                                                                   |
+
+
+## Privacy
+
+
+| Field                                                                          | Response                                                                           |
+| ------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------- |
+| **Generatable or reverse-engineerable personal data**                          | No 
+| **Personal data used to create this model**                                    | None Known                                                                         |      |
+| **Was consent obtained for any personal data used**                            | Not Applicable                                                                     |
+| **How often is the dataset reviewed**                                          | Before Release                                                                     |
+| **Was data from user interactions with the AI model used to train the model**  | No                                                                                 |
+| **Is there provenance for all datasets used in training**                      | Yes                                                                                |
+| **Does data labeling comply with privacy laws**                                | Yes                                                                                |
+| **Is the data compliant with data-subject requests for correction or removal** | No, not possible with externally-sourced data.                                     |
+| **Applicable NVIDIA Privacy Policy**                                           | [NVIDIA Privacy Policy](https://www.nvidia.com/en-us/about-nvidia/privacy-policy/) |

README.md

@@ -0,0 +1,161 @@
+---
+license: other
+license_name: nvidia-open-model-agreement
+license_link: LICENSE
+base_model: Qwen/Qwen3.8-Flash-Next
+base_model_relation: quantized
+pipeline_tag: image-text-to-text
+library_name: vllm
+tags:
+- qwen3.8
+- qwen3.8-flash-next
+- nvfp4
+- modelopt
+- mixture-of-experts
+- vllm
+- b12x
+- speculative-decoding
+- mtp
+- dgx-spark
+- gb10
+- blackwell
+---
+
+# Qwen3.8-Flash-Next FP4-Local for one DGX Spark
+
+NVIDIA's [Qwen3.8-Flash-Next-FP4-Local](https://huggingface.co/nvidia/Qwen3.8-Flash-Next-FP4-Local) checkpoint
+(unmodified), with a serving recipe for **a single DGX Spark** (GB10, 128 GB unified memory) on vLLM and the b12x
+kernels. The recipe is a pinned public image plus 22 patched or new vLLM / b12x files, all in `recipe/`.
+
+- **Agentic coding sessions: 63–64 tok/s decode** (54 ms per speculative step, 3.4–3.6 tokens accepted per step) and
+  0.76 s median time to first token on follow-up turns. The public 40-prompt category harness: 72.8–73.2 tok/s median.
+- **Prefill 2,130–2,160 tok/s** from 16K to 64K tokens (2,290 with the opt-in `A4_PREFILL=1`, at W4A4-level KL).
+- **KL divergence 0.0492 as served** against the BF16 model on AtomicChat's neutral set (top-1 91.7%, PPL +0.8%).
+- GSM8K (first 200) 96.5 %, HumanEval 95.7 % (greedy, thinking off); a needle is found at 29K, 96K and 125K tokens.
+
+Measured on one DGX Spark at a 0.70 GPU-memory fraction (131K context, FP8 KV cache, 4 MTP draft tokens, CUDA
+graphs). Changes over the stock launcher: 71.2 → 54.2 ms per decode step, 46.9 → 64 tok/s on the same sessions.
+
+## Fidelity
+
+Full-vocabulary KL(BF16 ‖ served model) in llama.cpp `--kl-divergence` semantics over AtomicChat's held-out neutral
+set: the original BF16 model's logits, 87 chunks of 4096 tokens, 178,089 scored positions, BF16 PPL 4.045
+([reference](https://huggingface.co/datasets/AtomicChat/Qwen3.8-Flash-Next-Uncensored-GGUF-metrics)). Every row is
+measured against that same reference.
+
+| Model | Size | Mean KLD | Top-1 | PPL ratio |
+|---|---:|---:|---:|---:|
+| **This repo, as served** (W4A16 experts) | 102 GB | **0.0492** | **91.7 %** | 1.008 |
+| Same checkpoint, vLLM default FP4 activations (W4A4) | 102 GB | 0.0665 | 90.4 % | |
+| This repo with `A4_PREFILL=1` (FP4 activations for prefill-size MoE calls) | 102 GB | 0.0665 | 90.3 % | 1.012 |
+| Vardaan97 Splash Quality-PLE8 GGUF ([card](https://huggingface.co/Vardaan97/Qwen3.8-Flash-Next-Splash-GGUF)) | 129–157 GB | 0.0126 | | 1.000 |
+| unsloth UD-Q4_K_XL + Q8 n-gram table (measured in the Splash card) | 141 GB | 0.0210 | 94.8 % | 1.005 |
+| AtomicChat AD-4.27bpw-Q4_K_M ([card](https://huggingface.co/AtomicChat/Qwen3.8-Flash-Next-GGUF)) | 95 GB | 0.0842 | 89.5 % | 1.026 |
+
+"As served" is the server's own forward pass through `recipe/patches/qwen4_exp/nvidia/logit_dump.py`, scored with
+`recipe/kl/atomic_kld.py`. NVIDIA's FP4 checkpoint is further from BF16 than the best GGUF builds; what it buys on the
+Spark is vLLM serving speed. Running the routed experts with bf16 activations (`MOE_A16`, the setting NVIDIA's own
+card recommends) is both closer to BF16 and faster than the FP4-activation default.
+
+Quality probes with thinking off and greedy decoding (`recipe/bench/quality-probe.py`):
+
+| | GSM8K-200 | HumanEval |
+|---|---:|---:|
+| This repo, as served | 96.5 % | 95.7 % |
+
+A needle planted in prompts of 28,842, 96,042 and 124,842 tokens is found at every length
+(`recipe/bench/needle-test.py`), and tool calls with `tool_choice=auto` parse (`recipe/bench/test-toolcall.py`).
+
+## Speed
+
+Test setup: one DGX Spark (GB10), `recipe/serve/serve.sh` with its defaults, no other clients, Spark idle otherwise.
+Throughput is in tokens per second.
+
+**Agentic coding sessions** (private session set, not published: multi-turn coding agents with tool definitions and
+tool calls in the history; 1,362 decode steps per run):
+
+| | Decode (weighted) | Decode (median per turn) | Time to first token (median) | Tokens per step |
+|---|---:|---:|---:|---:|
+| This release | 63.5–64.0 | 63.7–63.9 | 0.74–0.76 s | 3.40–3.56 |
+| Stock launcher (image + NVIDIA's flags) | 46.9 | | | |
+
+**Public 40-prompt category harness** (tonyd2wild's `bench_categories.py`, thinking off, temperature 0): 72.8–73.2 tok/s
+median decode, 0.28 s median time to first token, 53.5 ms per step.
+
+**Prefill** (cold, single request):
+
+| Prompt | Prefill |
+|---:|---:|
+| 2K tokens | 1,696 tok/s |
+| 16K tokens | 2,134–2,169 tok/s |
+| 64K tokens | 2,131 tok/s |
+
+`A4_PREFILL=1` runs prefill-size MoE calls with FP4 activations, as NVIDIA's card configures: 1,803 / 2,290 / 2,293 tok/s
+at 2K / 16K / 64K tokens (+6–8%), with decode unchanged, but the KL of prompt processing rises to the W4A4 level
+(table above), so it is off by default.
+
+**Synthetic prompts** (`recipe/bench/bench-sbs.py`: a random-word document plus a question, non-streaming; decode
+speed is the difference between a 400-token and a 16-token reply on a warm prefix cache). Random-word text is hard
+for the drafter, so decode here is below the agentic numbers above:
+
+| Benchmark | Prefill | Time to first token | Decode |
+|---|---:|---:|---:|
+| 2K prompt, 1 stream | 1,850 | 1.1 s | 51.7 |
+| 2K prompts, 4 streams | 2,230 aggregate | 4.1 s | 107.0 aggregate (26.7 per stream) |
+| 46K prompt, 1 stream | 2,259 | 20.3 s | 45.8 |
+| 46K prompts, 4 streams | ~2,190 aggregate (184K tokens in 84 s) | 84 s | 37.4 aggregate (9.4 per stream) |
+
+Four 46K-token contexts (184K tokens) exceed the KV cache, so vLLM preempts and recomputes; the configuration is
+sized for one long agentic session at a time.
+
+## How it fits in one Spark
+
+The FP4-Local checkpoint is 102 GB, including 51B parameters of n-gram (PLE) embedding tables. The launcher serves
+those tables from NVMe instead of memory (b12x reads them with io_uring each step, overlapped with the first decoder
+layer and kept warm with a keepalive read), so about 68 GiB of weights stay resident. vLLM gets 70% of the unified
+memory; the rest stays free for the OS and other services. The KV cache holds about 175–240K tokens depending on the
+boot, enough for one 131K-token request plus a smaller one.
+
+## What's in this repo
+
+| Files | Contents |
+|---|---|
+| `tensors/`, `config.json`, tokenizer, processor and generation files, `chat_template.jinja` | NVIDIA's `nvidia/Qwen3.8-Flash-Next-FP4-Local`, unmodified. |
+| `NVIDIA_MODEL_CARD.md` | NVIDIA's model card for those files. |
+| `recipe/` | Launcher, patches, draft vocabulary, benchmarks, KL client, raw results (`recipe/README.md`). |
+| `LICENSE`, `LICENSE.Qwen`, `NOTICE` | NVIDIA Open Model Agreement, Qwen Community License 1.0, attribution. |
+
+## Quick start
+
+You need one DGX Spark (or another GB10 system with 128 GB), Docker with the NVIDIA container toolkit, and about
+110 GB of free disk on fast NVMe.
+
+```bash
+hf download diffbot/Qwen3.8-Flash-Next-FP4-Local-DGX-Spark --local-dir ~/models/Qwen3.8-Flash-Next-FP4-Local-DGX-Spark
+cd ~/models/Qwen3.8-Flash-Next-FP4-Local-DGX-Spark/recipe
+API_KEY=change-me bash serve/serve.sh     # OpenAI-compatible API on :8000, model qwen3.8-flash-next
+```
+
+The first start compiles and autotunes the b12x kernels into `~/.cache/qwen38-flash-next-spark` and took 15.5 minutes
+from an empty cache on our Spark; later starts take about 5 minutes. Tool calling uses vLLM's `qwen3_xml` parser and reasoning the `qwen3` parser.
+
+## Limitations
+
+- Tested only on one DGX Spark with the pinned image; other GB10 systems should work but are untested.
+- b12x's MoE sums expert outputs with atomics, so greedy outputs can differ between runs. `MOE_DETERMINISTIC=1`
+  fixes this only for FP4-activation MoE.
+- `TAIL_CKPT` places the prompt checkpoint just before the prompt end; upstream vLLM's `--prefix-match-unit` does not
+  replace it on this image.
+- Up to 4 concurrent requests (`MAX_NUM_SEQS`); the configuration is tuned for single-user agentic work.
+- The quality evidence is the KL table, GSM8K-200, HumanEval and needle tests, not a full benchmark suite.
+
+## Credits and license
+
+Built on the work of NVIDIA (the FP4-Local checkpoint and Model Optimizer), the Qwen team (Qwen3.8-Flash-Next),
+local-inference-lab and the FlashInfer team (b12x), eugr (the Spark vLLM + b12x image), the vLLM team, AtomicChat
+(the KL reference) and Vardaan97 (the Splash measurements on the same reference).
+
+The weights are NVIDIA's, under the NVIDIA Open Model Agreement (`LICENSE`): "Licensed by NVIDIA Corporation under the
+NVIDIA Open Model Agreement." The base model is under the Qwen Community License 1.0 (`LICENSE.Qwen`), which requires a
+separate license from Qwen for commercial "Model as a Service" offerings. The recipe's own code is MIT; the files in
+`recipe/patches/` are Apache-2.0 (`recipe/THIRD_PARTY_NOTICES.md`).

chat_template.jinja

@@ -0,0 +1,170 @@
+{%- 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 %}
+{%- set reasoning_instructions = '' %}
+{%- if enable_thinking is undefined or enable_thinking is true %}
+    {%- set resolved_reasoning_effort = reasoning_effort|default('xhigh') %}
+    {%- if resolved_reasoning_effort not in ('xhigh', 'medium', 'low') %}
+        {{- raise_exception('Unexpected reasoning effort ' ~ reasoning_effort ~ '. Supported types are xhigh (default), medium, and low.') }}
+    {%- endif %}
+    {%- if resolved_reasoning_effort == 'xhigh' %}
+        {%- set reasoning_instructions = 'Reasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.' %}
+    {%- elif resolved_reasoning_effort == 'low' %}
+        {%- set reasoning_instructions = 'Reasoning effort is set to low. Keep your thinking brief and focused, moving directly to the conclusion without unnecessary elaboration.' %}
+    {%- endif %}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {%- if reasoning_instructions %}
+        {{- reasoning_instructions + '\n\n' }}
+    {%- endif %}
+    {{- "# 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 %}
+        {%- if content %}
+            {{- '<|im_start|>system\n' + (reasoning_instructions + '\n\n' if reasoning_instructions else '')  + content + '<|im_end|>\n' }}
+        {%- elif reasoning_instructions %}
+            {{- '<|im_start|>system\n' + reasoning_instructions + '<|im_end|>\n' }}
+        {%- endif %}
+    {%- elif reasoning_instructions %}
+        {{- '<|im_start|>system\n' + reasoning_instructions + '<|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 %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if preserve_thinking is undefined or preserve_thinking is true or 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 and tool_call.arguments != '' %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | string if args_value is string else args_value | tojson | safe %}
+                        {{- 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,2402 @@
+{
+  "architectures": [
+    "Qwen4ExpForConditionalGeneration"
+  ],
+  "image_token_id": 248056,
+  "language_model_only": false,
+  "model_type": "qwen4_exp",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 248044,
+    "dtype": "bfloat16",
+    "eos_token_id": 248044,
+    "full_attention_interval": 4,
+    "hc_count": 4,
+    "hc_lowrank": 320,
+    "head_dim": 256,
+    "heads_per_ngram": 8,
+    "hidden_act": "silu",
+    "hidden_size": 2560,
+    "indexer_budget": 2048,
+    "indexer_compress_ratio": 4,
+    "indexer_head_dim": 128,
+    "indexer_kv_heads": 1,
+    "indexer_n_heads": 4,
+    "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_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": 48,
+    "linear_value_head_dim": 128,
+    "make_ngram_vocab_size_divisible_by": 128,
+    "mamba_ssm_dtype": "float32",
+    "max_position_embeddings": 262144,
+    "model_type": "qwen4_exp_text",
+    "moe_intermediate_size": 640,
+    "mtp": {
+      "hybrid": true,
+      "layer_types": [
+        "full_attention"
+      ],
+      "mtp_use_hidden_state_from_layer": null,
+      "num_hidden_layers": 1,
+      "rope_theta": 10000000
+    },
+    "mtp_num_hidden_layers": 1,
+    "mtp_use_dedicated_embeddings": false,
+    "ngram_size": 3,
+    "ngram_vocab_size_base": 20000000,
+    "num_attention_heads": 24,
+    "num_experts": 512,
+    "num_experts_per_tok": 10,
+    "num_hidden_layers": 48,
+    "num_key_value_heads": 2,
+    "output_gate_type": "sigmoid",
+    "output_router_logits": false,
+    "pad_token_id": null,
+    "partial_rotary_factor": 0.25,
+    "ple_conv_kernel_size": 4,
+    "ple_embed_dim": 2560,
+    "ple_layer_ids": [
+      2
+    ],
+    "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": 640,
+    "split_ngram_parts": 128,
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "vocab_size": 248320,
+    "ple_embedding_dtype": "nvfp4"
+  },
+  "tie_word_embeddings": false,
+  "transformers_version": "5.8.0.dev0",
+  "video_token_id": 248057,
+  "vision_config": {
+    "deepstack_visual_indexes": [],
+    "depth": 27,
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "in_channels": 3,
+    "initializer_range": 0.02,
+    "intermediate_size": 4304,
+    "model_type": "qwen4_exp",
+    "num_heads": 16,
+    "num_position_embeddings": 2304,
+    "out_hidden_size": 2560,
+    "patch_size": 16,
+    "spatial_merge_size": 2,
+    "temporal_patch_size": 2
+  },
+  "vision_end_token_id": 248054,
+  "vision_start_token_id": 248053,
+  "quantization_config": {
+    "quant_method": "modelopt",
+    "quant_algo": "MIXED_PRECISION",
+    "kv_cache_quant_algo": null,
+    "ignore": [
+      "lm_head",
+      "mtp.fc_embedding",
+      "mtp.fc_hidden",
+      "mtp.pre_fc_norm_*",
+      "mtp.hyper_connection_mixer*",
+      "mtp.layers.0.self_attn*",
+      "mtp.layers.0.attn_hyper_connection*",
+      "mtp.layers.0.mlp_hyper_connection*",
+      "mtp.layers.0.mlp.gate",
+      "mtp.layers.0.mlp.shared_expert*"
+    ],
+    "quantized_layers": {
+      "model.language_model.layers.0.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.1.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.2.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.3.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.4.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.5.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.6.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.7.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.8.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.9.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.10.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.11.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.12.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.13.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.14.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.15.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.16.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.17.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.18.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.19.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.20.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.21.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.22.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.23.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.24.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.25.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.26.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.27.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.28.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.29.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.30.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.31.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.32.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.33.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.34.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.35.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.36.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.37.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.38.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.39.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.40.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.41.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.42.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.43.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.44.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.45.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.46.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.47.mlp.experts": {
+        "quant_algo": "NVFP4",
+        "group_size": 16,
+        "weight_scale_encoding": "csf"
+      },
+      "model.language_model.layers.0.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.0.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.0.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.1.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.1.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.1.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.2.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.2.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.2.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.3.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.3.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.3.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.4.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.4.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.4.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.5.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.5.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.5.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.6.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.6.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.6.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.7.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.7.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.7.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.8.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.8.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.8.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.9.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.9.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.9.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.10.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.10.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.10.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.11.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.11.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.11.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.12.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.12.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.12.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.13.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.13.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.13.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.14.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.14.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.14.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.15.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.15.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.15.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.16.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.16.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.16.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.17.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.17.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.17.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.18.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.18.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.18.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.19.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.19.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.19.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.20.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.20.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.20.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.21.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.21.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.21.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.22.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.22.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.22.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.23.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.23.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.23.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.24.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.24.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.24.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.25.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.25.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.25.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.26.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.26.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.26.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.27.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.27.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.27.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.28.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.28.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.28.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.29.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.29.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.29.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.30.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.30.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.30.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.31.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.31.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.31.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.32.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.32.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.32.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.33.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.33.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.33.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.34.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.34.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.34.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.35.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.35.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.35.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.36.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.36.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.36.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.37.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.37.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.37.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.38.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.38.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.38.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.39.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.39.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.39.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.40.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.40.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.40.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.41.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.41.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.41.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.42.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.42.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.42.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.43.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.43.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.43.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.44.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.44.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.44.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.45.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.45.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.45.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.46.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.46.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.46.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.47.mlp.shared_expert.gate_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.47.mlp.shared_expert.up_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.47.mlp.shared_expert.down_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.1.ple.ple_embedding.ngram_embedding": {
+        "quant_algo": "NVFP4",
+        "group_size": 16
+      },
+      "mtp.layers.0.mlp.experts": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.language_model.layers.0.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.0.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.0.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.0.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.0.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.1.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.1.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.1.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.1.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.1.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.10.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.10.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.10.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.10.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.10.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.11.self_attn.indexer.index_qk_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.11.self_attn.k_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.11.self_attn.o_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.11.self_attn.q_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.11.self_attn.v_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.12.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.12.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.12.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.12.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.12.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.13.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.13.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.13.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.13.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.13.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.14.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.14.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.14.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.14.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.14.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.15.self_attn.indexer.index_qk_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.15.self_attn.k_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.15.self_attn.o_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.15.self_attn.q_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.15.self_attn.v_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.16.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.16.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.16.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.16.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.16.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.17.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.17.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.17.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.17.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.17.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.18.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.18.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.18.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.18.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.18.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.19.self_attn.indexer.index_qk_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.19.self_attn.k_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.19.self_attn.o_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.19.self_attn.q_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.19.self_attn.v_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.2.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.2.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.2.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.2.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.2.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.20.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.20.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.20.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.20.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.20.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.21.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.21.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.21.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.21.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.21.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.22.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.22.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.22.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.22.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.22.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.23.self_attn.indexer.index_qk_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.23.self_attn.k_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.23.self_attn.o_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.23.self_attn.q_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.23.self_attn.v_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.24.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.24.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.24.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.24.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.24.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.25.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.25.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.25.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.25.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.25.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.26.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.26.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.26.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.26.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.26.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.27.self_attn.indexer.index_qk_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.27.self_attn.k_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.27.self_attn.o_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.27.self_attn.q_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.27.self_attn.v_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.28.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.28.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.28.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.28.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.28.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.29.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.29.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.29.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.29.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.29.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.3.self_attn.indexer.index_qk_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.3.self_attn.k_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.3.self_attn.o_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.3.self_attn.q_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.3.self_attn.v_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.30.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.30.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.30.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.30.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.30.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.31.self_attn.indexer.index_qk_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.31.self_attn.k_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.31.self_attn.o_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.31.self_attn.q_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.31.self_attn.v_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.32.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.32.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.32.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.32.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.32.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.33.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.33.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.33.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.33.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.33.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.34.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.34.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.34.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.34.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.34.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.35.self_attn.indexer.index_qk_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.35.self_attn.k_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.35.self_attn.o_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.35.self_attn.q_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.35.self_attn.v_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.36.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.36.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.36.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.36.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.36.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.37.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.37.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.37.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.37.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.37.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.38.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.38.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.38.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.38.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.38.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.39.self_attn.indexer.index_qk_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.39.self_attn.k_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.39.self_attn.o_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.39.self_attn.q_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.39.self_attn.v_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.4.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.4.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.4.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.4.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.4.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.40.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.40.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.40.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.40.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.40.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.41.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.41.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.41.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.41.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.41.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.42.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.42.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.42.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.42.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.42.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.43.self_attn.indexer.index_qk_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.43.self_attn.k_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.43.self_attn.o_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.43.self_attn.q_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.43.self_attn.v_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.44.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.44.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.44.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.44.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.44.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.45.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.45.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.45.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.45.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.45.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.46.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.46.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.46.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.46.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.46.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.47.self_attn.indexer.index_qk_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.47.self_attn.k_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.47.self_attn.o_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.47.self_attn.q_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.47.self_attn.v_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.5.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.5.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.5.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.5.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.5.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.6.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.6.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.6.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.6.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.6.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.7.self_attn.indexer.index_qk_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.7.self_attn.k_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.7.self_attn.o_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.7.self_attn.q_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.7.self_attn.v_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.8.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.8.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.8.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.8.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.8.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.9.linear_attn.in_proj_a": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.9.linear_attn.in_proj_b": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.9.linear_attn.in_proj_qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.9.linear_attn.in_proj_z": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.language_model.layers.9.linear_attn.out_proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "mtp.layers.48.mlp.experts": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.0.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.0.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.1.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.1.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.10.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.10.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.11.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.11.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.12.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.12.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.13.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.13.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.14.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.14.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.15.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.15.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.16.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.16.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.17.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.17.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.18.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.18.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.19.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.19.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.2.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.2.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.20.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.20.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.21.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.21.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.22.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.22.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.23.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.23.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.24.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.24.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.25.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.25.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.26.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.26.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.3.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.3.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.4.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.4.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.5.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.5.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.6.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.6.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.7.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.7.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.8.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.8.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.9.attn.proj": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.9.attn.qkv": {
+        "quant_algo": "MXFP8",
+        "group_size": 32
+      },
+      "model.visual.blocks.0.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.0.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.1.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.1.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.10.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.10.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.11.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.11.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.12.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.12.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.13.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.13.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.14.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.14.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.15.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.15.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.16.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.16.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.17.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.17.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.18.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.18.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.19.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.19.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.2.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.2.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.20.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.20.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.21.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.21.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.22.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.22.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.23.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.23.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.24.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.24.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.25.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.25.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.26.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.26.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.3.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.3.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.4.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.4.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.5.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.5.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.6.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.6.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.7.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.7.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.8.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.8.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.9.mlp.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.blocks.9.mlp.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.merger.linear_fc1": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      },
+      "model.visual.merger.linear_fc2": {
+        "quant_algo": "W4A16_NVFP4",
+        "group_size": 16
+      }
+    }
+  }
+}

generation_config.json

@@ -0,0 +1,12 @@
+{
+    "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
+}

preprocessor_config.json

@@ -0,0 +1,21 @@
+{
+    "size": {
+        "longest_edge": 16777216,
+        "shortest_edge": 65536
+    },
+    "patch_size": 16,
+    "temporal_patch_size": 2,
+    "merge_size": 2,
+    "image_mean": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "image_std": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "processor_class": "Qwen3VLProcessor",
+    "image_processor_type": "Qwen2VLImageProcessorFast"
+}

recipe/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Diffbot
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

recipe/README.md

@@ -0,0 +1,75 @@
+# Recipe: Qwen3.8-Flash-Next FP4-Local on one DGX Spark
+
+Everything needed to reproduce the published configuration: a launcher for a pinned public vLLM + b12x image, the
+files mounted over that image, the benchmark suite and the KL client.
+
+| Path | Contents |
+|---|---|
+| `serve/serve.sh` | The launcher. Every flag is documented inline and on by default; `FLAG=0` reverts that change. |
+| `serve/seccomp.json` | Docker's default seccomp profile plus `io_uring_setup/enter/register`, which b12x needs to read the n-gram (PLE) tables from disk. |
+| `serve/draft_vocab_64k.json` | 65,536 token ids for the MTP drafter's output head (see `DRAFT_VOCAB`). |
+| `patches/` | Modified and new vLLM / b12x files, mounted read-only over the image (Apache-2.0; `THIRD_PARTY_NOTICES.md`). |
+| `bench/` | The Diffbot release probe suite: tool calls, needle, GSM8K-200 + HumanEval, synthetic prefill/decode. |
+| `kl/` | `atomic_kld.py`: scores the server's prompt logits against AtomicChat's BF16 llama.cpp reference. |
+| `results/` | Raw outputs behind the numbers in the model card. |
+
+## The image
+
+`eugr/spark-vllm-b12x@sha256:03a956544e9f9b4b7ba91690469fbe7848bdbef29a7d2e0f5788368467bb3595` (2026-10-15 nightly:
+vLLM `0.1.dev21830+g8244d1c44` with the V2 model runner, flashinfer-python 0.7.1 with the experimental b12x port).
+It already contains everything NVIDIA's model card asks for; the launcher adds the flags the card's command omits
+(`--moe-backend b12x`, `--linear-backend b12x`, `VLLM_SSM_CONV_STATE_LAYOUT=DS`) and serves the 51B-parameter n-gram
+(PLE) tables from NVMe (`VLLM_PLE_TABLE_MEMORY=disk`), so about 68 GiB of weights are resident.
+
+## What the patches change
+
+Speeds are server inter-token latency per MTP verify step (4 draft tokens), measured on a private set of real
+multi-turn agentic coding sessions ("replay") and on the public 40-prompt category harness. Each step was A/B tested
+against the configuration before it, with tool calls and output checks; the numbers below are the cumulative effect.
+
+| Flag | Files | Change | Effect |
+|---|---|---|---|
+| `MXFP8_LM_HEAD` | (image) | b12x MXFP8 LM head with bf16 activations instead of a bf16 sm80 fallback | ~3% faster steps |
+| `SMALL_GEMV` | `ops/small_gemv.py`, `model.py` | split-K Triton GEMV for the MoE router at decode | |
+| `HC_FP8`, `HC_FUSE` | `ops/hc_fp8_gemv.py`, `hyperconnection.py` | rowwise-FP8 split-K GEMVs for the HyperConnection mix at decode; scaled SiLU and gate mix fused in (4 → 2 launches) | |
+| `MXFP8_GEMV` | `ops/mxfp8_gemv.py`, `vllm_kernels/linear/mxfp8/b12x.py` | split-K Triton GEMV over b12x's packed MXFP8 weights for narrow decode projections | |
+| `DRAFT_VOCAB` | `mtp.py`, `ops/mxfp8_gemv.py` | the MTP drafter scores a 64K-token subset through the target's MXFP8 LM head rows | |
+| `MTP_FP8`, `MTP_BF16_KV` | `mtp.py` | FP8 GEMVs for the drafter's bf16 projections; bf16 KV for the draft layer | |
+| `TARGET_FP8` | `model.py` | FP8 GEMVs for the target's remaining bf16 projections (PLE key/value) | |
+| `GEMMA_NORM` | `ops/gemma_norm.py` | one Triton launch per GemmaRMSNorm (stock: ~10 eager ops) | |
+| `SHARED_FUSE` | `model.py`, `ops/small_gemv.py` | shared-expert tail in one kernel at decode | |
+| `GDN_DEFER` | `gdn_defer.py`, `ops/gdn_deferred.py` | GDN verify writes per-token records; one commit after acceptance | replay 60.2 → 58.4 ms/step |
+| `PLE_OVERLAP`, `PLE_ASYNC`, `PLE_KEEPALIVE_US` | `b12x_ple.py`, `b12x/loader/*`, `disk_table.py` | n-gram table reads overlap decoder layer 0; reads enqueued as a stream host function; NVMe kept out of deep idle | ~1–2 ms/step |
+| `MOE_A16` | (image env) | routed experts with bf16 activations (W4A16) instead of the checkpoint default FP4 activations | KLD 0.0665 → 0.049, and faster |
+| `CSF_INLINE` | `b12x/moe/...` | the decode MoE kernel decodes NVFP4-CSF scales inline instead of a separate expansion pass; bit-identical | replay 57.5 → 54.2 ms/step |
+| `TAIL_CKPT` | `vllm_core/v1/core/boundary_checkpoint.py` | the prefix-cache checkpoint of a prompt is placed ≤31 tokens before its end instead of up to a whole prefill chunk earlier | agent follow-up TTFT 3.7 → 0.76 s at 4096-token chunks |
+
+From the stock launcher to this configuration: replay 71.2 → 54.2 ms per verify step, 46.9 → ~64 tok/s weighted
+decode.
+
+Measurement-only flags (off by default): `LOGIT_DUMP` (prompt logits for the KL client), `EXPERT_STATS`,
+`MOE_DETERMINISTIC` (b12x MoE sums top-k expert outputs with bf16 atomics, so outputs vary run to run; this switch makes
+FP4-activation MoE deterministic at +4% decode and +28% TTFT, and has no effect with `MOE_A16`).
+
+`--prefix-match-unit 16` (upstream vLLM) does not replace `TAIL_CKPT` here: it changes where a cache hit may land,
+not where the GDN state is saved; with it and `TAIL_CKPT=off`, the replay's median TTFT was 3.74 s.
+
+## Reproducing the numbers
+
+```bash
+cd recipe/bench && bash fetch-evaldata.sh
+export API_KEY=none SERVED_MODEL=qwen3.8-flash-next EVALDATA=$PWD/evaldata
+python3 test-toolcall.py
+for reps in 1800 6000 7800; do python3 needle-test.py $reps; done      # ~29K, ~97K, ~126K tokens
+python3 quality-probe.py                                               # GSM8K first 200 + HumanEval, greedy, thinking off
+for c in 2000 46000; do for n in 1 4; do CTX=$c CONC=$n python3 bench-sbs.py; done; done
+```
+
+KL (needs about 90 GB of disk for the reference): download `kld/base-neutral.*` from the dataset
+`AtomicChat/Qwen3.8-Flash-Next-Uncensored-GGUF-metrics` (revision `09e4a4c2`), start the server without `API_KEY` and
+with `LOGIT_DUMP=kld-fwd TAIL_CKPT=off` (one forward pass per 4096-token chunk), then:
+
+```bash
+python3 kl/atomic_kld.py http://127.0.0.1:8000 qwen3.8-flash-next kld/base-neutral.kld.00.part \
+    ~/.cache/qwen38-flash-next-spark/kld-fwd results/kl-neutral.json
+```

recipe/THIRD_PARTY_NOTICES.md

@@ -0,0 +1,50 @@
+# Third-party notices
+
+## Model weights
+
+| Component | License | Where |
+|---|---|---|
+| nvidia/Qwen3.8-Flash-Next-FP4-Local (NVIDIA ModelOpt QAD NVFP4 checkpoint), redistributed unmodified | NVIDIA Open Model Agreement | `../LICENSE`, `../NOTICE`, `../NVIDIA_MODEL_CARD.md` |
+| Qwen/Qwen3.8-Flash-Next (base model) | Qwen Community License 1.0 | `../LICENSE.Qwen` |
+
+The Qwen Community License requires a separate license from Qwen before commercial "Model as a Service" or
+"AI Work Assistant" use, and a visible model name in products above 100M monthly active users or US$20M monthly
+revenue. Read `../LICENSE.Qwen` before serving the model to third parties.
+
+## Code in `patches/` (Apache-2.0)
+
+Mounted over the files of the pinned image `eugr/spark-vllm-b12x@sha256:03a956544e9f9b4b7ba91690469fbe7848bdbef29a7d2e0f5788368467bb3595`
+(vLLM `0.1.dev21830+g8244d1c44`, flashinfer-python 0.7.1 with the experimental b12x port).
+
+Modified upstream files (originals: vLLM, Copyright the vLLM project, Apache-2.0; b12x, Copyright
+local-inference-lab and the FlashInfer contributors, Apache-2.0). Each is changed by Diffbot as described in
+`patches/README.md`:
+
+- `qwen4_exp/nvidia/model.py`, `hyperconnection.py`, `mtp.py`, `b12x_ple.py` (vLLM `models/qwen4_exp/nvidia/`)
+- `vllm_kernels/linear/mxfp8/b12x.py` (vLLM `model_executor/kernels/linear/mxfp8/b12x.py`)
+- `vllm_core/v1/core/boundary_checkpoint.py` (vLLM `v1/core/boundary_checkpoint.py`)
+- `b12x/moe/_shared/kernels/w4a16/kernel.py`, `b12x/moe/fused_moe/_impl.py`, `b12x/moe/fused_moe/_preparation.py`
+- `b12x/loader/_row_plan.h`, `b12x/loader/_ple_reader.c`, `b12x/loader/_storage.c`, `b12x/sequence/_shared/disk_table.py`
+
+New files by Diffbot, Apache-2.0 like the code they plug into:
+
+- `qwen4_exp/nvidia/gdn_defer.py`, `logit_dump.py`, `expert_stats.py`
+- `qwen4_exp/nvidia/ops/small_gemv.py`, `hc_fp8_gemv.py`, `mxfp8_gemv.py`, `nvfp4_gemv.py`, `gemma_norm.py`, `gdn_deferred.py`
+
+The inline NVFP4-CSF scale decode in `b12x/moe/_shared/kernels/w4a16/kernel.py` was written independently of,
+and does the same job as, upstream b12x PR #475 (merged 2026-10-06), which the pinned image's FlashInfer port
+does not include yet.
+
+## Other code in `recipe/` (MIT, `LICENSE`)
+
+`serve/`, `bench/` (the Diffbot release probe suite: `quality-probe.py`, `needle-test.py`, `test-toolcall.py`,
+`bench-sbs.py`, `fetch-evaldata.sh`) and `kl/`.
+
+## Data used by the benchmarks
+
+- GSM8K and HumanEval (OpenAI, MIT), downloaded by `bench/fetch-evaldata.sh`.
+- AtomicChat's llama.cpp KL reference: the original BF16 Qwen3.8-Flash-Next's logits over their held-out neutral
+  set, 87 chunks at 4096 context (`kld/base-neutral.kld.*`, dataset
+  https://huggingface.co/datasets/AtomicChat/Qwen3.8-Flash-Next-Uncensored-GGUF-metrics at revision
+  `09e4a4c2`), used by `kl/atomic_kld.py`.
+- The 40-prompt category benchmark is tonyd2wild's `bench_categories.py` (not redistributed here).

recipe/bench/bench-sbs.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python3
+"""Side-by-side on the REAL workload shape. Non-streaming (no client parse cap).
+  PREFILL: unique uncached prompt per stream, max_tokens=1 -> wall time is prefill.
+  DECODE : warm the prefix cache, then short(16)/long(400) pair -> subtraction cancels prefill."""
+import json, os, random, sys, threading, time, urllib.request
+K=os.environ["API_KEY"]
+M=os.environ.get("SERVED_MODEL","deepseek-v4.1-flash"); CTX=int(os.environ.get("CTX","46000"))
+CONC=int(os.environ.get("CONC","1")); URL=f"http://localhost:{os.environ.get('PORT','8000')}/v1/chat/completions"
+W="alpha beta gamma delta epsilon zeta eta theta iota kappa system network memory kernel buffer thread packet cluster tensor matrix".split()
+def prompt(seed):
+    r=random.Random(seed)
+    return ("Reference document:\n"+" ".join(r.choice(W) for _ in range(int(CTX*0.95)))+
+            "\n\nIgnore the document above. Write a detailed technical explanation of how speculative decoding works in LLM inference.")
+def call(p,nmax):
+    body=json.dumps({"model":M,"max_tokens":nmax,"temperature":0,"messages":[{"role":"user","content":p}]}).encode()
+    req=urllib.request.Request(URL,data=body,headers={"Content-Type":"application/json","Authorization":f"Bearer {K}"})
+    t0=time.time()
+    with urllib.request.urlopen(req,timeout=1800) as r: d=json.loads(r.read())
+    return time.time()-t0, d["usage"]["completion_tokens"], d["usage"]["prompt_tokens"]
+bar=threading.Barrier(CONC); res=[None]*CONC
+def worker(i):
+    seed=i*7919+int(time.time())
+    p=prompt(seed)
+    try:
+        bar.wait(); t_pf,_,ptok=call(p,1)          # uncached prefill
+        call(p,16)                                  # warm
+        bar.wait(); t_s,n_s,_=call(p,16); t_l,n_l,_=call(p,400)
+        dec=(n_l-n_s)/(t_l-t_s) if t_l>t_s else 0.0
+        res[i]=dict(prefill_s=t_pf, ptok=ptok, prefill_tps=ptok/t_pf, decode_tps=dec)
+    except Exception as e:
+        res[i]=dict(err=str(e)[:150])
+ts=[threading.Thread(target=worker,args=(i,)) for i in range(CONC)]
+[t.start() for t in ts]; [t.join() for t in ts]
+ok=[r for r in res if r and "err" not in r]; bad=[r for r in res if r and "err" in r]
+if bad: print(f"  ERRORS: {bad}")
+if ok:
+    print(f"  {M} ctx={CTX} conc={CONC}: prompt_tok={ok[0]['ptok']} | "
+          f"prefill {sum(o['prefill_tps'] for o in ok):8.0f} tok/s agg ({sum(o['prefill_tps'] for o in ok)/len(ok):7.0f}/stream, TTFT {max(o['prefill_s'] for o in ok):5.1f}s) | "
+          f"decode {sum(o['decode_tps'] for o in ok):6.1f} tok/s agg ({sum(o['decode_tps'] for o in ok)/len(ok):5.1f}/stream)")

recipe/bench/fetch-evaldata.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+# Download the quality-probe datasets into bench/evaldata: GSM8K test (answers reduced to the final number) and
+# HumanEval (164 problems). Both are MIT-licensed OpenAI releases.
+set -euo pipefail
+D=${EVALDATA:-$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/evaldata}; mkdir -p "$D"
+python3 - "$D" <<'PY'
+import gzip, json, sys, urllib.request
+d = sys.argv[1]
+g = urllib.request.urlopen("https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl").read().decode()
+with open(f"{d}/gsm8k-test.jsonl", "w") as f:
+    for line in g.splitlines():
+        r = json.loads(line)
+        f.write(json.dumps({"question": r["question"], "answer": r["answer"].split("####")[-1].strip().replace(",", "")}) + "\n")
+h = gzip.decompress(urllib.request.urlopen("https://raw.githubusercontent.com/openai/human-eval/master/data/HumanEval.jsonl.gz").read()).decode()
+with open(f"{d}/humaneval.jsonl", "w") as f:
+    for line in h.splitlines():
+        r = json.loads(line)
+        f.write(json.dumps({k: r[k] for k in ("task_id", "prompt", "test", "entry_point")}) + "\n")
+print("wrote", d)
+PY

recipe/bench/needle-test.py

@@ -0,0 +1,16 @@
+#!/usr/bin/env python3
+"""Long-context needle probe: ~29K-token prompt with a codename buried mid-document.
+Exercises the sparse-MLA PREFILL orchestrator (>64 tokens) + top-k decode. Prints prompt tokens, TTFT-ish latency, found?"""
+import os,re,json,time,urllib.request,sys
+key=os.environ["API_KEY"]
+H={"Content-Type":"application/json","Authorization":f"Bearer {key}"}
+reps=int(sys.argv[1]) if len(sys.argv)>1 else 1800
+filler=("The quarterly report discusses supply chains, logistics, and regional demand across several markets. ")*reps
+doc=filler[:len(filler)//2]+" The secret project codename is PELICAN-42. "+filler[len(filler)//2:]
+b={"model":os.environ.get("SERVED_MODEL","deepseek-v4.1-flash"),"messages":[{"role":"user","content":doc+"\n\nWhat is the secret project codename? Answer with just the codename."}],
+   "max_tokens":60,"temperature":0,"chat_template_kwargs":json.loads(os.environ.get("TEMPLATE_KWARGS",'{"enable_thinking": false}'))}
+t=time.time(); r=json.load(urllib.request.urlopen(urllib.request.Request(f"http://127.0.0.1:{os.environ.get('PORT','8000')}/v1/chat/completions",data=json.dumps(b).encode(),headers=H),timeout=900)); dt=time.time()-t
+m=r["choices"][0]["message"]; txt=(m.get("content") or "")+(m.get("reasoning") or m.get("reasoning_content") or "")
+ok="PELICAN-42" in txt
+print(f"prompt_tokens={r['usage']['prompt_tokens']:,} | {dt:.1f}s (~{r['usage']['prompt_tokens']/dt:,.0f} tok/s prefill) | found={ok} | {txt.strip()[:70]!r}")
+print("NEEDLE OK" if ok else "NEEDLE MISS")

recipe/bench/quality-probe.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+"""Quick quality probe for an OpenAI-compatible server: GSM8K subset + HumanEval pass@1.
+
+Greedy, thinking OFF by default (most sensitive to weight damage). Same prompts for every model.
+Env: PORT (8000), SERVED_MODEL, API_KEY (required), EVALDATA (bench/evaldata; see fetch-evaldata.sh), N_GSM (200), CONC (4),
+THINKING (0/1), MAX_TOKENS (4096), OUT (jsonl path), SKIP_HE=1 to skip HumanEval.
+Usage: API_KEY=... PORT=8000 python3 quality-probe.py
+"""
+import json, os, re, subprocess, sys, tempfile, time, urllib.request
+from concurrent.futures import ThreadPoolExecutor
+
+D = os.environ.get("EVALDATA", os.path.join(os.path.dirname(os.path.abspath(__file__)), "evaldata"))
+PORT = os.environ.get("PORT", "8000")
+MODEL = os.environ.get("SERVED_MODEL", "deepseek-v4.1-flash")
+KEY = os.environ["API_KEY"]
+N_GSM = int(os.environ.get("N_GSM", "200"))
+CONC = int(os.environ.get("CONC", "4"))
+THINK = os.environ.get("THINKING", "0") == "1"
+MAXT = int(os.environ.get("MAX_TOKENS", "4096"))
+OUT = os.environ.get("OUT", f"/tmp/quality-{MODEL.replace('/', '_')}-{int(time.time())}.jsonl")
+
+
+def chat(content):
+    body = {"model": MODEL, "messages": [{"role": "user", "content": content}],
+            "temperature": 0, "max_tokens": MAXT,
+            "chat_template_kwargs": {"enable_thinking": THINK, "thinking": THINK}}
+    req = urllib.request.Request(f"http://127.0.0.1:{PORT}/v1/chat/completions", json.dumps(body).encode(),
+                                 {"Content-Type": "application/json", "Authorization": f"Bearer {KEY}"})
+    for attempt in range(3):
+        try:
+            with urllib.request.urlopen(req, timeout=900) as r:
+                d = json.load(r)
+            m = d["choices"][0]["message"]
+            return (m.get("content") or ""), d["choices"][0].get("finish_reason"), d.get("usage", {})
+        except Exception as e:  # noqa: BLE001
+            err = str(e)
+            time.sleep(2)
+    return f"<<error {err}>>", "error", {}
+
+
+def num(s):
+    s = s.replace(",", "").replace("$", "")
+    m = re.findall(r"[Ff]inal answer[:\s]*\**\s*(-?\d+(?:\.\d+)?)", s) or re.findall(r"-?\d+(?:\.\d+)?", s)
+    return m[-1] if m else None
+
+
+def gsm_one(r):
+    c, fin, u = chat(r["question"] + "\n\nSolve step by step, then end with a line 'Final answer: <number>'.")
+    p = num(c)
+    ok = p is not None and abs(float(p) - float(r["answer"])) < 1e-6
+    return {"task": "gsm8k", "ok": ok, "pred": p, "ref": r["answer"], "finish": fin, "out_tokens": u.get("completion_tokens")}
+
+
+def he_one(r):
+    c, fin, u = chat("Complete this Python function. Reply with the complete function (including imports and the "
+                     "signature) in one ```python code block, no explanation.\n\n```python\n" + r["prompt"] + "```")
+    blocks = re.findall(r"```(?:python)?\n(.*?)```", c, re.S)
+    code = blocks[0] if blocks else c
+    if f"def {r['entry_point']}" not in code:
+        code = r["prompt"] + code
+    prog = code + "\n\n" + r["test"] + f"\n\ncheck({r['entry_point']})\n"
+    with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f:
+        f.write(prog)
+    try:
+        ok = subprocess.run([sys.executable, f.name], capture_output=True, timeout=15).returncode == 0
+    except subprocess.TimeoutExpired:
+        ok = False
+    os.unlink(f.name)
+    return {"task": "humaneval", "id": r["task_id"], "ok": ok, "finish": fin, "out_tokens": u.get("completion_tokens")}
+
+
+def run(fn, rows, name):
+    t = time.time()
+    with ThreadPoolExecutor(CONC) as ex:
+        res = list(ex.map(fn, rows))
+    acc = sum(r["ok"] for r in res) / len(res)
+    trunc = sum(r["finish"] == "length" for r in res); errs = sum(r["finish"] == "error" for r in res)
+    print(f"{MODEL} {name}: {sum(r['ok'] for r in res)}/{len(res)} = {acc:.1%}  (truncated {trunc}, errors {errs}, {time.time()-t:.0f}s)", flush=True)
+    return res
+
+
+if __name__ == "__main__":
+    gsm = [json.loads(l) for l in open(f"{D}/gsm8k-test.jsonl")][:N_GSM]
+    res = run(gsm_one, gsm, f"GSM8K first {len(gsm)}")
+    if os.environ.get("SKIP_HE") != "1":
+        he = [json.loads(l) for l in open(f"{D}/humaneval.jsonl")]
+        res += run(he_one, he, f"HumanEval {len(he)}")
+    with open(OUT, "w") as f:
+        for r in res:
+            f.write(json.dumps(r) + "\n")
+    print("wrote", OUT, "| thinking", THINK)

recipe/bench/test-toolcall.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+"""Verify tool_choice=auto + reasoning parser (what agent harnesses send)."""
+import json, os, re, urllib.request, sys
+key=os.environ["API_KEY"]
+H={"Content-Type":"application/json","Authorization":f"Bearer {key}"}
+tools=[{"type":"function","function":{"name":"web_search","description":"Search the web","parameters":{"type":"object","properties":{"query":{"type":"string"}},"required":["query"]}}}]
+def call(msgs, tool_choice="auto", n=300):
+    b={"model":os.environ.get("SERVED_MODEL","deepseek-v4.1-flash"),"messages":msgs,"tools":tools,"tool_choice":tool_choice,"max_tokens":n,"temperature":0}
+    req=urllib.request.Request(f"http://127.0.0.1:{os.environ.get('PORT','8000')}/v1/chat/completions",data=json.dumps(b).encode(),headers=H)
+    try:
+        r=json.load(urllib.request.urlopen(req,timeout=300))
+    except urllib.error.HTTPError as e:
+        print("HTTP", e.code, e.read().decode()[:300]); sys.exit(1)
+    m=r["choices"][0]["message"]
+    return m, r["choices"][0].get("finish_reason"), r["usage"]
+m,fr,u=call([{"role":"user","content":"Who won the 2026 Tour de France? Use the web_search tool."}])
+print("finish_reason:", fr, "| tool_calls:", json.dumps(m.get("tool_calls"))[:200] if m.get("tool_calls") else None)
+print("reasoning field present:", bool(m.get("reasoning") or m.get("reasoning_content")), "| content:", (m.get("content") or "")[:80].replace("\n"," "))
+m2,fr2,u2=call([{"role":"user","content":"What is 2+2? Answer briefly, no tools needed."}])
+print("no-tool case -> finish:", fr2, "| tool_calls:", bool(m2.get("tool_calls")), "| reasoning:", bool(m2.get("reasoning") or m2.get("reasoning_content")), "| content:", (m2.get("content") or "")[:60].replace("\n"," "))
+print("TOOLCALL TEST OK" if (m.get("tool_calls") and fr=="tool_calls") else "TOOLCALL TEST: no tool call emitted (check parser)")

recipe/kl/atomic_kld.py

@@ -0,0 +1,136 @@
+"""Replay a llama.cpp --kl-divergence-base file through the vLLM server (forward-side dump)
+and score it the way llama-perplexity --kl-divergence does.
+
+  atomic_kld.py BASE_URL MODEL BASE_FILE DUMP_DIR OUT.json [MAX_CHUNKS]
+
+The base file: "_logits_", n_ctx u32, n_vocab i32, n_chunk i32, tokens i32[n_ctx*n_chunk],
+then per chunk (n_ctx - 1 - n_ctx/2) rows of nv = 2*((n_vocab+1)//2) + 4 uint16: two floats
+(scale, min_log_prob) and n_vocab quantized log-probs, log p = scale*q + min_log_prob.
+Each chunk's n_ctx tokens are sent as one prompt (fresh cache_salt); positions n_ctx/2 ..
+n_ctx-2 are scored, KL summing base tokens with log p > -16 (llama.cpp's truncation).
+"""
+import json
+import os
+import sys
+import time
+import urllib.request
+
+import numpy as np
+import torch
+
+
+def post(url, body):
+    req = urllib.request.Request(url, json.dumps(body).encode(), {"Content-Type": "application/json"})
+    return json.load(urllib.request.urlopen(req, timeout=3600))
+
+
+class Base:
+    def __init__(self, paths):
+        self.paths = paths
+        self.sizes = [os.path.getsize(p) for p in paths]
+        head = self.read(0, 20)
+        assert head[:8] == b"_logits_", head[:8]
+        self.n_ctx, self.n_vocab, self.n_chunk = np.frombuffer(head[8:20], dtype=np.int32).tolist()
+        self.nv = 2 * ((self.n_vocab + 1) // 2) + 4
+        self.rows = self.n_ctx - 1 - self.n_ctx // 2
+        self.tok_off = 20
+        self.data_off = 20 + 4 * self.n_ctx * self.n_chunk
+        if self.n_chunk <= 0:  # header written before the chunk count was known
+            self.n_chunk = (sum(self.sizes) - 20) // (4 * self.n_ctx + 2 * self.nv * self.rows)
+            self.data_off = 20 + 4 * self.n_ctx * self.n_chunk
+        expect = self.data_off + 2 * self.nv * self.rows * self.n_chunk
+        assert expect == sum(self.sizes), (expect, sum(self.sizes))
+
+    def read(self, off, n):
+        out = bytearray()
+        for p, size in zip(self.paths, self.sizes):
+            if off >= size:
+                off -= size
+                continue
+            with open(p, "rb") as f:
+                f.seek(off)
+                chunk = f.read(min(n - len(out), size - off))
+            out += chunk
+            off = 0
+            if len(out) == n:
+                break
+        return bytes(out)
+
+    def tokens(self, c):
+        raw = self.read(self.tok_off + 4 * self.n_ctx * c, 4 * self.n_ctx)
+        return np.frombuffer(raw, dtype=np.int32).tolist()
+
+    def logprobs(self, c):
+        raw = self.read(self.data_off + 2 * self.nv * self.rows * c, 2 * self.nv * self.rows)
+        q = np.frombuffer(raw, dtype=np.uint16).reshape(self.rows, self.nv)
+        d = q[:, :4].copy().view(np.float32)  # [rows, 2]: scale, min_log_prob
+        lp = torch.from_numpy(q[:, 4:4 + self.n_vocab].astype(np.float32))
+        lp = lp * torch.from_numpy(d[:, :1]) + torch.from_numpy(d[:, 1:2])
+        return lp, q[:, 4:4 + self.n_vocab] == 0
+
+
+def score_chunk(base_lp, base_zero, ours, nxt):
+    lq = torch.log_softmax(ours.double(), -1)
+    keep = (base_lp > -16.0) & ~torch.from_numpy(base_zero)
+    p = torch.where(keep, base_lp.double().exp(), torch.zeros((), dtype=torch.float64))
+    kld = (p * (base_lp.double() - lq)).where(keep, torch.zeros((), dtype=torch.float64)).sum(-1)
+    same = base_lp.argmax(-1) == lq.argmax(-1)
+    nll_q = -lq.gather(1, nxt[:, None])[:, 0]
+    nll_b = -base_lp.double().gather(1, nxt[:, None])[:, 0]
+    dp = (-nll_q).exp() - (-nll_b).exp()
+    return kld, same, nll_q, nll_b, dp
+
+
+def main(base_url, model, base_file, dump_dir, out, max_chunks=None):
+    parts = sorted(p for p in [base_file] + [base_file.replace(".00.part", f".{i:02d}.part") for i in range(1, 8)]
+                   if os.path.exists(p))
+    base = Base(list(dict.fromkeys(parts)))
+    n = base.n_chunk if max_chunks is None else min(base.n_chunk, int(max_chunks))
+    print(f"base: n_ctx {base.n_ctx} vocab {base.n_vocab} chunks {base.n_chunk} (scoring {n}), {base.rows} rows/chunk", flush=True)
+    first = base.n_ctx // 2
+    acc = {k: [] for k in ("kld", "same", "nll_q", "nll_b", "dp")}
+    for c in range(n):
+        toks = base.tokens(c)
+        name = f"atomic-{c:03d}"
+        for ext in (".bin", ".idx"):
+            if os.path.exists(os.path.join(dump_dir, name + ext)):
+                os.remove(os.path.join(dump_dir, name + ext))
+        with open(os.path.join(dump_dir, "ARM"), "w") as f:
+            f.write(f"{name} {first}")
+        t0 = time.time()
+        try:
+            post(base_url + "/v1/completions", {"model": model, "prompt": toks, "max_tokens": 1,
+                                                "temperature": 0, "cache_salt": os.urandom(8).hex()})
+        finally:
+            os.remove(os.path.join(dump_dir, "ARM"))
+        idx = [tuple(map(int, l.split())) for l in open(os.path.join(dump_dir, name + ".idx"))]
+        assert idx[0][0] == first and sum(r for _, r, _ in idx) >= base.rows, idx
+        ours = torch.from_numpy(np.fromfile(os.path.join(dump_dir, name + ".bin"), dtype=np.float32)
+                                .reshape(-1, base.n_vocab)[: base.rows])
+        os.remove(os.path.join(dump_dir, name + ".bin"))
+        os.remove(os.path.join(dump_dir, name + ".idx"))
+        base_lp, base_zero = base.logprobs(c)
+        nxt = torch.tensor(toks[first + 1:first + 1 + base.rows], dtype=torch.long)
+        for k, v in zip(acc, score_chunk(base_lp, base_zero, ours, nxt)):
+            acc[k].append(v)
+        kl = torch.cat(acc["kld"])
+        print(f"chunk {c + 1}/{n} ({time.time() - t0:.1f}s): running mean KLD {kl.mean():.5f} "
+              f"± {kl.std() / kl.numel() ** 0.5:.5f}, same top-1 {100 * torch.cat(acc['same']).double().mean():.2f}%", flush=True)
+    kl, same, nq, nb, dp = (torch.cat(acc[k]) for k in acc)
+    q = torch.quantile(kl.float()[torch.randperm(kl.numel())[: min(kl.numel(), 1_000_000)]],
+                       torch.tensor([0.5, 0.9, 0.99, 0.999]))
+    res = {
+        "chunks": n, "positions": kl.numel(),
+        "mean_kld": kl.mean().item(), "kld_se": (kl.std() / kl.numel() ** 0.5).item(),
+        "median_kld": q[0].item(), "p90_kld": q[1].item(), "p99_kld": q[2].item(), "p999_kld": q[3].item(),
+        "same_top1_pct": 100 * same.double().mean().item(),
+        "ppl_q": nq.mean().exp().item(), "ppl_base": nb.mean().exp().item(),
+        "ln_ppl_ratio": (nq.mean() - nb.mean()).item(),
+        "rms_dp_pct": 100 * dp.pow(2).mean().sqrt().item(),
+    }
+    json.dump(res, open(out, "w"), indent=1)
+    print(json.dumps(res, indent=1))
+
+
+if __name__ == "__main__":
+    main(*sys.argv[1:7])

recipe/patches/b12x/loader/_ple_reader.c

@@ -0,0 +1,614 @@
+/* Included by _storage.c. File payloads are never mapped or retained between runs. */
+#include <fcntl.h>
+#include <time.h>
+#ifdef B12X_HAVE_LIBURING
+#include <liburing.h>
+#endif
+
+#include "_row_plan.h"
+#define PLE_CAPSULE "b12x.ple_reader"
+
+typedef struct ple_reader ple_reader_t;
+typedef struct {
+  char* buffer;
+  size_t job;
+} ple_slot_t;
+
+/* One stream-ordered read enqueued by ple_reader_run_async. */
+#define PLE_ASYNC_RING 8
+typedef struct {
+  ple_reader_t* reader;
+  const char* ids;
+  char *weights, *scales;
+  size_t count;
+  int pending; /* __atomic: set at enqueue, cleared by the host function */
+} ple_async_t;
+
+struct ple_reader {
+  ple_plan_t plan;
+  bool poisoned;
+  pthread_mutex_t api_mutex;
+  unsigned slots_count;
+  ple_slot_t* slots;
+  char* buffers;
+  char *weights, *scales;
+  ple_async_t async_jobs[PLE_ASYNC_RING];
+  unsigned async_next;
+  int async_failed; /* __atomic; async_failure is valid once set */
+  failure_t async_failure;
+  /* NVMe keep-alive (B12X_PLE_KEEPALIVE_US); see ple_keepalive_main. */
+  pthread_t keepalive_thread;
+  bool keepalive_started;
+  int keepalive_stop;        /* __atomic */
+  long long last_read_ns;    /* __atomic, CLOCK_MONOTONIC */
+  long keepalive_us;
+  char* keepalive_buffer;
+#ifdef B12X_HAVE_LIBURING
+  struct io_uring ring;
+  bool ring_ready;
+  bool files_registered;
+  int* registered_fds;
+  struct iovec* iovecs;
+  unsigned* free_slots;
+#endif
+};
+
+#ifdef B12X_HAVE_LIBURING
+static void ple_error(ple_reader_t* reader, const char* operation, int error) {
+  if (!reader->plan.failure.message[0])
+    snprintf(reader->plan.failure.message, sizeof(reader->plan.failure.message), "%s: %s",
+             operation, strerror(error));
+}
+
+static bool ple_scatter(ple_reader_t* reader, size_t index, char* buffer, int result) {
+  ple_job_t* job = &reader->plan.jobs[index];
+  reader->plan.read_calls++;
+  if (result > 0) reader->plan.read_bytes += (unsigned)result;
+  if (result < 0) {
+    ple_error(reader, "io_uring READ_FIXED", -result);
+    return false;
+  }
+  if ((unsigned)result != job->expected) {
+    if (!reader->plan.failure.message[0])
+      snprintf(reader->plan.failure.message, sizeof(reader->plan.failure.message),
+               "short PLE read at offset %lld: expected %u bytes, received %d; source changed or "
+               "was truncated",
+               (long long)job->offset, job->expected, result);
+    return false;
+  }
+  for (size_t i = job->begin; i < job->end; i++) {
+    ple_fragment_t* fragment = &reader->plan.fragments[i];
+    memcpy((fragment->scale ? reader->scales : reader->weights) + fragment->destination,
+           buffer + (fragment->offset - job->offset), fragment->length);
+  }
+  return true;
+}
+
+static void ple_release(ple_reader_t* reader) {
+  if (!reader) return;
+  if (reader->keepalive_started) {
+    __atomic_store_n(&reader->keepalive_stop, 1, __ATOMIC_RELEASE);
+    pthread_join(reader->keepalive_thread, NULL);
+  }
+  free(reader->keepalive_buffer);
+  if (reader->ring_ready) io_uring_queue_exit(&reader->ring);
+  free(reader->iovecs);
+  free(reader->free_slots);
+  free(reader->registered_fds);
+  for (size_t i = 0; i < reader->plan.file_count; i++) close(reader->plan.files[i].fd);
+  free(reader->plan.files);
+  free(reader->plan.sources);
+  free(reader->buffers);
+  free(reader->slots);
+  free(reader->plan.fragments);
+  free(reader->plan.jobs);
+  pthread_mutex_destroy(&reader->api_mutex);
+  free(reader);
+}
+
+static void ple_delete(PyObject* capsule) {
+  ple_reader_t* reader = PyCapsule_GetPointer(capsule, PLE_CAPSULE);
+  if (reader) {
+    Py_BEGIN_ALLOW_THREADS ple_release(reader);
+    Py_END_ALLOW_THREADS
+  }
+}
+#endif
+
+static PyObject* py_ple_reader(PyObject* self, PyObject* args) {
+  (void)self;
+  long long shard_rows, padded_rows, tp_start, tp_end, weight_bytes, scale_bytes, max_lookups;
+  int depth;
+  if (!PyArg_ParseTuple(args, "LLLLLLLi", &shard_rows, &padded_rows, &tp_start, &tp_end,
+                        &weight_bytes, &scale_bytes, &max_lookups, &depth))
+    return NULL;
+#ifndef B12X_HAVE_LIBURING
+  return PyErr_Format(PyExc_RuntimeError,
+                      "io_uring PLE support is unavailable: install liburing development headers "
+                      "and pkg-config, then rebuild the b12x native loader");
+#else
+  if (shard_rows <= 0 || padded_rows <= 0 || tp_start < 0 || tp_end < tp_start ||
+      tp_end > padded_rows || weight_bytes <= 0 || scale_bytes < 0 || max_lookups <= 0 ||
+      depth <= 0 || depth > 32768 || weight_bytes > PY_SSIZE_T_MAX / max_lookups ||
+      scale_bytes > PY_SSIZE_T_MAX / max_lookups || max_lookups > PY_SSIZE_T_MAX / 8)
+    return PyErr_Format(PyExc_ValueError, "invalid PLE reader geometry or I/O capacity");
+  /* At most ceil((row_bytes + block - 1) / block) fragments per plane. */
+  size_t per_row = (size_t)weight_bytes / PLE_BLOCK + 2;
+  if (scale_bytes) per_row += (size_t)scale_bytes / PLE_BLOCK + 2;
+  if ((size_t)max_lookups > SIZE_MAX / per_row ||
+      (size_t)max_lookups * per_row > SIZE_MAX / sizeof(ple_fragment_t) ||
+      (size_t)max_lookups * per_row > SIZE_MAX / sizeof(ple_job_t))
+    return PyErr_Format(PyExc_OverflowError, "PLE batch metadata is too large");
+  ple_reader_t* reader = calloc(1, sizeof(*reader));
+  if (!reader) return PyErr_NoMemory();
+  /* Only an initialized API mutex reaches ple_release. */
+  int error = pthread_mutex_init(&reader->api_mutex, NULL);
+  if (error) {
+    free(reader);
+    return PyErr_Format(PyExc_RuntimeError, "PLE mutex: %s", strerror(error));
+  }
+  reader->plan.shard_rows = shard_rows;
+  reader->plan.padded_rows = padded_rows;
+  reader->plan.tp_start = tp_start;
+  reader->plan.tp_end = tp_end;
+  reader->plan.weight_bytes = weight_bytes;
+  reader->plan.scale_bytes = scale_bytes;
+  reader->plan.max_lookups = max_lookups;
+  reader->plan.capacity = (size_t)max_lookups * per_row;
+  reader->slots_count = (unsigned)depth;
+  reader->slots = calloc(reader->slots_count, sizeof(*reader->slots));
+  reader->plan.fragments = calloc(reader->plan.capacity, sizeof(*reader->plan.fragments));
+  reader->plan.jobs = calloc(reader->plan.capacity, sizeof(*reader->plan.jobs));
+  error = posix_memalign((void**)&reader->buffers, PLE_BLOCK,
+                         (size_t)reader->slots_count * PLE_READ_MAX);
+  if (!reader->slots || !reader->plan.fragments || !reader->plan.jobs || error) {
+    ple_release(reader);
+    return PyErr_NoMemory();
+  }
+  for (unsigned i = 0; i < reader->slots_count; i++) {
+    reader->slots[i].buffer = reader->buffers + (size_t)i * PLE_READ_MAX;
+  }
+  error = io_uring_queue_init(reader->slots_count, &reader->ring, 0);
+  if (error < 0) {
+    ple_release(reader);
+    return PyErr_Format(
+        PyExc_RuntimeError,
+        "io_uring initialization failed: %s; enable io_uring in the kernel/container policy",
+        strerror(-error));
+  }
+  reader->ring_ready = true;
+  reader->iovecs = calloc(reader->slots_count, sizeof(*reader->iovecs));
+  reader->free_slots = calloc(reader->slots_count, sizeof(*reader->free_slots));
+  if (!reader->iovecs || !reader->free_slots) {
+    ple_release(reader);
+    return PyErr_NoMemory();
+  }
+  for (unsigned i = 0; i < reader->slots_count; i++) {
+    reader->iovecs[i].iov_base = reader->slots[i].buffer;
+    reader->iovecs[i].iov_len = PLE_READ_MAX;
+  }
+  error = io_uring_register_buffers(&reader->ring, reader->iovecs, reader->slots_count);
+  if (error < 0) {
+    ple_release(reader);
+    return PyErr_Format(PyExc_RuntimeError,
+                        "io_uring buffer registration failed: %s; raise the locked-memory limit or "
+                        "reduce queue_depth",
+                        strerror(-error));
+  }
+  PyObject* capsule = PyCapsule_New(reader, PLE_CAPSULE, ple_delete);
+  if (!capsule) ple_release(reader);
+  return capsule;
+#endif
+}
+
+static int ple_register_files(ple_plan_t* plan, void* owner) {
+  (void)plan;
+  ple_reader_t* reader = owner;
+#ifdef B12X_HAVE_LIBURING
+  int* fds = realloc(reader->registered_fds, reader->plan.file_count * sizeof(*fds));
+  if (!fds) return -ENOMEM;
+  reader->registered_fds = fds;
+  for (size_t i = 0; i < reader->plan.file_count; i++) fds[i] = reader->plan.files[i].fd;
+  if (reader->files_registered) {
+    int error = io_uring_unregister_files(&reader->ring);
+    if (error < 0) {
+      reader->poisoned = true;
+      return error;
+    }
+    reader->files_registered = false;
+  }
+  int error = io_uring_register_files(&reader->ring, fds, (unsigned)reader->plan.file_count);
+  if (!error)
+    reader->files_registered = true;
+  else
+    reader->poisoned = true;
+  return error;
+#else
+  (void)reader;
+  return 0;
+#endif
+}
+
+static PyObject* py_ple_reader_add(PyObject* self, PyObject* args) {
+  (void)self;
+  PyObject* capsule;
+  const char* path;
+  long long shard, offset;
+  int scale;
+  if (!PyArg_ParseTuple(args, "OLsLp", &capsule, &shard, &path, &offset, &scale)) return NULL;
+  ple_reader_t* reader = PyCapsule_GetPointer(capsule, PLE_CAPSULE);
+  if (!reader) return NULL;
+  failure_t failure = {{0}};
+  int error;
+  Py_BEGIN_ALLOW_THREADS pthread_mutex_lock(&reader->api_mutex);
+  reader->plan.failure.message[0] = 0;
+  error =
+      ple_add_source(&reader->plan, shard, path, offset, scale != 0, ple_register_files, reader);
+  failure = reader->plan.failure;
+  pthread_mutex_unlock(&reader->api_mutex);
+  Py_END_ALLOW_THREADS if (error == EINVAL) return PyErr_Format(
+      PyExc_ValueError, "invalid PLE source shard, offset, or scale plane");
+  if (error == EOVERFLOW)
+    return PyErr_Format(PyExc_OverflowError, "PLE source byte range overflows int64");
+  if (error)
+    return PyErr_Format(
+        PyExc_RuntimeError, "%s",
+        failure.message[0] ? failure.message : "invalid PLE source geometry or byte range");
+  Py_RETURN_NONE;
+}
+
+#ifdef B12X_HAVE_LIBURING
+/* Prepare a whole free-slot wave before entering the kernel. A short positive
+ * submit accounts for exactly that many SQEs; retry only the remaining entries.
+ * A failed submit may leave pending SQEs, so stop issuing, drain every submitted
+ * CQE, and retire the ring without submitting those pending entries. */
+static void ple_submit_wave(ple_reader_t* reader, unsigned free_count, size_t* next,
+                            unsigned* outstanding) {
+  size_t remaining = reader->plan.job_count - *next;
+  unsigned count = remaining < free_count ? (unsigned)remaining : free_count;
+  for (unsigned i = 0; i < count; i++) {
+    unsigned slot = reader->free_slots[i];
+    struct io_uring_sqe* sqe = io_uring_get_sqe(&reader->ring);
+    if (!sqe) {
+      ple_error(reader, "io_uring SQ capacity", ENOSPC);
+      reader->poisoned = true;
+      return;
+    }
+    ple_job_t* job = &reader->plan.jobs[*next];
+    reader->slots[slot].job = (*next)++;
+    io_uring_prep_read_fixed(sqe, (int)job->file, reader->slots[slot].buffer, job->length,
+                             job->offset, (int)slot);
+    sqe->flags |= IOSQE_FIXED_FILE;
+    io_uring_sqe_set_data64(sqe, slot);
+  }
+  unsigned pending = count;
+  while (pending) {
+    reader->plan.submit_calls++;
+    int result = io_uring_submit(&reader->ring);
+    if (result == -EINTR) continue;
+    if (result <= 0) {
+      ple_error(reader, "io_uring submit", result < 0 ? -result : EIO);
+      reader->poisoned = true;
+      return;
+    }
+    *outstanding += (unsigned)result;
+    pending -= (unsigned)result;
+  }
+}
+
+static void ple_uring_run(ple_reader_t* reader) {
+  unsigned outstanding = 0;
+  size_t next = 0;
+  for (unsigned i = 0; i < reader->slots_count; i++) reader->free_slots[i] = i;
+  ple_submit_wave(reader, reader->slots_count, &next, &outstanding);
+  while (outstanding) {
+    struct io_uring_cqe* cqe = NULL;
+    /* wait/peek do not submit pending SQEs when draining a failed submit. */
+    int result = io_uring_wait_cqe(&reader->ring, &cqe);
+    if (result < 0) {
+      if (result != -EINTR && result != -EAGAIN) {
+        ple_error(reader, "io_uring completion wait", -result);
+        reader->poisoned = true;
+      }
+      continue;
+    }
+    unsigned free_count = 0, head;
+    io_uring_for_each_cqe(&reader->ring, head, cqe) {
+      unsigned slot = (unsigned)io_uring_cqe_get_data64(cqe);
+      int status = cqe->res;
+      outstanding--;
+      ple_scatter(reader, reader->slots[slot].job, reader->slots[slot].buffer, status);
+      reader->free_slots[free_count++] = slot;
+      if (!outstanding) break;
+    }
+    io_uring_cq_advance(&reader->ring, free_count);
+    /* Do not prepare refill SQEs until every collected completion has been
+     * checked. An I/O error therefore leaves no speculative refill entries. */
+    if (!reader->plan.failure.message[0] && next < reader->plan.job_count)
+      ple_submit_wave(reader, free_count, &next, &outstanding);
+  }
+  if (reader->poisoned) {
+    io_uring_queue_exit(&reader->ring);
+    reader->ring_ready = false;
+  }
+}
+#endif
+
+static bool ple_overlap(const Py_buffer* a, size_t a_bytes, const Py_buffer* b, size_t b_bytes) {
+  uintptr_t x = (uintptr_t)a->buf, y = (uintptr_t)b->buf;
+  return a_bytes && b_bytes && (x <= y ? y - x < a_bytes : x - y < b_bytes);
+}
+
+/* Borrow and validate the caller's ID/destination buffers for count rows. */
+static int ple_buffers(ple_reader_t* reader, PyObject* ids_object, PyObject* weights_object,
+                       PyObject* scales_object, long long count, Py_buffer* ids,
+                       Py_buffer* weights, Py_buffer* scales) {
+  if (count < 0 || (uint64_t)count > reader->plan.max_lookups) {
+    PyErr_SetString(PyExc_ValueError, "PLE lookup count exceeds batch capacity");
+    return -1;
+  }
+  if (PyObject_GetBuffer(ids_object, ids, PyBUF_CONTIG_RO) < 0) return -1;
+  if (PyObject_GetBuffer(weights_object, weights, PyBUF_CONTIG) < 0) goto failed;
+  if (reader->plan.scale_bytes && PyObject_GetBuffer(scales_object, scales, PyBUF_CONTIG) < 0)
+    goto failed;
+  if (!reader->plan.scale_bytes && scales_object != Py_None) {
+    PyErr_SetString(PyExc_ValueError, "PLE reader has no scale plane; pass None");
+    goto failed;
+  }
+  size_t id_bytes = (size_t)count * 8;
+  size_t weight_bytes = (size_t)count * reader->plan.weight_bytes;
+  size_t scale_bytes = (size_t)count * reader->plan.scale_bytes;
+  if ((size_t)ids->len < id_bytes || (size_t)weights->len < weight_bytes ||
+      (size_t)scales->len < scale_bytes || !PyBuffer_IsContiguous(ids, 'C') ||
+      !PyBuffer_IsContiguous(weights, 'C') ||
+      (reader->plan.scale_bytes && !PyBuffer_IsContiguous(scales, 'C'))) {
+    PyErr_SetString(PyExc_ValueError,
+                    "PLE buffers must be C-contiguous and cover count rows (IDs are native signed "
+                    "int64 bytes)");
+    goto failed;
+  }
+  if (ple_overlap(ids, id_bytes, weights, weight_bytes) ||
+      ple_overlap(ids, id_bytes, scales, scale_bytes) ||
+      ple_overlap(weights, weight_bytes, scales, scale_bytes)) {
+    PyErr_SetString(PyExc_ValueError, "PLE IDs and destination byte ranges must not overlap");
+    goto failed;
+  }
+  return 0;
+failed:
+  if (ids->obj) PyBuffer_Release(ids);
+  if (weights->obj) PyBuffer_Release(weights);
+  if (scales->obj) PyBuffer_Release(scales);
+  return -1;
+}
+
+static void ple_release_buffers(Py_buffer* ids, Py_buffer* weights, Py_buffer* scales) {
+  PyBuffer_Release(ids);
+  PyBuffer_Release(weights);
+  if (scales->obj) PyBuffer_Release(scales);
+}
+
+/* Plan and read count rows; no Python API, so it also runs as a CUDA host function. */
+static failure_t ple_execute(ple_reader_t* reader, const char* ids, char* weights, char* scales,
+                             size_t count) {
+  failure_t failure = {{0}};
+  pthread_mutex_lock(&reader->api_mutex);
+  struct timespec start, planned_at, end;
+  clock_gettime(CLOCK_MONOTONIC, &start);
+  __atomic_store_n(&reader->last_read_ns, (long long)start.tv_sec * 1000000000LL + start.tv_nsec,
+                   __ATOMIC_RELEASE);
+  reader->plan.failure.message[0] = 0;
+  reader->plan.lookups = count;
+  reader->plan.requested_bytes = reader->plan.read_bytes = reader->plan.read_calls = 0;
+  reader->plan.unique_blocks = reader->plan.coalesced_reads = reader->plan.submit_calls = 0;
+  reader->plan.planning_seconds = 0;
+  if (reader->poisoned)
+    snprintf(reader->plan.failure.message, sizeof(reader->plan.failure.message),
+             "PLE io_uring reader is unusable after a submission/completion failure; create a new "
+             "reader");
+  else {
+    bool planned = ple_plan(&reader->plan, ids, weights, scales, count);
+    clock_gettime(CLOCK_MONOTONIC, &planned_at);
+    reader->plan.planning_seconds = (double)(planned_at.tv_sec - start.tv_sec) +
+                                    (double)(planned_at.tv_nsec - start.tv_nsec) * 1e-9;
+#ifdef B12X_HAVE_LIBURING
+    reader->weights = weights;
+    reader->scales = scales;
+    if (planned && reader->plan.job_count) ple_uring_run(reader);
+#else
+    (void)planned;
+#endif
+  }
+  clock_gettime(CLOCK_MONOTONIC, &end);
+  reader->plan.execution_seconds =
+      (double)(end.tv_sec - start.tv_sec) + (double)(end.tv_nsec - start.tv_nsec) * 1e-9;
+  failure = reader->plan.failure;
+  pthread_mutex_unlock(&reader->api_mutex);
+  return failure;
+}
+
+static PyObject* py_ple_reader_run(PyObject* self, PyObject* args) {
+  (void)self;
+  PyObject *capsule, *ids_object, *weights_object, *scales_object;
+  long long count;
+  if (!PyArg_ParseTuple(args, "OOOOL", &capsule, &ids_object, &weights_object, &scales_object,
+                        &count))
+    return NULL;
+  ple_reader_t* reader = PyCapsule_GetPointer(capsule, PLE_CAPSULE);
+  if (!reader) return NULL;
+  Py_buffer ids = {0}, weights = {0}, scales = {0};
+  if (ple_buffers(reader, ids_object, weights_object, scales_object, count, &ids, &weights,
+                  &scales) < 0)
+    return NULL;
+  failure_t failure;
+  Py_BEGIN_ALLOW_THREADS failure =
+      ple_execute(reader, ids.buf, weights.buf, scales.buf, (size_t)count);
+  Py_END_ALLOW_THREADS ple_release_buffers(&ids, &weights, &scales);
+  if (failure.message[0]) return PyErr_Format(PyExc_RuntimeError, "%s", failure.message);
+  Py_RETURN_NONE;
+}
+
+/* CUDA host function: the stream reaches it after the ID copy and before the
+ * decode kernel, so the CPU thread that enqueued it never waits on the read. */
+static void CUDART_CB ple_host_read(void* data) {
+  ple_async_t* job = data;
+  ple_reader_t* reader = job->reader;
+  failure_t failure = ple_execute(reader, job->ids, job->weights, job->scales, job->count);
+  if (failure.message[0] && !__atomic_load_n(&reader->async_failed, __ATOMIC_ACQUIRE)) {
+    reader->async_failure = failure;
+    __atomic_store_n(&reader->async_failed, 1, __ATOMIC_RELEASE);
+  }
+  /* Rows were copied into write-combined mapped memory; complete those stores
+   * before the driver releases the stream to the decode kernel. */
+#if defined(__aarch64__)
+  __asm__ __volatile__("dsb sy" ::: "memory");
+#else
+  __atomic_thread_fence(__ATOMIC_SEQ_CST);
+#endif
+  __atomic_store_n(&job->pending, 0, __ATOMIC_RELEASE);
+}
+
+#ifdef B12X_HAVE_LIBURING
+/* NVMe autonomous power-state transitions park the drive after tens of ms idle,
+ * and each decode step's burst of row reads (~60 ms apart) then pays the wake-up:
+ * fio on GB10 measured p99 2.2 ms per burst after 60 ms idle vs 0.44 ms with 1 ms
+ * gaps, and the GPU waits for the slowest read. While reads are recent (2 s), this
+ * thread keeps the drive busy with one aligned 4 KiB read every keepalive_us. */
+#define PLE_KEEPALIVE_WINDOW_NS 2000000000LL
+
+static long long ple_now_ns(void) {
+  struct timespec now;
+  clock_gettime(CLOCK_MONOTONIC, &now);
+  return (long long)now.tv_sec * 1000000000LL + now.tv_nsec;
+}
+
+static void* ple_keepalive_main(void* data) {
+  ple_reader_t* reader = data;
+  unsigned seed = 0x9e3779b9u;
+  while (!__atomic_load_n(&reader->keepalive_stop, __ATOMIC_ACQUIRE)) {
+    long long idle = ple_now_ns() - __atomic_load_n(&reader->last_read_ns, __ATOMIC_ACQUIRE);
+    long sleep_us = 20000;
+    if (idle < PLE_KEEPALIVE_WINDOW_NS && reader->plan.file_count) {
+      seed = seed * 1664525u + 1013904223u;
+      ple_file_t* file = &reader->plan.files[(seed >> 8) % reader->plan.file_count];
+      int64_t blocks = file->bytes / PLE_BLOCK;
+      if (blocks > 0) {
+        seed = seed * 1664525u + 1013904223u;
+        off_t offset = (off_t)((((uint64_t)seed << 16) ^ (uint64_t)idle) % (uint64_t)blocks) * PLE_BLOCK;
+        (void)!pread(file->fd, reader->keepalive_buffer, PLE_BLOCK, offset);
+      }
+      sleep_us = reader->keepalive_us;
+    }
+    struct timespec pause = {sleep_us / 1000000, (sleep_us % 1000000) * 1000};
+    nanosleep(&pause, NULL);
+  }
+  return NULL;
+}
+
+/* Start the keep-alive once rows are being read (sources are frozen by then). */
+static void ple_keepalive_maybe_start(ple_reader_t* reader) {
+  if (reader->keepalive_started || reader->keepalive_us < 0) return;
+  reader->keepalive_us = -1; /* one attempt; any failure leaves it off */
+  const char* setting = getenv("B12X_PLE_KEEPALIVE_US");
+  long interval = setting ? strtol(setting, NULL, 10) : 0;
+  if (interval <= 0) return;
+  if (posix_memalign((void**)&reader->keepalive_buffer, PLE_BLOCK, PLE_BLOCK) != 0) {
+    reader->keepalive_buffer = NULL;
+    return;
+  }
+  reader->keepalive_us = interval;
+  if (pthread_create(&reader->keepalive_thread, NULL, ple_keepalive_main, reader) != 0) {
+    reader->keepalive_us = -1;
+    return;
+  }
+  reader->keepalive_started = true;
+}
+#endif
+
+/* Enqueue the read on a CUDA stream. Buffers must stay valid until the stream
+ * passes this point; a failed read is raised by the next call. */
+static PyObject* py_ple_reader_run_async(PyObject* self, PyObject* args) {
+  (void)self;
+  PyObject *capsule, *ids_object, *weights_object, *scales_object;
+  long long count;
+  unsigned long long stream;
+  if (!PyArg_ParseTuple(args, "OOOOLK", &capsule, &ids_object, &weights_object, &scales_object,
+                        &count, &stream))
+    return NULL;
+  ple_reader_t* reader = PyCapsule_GetPointer(capsule, PLE_CAPSULE);
+  if (!reader) return NULL;
+#ifdef B12X_HAVE_LIBURING
+  ple_keepalive_maybe_start(reader);
+#endif
+  if (__atomic_load_n(&reader->async_failed, __ATOMIC_ACQUIRE)) {
+    failure_t failure = reader->async_failure;
+    __atomic_store_n(&reader->async_failed, 0, __ATOMIC_RELEASE);
+    return PyErr_Format(PyExc_RuntimeError, "asynchronous PLE read failed: %s",
+                        failure.message);
+  }
+  ple_async_t* job = &reader->async_jobs[reader->async_next % PLE_ASYNC_RING];
+  if (__atomic_load_n(&job->pending, __ATOMIC_ACQUIRE))
+    return PyErr_Format(PyExc_RuntimeError, "more than %d asynchronous PLE reads in flight",
+                        PLE_ASYNC_RING);
+  Py_buffer ids = {0}, weights = {0}, scales = {0};
+  if (ple_buffers(reader, ids_object, weights_object, scales_object, count, &ids, &weights,
+                  &scales) < 0)
+    return NULL;
+  reader->async_next++;
+  job->reader = reader;
+  job->ids = ids.buf;
+  job->weights = weights.buf;
+  job->scales = scales.buf;
+  job->count = (size_t)count;
+  __atomic_store_n(&job->pending, 1, __ATOMIC_RELEASE);
+  ple_release_buffers(&ids, &weights, &scales);
+  cudaError_t error =
+      cudaLaunchHostFunc((cudaStream_t)(uintptr_t)stream, ple_host_read, job);
+  if (error != cudaSuccess) {
+    __atomic_store_n(&job->pending, 0, __ATOMIC_RELEASE);
+    return PyErr_Format(PyExc_RuntimeError, "cudaLaunchHostFunc failed: %s",
+                        cudaGetErrorString(error));
+  }
+  Py_RETURN_NONE;
+}
+
+static PyObject* py_ple_reader_stats(PyObject* self, PyObject* capsule) {
+  (void)self;
+  ple_reader_t* reader = PyCapsule_GetPointer(capsule, PLE_CAPSULE);
+  if (!reader) return NULL;
+  uint64_t lookups, requested, bytes, calls, blocks, coalesced, submits, staging, metadata;
+  double seconds, planning;
+  Py_BEGIN_ALLOW_THREADS pthread_mutex_lock(&reader->api_mutex);
+  lookups = reader->plan.lookups;
+  requested = reader->plan.requested_bytes;
+  bytes = reader->plan.read_bytes;
+  calls = reader->plan.read_calls;
+  blocks = reader->plan.unique_blocks;
+  coalesced = reader->plan.coalesced_reads;
+  submits = reader->plan.submit_calls;
+  seconds = reader->plan.execution_seconds;
+  planning = reader->plan.planning_seconds;
+  staging = (uint64_t)reader->slots_count * PLE_READ_MAX;
+  metadata =
+      sizeof(*reader) +
+      reader->plan.capacity * (sizeof(*reader->plan.fragments) + sizeof(*reader->plan.jobs)) +
+      reader->slots_count * sizeof(*reader->slots) +
+      reader->plan.source_count * sizeof(*reader->plan.sources) +
+      reader->plan.file_count * sizeof(*reader->plan.files);
+#ifdef B12X_HAVE_LIBURING
+  metadata += reader->slots_count * (sizeof(*reader->iovecs) + sizeof(*reader->free_slots));
+  metadata += reader->plan.file_count * sizeof(*reader->registered_fds);
+  if (reader->ring_ready) {
+    metadata += reader->ring.sq.ring_sz;
+    if (reader->ring.cq.ring_ptr != reader->ring.sq.ring_ptr) metadata += reader->ring.cq.ring_sz;
+    metadata += reader->ring.sq.ring_entries * sizeof(struct io_uring_sqe);
+  }
+#endif
+  pthread_mutex_unlock(&reader->api_mutex);
+  Py_END_ALLOW_THREADS return Py_BuildValue(
+      "{s:K,s:K,s:K,s:K,s:K,s:K,s:K,s:K,s:K,s:d,s:d}", "lookups", (unsigned long long)lookups,
+      "requested_bytes", (unsigned long long)requested, "read_bytes", (unsigned long long)bytes,
+      "read_calls", (unsigned long long)calls, "unique_blocks", (unsigned long long)blocks,
+      "coalesced_reads", (unsigned long long)coalesced, "submit_calls", (unsigned long long)submits,
+      "staging_bytes", (unsigned long long)staging, "metadata_bytes", (unsigned long long)metadata,
+      "execution_seconds", seconds, "planning_seconds", planning);
+}

recipe/patches/b12x/loader/_row_plan.h

@@ -0,0 +1,314 @@
+/* Shared immutable row-source validation and allocation-free block planning. */
+#ifndef B12X_ROW_PLAN_H
+#define B12X_ROW_PLAN_H
+#include <fcntl.h>
+#include <stdlib.h>
+#include <time.h>
+#define PLE_BLOCK 4096u
+#define PLE_READ_MAX 65536u
+
+typedef struct {
+  int fd;
+  dev_t device;
+  ino_t inode;
+  int64_t bytes;
+} ple_file_t;
+
+typedef struct {
+  int64_t shard;
+  int64_t offset;
+  size_t file;
+  bool scale;
+} ple_source_t;
+
+typedef struct {
+  size_t file;
+  int64_t offset;
+  unsigned length;
+  size_t destination;
+  bool scale;
+} ple_fragment_t;
+
+typedef struct {
+  size_t file, begin, end;
+  int64_t offset;
+  unsigned length, expected;
+} ple_job_t;
+
+typedef struct {
+  int64_t shard_rows, padded_rows, tp_start, tp_end;
+  size_t weight_bytes, scale_bytes, max_lookups, capacity;
+  ple_fragment_t* fragments;
+  ple_job_t* jobs;
+  size_t fragment_count, job_count;
+  ple_file_t* files;
+  ple_source_t* sources;
+  size_t file_count, source_count;
+  failure_t failure;
+  uint64_t lookups, requested_bytes, read_bytes, read_calls;
+  uint64_t unique_blocks, coalesced_reads, submit_calls;
+  double execution_seconds, planning_seconds;
+} ple_plan_t;
+
+static inline int ple_source_compare(const void* left, const void* right) {
+  const ple_source_t *a = left, *b = right;
+  if (a->shard != b->shard) return a->shard < b->shard ? -1 : 1;
+  return (int)a->scale - (int)b->scale;
+}
+
+static inline void ple_plan_release(ple_plan_t* reader) {
+  for (size_t i = 0; i < reader->file_count; i++) close(reader->files[i].fd);
+  free(reader->files);
+  free(reader->sources);
+  free(reader->fragments);
+  free(reader->jobs);
+}
+
+static inline int ple_plan_init(ple_plan_t* reader, int64_t shard_rows, int64_t padded_rows,
+                                int64_t tp_start, int64_t tp_end, int64_t weight_bytes,
+                                int64_t scale_bytes, int64_t max_lookups) {
+  if (shard_rows <= 0 || padded_rows <= 0 || tp_start < 0 || tp_end < tp_start ||
+      tp_end > padded_rows || weight_bytes <= 0 || scale_bytes < 0 || max_lookups <= 0 ||
+      weight_bytes > PY_SSIZE_T_MAX / max_lookups || scale_bytes > PY_SSIZE_T_MAX / max_lookups ||
+      max_lookups > PY_SSIZE_T_MAX / 8)
+    return EINVAL;
+  size_t per_row = (size_t)weight_bytes / PLE_BLOCK + 2;
+  if (scale_bytes) per_row += (size_t)scale_bytes / PLE_BLOCK + 2;
+  if ((size_t)max_lookups > SIZE_MAX / per_row ||
+      (size_t)max_lookups * per_row > SIZE_MAX / sizeof(ple_fragment_t) ||
+      (size_t)max_lookups * per_row > SIZE_MAX / sizeof(ple_job_t))
+    return EOVERFLOW;
+  reader->shard_rows = shard_rows;
+  reader->padded_rows = padded_rows;
+  reader->tp_start = tp_start;
+  reader->tp_end = tp_end;
+  reader->weight_bytes = weight_bytes;
+  reader->scale_bytes = scale_bytes;
+  reader->max_lookups = max_lookups;
+  reader->capacity = (size_t)max_lookups * per_row;
+  reader->fragments = calloc(reader->capacity, sizeof(*reader->fragments));
+  reader->jobs = calloc(reader->capacity, sizeof(*reader->jobs));
+  return reader->fragments && reader->jobs ? 0 : ENOMEM;
+}
+
+static inline int ple_add_source(ple_plan_t* reader, int64_t shard, const char* path,
+                                 int64_t offset, bool scale,
+                                 int (*register_files)(ple_plan_t*, void*), void* owner) {
+  if (shard < 0 || shard > (reader->padded_rows - 1) / reader->shard_rows || offset < 0 ||
+      (scale && !reader->scale_bytes))
+    return EINVAL;
+  int64_t first = shard * reader->shard_rows;
+  int64_t rows = reader->padded_rows - first;
+  if (rows > reader->shard_rows) rows = reader->shard_rows;
+  size_t row_bytes = scale ? reader->scale_bytes : reader->weight_bytes;
+  if (rows > (INT64_MAX - offset) / (int64_t)row_bytes) return EOVERFLOW;
+  if (first >= reader->tp_end || first + rows <= reader->tp_start) return 0;
+  failure_t failure = {{0}};
+  ple_source_t key = {.shard = shard, .scale = scale};
+  if (reader->source_count &&
+      bsearch(&key, reader->sources, reader->source_count, sizeof(key), ple_source_compare)) {
+    snprintf(failure.message, sizeof(failure.message),
+             "PLE source shard/plane is already registered");
+  } else {
+    /* Reject FIFOs via fstat without blocking waiting for a writer. */
+    /* B12X_PLE_BUFFERED=1 reads through the page cache so hot n-gram rows
+     * stay resident in otherwise free RAM (evictable under memory pressure). */
+    const char* buffered = getenv("B12X_PLE_BUFFERED");
+    int direct = (buffered && buffered[0] == '1') ? 0 : O_DIRECT;
+    int fd = open(path, O_RDONLY | O_CLOEXEC | O_NONBLOCK | direct);
+    struct stat status;
+    if (fd < 0)
+      system_error(&failure, "open PLE source");
+    else if (fstat(fd, &status) != 0)
+      system_error(&failure, "fstat PLE source");
+    else if (!S_ISREG(status.st_mode) || offset + rows * (int64_t)row_bytes > status.st_size)
+      snprintf(failure.message, sizeof(failure.message),
+               "PLE source range exceeds a regular file's size");
+    else {
+      size_t file = 0;
+      while (file < reader->file_count && (reader->files[file].device != status.st_dev ||
+                                           reader->files[file].inode != status.st_ino))
+        file++;
+      ple_source_t* sources =
+          realloc(reader->sources, (reader->source_count + 1) * sizeof(*sources));
+      if (!sources)
+        snprintf(failure.message, sizeof(failure.message),
+                 "could not allocate PLE source descriptor");
+      else {
+        reader->sources = sources;
+        if (file == reader->file_count) {
+          ple_file_t* files = realloc(reader->files, (file + 1) * sizeof(*files));
+          if (!files)
+            snprintf(failure.message, sizeof(failure.message),
+                     "could not allocate PLE file descriptor");
+          else {
+            reader->files = files;
+            reader->files[file] = (ple_file_t){fd, status.st_dev, status.st_ino, status.st_size};
+            reader->file_count++;
+            fd = -1;
+            int error = register_files ? register_files(reader, owner) : 0;
+            if (error < 0) {
+              snprintf(failure.message, sizeof(failure.message), "row file registration failed: %s",
+                       strerror(-error));
+            }
+          }
+        } else if (reader->files[file].bytes != status.st_size) {
+          snprintf(failure.message, sizeof(failure.message),
+                   "PLE source file size changed during registration");
+        }
+        if (!failure.message[0]) {
+          reader->sources[reader->source_count++] = (ple_source_t){shard, offset, file, scale != 0};
+          qsort(reader->sources, reader->source_count, sizeof(*sources), ple_source_compare);
+        }
+      }
+    }
+    if (fd >= 0) close(fd);
+  }
+  if (failure.message[0] && !reader->failure.message[0]) reader->failure = failure;
+  return failure.message[0] ? EIO : 0;
+}
+
+static inline int ple_fragment_compare(const void* left, const void* right) {
+  const ple_fragment_t *a = left, *b = right;
+  if (a->file != b->file) return a->file < b->file ? -1 : 1;
+  return a->offset < b->offset ? -1 : a->offset > b->offset;
+}
+
+/* libc qsort may allocate a merge buffer on every run. Heap-sort in place to
+ * keep all batch metadata persistent, including for duplicate-heavy batches. */
+static inline void ple_sift(ple_fragment_t* items, size_t root, size_t count) {
+  ple_fragment_t value = items[root];
+  while (root < count / 2) {
+    size_t child = root * 2 + 1;
+    if (child + 1 < count && ple_fragment_compare(&items[child], &items[child + 1]) < 0) child++;
+    if (ple_fragment_compare(&value, &items[child]) >= 0) break;
+    items[root] = items[child];
+    root = child;
+  }
+  items[root] = value;
+}
+
+static inline void ple_sort(ple_plan_t* reader) {
+  ple_fragment_t* items = reader->fragments;
+  size_t count = reader->fragment_count;
+  if (count < 1024) {
+    for (size_t i = count / 2; i > 0; i--) ple_sift(items, i - 1, count);
+    for (size_t end = count; end > 1; end--) {
+      ple_fragment_t value = items[end - 1];
+      items[end - 1] = items[0];
+      items[0] = value;
+      ple_sift(items, 0, end - 1);
+    }
+    return;
+  }
+  /* Jobs are built only after sorting. Reuse that fixed-capacity allocation
+   * as radix scratch instead of allocating another prefill-sized buffer. */
+  _Static_assert(sizeof(ple_job_t) >= sizeof(ple_fragment_t),
+                 "job allocation must cover fragment sort scratch");
+  ple_fragment_t* scratch = (ple_fragment_t*)reader->jobs;
+  ple_fragment_t *source = items, *target = scratch;
+  uint64_t offset_variation = 0, file_variation = 0;
+  for (size_t i = 1; i < count; i++) {
+    offset_variation |= (uint64_t)(items[i].offset ^ items[0].offset);
+    file_variation |= (uint64_t)(items[i].file ^ items[0].file);
+  }
+  /* Stable LSD passes sort offset first, then the primary file key. Skip
+   * constant bytes, including the file key for a single-container table. */
+  for (int file = 0; file < 2; file++) {
+    uint64_t variation = file ? file_variation : offset_variation;
+    for (unsigned shift = 0; shift < 64; shift += 8) {
+      if (((variation >> shift) & 255) == 0) continue;
+      size_t positions[256] = {0};
+      for (size_t i = 0; i < count; i++) {
+        uint64_t key = file ? (uint64_t)source[i].file : (uint64_t)source[i].offset;
+        positions[(key >> shift) & 255]++;
+      }
+      size_t prefix = 0;
+      for (unsigned i = 0; i < 256; i++) {
+        size_t size = positions[i];
+        positions[i] = prefix;
+        prefix += size;
+      }
+      for (size_t i = 0; i < count; i++) {
+        uint64_t key = file ? (uint64_t)source[i].file : (uint64_t)source[i].offset;
+        target[positions[(key >> shift) & 255]++] = source[i];
+      }
+      ple_fragment_t* swap = source;
+      source = target;
+      target = swap;
+    }
+  }
+  if (source != items) memcpy(items, source, count * sizeof(*items));
+}
+
+static inline bool ple_plan(ple_plan_t* reader, const char* ids, char* weights, char* scales,
+                            size_t count) {
+  reader->fragment_count = reader->job_count = 0;
+  for (size_t i = 0; i < count; i++) {
+    int64_t id;
+    memcpy(&id, ids + i * sizeof(id), sizeof(id));
+    if (id < reader->tp_start || id >= reader->tp_end) {
+      if (weights) memset(weights + i * reader->weight_bytes, 0, reader->weight_bytes);
+      if (scales && reader->scale_bytes)
+        memset(scales + i * reader->scale_bytes, 0, reader->scale_bytes);
+      continue;
+    }
+    for (int plane = 0; plane < (reader->scale_bytes ? 2 : 1); plane++) {
+      ple_source_t key = {.shard = id / reader->shard_rows, .scale = plane != 0};
+      ple_source_t* source = reader->source_count
+                                 ? bsearch(&key, reader->sources, reader->source_count, sizeof(key),
+                                           ple_source_compare)
+                                 : NULL;
+      if (!source) {
+        snprintf(reader->failure.message, sizeof(reader->failure.message),
+                 "missing PLE %s source for shard %lld", plane ? "scale" : "weight",
+                 (long long)key.shard);
+        return false;
+      }
+      size_t bytes = plane ? reader->scale_bytes : reader->weight_bytes;
+      size_t destination = i * bytes;
+      int64_t offset = source->offset + (id % reader->shard_rows) * (int64_t)bytes;
+      reader->requested_bytes += bytes;
+      while (bytes) {
+        unsigned length = PLE_BLOCK - (uint64_t)offset % PLE_BLOCK;
+        if (bytes < length) length = (unsigned)bytes;
+        reader->fragments[reader->fragment_count++] =
+            (ple_fragment_t){source->file, offset, length, destination, plane != 0};
+        offset += length;
+        destination += length;
+        bytes -= length;
+      }
+    }
+  }
+  ple_sort(reader);
+  size_t cursor = 0;
+  while (cursor < reader->fragment_count) {
+    size_t begin = cursor;
+    ple_fragment_t* fragment = &reader->fragments[cursor];
+    size_t file = fragment->file;
+    int64_t first = fragment->offset & ~(int64_t)(PLE_BLOCK - 1);
+    int64_t last = first;
+    reader->unique_blocks++;
+    cursor++;
+    while (cursor < reader->fragment_count) {
+      fragment = &reader->fragments[cursor];
+      int64_t block = fragment->offset & ~(int64_t)(PLE_BLOCK - 1);
+      if (fragment->file != file ||
+          (block != last && (block - last != PLE_BLOCK || block - first >= PLE_READ_MAX)))
+        break;
+      if (block != last) {
+        reader->unique_blocks++;
+        last = block;
+      }
+      cursor++;
+    }
+    unsigned aligned = (unsigned)(last - first) + PLE_BLOCK;
+    int64_t available = reader->files[file].bytes - first;
+    unsigned expected = available < aligned ? (unsigned)available : aligned;
+    reader->jobs[reader->job_count++] = (ple_job_t){file, begin, cursor, first, aligned, expected};
+    if (last != first) reader->coalesced_reads++;
+  }
+  return true;
+}
+
+#endif

recipe/patches/b12x/loader/_storage.c

@@ -0,0 +1,107 @@
+#define _GNU_SOURCE
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+#include <cuda_runtime_api.h>
+#include <errno.h>
+#include <pthread.h>
+#include <stdbool.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/stat.h>
+#include <unistd.h>
+
+typedef char require_64_bit_offsets[(sizeof(off_t) == 8 && sizeof(size_t) == 8) ? 1 : -1];
+
+typedef struct {
+  char message[512];
+} failure_t;
+
+static bool cuda_ok(cudaError_t status, const char* operation, failure_t* failure) {
+  if (status == cudaSuccess) return true;
+  snprintf(failure->message, sizeof(failure->message), "%s: %s", operation,
+           cudaGetErrorString(status));
+  return false;
+}
+
+static void system_error(failure_t* failure, const char* operation) {
+  snprintf(failure->message, sizeof(failure->message), "%s: %s", operation, strerror(errno));
+}
+
+static PyObject* py_capabilities(PyObject* self, PyObject* args) {
+  (void)self;
+  int device;
+  if (!PyArg_ParseTuple(args, "i", &device)) return NULL;
+  static const struct {
+    const char* name;
+    enum cudaDeviceAttr attribute;
+  } keys[] = {
+      {"integrated", cudaDevAttrIntegrated},
+      {"can_map_host_memory", cudaDevAttrCanMapHostMemory},
+      {"managed_memory", cudaDevAttrManagedMemory},
+      {"concurrent_managed_access", cudaDevAttrConcurrentManagedAccess},
+      {"pageable_memory_access", cudaDevAttrPageableMemoryAccess},
+      {"host_page_tables", cudaDevAttrPageableMemoryAccessUsesHostPageTables},
+      {"host_register_supported", cudaDevAttrHostRegisterSupported},
+      {"registered_host_pointer", cudaDevAttrCanUseHostPointerForRegisteredMem},
+  };
+  PyObject* result = PyDict_New();
+  if (!result) return NULL;
+  for (size_t i = 0; i < sizeof(keys) / sizeof(keys[0]); i++) {
+    int value;
+    failure_t failure;
+    if (!cuda_ok(cudaDeviceGetAttribute(&value, keys[i].attribute, device),
+                 "cudaDeviceGetAttribute", &failure)) {
+      Py_DECREF(result);
+      return PyErr_Format(PyExc_RuntimeError, "%s", failure.message);
+    }
+    PyObject* number = PyLong_FromLong(value);
+    if (!number || PyDict_SetItemString(result, keys[i].name, number) != 0) {
+      Py_XDECREF(number);
+      Py_DECREF(result);
+      return NULL;
+    }
+    Py_DECREF(number);
+  }
+  return result;
+}
+
+// The batch reader uses definitions from the direct reader.
+// clang-format off
+#include "_direct.c"
+#include "_batch.c"
+#include "_ple_reader.c"
+// clang-format on
+
+static PyMethodDef methods[] = {
+    {"ple_reader", py_ple_reader, METH_VARARGS, NULL},
+    {"ple_reader_add", py_ple_reader_add, METH_VARARGS, NULL},
+    {"ple_reader_run", py_ple_reader_run, METH_VARARGS, NULL},
+    {"ple_reader_run_async", py_ple_reader_run_async, METH_VARARGS, NULL},
+    {"ple_reader_stats", py_ple_reader_stats, METH_O,
+     "Last-call counters; staging_bytes and metadata_bytes describe persistent allocations."},
+    {"batch_executor", py_batch_executor, METH_VARARGS, NULL},
+    {"batch_execute", py_batch_execute, METH_VARARGS, NULL},
+    {"batch_stats", py_batch_stats, METH_O, NULL},
+    {"direct_reader", py_direct_reader, METH_VARARGS, NULL},
+    {"direct_bytes", py_direct_bytes, METH_VARARGS, NULL},
+    {"direct_stats", py_direct_stats, METH_O, NULL},
+    {"capabilities", py_capabilities, METH_VARARGS, NULL},
+    {NULL, NULL, 0, NULL},
+};
+static PyModuleDef module = {
+    PyModuleDef_HEAD_INIT,
+    .m_name = "_b12x_loader_storage",
+    .m_size = -1,
+    .m_methods = methods,
+};
+
+PyMODINIT_FUNC PyInit__b12x_loader_storage(void) {
+  PyObject* result = PyModule_Create(&module);
+  if (result && PyModule_AddIntConstant(result, "ABI_VERSION", 1) < 0) {
+    Py_DECREF(result);
+    return NULL;
+  }
+  return result;
+}

recipe/patches/qwen4_exp/nvidia/b12x_ple.py

@@ -0,0 +1,1064 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""B12X-backed position-learning enhancement for Qwen4Exp."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import replace
+import os
+from typing import Any
+
+import torch
+from torch import nn
+
+import vllm.envs as envs
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_reduce,
+)
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.models.utils import AutoWeightsLoader
+from vllm.model_executor.weight_transfer import (
+    copy_weight,
+    flush_weight_transfers,
+    get_file_tensor_source,
+    materialize_weight,
+)
+from vllm.platforms import current_platform
+from vllm.utils.b12x import (
+    B12xPreparationUnit,
+    B12xWorkload,
+    get_b12x_ple,
+    get_b12x_ple_embedding,
+    set_b12x_preparation_provider,
+)
+from vllm.utils.torch_utils import direct_register_custom_op
+
+from ..config import Qwen4ExpTextConfig
+from .ple_shared_table import (
+    SharedTableConfig,
+    SharedTableGeometry,
+    checkpoint_identity,
+    open_shared_table,
+)
+
+logger = init_logger(__name__)
+
+_PLE_SPLITTING_OPS = (
+    "vllm::qwen4_exp_b12x_ple_embedding",
+    "vllm::qwen4_exp_b12x_ple",
+)
+
+
+
+# Patch knobs: io_uring queue depth for the disk PLE reader, and periodic
+# reader-stats logging every N prepare_disk calls (0 disables).
+_PLE_STATS_EVERY = int(os.environ.get("VLLM_PLE_STATS", "0"))
+# Overlap the disk PLE read with the first decoder layer: prepare_disk runs on a side
+# stream and the forward waits (cuStreamWaitValue32, also inside CUDA graphs) only where
+# the PLE layer consumes the embeddings. Flag protocol: the main stream clears the flag
+# in prepare_disk, the side stream sets it after the lookup kernel; forwards only wait
+# for 1 (never reset), so warm-up/capture forwards without a fresh prepare cannot hang.
+_PLE_OVERLAP = os.environ.get("VLLM_QWEN38_PLE_OVERLAP", "0") == "1"
+
+
+def _stream_value32(stream: int, ptr: int, value: int, *, wait: bool) -> None:
+    from cuda.bindings import driver as cu
+
+    if wait:
+        result = cu.cuStreamWaitValue32(
+            stream, ptr, value, cu.CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ
+        )
+    else:
+        result = cu.cuStreamWriteValue32(stream, ptr, value, 0)
+    if result[0] != cu.CUresult.CUDA_SUCCESS:
+        raise RuntimeError(f"PLE overlap stream memory op failed: {result[0]}")
+
+def _register_ple_compilation_context(
+    compilation_config: Any,
+    layer_name: str,
+    layer: nn.Module,
+) -> None:
+    """Keep request-dependent PLE transactions outside piecewise graphs.
+
+    Piecewise graph dispatch is keyed by padded token count, while PLE hashing
+    and recurrent-state routing also depend on the live request count and query
+    boundaries. These custom operators must therefore execute as partition
+    boundaries so a graph compiled for one request layout cannot replay stale
+    PLE metadata for another layout.
+    """
+    static_context = compilation_config.static_forward_context
+    if layer_name in static_context:
+        raise ValueError(f"duplicate layer name: {layer_name}")
+    static_context[layer_name] = layer
+
+    splitting_ops = compilation_config.splitting_ops
+    if splitting_ops is not None:
+        for op_name in _PLE_SPLITTING_OPS:
+            if op_name not in splitting_ops:
+                splitting_ops.append(op_name)
+
+
+def _b12x_module(name: str) -> Any:
+    api = {
+        "ple": get_b12x_ple,
+        "ple_embedding": get_b12x_ple_embedding,
+    }[name]()
+    if api is None:
+        raise ImportError(
+            f"Qwen4Exp requires b12x.sequence.{name}; install the b12x serving extra"
+        )
+    return api
+
+
+def _resolve_ple_table_memory(
+    additional_config: Any, embedding_dtype: str = "bfloat16"
+) -> str:
+    """Translate the public offload policy into a b12x storage mode."""
+    if isinstance(additional_config, dict) and "ple_table_memory" in additional_config:
+        table_memory = additional_config["ple_table_memory"]
+    else:
+        table_memory = envs.VLLM_PLE_TABLE_MEMORY
+        if table_memory is None:
+            if envs.is_set("VLLM_PLE_CPU_OFFLOAD"):
+                return "mapped_host" if envs.VLLM_PLE_CPU_OFFLOAD else "device"
+            return "io_uring" if embedding_dtype == "bfloat16" else "device"
+    if table_memory in ("ram", "shared"):
+        return "mapped_host"
+    if table_memory == "disk":
+        return "io_uring"
+    if table_memory == "device":
+        return "device"
+    raise ValueError(
+        "additional_config.ple_table_memory must be 'device', "
+        f"'ram', 'shared', or 'disk', got {table_memory!r}"
+    )
+
+
+def _resolve_ple_shared_table(vllm_config: Any) -> SharedTableConfig | None:
+    """Return the host-shared table settings when ``shared`` is selected."""
+    additional_config = vllm_config.additional_config
+    config = additional_config if isinstance(additional_config, dict) else {}
+    table_memory = config.get("ple_table_memory", envs.VLLM_PLE_TABLE_MEMORY)
+    if table_memory != "shared":
+        return None
+    model_config = vllm_config.model_config
+    identity = config.get(
+        "ple_shared_table_identity", envs.VLLM_PLE_SHARED_TABLE_IDENTITY
+    )
+    return SharedTableConfig(
+        directory=str(
+            config.get("ple_shared_table_dir", envs.VLLM_PLE_SHARED_TABLE_DIR)
+        ),
+        role=str(config.get("ple_shared_table_role", envs.VLLM_PLE_SHARED_TABLE_ROLE)),
+        lock_timeout_s=float(
+            config.get(
+                "ple_shared_table_lock_timeout_s",
+                envs.VLLM_PLE_SHARED_TABLE_LOCK_TIMEOUT_S,
+            )
+        ),
+        model_path=str(model_config.model),
+        revision=model_config.revision,
+        identity=str(identity)
+        if identity
+        else checkpoint_identity(str(model_config.model), model_config.revision),
+    )
+
+
+def _copy_embedding_shard(
+    destination: torch.Tensor,
+    loaded_weight: torch.Tensor,
+    *,
+    checkpoint_start: int,
+    tp_start: int,
+    tp_end: int,
+) -> torch.Tensor | None:
+    checkpoint_end = checkpoint_start + loaded_weight.shape[0]
+    overlap_start = max(checkpoint_start, tp_start)
+    overlap_end = min(checkpoint_end, tp_end)
+    if overlap_start >= overlap_end:
+        return None
+    rows = overlap_end - overlap_start
+    source = loaded_weight.narrow(0, overlap_start - checkpoint_start, rows)
+    target = destination.narrow(0, overlap_start - tp_start, rows)
+    with torch.no_grad():
+        if target.device.type == "cpu" and source.is_meta:
+            source = materialize_weight(source)
+        copy_weight(target, source)
+    return target
+
+
+class _NGramEmbeddingStorage(nn.Module):
+    def __init__(
+        self,
+        layout: Any,
+        shard_rows: int,
+        *,
+        shared_table: SharedTableConfig | None = None,
+        embedding_dim: int = 0,
+        dense_layer_ordinal: int = 0,
+    ) -> None:
+        super().__init__()
+        self.disk_table = None
+        self._table_storage = None
+        self._shared_table = None
+        if layout.caps.table_memory == "io_uring":
+            api = _b12x_module("ple_embedding")
+            self.disk_table = api.DiskTable(
+                layout,
+                shard_rows,
+                queue_depth=int(os.environ.get("VLLM_PLE_QUEUE_DEPTH", "64")),
+            )
+            tensors: dict[str, torch.Tensor | None] = {"weight": None}
+            for name in ("weight_scale", "weight_scale_2"):
+                shape = getattr(layout, f"{name}_shape")
+                tensors[name] = (
+                    torch.empty(
+                        shape,
+                        dtype=getattr(layout, f"{name}_dtype"),
+                        device=layout.caps.device,
+                    )
+                    if shape == (1,)
+                    else None
+                )
+        elif shared_table is not None:
+            geometry = SharedTableGeometry.from_layout(
+                layout,
+                model_path=shared_table.model_path,
+                revision=shared_table.revision,
+                identity=shared_table.identity,
+                embedding_dim=embedding_dim,
+                dense_layer_ordinal=dense_layer_ordinal,
+            )
+            self._shared_table = open_shared_table(
+                geometry, shared_table, device=layout.caps.device
+            )
+            try:
+                self._table_storage = layout.allocate_storage(
+                    host_allocator=self._shared_table.region
+                )
+            except BaseException:
+                self._shared_table.close()
+                raise
+            tensors = {
+                name: getattr(self._table_storage, name)
+                for name in ("weight", "weight_scale", "weight_scale_2")
+            }
+        else:
+            self._table_storage = layout.allocate_storage()
+            tensors = {
+                name: getattr(self._table_storage, name)
+                for name in ("weight", "weight_scale", "weight_scale_2")
+            }
+        for name, tensor in tensors.items():
+            self.register_parameter(
+                name,
+                nn.Parameter(tensor, requires_grad=False)
+                if tensor is not None
+                else None,
+            )
+
+    @property
+    def weight_load_view(self) -> torch.Tensor | None:
+        return self._table_storage.weight_load_view if self._table_storage else None
+
+    @property
+    def weight_scale_load_view(self) -> torch.Tensor | None:
+        if self._table_storage is None:
+            return self.weight_scale
+        return self._table_storage.weight_scale_load_view
+
+    @property
+    def weight_scale_2_load_view(self) -> torch.Tensor | None:
+        if self._table_storage is None:
+            return self.weight_scale_2
+        return self._table_storage.weight_scale_2_load_view
+
+    @property
+    def mapped_host_nbytes(self) -> int:
+        return int(self._table_storage.mapped_host_nbytes) if self._table_storage else 0
+
+    @property
+    def shared_table_attached(self) -> bool:
+        """Whether another process already filled every row of this table."""
+        return self._shared_table is not None and self._shared_table.attached
+
+    def abort_shared_table(self) -> None:
+        """Give up a shared table this process was populating."""
+        if self._shared_table is not None and self._shared_table.populating:
+            self._shared_table.abort()
+
+    def publish_shared_table(self) -> None:
+        """Mark a populated shared table complete for other processes."""
+        if self._shared_table is None or not self._shared_table.populating:
+            return
+        flush_weight_transfers()
+        torch.accelerator.synchronize()
+        self._shared_table.publish()
+
+
+class B12xNGramEmbedding(nn.Module):
+    """Prime-hashed learned n-gram embedding with fixed b12x storage."""
+
+    _STORAGE_MODES = {
+        "bfloat16": "bf16",
+        "float8_e4m3fn": "fp8_e4m3_per_tensor",
+        "float4_e2m1fn_x2": "nvfp4_group16",
+        "nvfp4": "nvfp4_group16",
+        "nvfp4_group16": "nvfp4_group16",
+        "uint8": "nvfp4_group16",
+    }
+
+    def __init__(
+        self,
+        config: Qwen4ExpTextConfig,
+        embedding_dim: int,
+        ple_dense_layer_id: int,
+        max_total_tokens: int,
+        max_num_reqs: int,
+        owner_prefix: str,
+        prefix: str,
+        dtype: torch.dtype,
+        table_memory: str,
+        shared_table: SharedTableConfig | None = None,
+    ) -> None:
+        super().__init__()
+        self.embedding_dim = int(embedding_dim)
+        self.max_total_tokens = int(max_total_tokens)
+        self.max_num_reqs = int(max_num_reqs)
+        self.ngram_size = int(config.ngram_size)
+        self.heads_per_ngram = int(config.heads_per_ngram)
+        self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram
+        self.requires_disk_preparation = table_memory == "io_uring"
+        self._disk_prepared_tokens = -1
+        self._disk_prepared = False
+        self.head_dim = self.embedding_dim // self.ngram_heads
+        self.eos_token_id = int(config.eos_token_id)
+        self.split_ngram_parts = int(getattr(config, "split_ngram_parts", 512))
+        self.owner_prefix = owner_prefix
+        self.embedding_storage_dtype = str(
+            getattr(config, "ple_embedding_dtype", "bfloat16")
+        )
+        if self.embedding_storage_dtype not in self._STORAGE_MODES:
+            raise NotImplementedError(
+                "Qwen4Exp PLE embedding storage dtype "
+                f"{self.embedding_storage_dtype!r} is unsupported"
+            )
+        self._quant_mode = self._STORAGE_MODES[self.embedding_storage_dtype]
+        self._embedding_load_ranges: set[tuple[int, int]] = set()
+        self._scale_load_ranges: set[tuple[int, int]] = set()
+        self._weight_scale_loaded = False
+        self._weight_scale_2_loaded = False
+        self._embedding_validated = False
+
+        device = torch.device(current_platform.current_device())
+        common_caps = {
+            "device": device,
+            "max_tokens": self.max_total_tokens,
+            "max_seqs": self.max_num_reqs,
+            "vocab_size": int(config.vocab_size),
+            "eos_token_id": self.eos_token_id,
+            "max_order": self.ngram_size,
+            "heads_per_order": self.heads_per_ngram,
+            "dense_layer_ordinal": int(ple_dense_layer_id),
+            "base_table_size": int(config.ngram_vocab_size_base),
+            "table_alignment": int(config.make_ngram_vocab_size_divisible_by),
+        }
+        api = _b12x_module("ple_embedding")
+        self._caps = api.Caps(
+            **common_caps,
+            embedding_dim=self.embedding_dim,
+            tp_size=get_tensor_model_parallel_world_size(),
+            tp_rank=get_tensor_model_parallel_rank(),
+            quant_mode=self._quant_mode,
+            table_memory=table_memory,
+            output_dtype=dtype,
+        )
+        self._geometry = api.compute_geometry(self._caps)
+        geometry_tensors = api.allocate_geometry(self._geometry, device=device)
+        self._table_layout = api.storage_layout(self._caps, geometry=self._geometry)
+        self.register_buffer("layer_multipliers", geometry_tensors.multipliers)
+        self.register_buffer("ngram_heads_offsets", geometry_tensors.table_offsets)
+        self.register_buffer("ngram_heads_vocab_sizes", geometry_tensors.prime_sizes)
+        self._plans: dict[int, object] = {}
+        shard_rows = (
+            self._table_layout.padded_vocab_size + self.split_ngram_parts - 1
+        ) // self.split_ngram_parts
+        self.ngram_embedding = _NGramEmbeddingStorage(
+            self._table_layout,
+            shard_rows,
+            shared_table=shared_table,
+            embedding_dim=self.embedding_dim,
+            dense_layer_ordinal=int(ple_dense_layer_id),
+        )
+        if self.ngram_embedding.mapped_host_nbytes:
+            logger.info(
+                "Using %.2f GiB of %s CUDA-mapped host memory for this TP rank's "
+                "PLE table",
+                self.ngram_embedding.mapped_host_nbytes / (1 << 30),
+                "shared" if shared_table is not None else "private",
+            )
+        scratch_spec = self._table_layout.scratch_specs()[0]
+        self.register_buffer(
+            "_scratch",
+            torch.empty(scratch_spec.shape, dtype=scratch_spec.dtype, device=device),
+            persistent=False,
+        )
+        self.register_buffer(
+            "_token_ids",
+            torch.empty(self.max_total_tokens, dtype=torch.int64, device=device),
+            persistent=False,
+        )
+        self.register_buffer(
+            "_query_start_loc",
+            torch.empty(self.max_num_reqs + 1, dtype=torch.int32, device=device),
+            persistent=False,
+        )
+        self.register_buffer(
+            "_committed_history",
+            torch.empty(
+                self.max_num_reqs,
+                self.ngram_size - 1,
+                dtype=torch.int64,
+                device=device,
+            ),
+            persistent=False,
+        )
+        self.register_buffer(
+            "_num_seqs",
+            torch.zeros(1, dtype=torch.int32, device=device),
+            persistent=False,
+        )
+        self.register_buffer(
+            "_num_tokens",
+            torch.zeros(1, dtype=torch.int32, device=device),
+            persistent=False,
+        )
+        self.register_buffer(
+            "_embedding_out",
+            torch.empty(
+                self._table_layout.output_shape,
+                dtype=self._table_layout.output_dtype,
+                device=device,
+            ),
+            persistent=False,
+        )
+        if not getattr(self, "b12x_preparation_suppressed", False):
+            set_b12x_preparation_provider(self, self)
+
+    def _declare_embedding_plan(self, token_count: int):
+        api = _b12x_module("ple_embedding")
+        return api.plan(
+            replace(self._caps, max_tokens=token_count),
+            geometry=self._geometry,
+            prime_sizes=self.ngram_heads_vocab_sizes,
+            table_offsets=self.ngram_heads_offsets,
+            multipliers=self.layer_multipliers,
+        )
+
+    def _capacity_for(self, token_count: int) -> int:
+        if not 0 <= token_count <= self.max_total_tokens:
+            raise ValueError("PLE embedding token count exceeds capacity")
+        if not self.requires_disk_preparation and token_count in self._plans:
+            return token_count
+        return self.max_total_tokens
+
+    def _plan_for(self, token_count: int):
+        """Resolve a live token count within the declared embedding capacity."""
+        capacity = self._capacity_for(token_count)
+        plan = self._plans.get(capacity)
+        if plan is None:
+            plan = self._declare_embedding_plan(capacity)
+            self._plans[capacity] = plan
+        return plan
+
+    def _bind_embedding(self, token_count: int):
+        capacity = self._capacity_for(token_count)
+        api = _b12x_module("ple_embedding")
+        return api.bind(
+            self._plan_for(token_count),
+            scratch=self._scratch,
+            weight=self.ngram_embedding.weight,
+            weight_scale=self.ngram_embedding.weight_scale,
+            weight_scale_2=self.ngram_embedding.weight_scale_2,
+            token_ids=self._token_ids[:capacity],
+            query_start_loc=self._query_start_loc,
+            committed_history=self._committed_history,
+            num_seqs=self._num_seqs,
+            num_tokens=self._num_tokens,
+            out=self._embedding_out[:capacity],
+            disk_table=self.ngram_embedding.disk_table,
+        )
+
+    def _bind_embedding_state(self, state: Any, token_count: int):
+        capacity = (
+            self.max_total_tokens if self.requires_disk_preparation else token_count
+        )
+        return state.bind(
+            scratch=self._scratch,
+            weight=self.ngram_embedding.weight,
+            weight_scale=self.ngram_embedding.weight_scale,
+            weight_scale_2=self.ngram_embedding.weight_scale_2,
+            token_ids=self._token_ids[:capacity],
+            query_start_loc=self._query_start_loc,
+            committed_history=self._committed_history,
+            num_seqs=self._num_seqs,
+            num_tokens=self._num_tokens,
+            out=self._embedding_out[:capacity],
+            disk_table=self.ngram_embedding.disk_table,
+        )
+
+    def _request_name(self, token_count: int) -> str:
+        return f"{self.owner_prefix}.ple_embedding.m{token_count}"
+
+    def get_b12x_preparation_units(
+        self, layer: torch.nn.Module, workload: B12xWorkload
+    ) -> tuple[B12xPreparationUnit, ...]:
+        if layer is not self:
+            raise ValueError("PLE embedding preparation owner mismatch")
+        self._validate_embedding_loaded()
+        requests = []
+        plans = self._plans
+        token_counts = (
+            (self.max_total_tokens,)
+            if self.requires_disk_preparation
+            else tuple(sorted({self.max_total_tokens, *workload.fixed_token_counts}))
+        )
+        # Plans are declared once per token count; a later collection reuses
+        # them so prepared state stays installed.
+        for token_count in token_counts:
+            plan = plans.get(token_count)
+            if plan is None:
+                plan = self._declare_embedding_plan(token_count)
+                plans[token_count] = plan
+            requests.append(
+                plan.request(
+                    name=self._request_name(token_count),
+                    prepare_call=self._embedding_prepare_call(token_count),
+                )
+            )
+        if not requests:
+            return ()
+        return (
+            B12xPreparationUnit(
+                name="PLE embedding",
+                key=self.owner_prefix,
+                requests=tuple(requests),
+                stage="weights",
+            ),
+        )
+
+    def _embedding_prepare_call(self, token_count: int):
+        def prepare(state):
+            from b12x.preparation import PreparedCall
+
+            # Prime one real row through the checkpoint-owned table and
+            # geometry.  These are the only mutable staging rows this callback
+            # touches; unlike a pool clone they are restored when the call is
+            # released and never replace checkpoint parameters.
+            capacity = (
+                self.max_total_tokens if self.requires_disk_preparation else token_count
+            )
+            token_ids = self._token_ids[:1].clone()
+            query_start = self._query_start_loc[:2].clone()
+            history = self._committed_history[:1].clone()
+            num_seqs = self._num_seqs.clone()
+            num_tokens = self._num_tokens.clone()
+            output = self._embedding_out[:1].clone()
+
+            def produce() -> None:
+                self._token_ids[0] = self.eos_token_id
+                self._query_start_loc[:2].zero_()
+                self._query_start_loc[1] = 1
+                self._committed_history[0].fill_(self.eos_token_id)
+                self._num_seqs.fill_(1)
+                self._num_tokens.fill_(1)
+
+            def restore() -> None:
+                self._token_ids[:1].copy_(token_ids)
+                self._query_start_loc[:2].copy_(query_start)
+                self._committed_history[:1].copy_(history)
+                self._num_seqs.copy_(num_seqs)
+                self._num_tokens.copy_(num_tokens)
+                self._embedding_out[:1].copy_(output)
+
+            binding = self._bind_embedding_state(state, capacity)
+            return PreparedCall(
+                run=lambda: state.run(binding, token_count=1),
+                produce=produce,
+                restore=restore,
+                owners=(binding,),
+            )
+
+        return prepare
+
+    def _prepare_inputs(
+        self,
+        input_ids: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        ngram_context: torch.Tensor,
+    ) -> int:
+        token_count = input_ids.numel()
+        num_seqs = query_start_loc.numel() - 1
+        if token_count > self.max_total_tokens or num_seqs > self.max_num_reqs:
+            raise ValueError(
+                "PLE hashing capacity exceeded: "
+                f"tokens={token_count}/{self.max_total_tokens}, "
+                f"requests={num_seqs}/{self.max_num_reqs}"
+            )
+        if tuple(ngram_context.shape) != (num_seqs, self.ngram_size - 1):
+            raise ValueError(
+                "ngram_context must have shape "
+                f"{(num_seqs, self.ngram_size - 1)}, got "
+                f"{tuple(ngram_context.shape)}"
+            )
+        self._token_ids.fill_(self.eos_token_id)
+        self._token_ids[:token_count].copy_(input_ids.reshape(-1).to(torch.int64))
+        self._query_start_loc.zero_()
+        self._query_start_loc[: num_seqs + 1].copy_(query_start_loc.to(torch.int32))
+        self._committed_history.fill_(self.eos_token_id)
+        self._committed_history[:num_seqs].copy_(ngram_context.to(torch.int64))
+        self._num_seqs.fill_(num_seqs)
+        self._num_tokens.copy_(query_start_loc[num_seqs : num_seqs + 1])
+        return token_count
+
+    def _run_embedding(
+        self,
+        input_ids: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        ngram_context: torch.Tensor,
+    ) -> None:
+        self._validate_embedding_loaded()
+        self._plan_for(input_ids.numel())
+        token_count = self._prepare_inputs(input_ids, query_start_loc, ngram_context)
+        _b12x_module("ple_embedding").run(
+            self._bind_embedding(token_count), token_count=token_count
+        )
+
+    def prepare_disk(
+        self,
+        input_ids: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        ngram_context: torch.Tensor,
+    ) -> None:
+        """Produce local embeddings before the model's graph is replayed."""
+        if not self.requires_disk_preparation:
+            raise RuntimeError("prepare_disk requires an io_uring PLE table")
+        if torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing():
+            raise RuntimeError("Disk PLE preparation must run outside CUDA graphs")
+        self._disk_prepared_tokens = -1
+        self._disk_prepared = False
+        if _PLE_OVERLAP:
+            ready = self._ple_ready_flag()
+            main = torch.cuda.current_stream()
+            side = self.__dict__.get("_ple_side_stream")
+            if side is None:
+                side = self.__dict__["_ple_side_stream"] = torch.cuda.Stream()
+            _stream_value32(main.cuda_stream, ready.data_ptr(), 0, wait=False)
+            side.wait_stream(main)
+            try:
+                with torch.cuda.stream(side):
+                    self._run_embedding(input_ids, query_start_loc, ngram_context)
+            finally:
+                # Release the forward even if the read failed; the error surfaces here.
+                _stream_value32(side.cuda_stream, ready.data_ptr(), 1, wait=False)
+        else:
+            self._run_embedding(input_ids, query_start_loc, ngram_context)
+        self._disk_prepared_tokens = input_ids.numel()
+        self._disk_prepared = True
+        if _PLE_STATS_EVERY:
+            self._ple_stats_calls = getattr(self, "_ple_stats_calls", 0) + 1
+            if self._ple_stats_calls % _PLE_STATS_EVERY == 0:
+                st = self.ngram_embedding.disk_table.stats()
+                logger.info(
+                    "PLE disk read: tokens=%d lookups=%d unique_blocks=%d read_calls=%d "
+                    "coalesced=%d submits=%d read_bytes=%d exec_ms=%.3f plan_ms=%.3f",
+                    input_ids.numel(), st["lookups"], st["unique_blocks"],
+                    st["read_calls"], st["coalesced_reads"], st["submit_calls"],
+                    st["read_bytes"], st["execution_seconds"] * 1e3,
+                    st["planning_seconds"] * 1e3,
+                )
+
+    def prepare_dummy_output(self, num_tokens: int) -> None:
+        """Initialize graph-visible profiling output without reading the table."""
+        if not self.requires_disk_preparation:
+            raise RuntimeError("Dummy disk output requires an io_uring table")
+        if torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing():
+            raise RuntimeError(
+                "Dummy disk PLE preparation must run outside CUDA graphs"
+            )
+        if not 0 <= num_tokens <= self.max_total_tokens:
+            raise ValueError("Dummy PLE output exceeds token capacity")
+        self._embedding_out[:num_tokens].zero_()
+        if _PLE_OVERLAP:
+            self._ple_ready_flag().fill_(1)
+        self._disk_prepared_tokens = num_tokens
+        self._disk_prepared = True
+
+    def _ple_ready_flag(self) -> torch.Tensor:
+        ready = self.__dict__.get("_ple_ready")
+        if ready is None:
+            ready = torch.ones(1, dtype=torch.int32, device=self._embedding_out.device)
+            self.__dict__["_ple_ready"] = ready
+            logger.info("Qwen3.8 PLE disk reads overlap the first decoder layer.")
+        return ready
+
+    def _validate_embedding_loaded(self) -> None:
+        if self._embedding_validated:
+            return
+        try:
+            _require_embedding_coverage(self)
+        except ValueError:
+            # Waiting replicas must not attach to a table this process cannot fill.
+            abort_shared_table = getattr(
+                self.ngram_embedding, "abort_shared_table", None
+            )
+            if abort_shared_table is not None:
+                abort_shared_table()
+            raise
+        self._embedding_validated = True
+        publish_shared_table = getattr(
+            self.ngram_embedding, "publish_shared_table", None
+        )
+        if publish_shared_table is not None:
+            publish_shared_table()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        ngram_context: torch.Tensor,
+    ) -> torch.Tensor:
+        input_ids = input_ids.reshape(-1)
+        if self.requires_disk_preparation:
+            # Batch counts must not become per-step Dynamo specialization guards.
+            if not self._disk_prepared or (
+                not torch.compiler.is_compiling()
+                and self._disk_prepared_tokens < input_ids.shape[0]
+            ):
+                raise RuntimeError(
+                    "Disk PLE output is not prepared; call prepare_disk before forward"
+                )
+        elif torch.compiler.is_compiling():
+            torch.ops.vllm.qwen4_exp_b12x_ple_embedding(
+                input_ids,
+                query_start_loc,
+                ngram_context,
+                self._embedding_out,
+                self.owner_prefix,
+            )
+        else:
+            self._run_embedding(input_ids, query_start_loc, ngram_context)
+        if _PLE_OVERLAP and self.requires_disk_preparation:
+            ready = self.__dict__.get("_ple_ready")
+            if ready is None:
+                if torch.cuda.is_current_stream_capturing():
+                    raise RuntimeError("PLE overlap flag must exist before graph capture")
+                ready = self._ple_ready_flag()
+            _stream_value32(
+                torch.cuda.current_stream().cuda_stream, ready.data_ptr(), 1, wait=True
+            )
+        embeddings = self._embedding_out[: input_ids.shape[0]]
+        if get_tensor_model_parallel_world_size() > 1:
+            embeddings = tensor_model_parallel_all_reduce(embeddings)
+        return embeddings
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        persistent_buffers = {
+            "layer_multipliers": self.layer_multipliers,
+            "ngram_heads_offsets": self.ngram_heads_offsets,
+            "ngram_heads_vocab_sizes": self.ngram_heads_vocab_sizes,
+        }
+        loaded: set[str] = set()
+        regular_weights: list[tuple[str, torch.Tensor]] = []
+        shard_prefix = "ngram_embedding.shard_"
+        embedding = self.ngram_embedding
+        org_vocab_size = self._table_layout.padded_vocab_size
+        tp_start = self._table_layout.shard_start
+        tp_end = self._table_layout.shard_end
+        shard_size = (
+            org_vocab_size + self.split_ngram_parts - 1
+        ) // self.split_ngram_parts
+
+        for name, loaded_weight in weights:
+            leaf_name = name.rsplit(".", 1)[-1]
+            if leaf_name.startswith("hashstats_") or leaf_name == "token_lookup":
+                continue
+            if name in persistent_buffers:
+                buffer = persistent_buffers[name]
+                if tuple(buffer.shape) != tuple(loaded_weight.shape):
+                    raise ValueError(
+                        f"shape mismatch for {name}: expected {tuple(buffer.shape)}, "
+                        f"got {tuple(loaded_weight.shape)}"
+                    )
+                checkpoint_value = loaded_weight.to(
+                    device=buffer.device, dtype=buffer.dtype
+                )
+                if not torch.equal(buffer, checkpoint_value):
+                    raise ValueError(
+                        f"checkpoint {name} does not match planned PLE geometry"
+                    )
+                loaded.add(name)
+                continue
+            if name == "ngram_embedding.weight_scale":
+                if self._quant_mode != "fp8_e4m3_per_tensor":
+                    regular_weights.append((name, loaded_weight))
+                    continue
+                expected_shape = self._table_layout.weight_scale_shape
+                if tuple(loaded_weight.shape) != expected_shape:
+                    raise ValueError(
+                        "shape mismatch for PLE embedding scale: expected "
+                        f"{expected_shape}, got {tuple(loaded_weight.shape)}"
+                    )
+                if loaded_weight.dtype != self._table_layout.weight_scale_dtype:
+                    raise TypeError(
+                        "PLE embedding weight_scale must have dtype "
+                        f"{self._table_layout.weight_scale_dtype}, "
+                        f"got {loaded_weight.dtype}"
+                    )
+                with torch.no_grad():
+                    target = getattr(
+                        embedding, "weight_scale_load_view", embedding.weight_scale
+                    )
+                    assert target is not None
+                    target.copy_(
+                        loaded_weight.to(
+                            device=target.device,
+                            dtype=target.dtype,
+                        )
+                    )
+                self._weight_scale_loaded = True
+                self._embedding_validated = False
+                loaded.add(name)
+                continue
+            if name == "ngram_embedding.weight_scale_2":
+                if self._quant_mode != "nvfp4_group16":
+                    regular_weights.append((name, loaded_weight))
+                    continue
+                expected_shape = self._table_layout.weight_scale_2_shape
+                if loaded_weight.ndim == 0 and expected_shape == (1,):
+                    loaded_weight = loaded_weight.reshape(1)
+                if tuple(loaded_weight.shape) != expected_shape:
+                    raise ValueError(
+                        "shape mismatch for PLE embedding weight_scale_2: expected "
+                        f"{expected_shape}, got {tuple(loaded_weight.shape)}"
+                    )
+                if loaded_weight.dtype != self._table_layout.weight_scale_2_dtype:
+                    raise TypeError(
+                        "PLE embedding weight_scale_2 must have dtype "
+                        f"{self._table_layout.weight_scale_2_dtype}, got "
+                        f"{loaded_weight.dtype}"
+                    )
+                with torch.no_grad():
+                    target = getattr(
+                        embedding,
+                        "weight_scale_2_load_view",
+                        embedding.weight_scale_2,
+                    )
+                    assert target is not None
+                    target.copy_(
+                        loaded_weight.to(device=target.device, dtype=target.dtype)
+                    )
+                self._weight_scale_2_loaded = True
+                self._embedding_validated = False
+                loaded.add(name)
+                continue
+            if name.startswith(shard_prefix):
+                shard_and_suffix = name[len(shard_prefix) :]
+                shard_text, separator, suffix = shard_and_suffix.partition(".")
+                if not shard_text.isdigit():
+                    regular_weights.append((name, loaded_weight))
+                    continue
+                shard_index = int(shard_text)
+                if shard_index >= self.split_ngram_parts:
+                    raise ValueError(
+                        f"PLE shard {shard_index} exceeds "
+                        f"split_ngram_parts={self.split_ngram_parts}"
+                    )
+                checkpoint_start = shard_index * shard_size
+                expected_rows = max(
+                    0,
+                    min(shard_size, org_vocab_size - checkpoint_start),
+                )
+                if separator != "." or suffix not in {"weight", "weight_scale"}:
+                    regular_weights.append((name, loaded_weight))
+                    continue
+                if suffix == "weight":
+                    expected_shape = (expected_rows, self._table_layout.weight_shape[1])
+                    expected_dtype = self._table_layout.weight_dtype
+                else:
+                    if self._quant_mode != "nvfp4_group16":
+                        regular_weights.append((name, loaded_weight))
+                        continue
+                    assert self._table_layout.weight_scale_shape is not None
+                    assert self._table_layout.weight_scale_dtype is not None
+                    expected_shape = (
+                        expected_rows,
+                        self._table_layout.weight_scale_shape[1],
+                    )
+                    expected_dtype = self._table_layout.weight_scale_dtype
+                if tuple(loaded_weight.shape) != expected_shape:
+                    raise ValueError(
+                        f"shape mismatch for PLE shard {shard_index} {suffix}: "
+                        f"expected {expected_shape}, got {tuple(loaded_weight.shape)}"
+                    )
+                if loaded_weight.dtype != expected_dtype:
+                    raise TypeError(
+                        f"PLE shard {shard_index} {suffix} must have dtype "
+                        f"{expected_dtype}, got {loaded_weight.dtype}"
+                    )
+                overlap_start = max(checkpoint_start, tp_start)
+                overlap_end = min(checkpoint_start + expected_rows, tp_end)
+                disk_table = getattr(embedding, "disk_table", None)
+                if disk_table is not None:
+                    source = get_file_tensor_source(loaded_weight)
+                    if source is None:
+                        raise ValueError(
+                            "File-backed PLE tables require safetensors weights "
+                            "from the b12x, instanttensor, fastsafetensors, or "
+                            "safetensors loader"
+                        )
+                    if source.shape != expected_shape or source.dtype != expected_dtype:
+                        raise ValueError(f"file source geometry does not match {name}")
+                    if overlap_start < overlap_end:
+                        disk_table.add_shard(
+                            shard_index,
+                            source.path,
+                            source.offset,
+                            scale=suffix == "weight_scale",
+                        )
+                elif getattr(embedding, "shared_table_attached", False):
+                    # The attached table already holds every row of this shard.
+                    pass
+                else:
+                    parameter = getattr(embedding, suffix)
+                    destination = getattr(
+                        embedding, f"{suffix}_load_view", parameter.data
+                    )
+                    assert destination is not None
+                    _copy_embedding_shard(
+                        destination,
+                        loaded_weight,
+                        checkpoint_start=checkpoint_start,
+                        tp_start=tp_start,
+                        tp_end=tp_end,
+                    )
+                if overlap_start < overlap_end:
+                    load_ranges = (
+                        self._embedding_load_ranges
+                        if suffix == "weight"
+                        else self._scale_load_ranges
+                    )
+                    load_ranges.add((overlap_start, overlap_end))
+                    self._embedding_validated = False
+                loaded.add(f"ngram_embedding.{suffix}")
+                continue
+            regular_weights.append((name, loaded_weight))
+
+        if regular_weights:
+            loaded.update(AutoWeightsLoader(self).load_weights(regular_weights))
+        return loaded
+
+
+def _require_embedding_coverage(embedding: B12xNGramEmbedding) -> None:
+    """Raise when the loaded shards or scales do not cover the local table."""
+    covered_until = embedding._table_layout.shard_start
+    for start, end in sorted(embedding._embedding_load_ranges):
+        if start > covered_until:
+            raise ValueError(
+                "PLE embedding shards do not cover the local table: "
+                f"expected row {covered_until}, got {start}"
+            )
+        covered_until = max(covered_until, end)
+    if covered_until != embedding._table_layout.shard_end:
+        raise ValueError(
+            "PLE embedding shards do not cover the local table: "
+            f"stopped at row {covered_until}, "
+            f"expected {embedding._table_layout.shard_end}"
+        )
+
+    if embedding._table_layout.weight_scale_shape is not None:
+        if embedding._quant_mode == "fp8_e4m3_per_tensor":
+            if not embedding._weight_scale_loaded:
+                raise ValueError("FP8 PLE embedding checkpoint is missing weight_scale")
+        else:
+            scale_covered_until = embedding._table_layout.shard_start
+            for start, end in sorted(embedding._scale_load_ranges):
+                if start > scale_covered_until:
+                    raise ValueError(
+                        "NVFP4 PLE scale shards do not cover the local table: "
+                        f"expected row {scale_covered_until}, got {start}"
+                    )
+                scale_covered_until = max(scale_covered_until, end)
+            if scale_covered_until != embedding._table_layout.shard_end:
+                raise ValueError(
+                    "NVFP4 PLE scale shards do not cover the local table: "
+                    f"stopped at row {scale_covered_until}, expected "
+                    f"{embedding._table_layout.shard_end}"
+                )
+    if (
+        embedding._table_layout.weight_scale_2_shape is not None
+        and not embedding._weight_scale_2_loaded
+    ):
+        raise ValueError("NVFP4 PLE embedding checkpoint is missing weight_scale_2")
+
+
+def _ple_embedding_op(
+    input_ids: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    ngram_context: torch.Tensor,
+    out: torch.Tensor,
+    layer_name: str,
+) -> None:
+    layer = get_forward_context().no_compile_layers[layer_name]
+    layer.ple_embedding._run_embedding(input_ids, query_start_loc, ngram_context)
+
+
+def _ple_embedding_fake(
+    input_ids: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    ngram_context: torch.Tensor,
+    out: torch.Tensor,
+    layer_name: str,
+) -> None:
+    return
+
+
+def _ple_op(
+    residual: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    out: torch.Tensor,
+    layer_name: str,
+) -> None:
+    layer = get_forward_context().no_compile_layers[layer_name]
+    layer._run_ple(residual, key, value, query_start_loc)
+
+
+def _ple_fake(
+    residual: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    out: torch.Tensor,
+    layer_name: str,
+) -> None:
+    return
+
+
+direct_register_custom_op(
+    op_name="qwen4_exp_b12x_ple_embedding",
+    op_func=_ple_embedding_op,
+    mutates_args=["out"],
+    fake_impl=_ple_embedding_fake,
+)
+direct_register_custom_op(
+    op_name="qwen4_exp_b12x_ple",
+    op_func=_ple_op,
+    mutates_args=["out"],
+    fake_impl=_ple_fake,
+)

recipe/patches/qwen4_exp/nvidia/expert_stats.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Distinct routed experts per small-batch MoE call (VLLM_QWEN38_EXPERT_STATS=<json path>).
+
+Measurement only. Each router gate also takes the top-k of its logits for calls of at
+most 16 tokens (decode / spec-verify shapes) and accumulates, on device and CUDA-graph
+safe, the number of distinct experts the call touched. Touch <path>.req and a daemon
+thread writes the running totals per gate to the JSON path (warmup and capture runs
+count too, so diff two snapshots around the measured workload).
+"""
+
+import json
+import os
+import threading
+import time
+
+import torch
+
+_MAX_TOKENS = 16
+_gates: list = []
+
+
+def install(gate: torch.nn.Module, name: str, top_k: int, num_experts: int, path: str) -> None:
+    stock = gate.forward
+    # Allocated here, not lazily: a first call inside CUDA-graph capture would put the
+    # buffer in the graph pool and re-zero it on every replay.
+    stats = {"buf": torch.zeros(_MAX_TOKENS + 1, 2, dtype=torch.int64, device="cuda")}
+
+    def forward(x):
+        out = stock(x)
+        logits = out[0] if isinstance(out, tuple) else out
+        t = logits.shape[0] if logits.dim() == 2 else 0
+        if 0 < t <= _MAX_TOKENS:
+            ids = logits.topk(top_k, dim=-1).indices.flatten()
+            seen = torch.zeros(num_experts, dtype=torch.int64, device=logits.device).scatter_(0, ids, 1)
+            stats["buf"][t, 0] += seen.sum()
+            stats["buf"][t, 1] += 1
+        return out
+
+    gate.forward = forward
+    _gates.append((name, stats))
+    if len(_gates) == 1:
+        threading.Thread(target=_writer, args=(path,), daemon=True).start()
+
+
+def _writer(path: str) -> None:
+    # Touches CUDA only when asked (<path>.req appears): a device read from this thread
+    # during the server's CUDA-graph capture would invalidate the capture.
+    while True:
+        time.sleep(1)
+        if not os.path.exists(path + ".req"):
+            continue
+        try:
+            os.remove(path + ".req")
+        except OSError:
+            pass
+        out = {}
+        for name, stats in _gates:
+            buf = stats.get("buf")
+            if buf is None:
+                continue
+            rows = buf.tolist()
+            out[name] = {str(t): {"distinct_sum": d, "calls": c} for t, (d, c) in enumerate(rows) if c}
+        try:
+            with open(path + ".tmp", "w") as f:
+                json.dump(out, f)
+            os.replace(path + ".tmp", path)
+        except OSError:
+            pass

recipe/patches/qwen4_exp/nvidia/gdn_defer.py

@@ -0,0 +1,402 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Deferred-commit GDN verify (VLLM_QWEN38_GDN_DEFER=1); kernels in ops/gdn_deferred.py.
+
+Verify: the uniform b12x spec-decode path keeps its metadata staging and output norm
+but runs gdn_verify, which writes per-token records instead of 5 full fp32 states per
+layer, and marks the request rows it deferred (a device flag; graph-replay safe).
+
+Commit: once acceptance is known, gdn_commit replays the accepted tokens of every
+deferred row in all GDN layers (one launch) and writes the state column the next step
+reads. It runs at the first of the runner's postprocess hooks that can read state
+columns: boundary-checkpoint capture (before it copies), model-state postprocess
+(before align block saves), or the end of postprocess_sampled (logits-only steps,
+committed with every verify column written exactly as stock verify would).
+
+Layers that ever ran stock verify inside a CUDA-graph capture (no records yet) stay on
+the stock path for good, so a layer is committed iff its graphs record.
+"""
+
+import threading
+
+import torch
+import triton
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_tls = threading.local()
+_layers: list = []
+_shared: dict = {"valid": None, "tables": None, "key": None}
+_DEBUG = __import__("os").environ.get("VLLM_QWEN38_GDN_DEBUG", "0") == "1"
+_FORCE_FULL = __import__("os").environ.get("VLLM_QWEN38_GDN_FULL", "0") == "1"
+_SYNC = __import__("os").environ.get("VLLM_QWEN38_GDN_SYNC", "0") == "1"
+# Ordinary rows commit on a side stream that overlaps the MTP draft phase; the main
+# stream joins it at the next execute_model. Rows that need every column (block
+# crossing, capture, logits-only) still commit on the main stream first.
+_ASYNC = __import__("os").environ.get("VLLM_QWEN38_GDN_COMMIT_ASYNC", "0") == "1"
+
+
+def _join():
+    """Make the current stream wait for an in-flight side-stream commit."""
+    ev = _shared.get("pending")
+    if ev is not None:
+        torch.cuda.current_stream().wait_event(ev)
+        _shared["pending"] = None
+
+
+def _snapshot(name, src):
+    """Copy a small per-step tensor into a persistent buffer (main stream), so the
+    side-stream commit reads stable values whatever the main stream does next."""
+    buf = _shared.get(name)
+    if buf is None or buf.numel() < src.numel() or buf.dtype != src.dtype:
+        buf = _shared[name] = torch.empty(max(src.numel(), 64), dtype=src.dtype, device=src.device)
+    out = buf[: src.numel()].view(src.shape)
+    out.copy_(src)
+    return out
+
+
+def _records(layer):
+    rec = layer.__dict__.get("_qwen38_gdn_records")
+    if rec is not None or layer.__dict__.get("_qwen38_gdn_stock"):
+        return rec
+    staging = getattr(layer, "_b12x_decode_staging", None)
+    if staging is None or torch.cuda.is_current_stream_capturing():
+        return None
+    dev = layer.kv_cache[1].device
+    cap = layer._b12x_max_seqs * layer._b12x_state_index_columns
+    vh, kh = layer._b12x_local_value_heads, layer._b12x_local_key_heads
+    rec = (
+        torch.zeros(cap, vh, device=dev, dtype=torch.float32),
+        torch.zeros(cap, vh, 128, device=dev, dtype=torch.float32),
+        torch.zeros(cap, kh, 128, device=dev, dtype=torch.float32),
+    )
+    layer.__dict__["_qwen38_gdn_records"] = rec
+    _layers.append(layer)
+    if _shared["valid"] is None:
+        _shared["valid"] = torch.zeros(layer._b12x_max_seqs, dtype=torch.int32, device=dev)
+        logger.info("Qwen3.8 GDN verify defers state writes to a post-acceptance commit.")
+    _shared["key"] = None
+    return rec
+
+
+def _tables():
+    key = tuple(
+        (layer.kv_cache[1].data_ptr(), layer._b12x_decode_staging.state_indices.data_ptr())
+        for layer in _layers
+    )
+    if key != _shared["key"]:
+        dev = _layers[0].kv_cache[1].device
+        cols = (
+            [layer.kv_cache[1].data_ptr() for layer in _layers],
+            [layer.__dict__["_qwen38_gdn_records"][0].data_ptr() for layer in _layers],
+            [layer.__dict__["_qwen38_gdn_records"][1].data_ptr() for layer in _layers],
+            [layer.__dict__["_qwen38_gdn_records"][2].data_ptr() for layer in _layers],
+            # each layer's own staged block ids (layers sit in different KV-cache groups)
+            [layer._b12x_decode_staging.state_indices.data_ptr() for layer in _layers],
+        )
+        ref = _layers[0]._b12x_decode_staging.state_indices
+        for layer in _layers:
+            si = layer._b12x_decode_staging.state_indices
+            if si.dtype != torch.int32 or si.shape != ref.shape or si.stride() != ref.stride():
+                raise RuntimeError("GDN defer needs identical int32 staged index layouts")
+        _shared["tables"] = tuple(torch.tensor(c, dtype=torch.int64, device=dev) for c in cols)
+        _shared["key"] = key
+        strides = {layer.kv_cache[1].stride(0) for layer in _layers}
+        if len(strides) != 1:
+            raise RuntimeError("GDN defer needs one state slot stride across layers")
+    return _shared["tables"]
+
+
+def _commit(capture_tokens=None, force_full=False):
+    from .ops.gdn_deferred import gdn_commit
+
+    ctx = _tls.ctx
+    ctx["done"] = True
+    if _SHADOW:
+        if _shared["valid"] is not None:
+            _shared["valid"].zero_()
+        return
+    valid = _shared["valid"]
+    num_sampled = ctx["num_sampled"]
+    if valid is None or not _layers or not isinstance(num_sampled, torch.Tensor):
+        return
+    runner = ctx["runner"]
+    layer0 = _layers[0]
+    st = layer0._b12x_decode_staging
+    cc = runner.cache_config
+    block = int(getattr(cc, "mamba_block_size", None) or cc.block_size)
+    if _DEBUG:
+        _shared["debug_commits"] = _shared.get("debug_commits", 0) + 1
+        k = _shared["debug_commits"]
+        if 15 <= k <= 55:
+            torch.cuda.synchronize()
+            logger.info(
+                "GDN defer commit #%d: valid=%s nseq=%s nacc=%s qsl=%s idx0=%s nsamp=%s map=%s cap=%s",
+                k, valid.tolist(), st.num_seqs.item(), st.num_accepted_tokens.tolist()[:2],
+                st.query_start_loc.tolist()[:3], st.state_indices[0].tolist(),
+                num_sampled.tolist(), ctx["idx_mapping"].tolist(),
+                None if capture_tokens is None else int(capture_tokens.max()),
+            )
+    _join()
+    args = (_tables(), st.state_indices, st.num_accepted_tokens, st.query_start_loc, st.num_seqs, valid)
+    kw = dict(key_heads=layer0._b12x_local_key_heads, value_heads=layer0._b12x_local_value_heads)
+    tail = (block, layer0.kv_cache[1].stride(0), layer0._b12x_max_seqs)
+    ncomp = runner.req_states.num_computed_tokens.gpu
+    full = force_full or _FORCE_FULL
+    if _ASYNC and capture_tokens is None and not full:
+        gdn_commit(*args, num_sampled, ctx["idx_mapping"], ncomp, None, *tail, row_mode=1, **kw)
+        ns = _snapshot("snap_ns", num_sampled)
+        rm = _snapshot("snap_map", ctx["idx_mapping"])
+        nc = _snapshot("snap_nc", ncomp)
+        main = torch.cuda.current_stream()
+        side = _shared.get("side")
+        if side is None:
+            side = _shared["side"] = torch.cuda.Stream(device=valid.device)
+            _shared["ev_fork"], _shared["ev_done"] = torch.cuda.Event(), torch.cuda.Event()
+        _shared["ev_fork"].record(main)
+        side.wait_event(_shared["ev_fork"])
+        with torch.cuda.stream(side):
+            gdn_commit(*args, ns, rm, nc, None, *tail, row_mode=2, **kw)
+            valid.zero_()
+            _shared["ev_done"].record(side)
+        _shared["pending"] = _shared["ev_done"]
+    else:
+        gdn_commit(*args, num_sampled, ctx["idx_mapping"], ncomp, capture_tokens, *tail,
+                   force_full=full, **kw)
+        valid.zero_()
+    if _SYNC:
+        torch.cuda.synchronize()
+    if _DEBUG and _shared.get('debug_streams', 0) < 3:
+        _shared['debug_streams'] = _shared.get('debug_streams', 0) + 1
+        logger.info('GDN defer commit stream %s', torch.cuda.current_stream())
+
+
+_SHADOW = __import__("os").environ.get("VLLM_QWEN38_GDN_SHADOW", "0") == "1"
+
+
+def _shadow_compare(layer, rec, stock_post, kwargs):
+    """Debug (eager only): run gdn_verify into a scratch output from the pristine source
+    state, then the stock verify (which drives the model and writes every column), and
+    compare outputs and torch-replayed per-token states against stock's columns."""
+    from .ops.gdn_deferred import gdn_verify
+
+    state = layer.kv_cache[1]
+    si, nacc, qsl = kwargs["state_indices"], kwargs["num_accepted_tokens"], kwargs["query_start_loc"]
+    nreq = int(kwargs["num_requests"])
+    staging = layer._b12x_decode_staging
+    # stage like stock (our kernel reads the staging copies)
+    import triton as _tr
+    from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn as qg
+
+    qg._stage_b12x_gdn_metadata_kernel[(1,)](
+        qsl, si, nacc, staging.query_start_loc, staging.state_indices,
+        staging.num_accepted_tokens, staging.num_seqs, staging.num_tokens, nreq,
+        qsl.stride(0), nacc.stride(0), si.stride(0), si.stride(1),
+        INPUT_COLUMNS=si.shape[1], STATE_COLUMNS=layer._b12x_state_index_columns,
+        MAX_SEQS=layer._b12x_max_seqs, HAS_ACCEPTED=True,
+        BLOCK=_tr.next_power_of_2(max(layer._b12x_max_seqs + 1, layer._b12x_max_seqs * layer._b12x_state_index_columns)),
+        num_warps=4,
+    )
+    src_slot = int(staging.state_indices[0, int(staging.num_accepted_tokens[0]) - 1])
+    src = state[src_slot].clone()
+    T = int(staging.query_start_loc[1] - staging.query_start_loc[0])
+    ours = kwargs["core_attn_out"].clone()
+    gdn_verify(
+        kwargs["mixed_qkv"], kwargs["a"], kwargs["b"], layer.A_log, layer.dt_bias, state,
+        staging.state_indices, staging.num_accepted_tokens, staging.query_start_loc,
+        staging.num_seqs, _shared["valid"], ours[: kwargs["mixed_qkv"].shape[0]], rec,
+        layer.head_k_dim ** -0.5, layer._b12x_max_seqs,
+        key_heads=layer._b12x_local_key_heads, value_heads=layer._b12x_local_value_heads,
+    )
+    from .ops.gdn_deferred import gdn_gated_rmsnorm
+
+    m = kwargs["mixed_qkv"].shape[0]
+    gdn_gated_rmsnorm(ours[:m], kwargs["output_gate"], layer.norm.weight, staging.num_tokens,
+                      layer.layer_norm_epsilon, layer._b12x_local_value_heads,
+                      layer.norm.activation == "sigmoid")
+    result = stock_post(layer, **kwargs)
+    stock_out = kwargs["core_attn_out"][:m].float()
+    out_err = ((ours[:m].float() - stock_out).abs().max() / stock_out.abs().max().clamp_min(1e-30)).item()
+    if layer is _layers[0] and _shared.setdefault("shadow_logs", 0) < 12:
+        _shared["shadow_logs"] += 1
+        g, d, kk = rec
+        S = src.double()
+        errs = []
+        for t in range(T):
+            kv = kk[t].double()[torch.arange(48, device=S.device) // 3]  # [48,128]
+            S = S * g[t].double()[:, None, None] + d[t].double()[:, :, None] * kv[:, None, :]
+            ref = state[int(staging.state_indices[0, t])].double()
+            errs.append(f"{((S - ref).abs().max() / ref.abs().max().clamp_min(1e-30)).item():.1e}")
+        logger.info(
+            "GDN shadow: layer0 T=%d src_slot=%d nacc=%d cols=%s | normed output rel err %.2e | "
+            "state replay rel err per column %s",
+            T, src_slot, int(staging.num_accepted_tokens[0]), staging.state_indices[0].tolist(), out_err, errs,
+        )
+    return result
+
+
+def install() -> None:
+    from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn as qg
+    from vllm.v1.worker.gpu import boundary_checkpoint as bc
+    from vllm.v1.worker.gpu import model_runner as mr
+    from vllm.v1.worker.gpu.model_states import mamba_hybrid as mh
+
+    from .ops.gdn_deferred import gdn_gated_rmsnorm, gdn_verify
+
+    Attn = qg.QwenGatedDeltaNetAttention
+    if getattr(Attn, "_qwen38_gdn_defer", False):
+        return
+    stock_fused = Attn._forward_core_fused_norm
+    stock_uniform = Attn._forward_core_decode_b12x_fused_norm
+    stock_post = Attn._run_b12x_gdn_decode_post_conv
+
+    def fused(self, mixed_qkv, b, a, output_gate, core_attn_out):
+        # Every forward starts from "nothing deferred"; the uniform path re-marks rows.
+        valid = _shared["valid"]
+        if valid is not None and _layers and self is _layers[0]:
+            valid.zero_()
+        return stock_fused(self, mixed_qkv, b, a, output_gate, core_attn_out)
+
+    def uniform(self, mixed_qkv, b, a, output_gate, core_attn_out, attn_metadata):
+        _tls.spec = attn_metadata.spec_sequence_masks is not None
+        try:
+            return stock_uniform(
+                self, mixed_qkv=mixed_qkv, b=b, a=a, output_gate=output_gate,
+                core_attn_out=core_attn_out, attn_metadata=attn_metadata,
+            )
+        finally:
+            _tls.spec = False
+
+    def post(self, *, mixed_qkv, b, a, output_gate, core_attn_out, state_indices,
+             query_start_loc, num_accepted_tokens, num_requests):
+        kwargs = dict(
+            mixed_qkv=mixed_qkv, b=b, a=a, output_gate=output_gate,
+            core_attn_out=core_attn_out, state_indices=state_indices,
+            query_start_loc=query_start_loc, num_accepted_tokens=num_accepted_tokens,
+            num_requests=num_requests,
+        )
+        rec = None
+        if (
+            getattr(_tls, "spec", False)
+            and num_accepted_tokens is not None
+            and state_indices.shape[1] > 1
+            and core_attn_out.dim() == 3
+            and core_attn_out.stride(2) == 1
+            and mixed_qkv.stride(1) == 1
+        ):
+            rec = _records(self)
+        if _DEBUG and torch.cuda.is_current_stream_capturing():
+            key = "debug_capture_deferred" if rec is not None else "debug_capture_stock"
+            _shared[key] = _shared.get(key, 0) + 1
+            if _shared[key] in (1, 36, 72):
+                logger.info("GDN defer capture: %s calls=%d (spec=%s cols=%d) stream %s", key, _shared[key],
+                            getattr(_tls, "spec", False), state_indices.shape[1], torch.cuda.current_stream())
+        if rec is not None and _SHADOW and not torch.cuda.is_current_stream_capturing():
+            return _shadow_compare(self, rec, stock_post, kwargs)
+        if rec is None:
+            if torch.cuda.is_current_stream_capturing() and "_qwen38_gdn_records" not in self.__dict__:
+                self.__dict__["_qwen38_gdn_stock"] = True
+            return stock_post(self, **kwargs)
+        staging = self._b12x_decode_staging
+        num_input_tokens = mixed_qkv.shape[0]
+        if (
+            num_input_tokens > self._b12x_max_tokens
+            or num_requests > self._b12x_max_seqs
+            or state_indices.shape[1] > self._b12x_state_index_columns
+            or num_input_tokens > rec[0].shape[0]
+        ):
+            raise ValueError("b12x GDN capacity exceeded (deferred verify)")
+        qg._stage_b12x_gdn_metadata_kernel[(1,)](
+            query_start_loc, state_indices, num_accepted_tokens,
+            staging.query_start_loc, staging.state_indices, staging.num_accepted_tokens,
+            staging.num_seqs, staging.num_tokens,
+            num_requests,
+            query_start_loc.stride(0),
+            num_accepted_tokens.stride(0),
+            state_indices.stride(0),
+            state_indices.stride(1),
+            INPUT_COLUMNS=state_indices.shape[1],
+            STATE_COLUMNS=self._b12x_state_index_columns,
+            MAX_SEQS=self._b12x_max_seqs,
+            HAS_ACCEPTED=True,
+            BLOCK=triton.next_power_of_2(
+                max(self._b12x_max_seqs + 1, self._b12x_max_seqs * self._b12x_state_index_columns)
+            ),
+            num_warps=4,
+        )
+        out = core_attn_out[:num_input_tokens]
+        gdn_verify(
+            mixed_qkv, a, b, self.A_log, self.dt_bias, self.kv_cache[1],
+            staging.state_indices, staging.num_accepted_tokens, staging.query_start_loc,
+            staging.num_seqs, _shared["valid"], out, rec, self.head_k_dim ** -0.5,
+            self._b12x_max_seqs, key_heads=self._b12x_local_key_heads,
+            value_heads=self._b12x_local_value_heads,
+        )
+        gdn_gated_rmsnorm(
+            out, output_gate, self.norm.weight, staging.num_tokens, self.layer_norm_epsilon,
+            self._b12x_local_value_heads, self.norm.activation == "sigmoid",
+        )
+
+    Attn._forward_core_fused_norm = fused
+    Attn._forward_core_decode_b12x_fused_norm = uniform
+    Attn._run_b12x_gdn_decode_post_conv = post
+    Attn._qwen38_gdn_defer = True
+
+    Runner = mr.GPUModelRunner
+    stock_pps = Runner.postprocess_sampled
+
+    def postprocess_sampled(self, idx_mapping, sampled_token_ids, num_sampled, num_rejected,
+                            query_start_loc=None, boundary_capture=None, logits_only=False):
+        if not _layers:
+            return stock_pps(self, idx_mapping, sampled_token_ids, num_sampled, num_rejected,
+                             query_start_loc, boundary_capture, logits_only=logits_only)
+        if _DEBUG and _shared.setdefault("debug_pps", 0) < 3:
+            _shared["debug_pps"] += 1
+            logger.info("GDN defer: postprocess_sampled hook (layers=%d, num_sampled type %s, capture %s, logits_only %s)", len(_layers), type(num_sampled).__name__, boundary_capture is not None, logits_only)
+        _tls.ctx = {"num_sampled": num_sampled, "idx_mapping": idx_mapping,
+                    "runner": self, "done": False}
+        try:
+            out = stock_pps(self, idx_mapping, sampled_token_ids, num_sampled, num_rejected,
+                            query_start_loc, boundary_capture, logits_only=logits_only)
+            if not _tls.ctx["done"]:
+                # postprocess_state was skipped (logits-only): reproduce stock columns.
+                _commit(force_full=True)
+            return out
+        finally:
+            _tls.ctx = None
+
+    Runner.postprocess_sampled = postprocess_sampled
+
+    if _ASYNC:
+        stock_exec = Runner.execute_model
+
+        def execute_model(self, *args, **kwargs):
+            # Before anything (state migration, staging, the forward) touches GDN state.
+            _join()
+            return stock_exec(self, *args, **kwargs)
+
+        Runner.execute_model = execute_model
+        logger.info("Qwen3.8 GDN commit runs on a side stream overlapping the draft phase.")
+
+    stock_capture = bc.BoundaryCheckpointState.capture_mamba
+
+    def capture_mamba(self, idx_mapping, capture, *, accepted_state_committed=False):
+        ctx = getattr(_tls, "ctx", None)
+        if ctx is not None and not ctx["done"]:
+            _commit(capture_tokens=capture[0])
+        return stock_capture(self, idx_mapping, capture,
+                             accepted_state_committed=accepted_state_committed)
+
+    bc.BoundaryCheckpointState.capture_mamba = capture_mamba
+
+    stock_pstate = mh.MambaHybridModelState.postprocess_state
+
+    def postprocess_state(self, idx_mapping, num_sampled, num_computed_tokens=None):
+        ctx = getattr(_tls, "ctx", None)
+        if ctx is not None and not ctx["done"]:
+            _commit()
+        return stock_pstate(self, idx_mapping, num_sampled, num_computed_tokens)
+
+    mh.MambaHybridModelState.postprocess_state = postprocess_state
+    logger.info("Qwen3.8 GDN deferred commit installed (verify records + post-acceptance commit).")

recipe/patches/qwen4_exp/nvidia/hyperconnection.py

@@ -0,0 +1,782 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HyperConnection (Gated Residual) utilities — NVIDIA model variant.
+
+Implements the HyperConnection residual scheme proposed in
+"HyperConnections" (https://arxiv.org/abs/2409.19606). This NVIDIA variant
+delays each HC combine to the following HC mix boundary. HC glue kernels,
+including fused combine+RMSNorm, live in ``ops/hc.py``; projections remain
+standard vLLM Linear modules.
+
+Hidden states between layers have shape ``[..., HC*HS]`` with HS inner
+(HC outer, HS inner — checkpoint-native layout).
+
+Typical usage inside a transformer decoder layer::
+
+    self.attn_hc = GatedResidual(hc_config)
+
+    hidden_states, block_input, injection = self.attn_hc.mix(hidden_states)
+    attention_output = attention(block_input)
+    hidden_states, block_input, injection = self.mlp_hc.combine_and_mix(
+        hidden_states, attention_output, injection
+    )
+"""
+
+from collections.abc import Iterable, Sequence
+from typing import Any
+
+import torch
+from torch import nn
+
+import vllm.envs as envs
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_gather,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    ReplicatedLinear,
+)
+from vllm.model_executor.models.utils import maybe_prefix
+from vllm.model_executor.weight_transfer import copy_weight
+from vllm.platforms import current_platform
+from vllm.utils.b12x import (
+    B12xPreparationUnit,
+    B12xWorkload,
+    PreparationResourceUnavailableError,
+    get_b12x_hyperconnection,
+    set_b12x_preparation_provider,
+)
+
+from ..common.hyperconnection import (
+    GroupedGemmaRMSNorm,
+    HyperConnectionConfig,
+)
+from .ops.hc import (
+    grouped_gemma_rmsnorm,
+    hc_combine,
+    hc_combine_norm,
+    hc_gate_mix,
+    hc_silu,
+)
+
+logger = init_logger(__name__)
+
+import os as _os
+
+# Rowwise-FP8 copies of the mix weights for decode-sized calls (ops/hc_fp8_gemv.py).
+_HC_FP8 = _os.environ.get("VLLM_QWEN38_HC_FP8", "0") == "1"
+# With _HC_FP8 at decode: scaled SiLU fused into the down projection's split-K reduce
+# and the gate mix fused into the up projection (two launches instead of four).
+_HC_FUSE = _HC_FP8 and _os.environ.get("VLLM_QWEN38_HC_FUSE", "0") == "1"
+
+
+# ---------------------------------------------------------------------------
+# Gated-residual variant
+# ---------------------------------------------------------------------------
+def _hyperconnection_api() -> Any:
+    api = get_b12x_hyperconnection()
+    if api is None:
+        raise ImportError(
+            "Qwen4Exp requires b12x.norm.hyperconnection; "
+            "install the b12x serving extra"
+        )
+    return api
+
+
+class _ShardedDownProjection(ReplicatedLinear):
+    """Local bottleneck rows with replicated injection rows in one GEMM."""
+
+    def __init__(
+        self,
+        config: HyperConnectionConfig,
+        tp_rank: int,
+        tp_size: int,
+        use_combine: bool,
+        prefix: str,
+    ) -> None:
+        self.hc_tp_rank = tp_rank
+        self.full_lowrank = config.hc_lowrank
+        self.local_lowrank = config.hc_lowrank // tp_size
+        self.injection_size = config.hc_count if use_combine else 0
+        rows = self.local_lowrank + self.injection_size
+        self.padding = -rows % 16
+        super().__init__(
+            config.hc_count * config.hidden_size,
+            rows + self.padding,
+            bias=False,
+            params_dtype=config.params_dtype,
+            return_bias=False,
+            disable_tp=True,
+            prefix=prefix,
+        )
+
+    def weight_loader(
+        self,
+        param: torch.Tensor,
+        loaded_weight: torch.Tensor,
+        shard_id: int | None = None,
+    ) -> None:
+        if shard_id is None and self.injection_size:
+            self.weight_loader(param, loaded_weight[: self.full_lowrank], 0)
+            self.weight_loader(
+                param,
+                loaded_weight[
+                    self.full_lowrank : self.full_lowrank + self.injection_size
+                ],
+                1,
+            )
+            return
+        if shard_id in (None, 0):
+            assert loaded_weight.shape == (self.full_lowrank, self.input_size)
+            source = loaded_weight.narrow(
+                0, self.hc_tp_rank * self.local_lowrank, self.local_lowrank
+            )
+            destination = param.data[: self.local_lowrank]
+        elif shard_id == 1 and self.injection_size:
+            assert loaded_weight.shape == (self.injection_size, self.input_size)
+            source = loaded_weight
+            destination = param.data[
+                self.local_lowrank : self.local_lowrank + self.injection_size
+            ]
+        else:
+            raise ValueError(f"Invalid HC projection shard: {shard_id}")
+        copy_weight(destination, source)
+        if self.padding and shard_id != 1:
+            copy_weight(
+                param.data[-self.padding :],
+                torch.zeros(
+                    self.padding, self.input_size, dtype=param.dtype, device="cpu"
+                ),
+            )
+
+    def load_weights(
+        self, weights: Iterable[tuple[str, torch.Tensor]]
+    ) -> Iterable[str]:
+        for name, weight in weights:
+            if name != "weight":
+                raise ValueError(f"Unexpected HC projection parameter: {name}")
+            self.weight_loader(self.weight, weight, getattr(weight, "shard_id", None))
+            yield name
+
+
+class _ShardedUpProjection(ReplicatedLinear):
+    """Shard features within every residual stream, preserving stream order."""
+
+    def __init__(
+        self, config: HyperConnectionConfig, tp_rank: int, tp_size: int, prefix: str
+    ) -> None:
+        self.hc_tp_rank = tp_rank
+        self.streams = config.hc_count
+        self.full_hidden = config.hidden_size
+        self.local_hidden = config.hidden_size // tp_size
+        super().__init__(
+            config.hc_lowrank,
+            self.streams * self.local_hidden,
+            bias=False,
+            params_dtype=config.params_dtype,
+            return_bias=False,
+            disable_tp=True,
+            prefix=prefix,
+        )
+
+    def weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+        assert loaded_weight.shape == (
+            self.streams * self.full_hidden,
+            self.input_size,
+        )
+        # Separate views retain b12x's file-backed source ranges during loading.
+        for stream in range(self.streams):
+            source = loaded_weight.narrow(
+                0,
+                stream * self.full_hidden + self.hc_tp_rank * self.local_hidden,
+                self.local_hidden,
+            )
+            destination = param.data.narrow(
+                0, stream * self.local_hidden, self.local_hidden
+            )
+            copy_weight(destination, source)
+
+
+class HyperConnectionWorkspace(nn.Module):
+    """Fixed-capacity storage shared by all HC modules in one model."""
+
+    def __init__(self, config: HyperConnectionConfig, max_tokens: int) -> None:
+        super().__init__()
+        if not config.hc_per_branch_norm:
+            raise NotImplementedError(
+                "Qwen4Exp requires one RMSNorm group per HC stream"
+            )
+        self.config = config
+        self.max_tokens = int(max_tokens)
+        self.device = torch.device(current_platform.current_device())
+        tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = (
+            tp_size
+            if envs.VLLM_QWEN3_8_FLASH_NEXT_HC_TP
+            and current_platform.is_cuda()
+            and config.params_dtype == torch.bfloat16
+            and config.hc_lowrank % tp_size == 0
+            and config.hidden_size % tp_size == 0
+            else 1
+        )
+        self.tp_rank = get_tensor_model_parallel_rank() if self.tp_size > 1 else 0
+        width = config.hc_count * config.hidden_size
+        factory = dict(device=self.device, dtype=config.params_dtype)
+        self.register_buffer(
+            "normalized", torch.empty(max_tokens, width, **factory), persistent=False
+        )
+        self.register_buffer(
+            "bottleneck",
+            torch.empty(max_tokens, config.hc_lowrank, **factory),
+            persistent=False,
+        )
+        self.register_buffer(
+            "block_input",
+            torch.empty(max_tokens, config.hidden_size, **factory),
+            persistent=False,
+        )
+        if self.tp_size > 1:
+            local_hidden = config.hidden_size // self.tp_size
+            local_lowrank = config.hc_lowrank // self.tp_size
+            logger.info_once(
+                "Sharding HyperConnection projections across %d ranks: "
+                "local bottleneck=%d, features/stream=%d; FP32 down outputs.",
+                self.tp_size,
+                local_lowrank,
+                local_hidden,
+            )
+            down_width = local_lowrank + config.hc_count
+            down_width += -down_width % 16
+            for name, columns in (
+                ("local_normalized", config.hc_count * local_hidden),
+                ("local_bottleneck", local_lowrank),
+                ("local_block_input", local_hidden),
+                ("local_gates", config.hc_count * local_hidden),
+                ("local_down", down_width),
+            ):
+                self.register_buffer(
+                    name, torch.empty(max_tokens, columns, **factory), persistent=False
+                )
+            self.register_buffer(
+                "local_down_fp32",
+                torch.empty(
+                    max_tokens, down_width, device=self.device, dtype=torch.float32
+                ),
+                persistent=False,
+            )
+
+    def caps(self, max_tokens: int, *, local: bool = False):
+        api = _hyperconnection_api()
+        partitions = self.tp_size if local else 1
+        return api.Caps(
+            device=self.device,
+            max_tokens=max_tokens,
+            hidden_size=self.config.hidden_size // partitions,
+            streams=self.config.hc_count,
+            lowrank=self.config.hc_lowrank // partitions,
+            dtype=self.config.params_dtype,
+        )
+
+    def bind(self, plan, tokens: int, *, local: bool = False):
+        return _hyperconnection_api().bind(
+            plan,
+            normalized=self.local_normalized if local else self.normalized,
+            bottleneck=self.local_bottleneck if local else self.bottleneck,
+            block_input=self.local_block_input if local else self.block_input,
+            tokens=tokens,
+        )
+
+
+class GatedResidual(nn.Module):
+    """Gated HyperConnection with learnable low-rank mixing and injection.
+
+    ``combine_and_mix()`` runs the pre pipeline (grouped GemmaRMSNorm -> merged
+    low-rank down+inject GEMM -> silu -> up GEMM -> sigmoid -> gated mean
+    over the HC streams). When passed a pending block output, it fuses its
+    residual combine with the RMSNorm. A missing injection selects unit-weight
+    combine. Final mixers use ``use_combine=False`` and do not produce a new
+    injection.
+
+    TP shards bottleneck rows and output features within each stream, with a
+    gather at each boundary. The down GEMM returns FP32 before its BF16 cast.
+    Norm and residual state remain replicated; TP1 uses standard Linear dispatch.
+    """
+
+    def __init__(
+        self,
+        config: HyperConnectionConfig,
+        use_combine: bool = True,
+        prefix: str = "",
+        *,
+        workspace: HyperConnectionWorkspace | None = None,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.tp_size = workspace.tp_size if workspace is not None else 1
+        self.tp_rank = workspace.tp_rank if workspace is not None else 0
+        self.lora_rank = config.hc_lowrank // self.tp_size
+        self.hc_count = config.hc_count
+        self.hidden_size = config.hidden_size
+        self.use_combine = use_combine
+
+        norm_size = (
+            self.hyper_hidden_size if config.hc_per_branch_norm else config.hidden_size
+        )
+        group_size = config.hidden_size if config.hc_per_branch_norm else None
+        # Normalize each H-sized HC stream independently while retaining a
+        # separate affine weight for every element of the HC*H layout.
+        self.hc_norm = GroupedGemmaRMSNorm(
+            norm_size,
+            eps=config.rms_norm_eps,
+            group_size=group_size,
+            dtype=config.params_dtype,
+        )
+
+        # -- vLLM Linear weights --------------------------------------------
+        # The merged skinny-GEMM shape is physically padded to 16 rows to ensure
+        # good alignment and performant implementation chosen by CuBLAS heuristics.
+        self.pad_size = (-(self.lora_rank + self.hc_count)) % 16 if use_combine else 0
+        if self.tp_size > 1:
+            name = (
+                "input_mix_weight_down_block_inject"
+                if use_combine
+                else "input_mix_weight_down"
+            )
+            setattr(
+                self,
+                name,
+                _ShardedDownProjection(
+                    config,
+                    self.tp_rank,
+                    self.tp_size,
+                    use_combine,
+                    maybe_prefix(prefix, name),
+                ),
+            )
+        elif use_combine:
+            self.input_mix_weight_down_block_inject = MergedColumnParallelLinear(
+                self.hyper_hidden_size,
+                [self.lora_rank, self.hc_count]
+                + ([self.pad_size] if self.pad_size else []),
+                bias=False,
+                params_dtype=config.params_dtype,
+                quant_config=None,
+                prefix=maybe_prefix(prefix, "input_mix_weight_down_block_inject"),
+                return_bias=False,
+                disable_tp=True,
+            )
+        else:
+            self.input_mix_weight_down = ReplicatedLinear(
+                self.hyper_hidden_size,
+                self.lora_rank,
+                bias=False,
+                params_dtype=config.params_dtype,
+                quant_config=None,
+                prefix=maybe_prefix(prefix, "input_mix_weight_down"),
+                return_bias=False,
+            )
+        self.input_mix_weight_up = (
+            _ShardedUpProjection(
+                config,
+                self.tp_rank,
+                self.tp_size,
+                maybe_prefix(prefix, "input_mix_weight_up"),
+            )
+            if self.tp_size > 1
+            else ReplicatedLinear(
+                self.lora_rank,
+                self.hyper_hidden_size,
+                bias=False,
+                params_dtype=config.params_dtype,
+                quant_config=None,
+                prefix=maybe_prefix(prefix, "input_mix_weight_up"),
+                return_bias=False,
+            )
+        )
+        object.__setattr__(self, "_workspace", workspace)
+        self._preparation_prefix = prefix or "qwen4_exp.hyperconnection"
+        self._plans: dict[str, object] = {}
+        if workspace is not None and not getattr(
+            self, "b12x_preparation_suppressed", False
+        ):
+            set_b12x_preparation_provider(self, self)
+
+    def mix(self, hidden_states: torch.Tensor):
+        if self.workspace is not None:
+            normalized = _hyperconnection_api().run_grouped_rmsnorm(
+                hidden_states,
+                self.hc_norm.weight,
+                eps=self.config.rms_norm_eps,
+                binding=self._binding(hidden_states, "grouped_rmsnorm"),
+            )
+        else:
+            normalized = grouped_gemma_rmsnorm(
+                hidden_states,
+                self.hc_norm.weight,
+                self.config.rms_norm_eps,
+                self.hc_count,
+            )
+        block_input, injection = self._mix_normalized(normalized)
+        return hidden_states, block_input, injection
+
+    def combine_and_mix(self, hidden_states, prev_block_output, prev_injection):
+        if prev_block_output is None:
+            return self.mix(hidden_states)
+        if self.workspace is not None and prev_injection is not None:
+            combined, normalized = _hyperconnection_api().run_combine_norm(
+                hidden_states,
+                prev_block_output,
+                prev_injection,
+                self.hc_norm.weight,
+                eps=self.config.rms_norm_eps,
+                plan=self._plan_for("combine_norm"),
+            )
+        else:
+            combined, normalized = hc_combine_norm(
+                hidden_states,
+                prev_block_output,
+                prev_injection,
+                self.hc_norm.weight,
+                self.config.rms_norm_eps,
+                self.hc_count,
+            )
+        block_input, injection = self._mix_normalized(normalized)
+        return combined, block_input, injection
+
+    def combine(self, hidden_states, block_output, injection):
+        if self.workspace is not None and injection is not None:
+            return _hyperconnection_api().run_combine(
+                hidden_states,
+                block_output,
+                injection,
+                plan=self._plan_for("combine"),
+            )
+        return hc_combine(hidden_states, block_output, injection, self.hc_count)
+
+    @property
+    def hyper_hidden_size(self) -> int:
+        return self.hc_count * self.hidden_size
+
+    def _request_name(self, operation: str, tokens: int) -> str:
+        return f"{self._preparation_prefix}.hc.{operation}.m{tokens}"
+
+    def get_b12x_preparation_units(
+        self, layer: torch.nn.Module, workload: B12xWorkload
+    ) -> Sequence[B12xPreparationUnit]:
+        if layer is not self:
+            raise ValueError("HC preparation owner mismatch")
+        if self.workspace is None or workload.stage != "weights":
+            return ()
+        if self.hc_norm.weight.is_meta:
+            return ()
+        if workload.max_tokens > self.workspace.max_tokens:
+            raise PreparationResourceUnavailableError(
+                f"{self._preparation_prefix} HC workspace capacity "
+                f"{self.workspace.max_tokens} cannot serve {workload.max_tokens}"
+            )
+        api = _hyperconnection_api()
+        operations = (
+            "grouped_rmsnorm",
+            "scaled_silu",
+            "gate_mean",
+            "combine",
+            "combine_norm",
+        )
+        requests = []
+        tokens = workload.max_tokens
+        plans = {}
+        for operation in operations:
+            local = self.tp_size > 1 and operation in ("scaled_silu", "gate_mean")
+            plan = api.plan(
+                self.workspace.caps(tokens, local=True)
+                if local
+                else self.workspace.caps(tokens),
+                invocation={"operation": operation, "eps": self.config.rms_norm_eps},
+            )
+            plans[operation] = plan
+            requests.append(
+                plan.request(
+                    name=self._request_name(operation, tokens),
+                    prepare_call=self._prepare_call(operation, tokens),
+                    benchmark_call=self._benchmark_call(operation, tokens),
+                )
+            )
+        self._plans = plans
+        return (
+            B12xPreparationUnit(
+                name="HYPERCONNECTION",
+                key=(self._preparation_prefix, tokens),
+                requests=tuple(requests),
+                stage="weights",
+            ),
+        )
+
+    def _prepare_call(self, operation: str, tokens: int):
+        """Prime the installed operation against its durable workspace."""
+        return self._call_factory(operation, tokens, benchmark=False)
+
+    def _benchmark_call(self, operation: str, tokens: int):
+        """Measure an isolated binding; it is never retained for serving."""
+        return self._call_factory(operation, tokens, benchmark=True)
+
+    def _call_factory(self, operation: str, tokens: int, *, benchmark: bool):
+        def prepare(state):
+            from b12x.norm.hyperconnection import _impl
+            from b12x.preparation import PreparedCall
+
+            factory = dict(device=self.workspace.device, dtype=self.config.params_dtype)
+            local = self.tp_size > 1 and operation in ("scaled_silu", "gate_mean")
+            hidden_size = (
+                self.hidden_size // self.tp_size if local else self.hidden_size
+            )
+            width = self.hc_count * hidden_size
+            activation_inputs = []
+            activation_owners: list[torch.Tensor] = []
+            owners: tuple[torch.Tensor, ...]
+
+            def activation(shape):
+                value = torch.empty(shape, **factory)
+                template = torch.arange(value.numel(), **factory).reshape(shape)
+                template.div_(max(value.numel(), 1))
+                activation_inputs.append((value, template))
+                activation_owners.extend((value, template))
+                return value
+
+            def produce():
+                for value, template in activation_inputs:
+                    value.copy_(template)
+
+            # Serving priming writes its owned workspace. Trials instead own
+            # independent output storage so no measured binding can escape.
+            def output(shape, serving):
+                return torch.empty(shape, **factory) if benchmark else serving
+
+            if operation == "grouped_rmsnorm":
+                source = activation((tokens, width))
+                out = output((tokens, width), self.workspace.normalized)
+                run = lambda: _impl.run_grouped_rmsnorm_impl(
+                    source,
+                    self.hc_norm.weight,
+                    eps=self.config.rms_norm_eps,
+                    plan=state,
+                    out=out,
+                )
+                owners = (out,)
+            elif operation == "scaled_silu":
+                source = activation((tokens, self.lora_rank))
+                out = output(
+                    (tokens, self.lora_rank),
+                    self.workspace.local_bottleneck
+                    if local
+                    else self.workspace.bottleneck,
+                )
+                run = lambda: _impl.run_scaled_silu_impl(source, plan=state, out=out)
+                owners = (out,)
+            elif operation == "gate_mean":
+                source = activation((tokens, width))
+                gates = activation((tokens, width))
+                out = output(
+                    (tokens, hidden_size),
+                    self.workspace.local_block_input
+                    if local
+                    else self.workspace.block_input,
+                )
+                run = lambda: _impl.run_gate_mean_impl(
+                    source, gates, plan=state, out=out
+                )
+                owners = (out,)
+            else:
+                hidden = activation((tokens, width))
+                block = activation((tokens, self.hidden_size))
+                injection = activation((tokens, self.hc_count))
+                if operation == "combine":
+                    run = lambda: _impl.run_combine_impl(
+                        hidden,
+                        block,
+                        injection,
+                        plan=state,
+                    )
+                else:
+                    run = lambda: _impl.run_combine_norm_impl(
+                        hidden,
+                        block,
+                        injection,
+                        self.hc_norm.weight,
+                        eps=self.config.rms_norm_eps,
+                        plan=state,
+                    )
+                owners = ()
+            return PreparedCall(
+                run=run,
+                produce=produce,
+                owners=(*activation_owners, *owners),
+            )
+
+        return prepare
+
+    def _plan_for(self, operation: str):
+        try:
+            return self._plans[operation]
+        except KeyError:
+            raise PreparationResourceUnavailableError(
+                f"{self._preparation_prefix} lacks a declared {operation} plan"
+            ) from None
+
+    @property
+    def workspace(self) -> HyperConnectionWorkspace:
+        return self._workspace
+
+    def _binding(self, hidden_states: torch.Tensor, operation: str):
+        if self.tp_size > 1 and operation in ("scaled_silu", "gate_mean"):
+            return self.workspace.bind(
+                self._plan_for(operation), hidden_states.shape[0], local=True
+            )
+        return self.workspace.bind(self._plan_for(operation), hidden_states.shape[0])
+
+    def _hc_fp8(self):
+        """Rowwise-FP8 (down_w8, down_scale, up_w8, up_scale), or None.
+
+        Built on the first call outside CUDA graph capture (the profile run),
+        so the copies are counted before the KV cache is sized.
+        """
+        cache = self.__dict__.get("_hc_fp8_cache")
+        if cache is not None or not _HC_FP8 or torch.cuda.is_current_stream_capturing():
+            return cache
+        from .ops.hc_fp8_gemv import quantize_rowwise_fp8
+
+        down = (
+            self.input_mix_weight_down_block_inject
+            if self.use_combine
+            else self.input_mix_weight_down
+        )
+        cache = (
+            *quantize_rowwise_fp8(down.weight),
+            *quantize_rowwise_fp8(self.input_mix_weight_up.weight),
+        )
+        self.__dict__["_hc_fp8_cache"] = cache
+        logger.info_once("Qwen3.8 HyperConnection mix uses rowwise-FP8 decode GEMVs.")
+        return cache
+
+    def _mix_normalized(self, normalized: torch.Tensor):
+        if self.tp_size > 1:
+            return self._mix_sharded(normalized)
+        api = _hyperconnection_api() if self.workspace is not None else None
+        fp8 = self._hc_fp8()
+        if fp8 is not None and normalized.shape[0] <= 16:
+            from .ops.hc_fp8_gemv import fp8_gemv
+        else:
+            fp8 = None
+        if fp8 is not None and _HC_FUSE and normalized.stride(1) == 1:
+            from .ops.hc_fp8_gemv import (
+                _FUSED_DOWN,
+                _MIX_CFG,
+                fp8_gemv_down_silu,
+                fp8_gemv_down_silu_fused,
+                fp8_gemv_gate_mix,
+            )
+
+            down = fp8_gemv_down_silu_fused if _FUSED_DOWN else fp8_gemv_down_silu
+            projected, bottleneck = down(
+                normalized, fp8[0], fp8[1], self.lora_rank, self.hc_count
+            )
+            injection = (
+                projected[:, self.lora_rank : self.lora_rank + self.hc_count]
+                if self.use_combine
+                else None
+            )
+            block_input = fp8_gemv_gate_mix(
+                bottleneck, fp8[2], fp8[3], normalized, self.hc_count, *_MIX_CFG
+            )
+            return block_input, injection
+        if fp8 is not None and self.use_combine:
+            down_and_injection = fp8_gemv(normalized, fp8[0], fp8[1])
+            projected_down = down_and_injection[:, : self.lora_rank]
+            injection = down_and_injection[
+                :, self.lora_rank : self.lora_rank + self.hc_count
+            ]
+        elif fp8 is not None:
+            projected_down = fp8_gemv(normalized, fp8[0], fp8[1])
+            injection = None
+        elif self.use_combine:
+            down_and_injection = self.input_mix_weight_down_block_inject(normalized)
+            projected_down = down_and_injection[:, : self.lora_rank]
+            injection_start = self.lora_rank
+            # The projection owner stays live through the downstream residual
+            # combine; readers consume row-strided slices without staging.
+            injection = down_and_injection[
+                :, injection_start : injection_start + self.hc_count
+            ]
+        else:
+            projected_down = self.input_mix_weight_down(normalized)
+            injection = None
+
+        bottleneck = (
+            api.run_scaled_silu(
+                projected_down, binding=self._binding(normalized, "scaled_silu")
+            )
+            if api is not None
+            else hc_silu(projected_down, self.hc_count)
+        )
+        gate_logits = (
+            fp8_gemv(bottleneck, fp8[2], fp8[3])
+            if fp8 is not None
+            else self.input_mix_weight_up(bottleneck)
+        )
+        block_input = (
+            api.run_gate_mean(
+                normalized, gate_logits, binding=self._binding(normalized, "gate_mean")
+            )
+            if api is not None
+            else hc_gate_mix(normalized, gate_logits, self.hc_count)
+        )
+        return block_input, injection
+
+    def _mix_sharded(self, normalized: torch.Tensor):
+        api = _hyperconnection_api()
+        workspace = self.workspace
+        tokens = normalized.shape[0]
+        down = (
+            self.input_mix_weight_down_block_inject
+            if self.use_combine
+            else self.input_mix_weight_down
+        )
+        width = down.weight.shape[0]
+        projected_fp32 = workspace.local_down_fp32.view(-1)[: tokens * width].view(
+            tokens, width
+        )
+        torch.mm(normalized, down.weight.T, out=projected_fp32, out_dtype=torch.float32)
+        projected = workspace.local_down.view(-1)[: tokens * width].view(tokens, width)
+        projected.copy_(projected_fp32)
+        injection = (
+            projected[:, self.lora_rank : self.lora_rank + self.hc_count]
+            if self.use_combine
+            else None
+        )
+        bottleneck = api.run_scaled_silu(
+            projected[:, : self.lora_rank],
+            binding=self._binding(normalized, "scaled_silu"),
+        )
+        bottleneck = tensor_model_parallel_all_gather(bottleneck, dim=-1)
+        logits = workspace.local_gates[:tokens]
+        torch.mm(bottleneck, self.input_mix_weight_up.weight.T, out=logits)
+        local_hidden = self.hidden_size // self.tp_size
+        local_normalized = workspace.local_normalized[:tokens]
+        local_normalized.view(tokens, self.hc_count, local_hidden).copy_(
+            normalized.view(tokens, self.hc_count, self.hidden_size)[
+                :, :, self.tp_rank * local_hidden : (self.tp_rank + 1) * local_hidden
+            ]
+        )
+        block_input = api.run_gate_mean(
+            local_normalized, logits, binding=self._binding(normalized, "gate_mean")
+        )
+        return tensor_model_parallel_all_gather(block_input, dim=-1), injection
+
+
+__all__ = [
+    "GatedResidual",
+    "GroupedGemmaRMSNorm",
+    "HyperConnectionConfig",
+]

recipe/patches/qwen4_exp/nvidia/logit_dump.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Full-vocabulary prompt-logit dump for KLD measurement (VLLM_QWEN38_LOGIT_DUMP=<dir>).
+
+Use the forward-side dump (maybe_dump_forward, armed by <dir>/ARM) for measurements:
+on this stack a request that asks for prompt logprobs gets a different, run-to-run
+non-deterministic forward, so the prompt-logprob hook below is kept for diagnosis only.
+
+For every request that asks for prompt logprobs, each prefill chunk's logits (the
+same ``compute_logits`` the prompt-logprob path uses) are appended as float32 rows to
+``<dir>/<req_id>.bin``, with ``<first position> <rows> <vocab>`` per chunk in
+``<req_id>.idx``. Row j of a chunk holds the logits at prompt position first + j
+(the distribution over the token at first + j + 1). Measurement only: it copies
+every row to the host.
+"""
+
+import os
+import re
+
+import numpy as np
+import torch
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_ROWS_PER_CALL = 512
+_DIAG = os.environ.get("VLLM_QWEN38_LOGIT_DUMP_DIAG", "0") == "1"
+_last_forward: dict = {}
+
+
+def note_forward_output(hidden: torch.Tensor) -> None:
+    """Diagnostic: snapshot the model's final hidden states as the forward returns them."""
+    if _DIAG and not torch.cuda.is_current_stream_capturing() and hidden.shape[0] > 16:
+        _last_forward["hs"] = hidden.detach().clone()
+
+
+def _diagnose(logits_fn, hidden_states, input_batch) -> None:
+    n = int(input_batch.num_tokens)
+    hs = hidden_states[:n]
+    snap = _last_forward.get("hs")
+    msg = [f"rows {n}"]
+    if snap is not None and snap.shape[0] >= n:
+        d = (hs.float() - snap[:n].float()).abs()
+        msg.append(f"hidden vs forward-end snapshot: max {d.max().item():.3e} rows changed "
+                   f"{int((d.amax(-1) > 0).sum())}/{n}")
+    with torch.inference_mode():
+        one_a = logits_fn(hs[n - 1:n]).float()
+        one_b = logits_fn(hs[n - 1:n]).float()
+        chunk = logits_fn(hs[max(0, n - 512):n]).float()[-1:]
+        msg.append(f"last row: M=1 repeat max|d| {(one_a - one_b).abs().max().item():.3e}, "
+                   f"M=1 vs M={min(n, 512)} max|d| {(one_a - chunk).abs().max().item():.3e}")
+        if snap is not None and snap.shape[0] >= n:
+            one_s = logits_fn(snap[n - 1:n]).float()
+            msg.append(f"M=1 on snapshot vs current max|d| {(one_a - one_s).abs().max().item():.3e}")
+    logger.info("LOGIT DIAG: %s", " | ".join(msg))
+
+
+def maybe_dump_forward(out_dir, compute_logits, hidden: torch.Tensor, positions: torch.Tensor) -> None:
+    """Forward-side dump, independent of prompt logprobs: while ``<out_dir>/ARM`` exists
+    (its text: dump name, optional first position to keep), every eager forward of more than 16 rows appends its
+    full-vocab float32 logits to ``<out_dir>/<name>.bin`` and ``<first position> <rows>
+    <vocab>`` to ``<name>.idx``. Arm only around one request at a time."""
+    if hidden.shape[0] <= 16 or torch.cuda.is_current_stream_capturing():
+        return
+    arm = os.path.join(out_dir, "ARM")
+    if not os.path.exists(arm):
+        return
+    words = open(arm).read().split()
+    name = re.sub(r"[^A-Za-z0-9_.-]", "_", words[0] if words else "") or "dump"
+    min_pos = int(words[1]) if len(words) > 1 else 0
+    path = os.path.join(out_dir, name)
+    pos = positions[0] if positions.dim() == 2 else positions
+    first = int(pos[0].item())
+    skip = min(hidden.shape[0], max(0, min_pos - first))
+    if skip == hidden.shape[0]:
+        return
+    hidden, first = hidden[skip:], first + skip
+    rows = hidden.shape[0]
+    vocab = 0
+    with torch.inference_mode(), open(path + ".bin", "ab") as f:
+        for s in range(0, rows, _ROWS_PER_CALL):
+            logits = compute_logits(hidden[s:s + _ROWS_PER_CALL]).float()
+            vocab = logits.shape[-1]
+            logits.cpu().numpy().tofile(f)
+    with open(path + ".idx", "a") as f:
+        f.write(f"{first} {rows} {vocab}\n")
+
+
+def install(out_dir: str) -> None:
+    from vllm.v1.worker.gpu.sample import prompt_logprob as pl
+
+    Worker = pl.PromptLogprobsWorker
+    if getattr(Worker, "_qwen38_logit_dump", False):
+        return
+    os.makedirs(out_dir, exist_ok=True)
+    stock = Worker.compute_prompt_logprobs
+
+    def compute_prompt_logprobs(self, logits_fn, hidden_states, input_batch, all_token_ids,
+                                num_computed_tokens, prompt_lens):
+        need = self.uses_prompt_logprobs[input_batch.idx_mapping_np]
+        if np.any(need) and _DIAG:
+            _diagnose(logits_fn, hidden_states, input_batch)
+        if np.any(need):
+            qsl = input_batch.query_start_loc_np
+            first = input_batch.num_computed_prefill_tokens_np
+            for i, req_id in enumerate(input_batch.req_ids):
+                start, end = int(qsl[i]), int(qsl[i + 1])
+                if not need[i] or start >= end:
+                    continue
+                path = os.path.join(out_dir, re.sub(r"[^A-Za-z0-9_.-]", "_", req_id))
+                with torch.inference_mode(), open(path + ".bin", "ab") as f:
+                    for s in range(start, end, _ROWS_PER_CALL):
+                        e = min(end, s + _ROWS_PER_CALL)
+                        logits_fn(hidden_states[s:e]).float().cpu().numpy().tofile(f)
+                vocab = int(logits_fn(hidden_states[start:start + 1]).shape[-1])
+                with open(path + ".idx", "a") as f:
+                    f.write(f"{int(first[i])} {end - start} {vocab}\n")
+        return stock(self, logits_fn, hidden_states, input_batch, all_token_ids,
+                     num_computed_tokens, prompt_lens)
+
+    Worker.compute_prompt_logprobs = compute_prompt_logprobs
+    Worker._qwen38_logit_dump = True
+    logger.info("Qwen3.8 prompt-logit dump to %s (measurement only).", out_dir)

recipe/patches/qwen4_exp/nvidia/model.py

@@ -0,0 +1,1310 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inference-only Qwen4Exp model."""
+
+from collections.abc import Callable, Iterable
+from itertools import islice
+
+import torch
+from torch import nn
+
+from vllm import envs
+from vllm.config import VllmConfig
+from vllm.distributed import get_pp_group
+from vllm.model_executor.layers.fused_moe.utils import (
+    is_model_fused_shared_expert_compatible,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+import os as _os
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_FP8_LM_HEAD = _os.environ.get("VLLM_QWEN38_FP8_LM_HEAD", "0") == "1"
+_SMALL_GEMV = _os.environ.get("VLLM_QWEN38_SMALL_GEMV", "0") == "1"
+# Rowwise-FP8 decode GEMVs for the target's remaining plain bf16 projections
+# (PLE key/value projections; ~65 MB per step through sm80 WMMA fallbacks).
+_TARGET_FP8 = _os.environ.get("VLLM_QWEN38_TARGET_FP8", "0") == "1"
+# Fused shared-expert tail at decode: SiLU*mul prologue + MXFP8 down GEMV +
+# sigmoid(shared_expert_gate) epilogue in one kernel (stock: 5-6 small kernels).
+_SHARED_FUSE = _os.environ.get("VLLM_QWEN38_SHARED_FUSE", "0") == "1"
+_LOGIT_DUMP_DIAG = _os.environ.get("VLLM_QWEN38_LOGIT_DUMP_DIAG", "0") == "1"
+_LOGIT_DUMP_DIR = _os.environ.get("VLLM_QWEN38_LOGIT_DUMP", "")
+# One-launch GemmaRMSNorm (QSA q_norm/k_norm, MTP norms): the stock residual-free
+# path runs ~10 eager torch ops per call with compilation off.
+if _os.environ.get("VLLM_QWEN38_GEMMA_NORM", "0") == "1":
+    from .ops.gemma_norm import install_gemma_norm
+
+    install_gemma_norm()
+# Deferred-commit GDN verify: per-token records at verify, one post-acceptance commit.
+if _os.environ.get("VLLM_QWEN38_GDN_DEFER", "0") == "1":
+    from .gdn_defer import install as _install_gdn_defer
+
+    _install_gdn_defer()
+# Measurement only: full-vocab prompt logits to disk for KLD against a reference.
+if _os.environ.get("VLLM_QWEN38_LOGIT_DUMP"):
+    from .logit_dump import install as _install_logit_dump
+
+    _install_logit_dump(_os.environ["VLLM_QWEN38_LOGIT_DUMP"])
+from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
+    QwenGatedDeltaNetAttention,
+)
+from vllm.model_executor.layers.mamba.mamba_utils import (
+    MambaStateCopyFunc,
+    MambaStateCopyFuncCalculator,
+    MambaStateCopyFuncsByType,
+    MambaStateDtypeCalculator,
+    MambaStateShapeCalculator,
+)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.models.interfaces import (
+    HasInnerState,
+    IsHybrid,
+    MixtureOfExperts,
+    MultiModalEmbeddings,
+    SupportsLoRA,
+    SupportsMRoPE,
+    SupportsPP,
+    _require_is_multimodal,
+)
+from vllm.model_executor.models.qwen3_5 import (
+    Qwen3_5ForConditionalGeneration,
+    Qwen3_5Model,
+)
+from vllm.model_executor.models.qwen3_next import (
+    Qwen3NextAttention,
+    Qwen3NextMLP,
+    Qwen3NextSparseMoeBlock,
+)
+from vllm.model_executor.models.qwen3_vl import (
+    Qwen3_VisionTransformer,
+    Qwen3VLDummyInputsBuilder,
+    Qwen3VLMultiModalProcessor,
+    Qwen3VLProcessingInfo,
+)
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    StageMissingLayer,
+    WeightsMapper,
+    _merge_multimodal_embeddings,
+    extract_layer_index,
+    make_empty_intermediate_tensors_factory,
+    make_layers,
+    maybe_fuse_shared_experts,
+    maybe_prefix,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import MultiModalFeatureSpec
+from vllm.sequence import IntermediateTensors
+from vllm.tokenizers.registry import cached_tokenizer_from_config
+from vllm.transformers_utils.configs.qwen4_exp import (
+    Qwen4ExpTextConfig,
+)
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
+from vllm.v1.kv_cache_interface import MambaSpec
+
+from ..config import Qwen4ExpConfig
+from .b12x_ple import B12xNGramEmbedding, _resolve_ple_table_memory
+from .b12x_qsa import Qwen4ExpQSAAttention as B12xQSAAttention
+from .backend import uses_b12x
+from .hyperconnection import (
+    GatedResidual,
+    HyperConnectionConfig,
+    HyperConnectionWorkspace,
+)
+from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
+from .ple_layer import Qwen4ExpPLELayer
+from .qsa import Qwen4ExpQSAAttention
+
+
+def _is_file_backed_ple_weight(name: str) -> bool:
+    _, marker, shard_suffix = name.rpartition(
+        ".ple.ple_embedding.ngram_embedding.shard_"
+    )
+    shard_index, separator, suffix = shard_suffix.partition(".")
+    return bool(
+        marker
+        and separator
+        and shard_index.isdigit()
+        and suffix in {"weight", "weight_scale"}
+    )
+
+
+def without_modelopt_fp4(
+    quant_config: QuantizationConfig | None,
+) -> QuantizationConfig | None:
+    """Return ``None`` for weights excluded from Qwen4Exp ModelOpt-FP4."""
+
+    if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
+        return None
+    return quant_config
+
+
+def _remap_qsa_cache_scale_name(
+    name: str,
+    qsa_layer_ids: frozenset[int],
+) -> str:
+    """Map serialized main-cache scales onto the merged QSA owner.
+
+    Regular attention keeps cache scales below its ``attn`` child. QSA owns
+    that cache directly, so only QSA layers need the final path component
+    moved to the owner's persistent ``_k_scale``/``_v_scale`` buffers.
+    """
+
+    scale_suffixes = {
+        "k_proj.k_scale": "_k_scale",
+        "k_proj.output_scale": "_k_scale",
+        "attn.k_scale": "_k_scale",
+        "attn._k_scale": "_k_scale",
+        "k_scale": "_k_scale",
+        "_k_scale": "_k_scale",
+        "v_proj.v_scale": "_v_scale",
+        "v_proj.output_scale": "_v_scale",
+        "attn.v_scale": "_v_scale",
+        "attn._v_scale": "_v_scale",
+        "v_scale": "_v_scale",
+        "_v_scale": "_v_scale",
+    }
+    for layer_id in qsa_layer_ids:
+        marker = f"layers.{layer_id}.self_attn."
+        marker_start = name.find(marker)
+        if marker_start < 0 or (marker_start > 0 and name[marker_start - 1] != "."):
+            continue
+        suffix = name[marker_start + len(marker) :]
+        mapped_suffix = scale_suffixes.get(suffix)
+        if mapped_suffix is not None:
+            return f"{name[: marker_start + len(marker)]}{mapped_suffix}"
+    return name
+
+
+_QWEN4_EXP_IGNORED_MISSING_SUFFIXES = [
+    ".bias",
+    "_bias",
+    ".k_scale",
+    "_k_scale",
+    ".v_scale",
+    "_v_scale",
+    "_weight_scale",
+    "_input_scale",
+]
+
+# The checkpoint stores these projections separately; runtime packs each group
+# into adjacent logical shards of a MergedColumnParallelLinear.
+_EXTRA_WEIGHTS_MAPPER = WeightsMapper(
+    orig_to_new_stacked={
+        "hyper_connection.input_mix_weight_down.weight": (
+            "hyper_connection.input_mix_weight_down_block_inject.weight",
+            0,
+        ),
+        "hyper_connection.block_inject_weight.weight": (
+            "hyper_connection.input_mix_weight_down_block_inject.weight",
+            1,
+        ),
+        "ple.key_proj": ("ple.kv_proj", 0),
+        "ple.value_proj": ("ple.kv_proj", 1),
+    }
+)
+
+
+def _install_small_gemv_gate(gate: torch.nn.Module) -> None:
+    """Router GEMM fast path (VLLM_QWEN38_SMALL_GEMV=1). GB10 (sm_121) runs the
+    bf16 [tokens, 2560] x [512, 2560]^T router on an sm80 WMMA fallback at
+    ~27 us against a ~10 us bandwidth floor; decode-sized batches use the
+    split-K Triton GEMV instead. Overrides only this instance's forward, so
+    parameter names (gate.weight) and checkpoint loading are unchanged."""
+    if getattr(gate, "bias", None) is not None or not hasattr(gate, "weight"):
+        return
+    from vllm.models.qwen4_exp.nvidia.ops.small_gemv import small_mm
+
+    stock_forward = gate.forward
+
+    def forward(x: torch.Tensor):
+        if x.dim() == 2:
+            out = small_mm(x, gate.weight, x.dtype)
+            return out if getattr(gate, "return_bias", True) is False else (out, None)
+        return stock_forward(x)
+
+    gate.forward = forward
+
+
+def _install_fused_shared_expert(mlp: torch.nn.Module) -> None:
+    """Decode rows (<= 16) of a shared-expert MLP with an MXFP8 b12x down_proj
+    and a sigmoid expert gate run gate_up + one fused kernel
+    (ops.mxfp8_gemv.mxfp8_gemv_swiglu_gate); anything else keeps the stock path."""
+    from vllm.models.qwen4_exp.nvidia.ops.mxfp8_gemv import mxfp8_gemv_swiglu_gate
+
+    stock = mlp.forward
+
+    def forward(x):
+        down = mlp.down_proj
+        packed = getattr(down, "b12x_mxfp8_packed_weight", None)
+        gate = getattr(mlp, "expert_gate", None)
+        if (
+            packed is None
+            or gate is None
+            or x.dim() != 2
+            or x.shape[0] > 16
+            or x.dtype != torch.bfloat16
+            or packed.weight.values.is_meta
+            or getattr(down, "bias", None) is not None
+        ):
+            return stock(x)
+        gate_up, _ = mlp.gate_up_proj(x)
+        return mxfp8_gemv_swiglu_gate(
+            gate_up,
+            x if x.stride(1) == 1 else x.contiguous(),
+            gate.weight.view(-1),
+            packed.weight.values,
+            packed.weight.scale_rows[0],
+            down.b12x_mxfp8_output_size,
+            packed.in_features,
+        )
+
+    mlp.forward = forward
+
+
+class Qwen4ExpSparseMoeBlock(Qwen3NextSparseMoeBlock):
+    """Qwen3Next MoE with Qwen4Exp HC validation."""
+
+    def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
+        parallel_config = vllm_config.parallel_config
+        if parallel_config.use_sequence_parallel_moe:
+            raise NotImplementedError(
+                "Qwen4Exp HC does not support sequence-parallel MoE"
+            )
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+        config = vllm_config.model_config.hf_text_config
+        self.n_shared_experts = int(config.shared_expert_intermediate_size > 0)
+        if _SMALL_GEMV:
+            _install_small_gemv_gate(self.gate)
+        if _SHARED_FUSE and getattr(self, "shared_expert", None) is not None:
+            _install_fused_shared_expert(self.shared_expert)
+        if _os.environ.get("VLLM_QWEN38_EXPERT_STATS"):
+            from .expert_stats import install as _install_expert_stats
+
+            _install_expert_stats(self.gate, prefix, config.num_experts_per_tok,
+                                  config.num_experts, _os.environ["VLLM_QWEN38_EXPERT_STATS"])
+
+
+class Qwen4ExpDecoderLayer(nn.Module):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        layer_type: str,
+        prefix: str = "",
+        *,
+        workspace: HyperConnectionWorkspace | None = None,
+    ) -> None:
+        super().__init__()
+        config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
+        model_config = vllm_config.model_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+
+        self.config = config
+        self.layer_type = layer_type
+        self.layer_idx = extract_layer_index(prefix)
+        if vllm_config.parallel_config.use_sequence_parallel_moe:
+            raise NotImplementedError(
+                "Qwen4Exp HC does not support sequence-parallel MoE"
+            )
+        self.ple: Qwen4ExpPLELayer | None = None
+        ple_layer_ids = config.ple_layer_ids
+        if (self.layer_idx + 1) in ple_layer_ids:
+            ple_layer_ids_sorted = sorted(set(ple_layer_ids))
+            ple_dense_layer_id_map = {
+                abs_id: idx for idx, abs_id in enumerate(ple_layer_ids_sorted)
+            }
+            ple_dense_layer_id = ple_dense_layer_id_map[self.layer_idx + 1]
+            self.ple = Qwen4ExpPLELayer(
+                config,
+                vllm_config=vllm_config,
+                layer_idx=self.layer_idx,
+                ple_dense_layer_id=ple_dense_layer_id,
+                prefix=f"{prefix}.ple",
+            )
+
+        if layer_type == "linear_attention":
+            self.linear_attn = QwenGatedDeltaNetAttention(
+                config,
+                vllm_config=vllm_config,
+                prefix=f"{prefix}.linear_attn",
+                gqa_interleaved_layout=False,
+                overlap_input_projections=uses_b12x(vllm_config)
+                and envs.VLLM_QWEN3_8_FLASH_NEXT_OVERLAP,
+            )
+        elif layer_type == "full_attention":
+            use_qsa = getattr(config, "indexer_n_heads", None) is not None
+            if not use_qsa:
+                self.self_attn = Qwen3NextAttention(
+                    config,
+                    model_config=model_config,
+                    cache_config=cache_config,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.self_attn",
+                )
+            else:
+                attention_cls = (
+                    B12xQSAAttention if uses_b12x(vllm_config) else Qwen4ExpQSAAttention
+                )
+                self.self_attn = attention_cls(
+                    vllm_config=vllm_config,
+                    config=config,
+                    layer_id=self.layer_idx,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.self_attn",
+                )
+        else:
+            raise ValueError(f"Invalid layer_type {layer_type}")
+
+        mlp_only_layers = getattr(config, "mlp_only_layers", [])
+        num_experts = getattr(config, "num_experts", 0) or 0
+        absolute_layer_id = self.layer_idx + 1
+        is_moe_layer = self.layer_idx not in mlp_only_layers and (
+            num_experts > 0 and absolute_layer_id % config.decoder_sparse_step == 0
+        )
+        if is_moe_layer:
+            self.mlp = Qwen4ExpSparseMoeBlock(
+                vllm_config=vllm_config, prefix=f"{prefix}.mlp"
+            )
+        else:
+            self.mlp = Qwen3NextMLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+
+        hc_config = HyperConnectionConfig(
+            hc_count=config.hc_count,
+            hidden_size=config.hidden_size,
+            params_dtype=torch.bfloat16,
+            hc_lowrank=config.hc_lowrank,
+            rms_norm_eps=config.rms_norm_eps,
+            hc_per_branch_norm=True,
+        )
+        self.attn_hyper_connection = GatedResidual(
+            hc_config,
+            workspace=workspace,
+            prefix=maybe_prefix(prefix, "attn_hyper_connection"),
+        )
+        self.mlp_hyper_connection = GatedResidual(
+            hc_config,
+            workspace=workspace,
+            prefix=maybe_prefix(prefix, "mlp_hyper_connection"),
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        prev_block_output: torch.Tensor | None,
+        prev_injection: torch.Tensor | None,
+        positions: torch.Tensor,
+        *,
+        input_ids: torch.Tensor | None,
+        query_start_loc: torch.Tensor | None,
+        ngram_context: torch.Tensor | None,
+        output_indices: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if prev_block_output is None:
+            assert prev_injection is None
+        attn_hc = self.attn_hyper_connection
+        if self.ple is not None:
+            # PLE adds directly to the multi-stream state, so pending HC state
+            # must be materialized before the addition.
+            if prev_block_output is not None:
+                hidden_states = attn_hc.combine(
+                    hidden_states, prev_block_output, prev_injection
+                )
+                prev_block_output = prev_injection = None
+
+            if input_ids is None or query_start_loc is None or ngram_context is None:
+                raise RuntimeError("PLE inputs were not prepared")
+            hidden_states = self.ple(
+                hidden_states,
+                input_ids,
+                query_start_loc,
+                ngram_context,
+            )
+
+        # Fuse a pending combine with this HC module's mix when possible.
+        if prev_block_output is not None:
+            hidden_states, block_input, injection = attn_hc.combine_and_mix(
+                hidden_states, prev_block_output, prev_injection
+            )
+        else:
+            hidden_states, block_input, injection = attn_hc.mix(hidden_states)
+
+        if self.layer_type == "linear_attention":
+            attn_out = self.linear_attn(hidden_states=block_input)
+        elif self.layer_type == "full_attention":
+            attn_out = self.self_attn(
+                hidden_states=block_input,
+                positions=positions,
+            )
+        else:
+            raise ValueError("Invalid layer_type")
+
+        if output_indices is not None:
+            hidden_states = hidden_states[output_indices]
+            attn_out = attn_out[output_indices]
+            assert injection is not None
+            injection = injection[output_indices]
+        mlp_hc = self.mlp_hyper_connection
+        hidden_states, block_input, injection = mlp_hc.combine_and_mix(
+            hidden_states, attn_out, injection
+        )
+        mlp_out = self.mlp(block_input)
+        return hidden_states, mlp_out, injection
+
+
+class Qwen4ExpMixtureOfExperts(MixtureOfExperts):
+    """Expose Qwen4Exp routed experts through vLLM's EPLB protocol."""
+
+    def set_moe_parameters(self, layers: Iterable[nn.Module]) -> None:
+        self.moe_layers = []
+        self.moe_mlp_layers = []
+        example_moe = None
+        for layer in layers:
+            if isinstance(layer, Qwen4ExpDecoderLayer) and isinstance(
+                layer.mlp, Qwen4ExpSparseMoeBlock
+            ):
+                example_moe = layer.mlp
+                self.moe_mlp_layers.append(layer.mlp)
+                self.moe_layers.append(layer.mlp.experts)
+
+        self.num_moe_layers = len(self.moe_layers)
+        if example_moe is None:
+            self.num_expert_groups = 0
+            self.num_shared_experts = 0
+            self.num_logical_experts = 0
+            self.num_physical_experts = 0
+            self.num_local_physical_experts = 0
+            self.num_routed_experts = 0
+            self.num_redundant_experts = 0
+            return
+
+        self.num_expert_groups = 1
+        self.num_shared_experts = example_moe.n_shared_experts
+        self.num_logical_experts = example_moe.n_logical_experts
+        self.num_physical_experts = example_moe.n_physical_experts
+        self.num_local_physical_experts = example_moe.n_local_physical_experts
+        self.num_routed_experts = example_moe.n_routed_experts
+        self.num_redundant_experts = example_moe.n_redundant_experts
+
+    def update_physical_experts_metadata(
+        self,
+        num_physical_experts: int,
+        num_local_physical_experts: int,
+    ) -> None:
+        assert self.num_local_physical_experts == num_local_physical_experts
+        self.num_physical_experts = num_physical_experts
+        self.num_local_physical_experts = num_local_physical_experts
+        self.num_redundant_experts = num_physical_experts - self.num_logical_experts
+        for moe in self.moe_mlp_layers:
+            moe.n_physical_experts = num_physical_experts
+            moe.n_local_physical_experts = num_local_physical_experts
+            moe.n_redundant_experts = self.num_redundant_experts
+            moe.experts.update_expert_map()
+
+
+class Qwen4ExpModel(nn.Module):
+    hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
+        self.config = config
+        self._use_b12x = uses_b12x(vllm_config)
+        self.num_redundant_experts = (
+            vllm_config.parallel_config.eplb_config.num_redundant_experts
+        )
+        self.vocab_size = config.vocab_size
+        self._qsa_layer_ids = frozenset(
+            layer_idx
+            for layer_idx, layer_type in enumerate(config.layer_types)
+            if layer_type == "full_attention"
+            and getattr(config, "indexer_n_heads", None) is not None
+        )
+        self.embed_tokens = VocabParallelEmbedding(
+            self.vocab_size,
+            config.hidden_size,
+            prefix=maybe_prefix(prefix, "embed_tokens"),
+        )
+        hc_config = HyperConnectionConfig(
+            hc_count=config.hc_count,
+            hidden_size=config.hidden_size,
+            params_dtype=torch.bfloat16,
+            hc_lowrank=config.hc_lowrank,
+            rms_norm_eps=config.rms_norm_eps,
+            hc_per_branch_norm=True,
+        )
+        self.hyper_connection_workspace = (
+            HyperConnectionWorkspace(
+                hc_config, vllm_config.scheduler_config.max_num_batched_tokens
+            )
+            if self._use_b12x
+            else None
+        )
+
+        def get_layer(prefix: str) -> Qwen4ExpDecoderLayer:
+            layer_idx = extract_layer_index(prefix)
+            return Qwen4ExpDecoderLayer(
+                vllm_config,
+                layer_type=config.layer_types[layer_idx],
+                workspace=self.hyper_connection_workspace,
+                prefix=prefix,
+            )
+
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
+        )
+        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
+            self.layers,
+            Qwen4ExpSparseMoeBlock,
+            "mlp",
+        )
+        intermediate_size = config.hidden_size * config.hc_count
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states"], intermediate_size
+        )
+
+        self.hyper_connection_mixer: GatedResidual | None
+        if get_pp_group().is_last_rank:
+            hc_config = HyperConnectionConfig(
+                hc_count=config.hc_count,
+                hidden_size=config.hidden_size,
+                params_dtype=torch.bfloat16,
+                hc_lowrank=config.hc_lowrank,
+                rms_norm_eps=config.rms_norm_eps,
+                hc_per_branch_norm=True,
+            )
+            self.hyper_connection_mixer = GatedResidual(
+                hc_config,
+                use_combine=False,
+                workspace=self.hyper_connection_workspace,
+                prefix=maybe_prefix(prefix, "hyper_connection_mixer"),
+            )
+        else:
+            self.hyper_connection_mixer = None
+
+        spec_config = vllm_config.speculative_config
+        # MTP HC multi-stream outputs: when speculative method=="mtp" and the
+        # model uses HC with hc_count>1, retain the pre-final-mixer multi-stream
+        # hidden state [T, hc_count*H] so the MTP drafter can feed a real
+        # multi-stream backbone hidden on its first step (scheme A). Derived
+        # purely from config (NOT node identity) so P/D nodes stay consistent.
+        needs_mtp_hidden = (
+            spec_config is not None
+            and getattr(spec_config, "method", None) == "mtp"
+            and get_pp_group().is_last_rank
+        )
+        if needs_mtp_hidden:
+            self.register_buffer(
+                "_mtp_hidden_buffer",
+                torch.empty(
+                    vllm_config.scheduler_config.max_num_batched_tokens,
+                    config.hc_count * config.hidden_size,
+                    dtype=vllm_config.model_config.dtype,
+                ),
+                persistent=False,
+            )
+        else:
+            self.register_buffer("_mtp_hidden_buffer", None, persistent=False)
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        query_start_loc: torch.Tensor | None = None,
+        ngram_context: torch.Tensor | None = None,
+        deepstack_input_embeds: IntermediateTensors | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                if input_ids is None:
+                    raise ValueError("input_ids or inputs_embeds is required")
+                hidden_states = self.embed_input_ids(input_ids)
+            hidden_states = hidden_states.repeat(1, self.config.hc_count)
+        else:
+            if intermediate_tensors is None:
+                raise ValueError("pipeline stage requires intermediate tensors")
+            hidden_states = intermediate_tensors["hidden_states"]
+
+        block_output = None
+        injection = None
+        last_layer = None
+        for layer_idx, layer in islice(
+            enumerate(self.layers), self.start_layer, self.end_layer
+        ):
+            last_layer = layer
+            hidden_states, block_output, injection = layer(
+                hidden_states=hidden_states,
+                prev_block_output=block_output,
+                prev_injection=injection,
+                positions=positions,
+                input_ids=input_ids,
+                query_start_loc=query_start_loc,
+                ngram_context=ngram_context,
+            )
+            if deepstack_input_embeds is not None and layer_idx < len(
+                deepstack_input_embeds
+            ):
+                deepstack_embed = deepstack_input_embeds[
+                    f"deepstack_input_embeds_{layer_idx}"
+                ]
+                deepstack_embed = (
+                    deepstack_embed.unsqueeze(-2)
+                    .expand(
+                        *deepstack_embed.shape[:-1],
+                        self.config.hc_count,
+                        self.config.hidden_size,
+                    )
+                    .flatten(-2)
+                )
+                # Deepstack is an external addition to the materialized
+                # multi-stream state and therefore terminates delayed combine.
+                hidden_states = layer.mlp_hyper_connection.combine(
+                    hidden_states, block_output, injection
+                )
+                block_output = None
+                injection = None
+                hidden_states = hidden_states + deepstack_embed
+
+        if not get_pp_group().is_last_rank:
+            # PP transports one tensor, not the delayed HC tuple. Materialize
+            # with the HC module that produced the pending injection.
+            if last_layer is not None and block_output is not None:
+                hidden_states = last_layer.mlp_hyper_connection.combine(
+                    hidden_states, block_output, injection
+                )
+            return IntermediateTensors({"hidden_states": hidden_states})
+
+        # The final mixer consumes the last pending combine and returns both
+        # the sampled single stream and the materialized multi-stream state.
+        final_mixer = self.hyper_connection_mixer
+        assert final_mixer is not None
+        multi_hidden, sample_hidden_states, _ = final_mixer.combine_and_mix(
+            hidden_states, block_output, injection
+        )
+        if self._mtp_hidden_buffer is not None:
+            # Capture the pre-final-mixer multi-stream hidden state
+            # [T, hc_count*H] for the MTP drafter (zero extra compute:
+            # this tensor is needed by the final mixer regardless).
+            num_tokens = multi_hidden.shape[0]
+            self._mtp_hidden_buffer[:num_tokens].copy_(multi_hidden)
+        if _LOGIT_DUMP_DIAG:
+            from .logit_dump import note_forward_output
+
+            note_forward_output(sample_hidden_states)
+        return sample_hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        weights = (
+            (
+                _remap_qsa_cache_scale_name(name, self._qsa_layer_ids),
+                weight,
+            )
+            for name, weight in weights
+        )
+        weights = maybe_fuse_shared_experts(
+            weights,
+            enabled=self.is_fused_shared_expert_enabled,
+            n_routed_experts=getattr(self.config, "num_experts", 0) or 0,
+            n_shared_experts=1,
+            ckpt_prefix="mlp.shared_expert",
+        )
+        # Non-persistent PLE state rebuilt in __init__; skip any ckpt
+        # column for them.
+        skip_substrs = (
+            "hashstats_",
+            "token_lookup",
+            "hyper_connection_mixer.block_inject_weight",
+        )
+        mapper = self.hf_to_vllm_mapper | WeightsMapper(
+            orig_to_new_substr={substr: None for substr in skip_substrs}
+        )
+        # The final HC mixer only exists on the last PP rank; earlier ranks
+        # must drop its checkpoint weights instead of failing to place them.
+        ignore_prefixes = (
+            None
+            if self.hyper_connection_mixer is not None
+            else ["hyper_connection_mixer."]
+        )
+        loader = AutoWeightsLoader(
+            self,
+            ignore_unexpected_prefixes=ignore_prefixes,
+            ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
+        )
+        loaded = loader.load_weights(
+            weights,
+            mapper=mapper,
+        )
+        if _TARGET_FP8:
+            from vllm.logger import init_logger
+
+            from .ops.hc_fp8_gemv import install_fp8_decode_linears
+
+            names = install_fp8_decode_linears(self)
+            init_logger(__name__).info(
+                "Qwen3.8 target: %d bf16 projections use rowwise-FP8 decode GEMVs: %s",
+                len(names), ", ".join(names),
+            )
+        return loaded
+
+
+class Qwen4ExpForCausalLM(
+    nn.Module,
+    HasInnerState,
+    SupportsLoRA,
+    SupportsMRoPE,
+    SupportsPP,
+    Qwen4ExpMixtureOfExperts,
+    IsHybrid,
+):
+    packed_modules_mapping = {
+        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+        "kv_proj": ["key_proj", "value_proj"],
+        "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
+        "in_proj_ba": ["in_proj_b", "in_proj_a"],
+        "input_mix_weight_down_block_inject": [
+            "input_mix_weight_down",
+            "block_inject_weight",
+            "_input_mix_padding",
+        ],
+    }
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={"model.language_model.": "model."}
+    )
+    requires_raw_input_tokens = True
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
+        self.vllm_config = vllm_config
+        self.model_config = vllm_config.model_config
+        self.quant_config = vllm_config.quant_config
+        self.config = config
+        self.scheduler_config = vllm_config.scheduler_config
+        if vllm_config.cache_config.mamba_cache_mode == "all":
+            raise NotImplementedError(
+                "Qwen4Exp currently does not support 'all' prefix caching, "
+                "please use '--mamba-cache-mode=align' instead"
+            )
+        self.model = Qwen4ExpModel(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+        self.lm_head = ParallelLMHead(
+            config.vocab_size,
+            config.hidden_size,
+            prefix=maybe_prefix(prefix, "lm_head"),
+        )
+        self.logits_processor = LogitsProcessor(config.vocab_size, lm_head=self.lm_head)
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors
+        )
+        self.set_moe_parameters(self.model.layers)
+        if not uses_b12x(vllm_config):
+            enable_qwen4_exp_low_latency_gemm(self, self.model_config.dtype)
+
+    @staticmethod
+    def get_model_state_cls():
+        from .model_state import Qwen4ExpModelState
+
+        return Qwen4ExpModelState
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors:
+        # Forward kwargs unchanged so the runner's _maybe_add_ngram_kwargs
+        # path (query_start_loc / ngram_context) reaches Qwen4ExpModel.
+        return self.model(
+            input_ids,
+            positions,
+            intermediate_tensors,
+            inputs_embeds,
+            **kwargs,
+        )
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    @classmethod
+    def get_ple_mamba_state_dtype_from_config(
+        cls,
+        vllm_config: VllmConfig,
+    ) -> tuple[torch.dtype, ...]:
+        return MambaStateDtypeCalculator.short_conv_state_dtype(
+            vllm_config.model_config.dtype,
+            vllm_config.cache_config.mamba_cache_dtype,
+        )
+
+    @classmethod
+    def get_ple_mamba_state_shape_from_config(
+        cls, vllm_config: VllmConfig
+    ) -> tuple[tuple[int, int]]:
+        hf_config = vllm_config.model_config.hf_text_config
+        conv_kernel_size = hf_config.ple_conv_kernel_size
+        short_conv_dilation = hf_config.ngram_size
+        conv_state_len = (conv_kernel_size - 1) * short_conv_dilation
+        num_spec = (
+            vllm_config.speculative_config.num_speculative_tokens
+            if vllm_config.speculative_config
+            else 0
+        )
+        hc_count = hf_config.hc_count
+        hc_hidden_size = hf_config.hidden_size * hc_count
+        return MambaStateShapeCalculator.short_conv_state_shape(
+            tp_world_size=1,
+            intermediate_size=hc_hidden_size,
+            conv_kernel=conv_state_len + 1,
+            num_spec=num_spec,
+        )
+
+    @classmethod
+    def get_gdn_mamba_state_dtype_from_config(
+        cls, vllm_config: VllmConfig
+    ) -> tuple[torch.dtype, torch.dtype]:
+        return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
+            vllm_config.model_config.dtype,
+            vllm_config.cache_config.mamba_cache_dtype,
+            vllm_config.cache_config.mamba_ssm_cache_dtype,
+        )
+
+    @classmethod
+    def get_gdn_mamba_state_shape_from_config(
+        cls, vllm_config: VllmConfig
+    ) -> tuple[tuple[int, int], tuple[int, int]]:
+        parallel_config = vllm_config.parallel_config
+        hf_config = vllm_config.model_config.hf_text_config
+        tp_size = parallel_config.tensor_parallel_size
+        num_spec = (
+            vllm_config.speculative_config.num_speculative_tokens
+            if vllm_config.speculative_config
+            else 0
+        )
+        return MambaStateShapeCalculator.gated_delta_net_state_shape(
+            tp_size,
+            hf_config.linear_num_key_heads,
+            hf_config.linear_num_value_heads,
+            hf_config.linear_key_head_dim,
+            hf_config.linear_value_head_dim,
+            hf_config.linear_conv_kernel_dim,
+            num_spec,
+        )
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(
+        cls,
+        vllm_config: VllmConfig,
+    ) -> tuple[torch.dtype, torch.dtype]:
+        return cls.get_gdn_mamba_state_dtype_from_config(vllm_config)
+
+    @classmethod
+    def get_mamba_state_shape_from_config(
+        cls, vllm_config: VllmConfig
+    ) -> tuple[tuple[int, int], tuple[int, int]]:
+        return cls.get_gdn_mamba_state_shape_from_config(vllm_config)
+
+    @classmethod
+    def get_mamba_state_copy_func(
+        cls,
+    ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
+        return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()
+
+    @classmethod
+    def get_mamba_state_copy_funcs(
+        cls,
+        mamba_types: set[MambaAttentionBackendEnum],
+    ) -> MambaStateCopyFuncsByType:
+        copy_funcs_by_type = {
+            MambaAttentionBackendEnum.GDN_ATTN: cls.get_mamba_state_copy_func(),
+            MambaAttentionBackendEnum.SHORT_CONV: (
+                MambaStateCopyFuncCalculator.short_conv_state_copy_func()
+            ),
+        }
+        missing_types = mamba_types - copy_funcs_by_type.keys()
+        assert not missing_types, f"missing state copy funcs for {missing_types}"
+        return {
+            mamba_type: copy_funcs_by_type[mamba_type] for mamba_type in mamba_types
+        }
+
+    @classmethod
+    def get_mamba_specs_from_config(
+        cls, vllm_config: VllmConfig
+    ) -> tuple[MambaSpec, ...]:
+        """Return all MambaSpecs for this model (GDN layers + PLE layer).
+
+        The PLE layer uses a separate short_conv MambaSpec whose page_size_bytes
+        may exceed the GDN spec; callers should take the maximum.
+        """
+        return (
+            MambaSpec(
+                shapes=cls.get_gdn_mamba_state_shape_from_config(vllm_config),
+                dtypes=cls.get_gdn_mamba_state_dtype_from_config(vllm_config),
+                block_size=-1,
+            ),
+            MambaSpec(
+                shapes=cls.get_ple_mamba_state_shape_from_config(vllm_config),
+                dtypes=cls.get_ple_mamba_state_dtype_from_config(vllm_config),
+                block_size=-1,
+                tp_replicated=True,
+            ),
+        )
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        if _FP8_LM_HEAD and hidden_states.dim() == 2:
+            head = self._fp8_lm_head()
+            if head is not None:
+                return self._fp8_logits(hidden_states, head)
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    # --- Opt-in rowwise-fp8 target lm_head (VLLM_QWEN38_FP8_LM_HEAD=1) ---------
+    # The bf16 head is a 248320 x 2560 (1.27 GB) read on every decode step and the
+    # GB10 decode step is bandwidth-bound (~5.4 ms of a ~72 ms step); fp8 halves it.
+    # Changes verify logits slightly, so it is not bitwise lossless.
+    def _fp8_lm_head(self):
+        head = getattr(self, "_fp8_head_cache", None)
+        if head is not None:
+            return head
+        if torch.cuda.is_current_stream_capturing():
+            return None  # never quantize inside a graph capture; bf16 path instead
+        lp = self.logits_processor
+        if lp.soft_cap is not None or lp.scale != 1.0 or lp.logits_as_input:
+            return None
+        with torch.no_grad():
+            w = self.lm_head.weight.detach()
+            row_max = w.abs().amax(dim=1, keepdim=True).float().clamp(min=1e-6)
+            w8 = (w.float() * (448.0 / row_max)).to(torch.float8_e4m3fn)
+            row_scale = (row_max / 448.0).to(w.dtype).reshape(1, -1)
+            unit = torch.ones(1, dtype=torch.float32, device=w.device)
+        self._fp8_head_cache = (w8, row_scale, unit)
+        logger.info("Qwen4Exp target lm_head uses a rowwise-fp8 copy.")
+        return self._fp8_head_cache
+
+    def _fp8_logits(self, hidden_states: torch.Tensor, head) -> torch.Tensor | None:
+        w8, row_scale, unit = head
+        act_max = hidden_states.abs().amax(dim=-1, keepdim=True).clamp(min=1e-6)
+        a8 = (hidden_states * (448.0 / act_max)).to(torch.float8_e4m3fn)
+        logits = torch._scaled_mm(
+            a8, w8.t(), scale_a=unit, scale_b=unit, out_dtype=hidden_states.dtype
+        )
+        logits = logits * row_scale * (act_max / 448.0).to(hidden_states.dtype)
+        lp = self.logits_processor
+        logits = lp._gather_logits(logits)
+        if logits is not None:
+            logits = logits[..., : lp.org_vocab_size]
+        return logits
+
+    def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
+        return self.model._mtp_hidden_buffer
+
+    def get_mrope_input_positions(
+        self,
+        input_tokens: list[int],
+        mm_features: list[MultiModalFeatureSpec],
+    ) -> tuple[torch.Tensor, int]:
+        positions = torch.arange(len(input_tokens), dtype=torch.long)
+        return positions.unsqueeze(0).expand(3, -1), 0
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        mapper = self.hf_to_vllm_mapper | WeightsMapper(
+            orig_to_new_substr={"mtp.": None}
+        )
+        loader = AutoWeightsLoader(
+            self,
+            ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
+        )
+        return loader.load_weights(weights, mapper=mapper)
+
+    def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.logits_processor(self.lm_head, hidden_states, skip_gather=True)
+
+    @property
+    def checkpoint_file_weight_filter(self) -> Callable[[str], bool] | None:
+        if not (uses_b12x(self.vllm_config) and self.config.ple_layer_ids):
+            return None
+        if (
+            _resolve_ple_table_memory(
+                self.vllm_config.additional_config, self.config.ple_embedding_dtype
+            )
+            == "io_uring"
+        ):
+            return _is_file_backed_ple_weight
+        embeddings = [
+            module
+            for module in self.modules()
+            if isinstance(module, B12xNGramEmbedding)
+        ]
+        # Attached shared tables already hold every shard row: skip reading them.
+        if embeddings and all(
+            embedding.ngram_embedding.shared_table_attached for embedding in embeddings
+        ):
+            return _is_file_backed_ple_weight
+        return None
+
+
+class Qwen4ExpProcessingInfo(Qwen3VLProcessingInfo):
+    def get_hf_config(self) -> Qwen4ExpConfig:
+        return self.ctx.get_hf_config(Qwen4ExpConfig)
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Qwen3VLMultiModalProcessor,
+    info=Qwen4ExpProcessingInfo,
+    dummy_inputs=Qwen3VLDummyInputsBuilder,
+)
+class Qwen4ExpForConditionalGeneration(
+    Qwen3_5ForConditionalGeneration,
+    HasInnerState,
+    Qwen4ExpMixtureOfExperts,
+):
+    """Qwen3-VL vision tower backed by the Qwen4Exp language model."""
+
+    requires_raw_input_tokens = True
+
+    packed_modules_mapping = Qwen3_5ForConditionalGeneration.packed_modules_mapping | {
+        "kv_proj": ["key_proj", "value_proj"],
+        "input_mix_weight_down_block_inject": [
+            "input_mix_weight_down",
+            "block_inject_weight",
+            "_input_mix_padding",
+        ],
+    }
+
+    @staticmethod
+    def get_model_state_cls():
+        from .model_state import Qwen4ExpModelState
+
+        return Qwen4ExpModelState
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model") -> None:
+        nn.Module.__init__(self)
+        config: Qwen4ExpConfig = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        multimodal_config = vllm_config.model_config.multimodal_config
+        if multimodal_config is None:
+            raise ValueError(
+                "Qwen4ExpForConditionalGeneration requires multimodal_config"
+            )
+
+        self.config = config
+        self.model_config = vllm_config.model_config
+        self.multimodal_config = multimodal_config
+        self.language_model_only = multimodal_config.language_model_only
+        if self.language_model_only:
+            self.use_data_parallel = False
+            self.is_multimodal_pruning_enabled = False
+            self.video_pruning_method = None
+            self.video_pruning_rate = 0.0
+            self._tokenizer = None
+            self.visual = StageMissingLayer("vision_tower")
+            self._tower_model_names = []
+        else:
+            self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
+            self._init_video_pruning(multimodal_config)
+            self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config)
+
+            with self._mark_tower_model(vllm_config, {"image", "video"}):
+                self.visual = Qwen3_VisionTransformer(
+                    config.vision_config,
+                    norm_eps=config.text_config.rms_norm_eps,
+                    quant_config=quant_config,
+                    prefix=maybe_prefix(prefix, "visual"),
+                )
+
+        self.use_deepstack = (
+            not self.language_model_only
+            and bool(config.vision_config.deepstack_visual_indexes)
+            and not isinstance(self.visual, StageMissingLayer)
+        )
+        self.deepstack_num_level = (
+            len(config.vision_config.deepstack_visual_indexes)
+            if self.use_deepstack
+            else 0
+        )
+        self.visual_dim = config.vision_config.out_hidden_size
+        self.multiscale_dim = self.visual_dim * self.deepstack_num_level
+
+        if self.use_deepstack:
+            self.deepstack_input_embeds = [
+                torch.zeros(
+                    vllm_config.scheduler_config.max_num_batched_tokens,
+                    config.text_config.hidden_size,
+                )
+                for _ in range(self.deepstack_num_level)
+            ]
+            self.deepstack_input_embeds_num_tokens = 0
+
+        with self._mark_language_model(vllm_config):
+            self.language_model = Qwen4ExpForCausalLM(
+                vllm_config=vllm_config,
+                prefix=maybe_prefix(prefix, "language_model"),
+            )
+
+        self.make_empty_intermediate_tensors = (
+            self.language_model.make_empty_intermediate_tensors
+        )
+        if not get_pp_group().is_first_rank and self.use_deepstack:
+            assert self.language_model.model.start_layer >= len(
+                config.vision_config.deepstack_visual_indexes
+            ), (
+                "start_layer should be greater than or equal to "
+                "len(deepstack_visual_indexes)"
+            )
+        self.set_moe_parameters(self.language_model.model.layers)
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: MultiModalEmbeddings | None = None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        inputs_embeds = self._embed_text_input_ids(
+            input_ids,
+            self.language_model.embed_input_ids,
+            is_multimodal=is_multimodal,
+        )
+        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
+            return inputs_embeds
+        if self.language_model_only:
+            raise ValueError(
+                "Qwen4Exp language_model_only does not accept multimodal embeddings"
+            )
+
+        is_multimodal = _require_is_multimodal(is_multimodal)
+        if self.use_deepstack:
+            deepstack_input_embeds, multimodal_embeddings = (
+                self._compute_deepstack_embeds(
+                    inputs_embeds=inputs_embeds,
+                    multimodal_embeddings=multimodal_embeddings,
+                    is_multimodal=is_multimodal,
+                )
+            )
+        else:
+            deepstack_input_embeds = None
+
+        inputs_embeds = _merge_multimodal_embeddings(
+            inputs_embeds=inputs_embeds,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=is_multimodal,
+        )
+        if deepstack_input_embeds is not None:
+            self._set_deepstack_input_embeds(deepstack_input_embeds)
+        return inputs_embeds
+
+    def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
+        return self.language_model.get_mtp_target_hidden_states()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor | IntermediateTensors:
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+        if inputs_embeds is not None and get_pp_group().is_first_rank:
+            deepstack_input_embeds = self._get_deepstack_input_embeds(
+                inputs_embeds.size(0)
+            )
+        else:
+            deepstack_input_embeds = None
+
+        hidden_states = self.language_model.model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            query_start_loc=kwargs.get("query_start_loc"),
+            ngram_context=kwargs.get("ngram_context"),
+            deepstack_input_embeds=deepstack_input_embeds,
+        )
+        if inputs_embeds is not None and get_pp_group().is_first_rank:
+            self._clear_deepstack_input_embeds(inputs_embeds.size(0))
+        if _LOGIT_DUMP_DIR and isinstance(hidden_states, torch.Tensor):
+            from .logit_dump import maybe_dump_forward
+
+            maybe_dump_forward(_LOGIT_DUMP_DIR, self.language_model.compute_logits,
+                               hidden_states, positions)
+        return hidden_states
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        mapper = self.hf_to_vllm_mapper | WeightsMapper(
+            orig_to_new_substr={"mtp.": None},
+            orig_to_new_prefix={"visual.": None} if self.language_model_only else {},
+        )
+        loader = AutoWeightsLoader(
+            self,
+            ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
+        )
+        return loader.load_weights(weights, mapper=mapper)
+
+    @classmethod
+    def get_mamba_state_dtype_from_config(
+        cls,
+        vllm_config: VllmConfig,
+    ) -> tuple[torch.dtype, torch.dtype]:
+        return Qwen4ExpForCausalLM.get_mamba_state_dtype_from_config(vllm_config)
+
+    @classmethod
+    def get_mamba_state_shape_from_config(
+        cls,
+        vllm_config: VllmConfig,
+    ) -> tuple[tuple[int, int], tuple[int, int]]:
+        return Qwen4ExpForCausalLM.get_mamba_state_shape_from_config(vllm_config)
+
+    @classmethod
+    def get_mamba_state_copy_func(
+        cls,
+    ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
+        return Qwen4ExpForCausalLM.get_mamba_state_copy_func()
+
+    @classmethod
+    def get_mamba_state_copy_funcs(
+        cls,
+        mamba_types: set[MambaAttentionBackendEnum],
+    ) -> MambaStateCopyFuncsByType:
+        return Qwen4ExpForCausalLM.get_mamba_state_copy_funcs(mamba_types)
+
+    @classmethod
+    def get_mamba_specs_from_config(
+        cls, vllm_config: VllmConfig
+    ) -> tuple[MambaSpec, ...]:
+        return Qwen4ExpForCausalLM.get_mamba_specs_from_config(vllm_config)
+
+    @property
+    def checkpoint_file_weight_filter(self) -> Callable[[str], bool] | None:
+        return self.language_model.checkpoint_file_weight_filter
+
+
+__all__ = [
+    "Qwen4ExpDecoderLayer",
+    "Qwen4ExpForCausalLM",
+    "Qwen4ExpForConditionalGeneration",
+    "Qwen4ExpMixtureOfExperts",
+    "Qwen4ExpModel",
+    "Qwen4ExpSparseMoeBlock",
+]

recipe/patches/qwen4_exp/nvidia/mtp.py

@@ -0,0 +1,968 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inference-only Qwen4Exp multi-token predictor."""
+
+from __future__ import annotations
+
+import json
+import os
+import weakref
+from collections.abc import Iterable, Sequence
+from typing import Any
+
+import regex as re
+import torch
+from torch import nn
+
+import vllm.envs as envs
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig, replace, set_current_vllm_config
+from vllm.distributed import get_pp_group
+from vllm.model_executor.layers.fused_moe.utils import (
+    is_model_fused_shared_expert_compatible,
+)
+from vllm.model_executor.layers.linear import ColumnParallelLinear, ReplicatedLinear
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.utils import configure_quant_config
+from vllm.model_executor.models.interfaces import LocalArgmaxMixin, SupportsPP
+from vllm.model_executor.models.qwen3_5 import Qwen3_5Model
+from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    WeightsMapper,
+    get_draft_quant_config,
+    make_empty_intermediate_tensors_factory,
+    maybe_fuse_shared_experts,
+    maybe_prefix,
+)
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.sequence import IntermediateTensors
+from vllm.utils.b12x import (
+    B12xPreparationUnit,
+    B12xWorkload,
+    PreparationResourceUnavailableError,
+    b12x_layer,
+    b12x_layer_prefix,
+    get_b12x_mtp_feedback,
+    register_b12x_layer,
+    set_b12x_preparation_provider,
+)
+from vllm.utils.torch_utils import (
+    LayerNameType,
+    _encode_layer_name,
+    _resolve_layer_name,
+    direct_register_custom_op,
+)
+
+from ..config import Qwen4ExpTextConfig
+from .backend import uses_b12x
+from .hyperconnection import (
+    GatedResidual,
+    GroupedGemmaRMSNorm,
+    HyperConnectionConfig,
+    HyperConnectionWorkspace,
+)
+from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
+from .model import (
+    _EXTRA_WEIGHTS_MAPPER,
+    _QWEN4_EXP_IGNORED_MISSING_SUFFIXES,
+    Qwen4ExpDecoderLayer,
+    Qwen4ExpMixtureOfExperts,
+    Qwen4ExpSparseMoeBlock,
+    _remap_qsa_cache_scale_name,
+)
+
+
+def _mtp_api() -> Any:
+    api = get_b12x_mtp_feedback()
+    if api is None:
+        raise ImportError(
+            "Qwen4Exp MTP requires b12x.sequence.mtp_feedback; "
+            "install the b12x serving extra"
+        )
+    return api
+
+
+def _remap_ignored_layers(
+    ignored_layers: list[str],
+    mtp_start_layer_idx: int,
+) -> list[str]:
+    return [_remap_mtp_layer_name(name, mtp_start_layer_idx) for name in ignored_layers]
+
+
+def _remap_mtp_layer_name(name: str, mtp_start_layer_idx: int) -> str:
+    if not name.startswith("mtp."):
+        return name
+    return re.sub(
+        r"(?<=\.layers\.)\d+",
+        lambda match: str(mtp_start_layer_idx + int(match.group(0))),
+        name,
+    )
+
+
+def _remap_mtp_quantized_layers(
+    quantized_layers: dict[str, dict[str, Any]],
+    mtp_start_layer_idx: int,
+) -> dict[str, dict[str, Any]]:
+    return {
+        _remap_mtp_layer_name(name, mtp_start_layer_idx): layer_config
+        for name, layer_config in quantized_layers.items()
+    }
+
+
+_remap_quantized_layers = _remap_mtp_quantized_layers
+
+
+def _remap_mtp_weight_name(name: str) -> str | None:
+    """Map target-checkpoint names into the standalone draft model."""
+    for checkpoint_prefix in ("model.language_model.", "language_model."):
+        if name.startswith(checkpoint_prefix):
+            name = name.removeprefix(checkpoint_prefix)
+            break
+
+    if name.startswith("embed_tokens."):
+        name = f"model.{name}"
+    if name.startswith("model.mtp."):
+        name = name.removeprefix("model.")
+    for shared_head_prefix in (
+        "mtp.shared_head.head.",
+        "model.shared_head.head.",
+        "shared_head.head.",
+    ):
+        if name.startswith(shared_head_prefix):
+            return name.replace(shared_head_prefix, "lm_head.", 1)
+    if name.startswith("model.lm_head."):
+        return name.removeprefix("model.")
+    if name.startswith("mtp."):
+        return name.replace("mtp.", "model.", 1)
+    if name.startswith("model.embed_tokens.") or name.startswith("lm_head."):
+        return name
+    return None
+
+
+def _make_draft_vllm_config(
+    vllm_config: VllmConfig,
+    mtp_start_layer_idx: int,
+) -> VllmConfig:
+    speculative_config = vllm_config.speculative_config
+    if speculative_config is None or speculative_config.draft_model_config is None:
+        raise ValueError("speculative_config.draft_model_config must be set")
+
+    draft_quant_config = get_draft_quant_config(vllm_config)
+    if draft_quant_config is not None:
+        configure_quant_config(draft_quant_config, Qwen4ExpMTP)
+        quantized_layers = getattr(draft_quant_config, "quantized_layers", None)
+        if quantized_layers:
+            draft_quant_config.quantized_layers = (  # type: ignore[attr-defined]
+                _remap_mtp_quantized_layers(
+                    quantized_layers,
+                    mtp_start_layer_idx,
+                )
+            )
+        for attribute in ("ignored_layers", "exclude_modules"):
+            names = getattr(draft_quant_config, attribute, None)
+            if names:
+                setattr(
+                    draft_quant_config,
+                    attribute,
+                    _remap_ignored_layers(names, mtp_start_layer_idx),
+                )
+
+    cache_config = vllm_config.cache_config
+    if _MTP_BF16_KV and cache_config.cache_dtype in ("fp8", "fp8_e4m3"):
+        cache_config = replace(cache_config, cache_dtype="bfloat16")
+    draft_vllm_config = replace(
+        vllm_config,
+        model_config=speculative_config.draft_model_config,
+        cache_config=cache_config,
+    )
+    draft_vllm_config.quant_config = draft_quant_config
+    return draft_vllm_config
+
+
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        "positions": -1,
+        "intermediate_tensors": 0,
+        "inputs_embeds": 0,
+        "hidden_states": 0,
+    }
+)
+class Qwen4ExpMultiTokenPredictor(nn.Module):
+    """Shared draft stack with optional fused token/multi-stream feedback."""
+
+    hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
+        self.config = config
+        self._use_b12x = uses_b12x(vllm_config)
+        self.vocab_size = config.vocab_size
+        self.hidden_size = config.hidden_size
+        self.hc_count = config.hc_count
+        self.mtp_start_layer_idx = config.num_hidden_layers
+        self.num_mtp_layers = int(getattr(config, "mtp_num_hidden_layers", 1))
+        self.supports_mtp_prefill_compaction = (
+            self._use_b12x
+            and self.num_mtp_layers == 1
+            and envs.VLLM_QWEN3_8_FLASH_NEXT_MTP_COMPACT
+            and vllm_config.scheduler_config.max_num_seqs == 1
+        )
+        self._prefill_output_indices: torch.Tensor | None = None
+        if self.num_mtp_layers < 1:
+            raise ValueError("MTP requires at least one predictor layer")
+
+        self.embed_tokens = VocabParallelEmbedding(
+            self.vocab_size,
+            self.hidden_size,
+            prefix=maybe_prefix(prefix, "embed_tokens"),
+        )
+        draft_vllm_config = _make_draft_vllm_config(
+            vllm_config,
+            self.mtp_start_layer_idx,
+        )
+        with set_current_vllm_config(draft_vllm_config, prefix=prefix):
+            projection_cls = (
+                ReplicatedLinear if self._use_b12x else ColumnParallelLinear
+            )
+            projection_kwargs = {} if self._use_b12x else {"gather_output": True}
+            self.fc_embedding = projection_cls(
+                self.hidden_size,
+                self.hidden_size,
+                bias=False,
+                params_dtype=torch.bfloat16,
+                quant_config=None if self._use_b12x else draft_vllm_config.quant_config,
+                **projection_kwargs,
+                prefix=maybe_prefix(prefix, "fc_embedding"),
+                return_bias=False,
+            )
+            self.fc_hidden = projection_cls(
+                self.hidden_size,
+                self.hidden_size,
+                bias=False,
+                params_dtype=torch.bfloat16,
+                quant_config=None if self._use_b12x else draft_vllm_config.quant_config,
+                **projection_kwargs,
+                prefix=maybe_prefix(prefix, "fc_hidden"),
+                return_bias=False,
+            )
+            hc_config = HyperConnectionConfig(
+                hc_count=self.hc_count,
+                hidden_size=self.hidden_size,
+                params_dtype=torch.bfloat16,
+                hc_lowrank=config.hc_lowrank,
+                rms_norm_eps=config.rms_norm_eps,
+                hc_per_branch_norm=True,
+            )
+            max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+            self.hyper_connection_workspace = (
+                HyperConnectionWorkspace(hc_config, max_tokens)
+                if self._use_b12x
+                else None
+            )
+            self.layers = nn.ModuleList(
+                [
+                    Qwen4ExpDecoderLayer(
+                        draft_vllm_config,
+                        layer_type="full_attention",
+                        workspace=self.hyper_connection_workspace,
+                        prefix=(
+                            f"{prefix}.layers.{self.mtp_start_layer_idx + idx}"
+                            if prefix
+                            else f"layers.{self.mtp_start_layer_idx + idx}"
+                        ),
+                    )
+                    for idx in range(self.num_mtp_layers)
+                ]
+            )
+            self.hyper_connection_mixer = GatedResidual(
+                hc_config,
+                workspace=self.hyper_connection_workspace,
+                use_combine=False,
+                prefix=maybe_prefix(prefix, "hyper_connection_mixer"),
+            )
+
+        self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
+            self.layers, Qwen4ExpSparseMoeBlock, "mlp"
+        )
+        self.pre_fc_norm_embedding = GroupedGemmaRMSNorm(
+            self.hidden_size,
+            eps=config.rms_norm_eps,
+            group_size=None,
+            dtype=torch.bfloat16,
+        )
+        self.pre_fc_norm_hidden = GroupedGemmaRMSNorm(
+            self.hidden_size * self.hc_count,
+            eps=config.rms_norm_eps,
+            group_size=None,
+            dtype=torch.bfloat16,
+        )
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states"], self.hidden_size * self.hc_count
+        )
+
+        device = torch.device(current_platform.current_device())
+        self.register_buffer(
+            "_decode_output_indices",
+            torch.zeros(1, dtype=torch.int64, device=device),
+            persistent=False,
+        )
+        if self.supports_mtp_prefill_compaction:
+            self._prefill_output_indices = self._decode_output_indices
+        if not self._use_b12x:
+            return
+        self._feedback_caps = _mtp_api().Caps(
+            device=device,
+            max_tokens=max_tokens,
+            hidden_size=self.hidden_size,
+            streams=self.hc_count,
+            dtype=torch.bfloat16,
+        )
+        self._plan = None
+        self._feedback_scratch_spec: tuple[tuple[int, ...], torch.dtype] | None = None
+        self.register_buffer("_feedback_token_embedding", None, persistent=False)
+        self.register_buffer("_feedback_multi_state", None, persistent=False)
+        self.register_buffer("_feedback_output", None, persistent=False)
+        compilation_config = vllm_config.compilation_config
+        if prefix in compilation_config.static_forward_context:
+            raise ValueError(f"duplicate layer name: {prefix}")
+        compilation_config.static_forward_context[prefix] = self
+        self.prefix = prefix
+        b12x_name = b12x_layer_prefix(self)
+        self._b12x_layer_name = _encode_layer_name(b12x_name)
+        register_b12x_layer(b12x_name, self)
+        if not getattr(self, "b12x_preparation_suppressed", False):
+            set_b12x_preparation_provider(self, self)
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def set_prefill_output_indices(self, output_indices: torch.Tensor | None) -> None:
+        """Select request-tail outputs after populating the attention cache."""
+        # AOT reuses the same compiled callable for first and subsequent drafts.
+        # Keep the one-row selector's tensor contract stable in both phases.
+        self._prefill_output_indices = (
+            self._decode_output_indices if output_indices is None else output_indices
+        )
+
+    def snapshot_qsa_interval_starts(self) -> None:
+        for layer in self.layers:
+            attention = getattr(layer, "self_attn", None)
+            snapshot = getattr(attention, "snapshot_speculative_interval_starts", None)
+            if snapshot is not None:
+                snapshot()
+
+    def restore_qsa_interval_starts(self) -> None:
+        for layer in self.layers:
+            attention = getattr(layer, "self_attn", None)
+            restore = getattr(attention, "restore_speculative_interval_starts", None)
+            if restore is not None:
+                restore()
+
+    def set_skip_topk(self, skip: bool) -> None:
+        """Reuse each draft attention layer's anchor selection within a round."""
+        for layer in self.layers:
+            attention = getattr(layer, "self_attn", None)
+            setter = getattr(attention, "set_skip_topk", None)
+            if setter is not None:
+                setter(skip and self.num_mtp_layers == 1)
+            elif (
+                attention is not None
+                and getattr(attention, "indexer", None) is not None
+            ):
+                attention.indexer.skip_topk = skip and self.num_mtp_layers == 1
+
+    def compact_topk_indices(self, source_rows: torch.Tensor) -> None:
+        """Capture accepted-token-aligned draft selections in request order."""
+        for layer in self.layers:
+            attention = getattr(layer, "self_attn", None)
+            compact = getattr(attention, "compact_topk_indices", None)
+            if compact is not None:
+                compact(source_rows)
+            elif (
+                attention is not None
+                and getattr(attention, "indexer", None) is not None
+            ):
+                buffer = attention.topk_indices_buffer
+                buffer[: source_rows.numel()].copy_(buffer.index_select(0, source_rows))
+
+    @property
+    def _feedback_request_name(self) -> str:
+        return f"{self.prefix}.mtp.feedback"
+
+    @staticmethod
+    def _alignment(tensor: torch.Tensor) -> int:
+        pointer = tensor.data_ptr()
+        return min(16, pointer & -pointer) if pointer else 16
+
+    def _feedback_plan(self):
+        return _mtp_api().plan(
+            self._feedback_caps,
+            invocation={
+                "input_alignments": (
+                    16,
+                    16,
+                    self._alignment(self.pre_fc_norm_embedding.weight),
+                    self._alignment(self.pre_fc_norm_hidden.weight),
+                )
+            },
+        )
+
+    def _feedback_resources_ready(self) -> bool:
+        weights = (
+            self.pre_fc_norm_embedding.weight,
+            self.pre_fc_norm_hidden.weight,
+            self.fc_embedding.weight,
+            self.fc_hidden.weight,
+        )
+        return not any(weight.is_meta for weight in weights)
+
+    def get_b12x_preparation_units(
+        self, layer: torch.nn.Module, workload: B12xWorkload
+    ) -> Sequence[B12xPreparationUnit]:
+        if layer is not self:
+            raise ValueError("MTP preparation owner mismatch")
+        if not self._use_b12x or workload.stage != "weights":
+            return ()
+        if not self._feedback_resources_ready():
+            return ()
+        if workload.max_tokens > self._feedback_caps.max_tokens:
+            raise PreparationResourceUnavailableError(
+                f"{self.prefix} MTP feedback capacity "
+                f"{self._feedback_caps.max_tokens} cannot serve "
+                f"{workload.max_tokens} tokens"
+            )
+        plan = self._feedback_plan()
+        self._plan = plan
+        request = plan.request(
+            name=self._feedback_request_name,
+            prepare_call=self._prepare_feedback_call,
+            benchmark_call=self._feedback_benchmark_factory(),
+        )
+        return (
+            B12xPreparationUnit(
+                name="MTP_FEEDBACK",
+                key=(self.prefix, workload.max_tokens),
+                requests=(request,),
+                stage="weights",
+            ),
+        )
+
+    def _publish_feedback_storage(self, state) -> None:
+        """Publish fixed feedback storage while keeping arena views transient."""
+        (spec,) = state.layout.scratch_specs()
+        self._feedback_scratch_spec = (spec.shape, spec.dtype)
+        factory = dict(device=spec.device, dtype=torch.bfloat16)
+        tokens = state.layout.caps.max_tokens
+        self._feedback_token_embedding = torch.empty(
+            (tokens, self.hidden_size), **factory
+        )
+        self._feedback_multi_state = torch.empty(
+            (tokens, self.hc_count, self.hidden_size), **factory
+        )
+        self._feedback_output = torch.empty(
+            state.layout.output_storage_shape(), **factory
+        )
+
+    def _prepare_feedback_call(self, state):
+        self._publish_feedback_storage(state)
+        token_embedding = self._feedback_token_embedding
+        multi_state = self._feedback_multi_state
+        output = self._feedback_output
+        assert token_embedding is not None
+        assert multi_state is not None
+        assert output is not None
+        assert self._feedback_scratch_spec is not None
+        # Trial scratch belongs to the trial; the serving workspace is drawn
+        # from only inside the feedback op body.
+        shape, dtype = self._feedback_scratch_spec
+        scratch = torch.empty(shape, dtype=dtype, device=token_embedding.device)
+        return self._feedback_call(
+            state,
+            scratch=scratch,
+            token_embedding=token_embedding,
+            multi_state=multi_state,
+            output=output,
+        )
+
+    def _feedback_benchmark_factory(self):
+        shared = None
+
+        def factory(state):
+            nonlocal shared
+            tensors = None if shared is None else tuple(ref() for ref in shared)
+            if tensors is None or any(tensor is None for tensor in tensors):
+                (spec,) = state.layout.scratch_specs()
+                scratch = torch.empty(spec.shape, dtype=spec.dtype, device=spec.device)
+                tensor_factory = dict(device=spec.device, dtype=torch.bfloat16)
+                tokens = state.layout.caps.max_tokens
+                token_embedding = torch.empty(
+                    (tokens, self.hidden_size), **tensor_factory
+                )
+                multi_state = torch.empty(
+                    (tokens, self.hc_count, self.hidden_size), **tensor_factory
+                )
+                output = torch.empty(
+                    state.layout.output_storage_shape(), **tensor_factory
+                )
+                tensors = scratch, token_embedding, multi_state, output
+                shared = tuple(weakref.ref(tensor) for tensor in tensors)
+            scratch, token_embedding, multi_state, output = tensors
+            return self._feedback_call(
+                state,
+                scratch=scratch,
+                token_embedding=token_embedding,
+                multi_state=multi_state,
+                output=output,
+            )
+
+        return factory
+
+    def _feedback_call(
+        self,
+        state,
+        *,
+        scratch,
+        token_embedding,
+        multi_state,
+        output,
+    ):
+        from b12x.preparation import PreparedCall
+
+        def produce() -> None:
+            # Prime the installed fixed route with real-shaped, nonzero
+            # activations while borrowing the loaded norm and FC weights.
+            token_embedding.fill_(0.125)
+            multi_state.fill_(0.25)
+
+        def reset() -> None:
+            scratch.zero_()
+            output.zero_()
+
+        def run():
+            return state.run_tensors(
+                token_embedding,
+                multi_state,
+                self.pre_fc_norm_embedding.weight,
+                self.pre_fc_norm_hidden.weight,
+                self.fc_embedding.weight,
+                self.fc_hidden.weight,
+                scratch,
+                output,
+                eps=self.config.rms_norm_eps,
+            )
+
+        return PreparedCall(
+            run=run,
+            produce=produce,
+            reset=reset,
+            capture_safe=False,
+        )
+
+    def _run_feedback(
+        self,
+        token_embedding: torch.Tensor,
+        multi_state: torch.Tensor,
+    ) -> None:
+        plan = self._plan
+        if plan is None or self._feedback_scratch_spec is None:
+            raise PreparationResourceUnavailableError(
+                f"{self.prefix} MTP feedback is not prepared"
+            )
+        from vllm.v1.worker.workspace import current_workspace_manager
+
+        (scratch,) = current_workspace_manager().get_simultaneous(
+            self._feedback_scratch_spec
+        )
+        num_tokens = token_embedding.shape[0]
+        if num_tokens > self._feedback_caps.max_tokens:
+            raise ValueError(
+                "Qwen4Exp MTP feedback capacity exceeded: "
+                f"{num_tokens}/{self._feedback_caps.max_tokens} tokens"
+            )
+        multi_state = multi_state.reshape(num_tokens, self.hc_count, self.hidden_size)
+        self._feedback_token_embedding[:num_tokens].copy_(token_embedding)
+        self._feedback_multi_state[:num_tokens].copy_(multi_state)
+        binding = _mtp_api().bind(
+            plan,
+            scratch=scratch,
+            token_embedding=self._feedback_token_embedding,
+            multi_state=self._feedback_multi_state,
+            token_norm_weight=self.pre_fc_norm_embedding.weight,
+            state_norm_weight=self.pre_fc_norm_hidden.weight,
+            embedding_fc_weight=self.fc_embedding.weight.data,
+            hidden_fc_weight=self.fc_hidden.weight.data,
+            output=self._feedback_output,
+            tokens=num_tokens,
+        )
+        _mtp_api().run(binding, eps=self.config.rms_norm_eps)
+
+    def _prepare_feedback(
+        self,
+        token_embedding: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        if not self._use_b12x:
+            embedding = self.fc_embedding(self.pre_fc_norm_embedding(token_embedding))
+            multi_state = self.pre_fc_norm_hidden(hidden_states).reshape(
+                -1, self.hc_count, self.hidden_size
+            )
+            return (self.fc_hidden(multi_state) + embedding.unsqueeze(1)).flatten(-2)
+        if self._plan is None or self._feedback_output is None:
+            raise PreparationResourceUnavailableError(
+                f"{self.prefix} MTP feedback is not prepared"
+            )
+        if torch.compiler.is_compiling():
+            torch.ops.vllm.qwen4_exp_b12x_mtp_feedback(
+                token_embedding,
+                hidden_states,
+                self._feedback_output,
+                self._b12x_layer_name,
+            )
+        else:
+            self._run_feedback(token_embedding, hidden_states)
+        return self._feedback_output[: token_embedding.shape[0]].flatten(-2)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor | None = None,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | IntermediateTensors:
+        if get_pp_group().is_first_rank:
+            if hidden_states is None:
+                raise ValueError("MTP requires target-model hidden states")
+            if inputs_embeds is None:
+                if input_ids is None:
+                    raise ValueError("input_ids or inputs_embeds is required")
+                inputs_embeds = self.embed_input_ids(input_ids)
+            hidden_states = self._prepare_feedback(inputs_embeds, hidden_states)
+        else:
+            if intermediate_tensors is None:
+                raise ValueError("pipeline stage requires intermediate tensors")
+            hidden_states = intermediate_tensors["hidden_states"]
+
+        layer = self.layers[spec_step_idx % self.num_mtp_layers]
+        hidden_states, block_output, injection = layer(
+            hidden_states=hidden_states,
+            prev_block_output=None,
+            prev_injection=None,
+            positions=positions,
+            input_ids=None,
+            query_start_loc=None,
+            ngram_context=None,
+            output_indices=self._prefill_output_indices,
+        )
+        if not get_pp_group().is_last_rank:
+            hidden_states = layer.mlp_hyper_connection.combine(
+                hidden_states, block_output, injection
+            )
+            return IntermediateTensors({"hidden_states": hidden_states})
+
+        multi_hidden, sample_hidden_states, _ = (
+            self.hyper_connection_mixer.combine_and_mix(
+                hidden_states, block_output, injection
+            )
+        )
+        return sample_hidden_states, multi_hidden
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        weights = maybe_fuse_shared_experts(
+            weights,
+            enabled=self.is_fused_shared_expert_enabled,
+            n_routed_experts=getattr(self.config, "num_experts", 0) or 0,
+            n_shared_experts=1,
+            ckpt_prefix="mlp.shared_expert",
+        )
+        loader = AutoWeightsLoader(
+            self,
+            ignore_unexpected_suffixes=(_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy()),
+        )
+        mapper = self.hf_to_vllm_mapper | WeightsMapper(
+            orig_to_new_substr={"hyper_connection_mixer.block_inject_weight": None}
+        )
+        return loader.load_weights(weights, mapper=mapper)
+
+
+def _mtp_feedback_op(
+    token_embedding: torch.Tensor,
+    multi_state: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: LayerNameType,
+) -> None:
+    layer = b12x_layer(_resolve_layer_name(layer_name))
+    layer._run_feedback(token_embedding, multi_state)
+
+
+def _mtp_feedback_fake(
+    token_embedding: torch.Tensor,
+    multi_state: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: LayerNameType,
+) -> None:
+    return
+
+
+direct_register_custom_op(
+    op_name="qwen4_exp_b12x_mtp_feedback",
+    op_func=_mtp_feedback_op,
+    mutates_args=["output"],
+    fake_impl=_mtp_feedback_fake,
+)
+
+
+logger = init_logger(__name__)
+
+# Reduced draft vocabulary: path to a JSON list of target token ids. The NVFP4
+# draft head then holds only those rows (the 4 draft steps read ~1.4 ms of
+# full-vocab head per step otherwise) and logits are scattered back to the
+# full vocabulary with -inf elsewhere, so rejection sampling stays exact.
+_DRAFT_VOCAB = os.environ.get("VLLM_QWEN38_DRAFT_VOCAB")
+# Rowwise-FP8 decode GEMVs for the draft layer's unquantized bf16 projections
+# (self-attention, shared expert, fc_embedding/fc_hidden; ~139 MB per draft step).
+_MTP_FP8 = os.environ.get("VLLM_QWEN38_MTP_FP8", "0") == "1"
+# Store the draft-vocabulary rows as row-major NVFP4 (0.56 vs 1.03 bytes per
+# weight), requantized from the shared MXFP8 head; the stock draft head is NVFP4.
+_DRAFT_NVFP4 = os.environ.get("VLLM_QWEN38_DRAFT_NVFP4", "0") == "1"
+# Keep the draft layer's KV cache in bf16 while the target's stays FP8: the draft
+# attends over its own cache at every draft step, so FP8 rounding there costs
+# acceptance (reported ~+7.5% greedy decode on the NVFP4 checkpoint).
+_MTP_BF16_KV = os.environ.get("VLLM_QWEN38_MTP_BF16_KV", "0") == "1"
+
+
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        "positions": -1,
+        "intermediate_tensors": 0,
+        "inputs_embeds": 0,
+        "hidden_states": 0,
+    }
+)
+class Qwen4ExpMTP(
+    LocalArgmaxMixin,
+    nn.Module,
+    SupportsPP,
+    Qwen4ExpMixtureOfExperts,
+):
+    checkpoint_weight_name_prefixes = tuple(
+        checkpoint_prefix + weight_prefix
+        for checkpoint_prefix in ("", "model.language_model.", "language_model.")
+        for weight_prefix in (
+            "mtp.",
+            "model.mtp.",
+            "embed_tokens.",
+            "model.embed_tokens.",
+            "lm_head.",
+            "model.lm_head.",
+            "shared_head.head.",
+            "model.shared_head.head.",
+        )
+    )
+
+    packed_modules_mapping = {
+        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+        "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
+        "in_proj_ba": ["in_proj_b", "in_proj_a"],
+        "input_mix_weight_down_block_inject": [
+            "input_mix_weight_down",
+            "block_inject_weight",
+            "_input_mix_padding",
+        ],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
+        self.vllm_config = vllm_config
+        if vllm_config.cache_config.mamba_cache_mode == "all":
+            raise NotImplementedError("Qwen4Exp MTP requires --mamba-cache-mode=align")
+        self.quant_config = vllm_config.quant_config
+        super().__init__()
+        self.has_own_lm_head = envs.VLLM_MTP_NVFP4_LM_HEAD
+        if (
+            self.has_own_lm_head
+            and envs.is_set("VLLM_MTP_NVFP4_LM_HEAD")
+            and config.tie_word_embeddings
+        ):
+            raise ValueError("NVFP4 draft head requires untied word embeddings")
+        self.config = config
+        self.model = Qwen4ExpMultiTokenPredictor(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "mtp"),
+        )
+
+        if get_pp_group().is_last_rank:
+            self.lm_head = ParallelLMHead(
+                config.vocab_size,
+                config.hidden_size,
+                quant_config=self.quant_config,
+                prefix=maybe_prefix(prefix, "lm_head"),
+                lm_head_quantization="nvfp4" if self.has_own_lm_head else None,
+            )
+            self.has_own_lm_head = self.lm_head.runtime_lm_head_quantization == "nvfp4"
+            if config.tie_word_embeddings:
+                self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
+        else:
+            self.lm_head = PPMissingLayer()
+
+        self.logits_processor = LogitsProcessor(
+            config.vocab_size,
+            lm_head=self.lm_head,
+        )
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors
+        )
+        self.set_moe_parameters(self.model.layers)
+        if not uses_b12x(vllm_config):
+            enable_qwen4_exp_low_latency_gemm(self, vllm_config.model_config.dtype)
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor | None = None,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_idx: int = 0,
+        **kwargs: object,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | IntermediateTensors:
+        return self.model(
+            input_ids,
+            positions,
+            hidden_states,
+            intermediate_tensors,
+            inputs_embeds,
+            spec_step_idx=spec_step_idx,
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor | None:
+        subset = self._draft_subset()
+        if subset is not None and hidden_states.dim() == 2 and hidden_states.shape[0] <= 16:
+            ids, *weights = subset
+            if len(weights) == 3:
+                from .ops.nvfp4_gemv import nvfp4_gemv
+
+                logits = nvfp4_gemv(hidden_states.contiguous(), *weights)
+            else:
+                from .ops.mxfp8_gemv import mxfp8_gemv
+
+                logits = mxfp8_gemv(
+                    hidden_states.contiguous(), *weights, ids.numel(), self.config.hidden_size
+                )
+            full = logits.new_full(
+                (logits.shape[0], self.config.vocab_size), float("-inf"), dtype=torch.float32
+            )
+            full.index_copy_(1, ids, logits.float())
+            return full
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    def _draft_subset(self):
+        """(ids, MXFP8 values, E8M0 scale rows) of the draft-vocabulary head rows.
+
+        Gathered once, outside graph capture, from the shared target lm_head's
+        b12x-packed MXFP8 weight (VLLM_MTP_NVFP4_LM_HEAD=0 makes the draft share
+        it). None when disabled or the head is not MXFP8-packed.
+        """
+        cached = self.__dict__.get("_draft_subset_cache")
+        if cached is not None or not _DRAFT_VOCAB:
+            return cached or None
+        if torch.cuda.is_current_stream_capturing():
+            return None
+        packed = getattr(self.lm_head, "b12x_mxfp8_packed_weight", None)
+        if packed is None:
+            logger.warning_once(
+                "VLLM_QWEN38_DRAFT_VOCAB needs a b12x MXFP8 lm_head "
+                "(VLLM_MXFP8_LM_HEAD=1, VLLM_MTP_NVFP4_LM_HEAD=0); using the full head."
+            )
+            self.__dict__["_draft_subset_cache"] = ()
+            return None
+        with open(_DRAFT_VOCAB) as f:
+            ids = sorted({int(t) for t in json.load(f)})
+        device = packed.weight.values.device
+        ids = torch.tensor(ids, dtype=torch.long, device=device)
+        scale_rows = packed.weight.scale_rows[0]
+        if scale_rows.dtype != torch.uint8:
+            scale_rows = scale_rows.view(torch.uint8)
+        if _DRAFT_NVFP4:
+            from .ops.nvfp4_gemv import quantize_nvfp4_rows
+
+            k = self.config.hidden_size
+            n = ids.numel()
+            nv = (
+                torch.empty((n, k // 2), dtype=torch.uint8, device=device),
+                torch.empty((n, k // 16), dtype=torch.uint8, device=device),
+                torch.empty(n, dtype=torch.float32, device=device),
+            )
+            for r in range(0, n, 2048):  # small temporaries: this runs in the profile run
+                rows = ids[r : r + 2048]
+                v = packed.weight.values.index_select(0, rows)[:, :k].float()
+                e = scale_rows.index_select(0, rows)[:, : k // 32].to(torch.int32)
+                e = (e << 23).view(torch.float32)  # E8M0 -> 2^(e-127)
+                w = (v.view(v.shape[0], -1, 32) * e[..., None]).view(v.shape[0], k)
+                for dst, src in zip(nv, quantize_nvfp4_rows(w)):
+                    dst[r : r + 2048] = src
+            cache = (ids, *nv)
+            kind = "NVFP4 requantized from"
+        else:
+            cache = (
+                ids,
+                packed.weight.values.index_select(0, ids).contiguous(),
+                scale_rows.index_select(0, ids).contiguous(),
+            )
+            kind = "MXFP8 rows of"
+        self.__dict__["_draft_subset_cache"] = cache
+        logger.info(
+            "Qwen3.8 MTP draft head: %d-token vocabulary subset (%s the shared lm_head).",
+            ids.numel(), kind,
+        )
+        return cache
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        qsa_layer_ids = frozenset(range(self.model.num_mtp_layers))
+
+        def remap_weight_names():
+            for name, weight in weights:
+                remapped_name = _remap_mtp_weight_name(name)
+                if remapped_name is not None:
+                    yield (
+                        _remap_qsa_cache_scale_name(remapped_name, qsa_layer_ids),
+                        weight,
+                    )
+
+        loader = AutoWeightsLoader(
+            self,
+            ignore_unexpected_suffixes=(_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy()),
+        )
+        mapper = WeightsMapper(
+            orig_to_new_substr={"hyper_connection_mixer.block_inject_weight": None}
+        )
+        loaded = loader.load_weights(remap_weight_names(), mapper=mapper)
+        if _MTP_FP8:
+            from .ops.hc_fp8_gemv import install_fp8_decode_linears
+
+            names = install_fp8_decode_linears(self)
+            logger.info(
+                "Qwen3.8 MTP: %d bf16 projections use rowwise-FP8 decode GEMVs.", len(names)
+            )
+        return loaded
+
+
+__all__ = ["Qwen4ExpMTP", "Qwen4ExpMultiTokenPredictor"]

recipe/patches/qwen4_exp/nvidia/ops/gdn_deferred.py

@@ -0,0 +1,330 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Deferred-commit GDN decode for speculative verify steps (Qwen3.8 GatedDeltaNet).
+
+Stock b12x verify reads the committed fp32 state once and writes one full state per
+verify token (5 x 3.1 MB per layer at MTP-4); the next step reads only the column of
+the last accepted token. gdn_verify keeps the recurrence in registers and writes small
+per-token records instead (decay g, delta per value row, normalized k). Once
+acceptance is known, gdn_commit replays the accepted tokens from the same source state
+and writes column n-1 only. Rows that end a request, cross a mamba block or take a
+logits-only step get every verify column written, exactly as stock verify would have,
+so prefix-cache checkpoint/align copies always find the columns they read.
+
+Math (b12x _PackedRecurrentQwenKernel, ungrouped path), per value head h, key head
+h // (VH/KH), per token t:
+  q^ = q * rsqrt(|q|^2 + 1e-6) * scale,  k^ = k * rsqrt(|k|^2 + 1e-6)
+  g = exp(-exp(A_log[h]) * softplus(a + dt_bias[h]))   (softplus linear above 20)
+  beta = bf16(sigmoid(b))
+  S = S * g ;  delta = (v - S k^) * beta ;  S = S + delta k^T ;  o = S q^
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+KEY_DIM = 128
+VALUE_DIM = 128
+ROWS = 32  # value rows per program (b12x uses 32 per CTA)
+
+
+@triton.jit
+def _gdn_verify_kernel(
+    qkv_ptr, a_ptr, b_ptr, alog_ptr, dtb_ptr,
+    state_ptr, idx_ptr, nacc_ptr, qsl_ptr, nseq_ptr, valid_ptr,
+    out_ptr, g_rec_ptr, d_rec_ptr, k_rec_ptr,
+    stride_qkv, stride_a, stride_b, stride_ot, stride_oh,
+    slot_stride, idx_stride_r, idx_stride_c,
+    scale,
+    KEY_HEADS: tl.constexpr, VALUE_HEADS: tl.constexpr, COLUMNS: tl.constexpr,
+    ROWS: tl.constexpr, KD: tl.constexpr, VD: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    tiles: tl.constexpr = VD // ROWS
+    r = pid // (VALUE_HEADS * tiles)
+    vh = (pid // tiles) % VALUE_HEADS
+    tile = pid % tiles
+    if r >= tl.load(nseq_ptr):
+        return
+    kh = vh // (VALUE_HEADS // KEY_HEADS)
+    start = tl.load(qsl_ptr + r)
+    end = tl.load(qsl_ptr + r + 1)
+    col = tl.load(nacc_ptr + r) - 1
+    src = tl.load(idx_ptr + r * idx_stride_r + col * idx_stride_c).to(tl.int64)
+    rows = tile * ROWS + tl.arange(0, ROWS)
+    cols = tl.arange(0, KD)
+    S = tl.load(state_ptr + src * slot_stride + vh * (VD * KD) + rows[:, None] * KD + cols[None, :])
+    a_scale = tl.exp(tl.load(alog_ptr + vh).to(tl.float32))
+    dtb = tl.load(dtb_ptr + vh).to(tl.float32)
+    for t in tl.static_range(COLUMNS):
+        tok = start + t
+        if tok < end:
+            base = qkv_ptr + tok.to(tl.int64) * stride_qkv
+            q = tl.load(base + kh * KD + cols).to(tl.float32)
+            k = tl.load(base + KEY_HEADS * KD + kh * KD + cols).to(tl.float32)
+            v = tl.load(base + 2 * KEY_HEADS * KD + vh * VD + rows).to(tl.float32)
+            qn = q * (1.0 / tl.sqrt(tl.sum(q * q, axis=0) + 1e-6)) * scale
+            kn = k * (1.0 / tl.sqrt(tl.sum(k * k, axis=0) + 1e-6))
+            sp = tl.load(a_ptr + tok * stride_a + vh).to(tl.float32) + dtb
+            sp = tl.where(sp <= 20.0, tl.log(1.0 + tl.exp(sp)), sp)
+            g = tl.exp(-a_scale * sp)
+            beta = 1.0 / (1.0 + tl.exp(-tl.load(b_ptr + tok * stride_b + vh).to(tl.float32)))
+            beta = beta.to(tl.bfloat16).to(tl.float32)
+            S = S * g
+            delta = (v - tl.sum(S * kn[None, :], axis=1)) * beta
+            S = S + delta[:, None] * kn[None, :]
+            o = tl.sum(S * qn[None, :], axis=1)
+            tl.store(out_ptr + tok * stride_ot + vh * stride_oh + rows, o.to(tl.bfloat16))
+            tl.store(d_rec_ptr + (tok * VALUE_HEADS + vh) * VD + rows, delta)
+            if tile == 0:
+                tl.store(g_rec_ptr + tok * VALUE_HEADS + vh, g)
+                if vh % (VALUE_HEADS // KEY_HEADS) == 0:
+                    tl.store(k_rec_ptr + (tok * KEY_HEADS + kh) * KD + cols, kn)
+    if (vh == 0) & (tile == 0):
+        tl.store(valid_ptr + r, 1)
+
+
+@triton.jit
+def _gdn_verify_flat_kernel(
+    qkv_ptr, a_ptr, b_ptr, alog_ptr, dtb_ptr,
+    state_ptr, idx_ptr, nacc_ptr, qsl_ptr, nseq_ptr, valid_ptr,
+    out_ptr, g_rec_ptr, d_rec_ptr, k_rec_ptr,
+    stride_qkv, stride_a, stride_b, stride_ot, stride_oh,
+    slot_stride, idx_stride_r, idx_stride_c,
+    scale,
+    KEY_HEADS: tl.constexpr, VALUE_HEADS: tl.constexpr, COLUMNS: tl.constexpr,
+    ROWS: tl.constexpr, KD: tl.constexpr, VD: tl.constexpr,
+):
+    """_gdn_verify_kernel without per-token branches: the state load is issued first,
+    every token's inputs load up front (masked), and each step's S.k and S.q are two
+    independent reductions over the pre-decay state:
+      delta = (v - g (S k^)) beta ;  o = g (S q^) + delta (k^.q^) ;  S = g S + delta k^T
+    """
+    pid = tl.program_id(0)
+    tiles: tl.constexpr = VD // ROWS
+    r = pid // (VALUE_HEADS * tiles)
+    vh = (pid // tiles) % VALUE_HEADS
+    tile = pid % tiles
+    if r >= tl.load(nseq_ptr):
+        return
+    kh = vh // (VALUE_HEADS // KEY_HEADS)
+    start = tl.load(qsl_ptr + r)
+    end = tl.load(qsl_ptr + r + 1)
+    col = tl.load(nacc_ptr + r) - 1
+    src = tl.load(idx_ptr + r * idx_stride_r + col * idx_stride_c).to(tl.int64)
+    rows = tile * ROWS + tl.arange(0, ROWS)
+    cols = tl.arange(0, KD)
+    S = tl.load(state_ptr + src * slot_stride + vh * (VD * KD) + rows[:, None] * KD + cols[None, :])
+    a_scale = tl.exp(tl.load(alog_ptr + vh).to(tl.float32))
+    dtb = tl.load(dtb_ptr + vh).to(tl.float32)
+    for t in tl.static_range(COLUMNS):
+        tok = start + t
+        live = tok < end
+        base = qkv_ptr + tok.to(tl.int64) * stride_qkv
+        q = tl.load(base + kh * KD + cols, mask=live & (cols < KD), other=0.0).to(tl.float32)
+        k = tl.load(base + KEY_HEADS * KD + kh * KD + cols, mask=live & (cols < KD), other=0.0).to(tl.float32)
+        v = tl.load(base + 2 * KEY_HEADS * KD + vh * VD + rows, mask=live & (rows < VD), other=0.0).to(tl.float32)
+        qn = q * (1.0 / tl.sqrt(tl.sum(q * q, axis=0) + 1e-6)) * scale
+        kn = k * (1.0 / tl.sqrt(tl.sum(k * k, axis=0) + 1e-6))
+        sp = tl.load(a_ptr + tok * stride_a + vh, mask=live, other=0.0).to(tl.float32) + dtb
+        sp = tl.where(sp <= 20.0, tl.log(1.0 + tl.exp(sp)), sp)
+        g = tl.exp(-a_scale * sp)
+        beta = 1.0 / (1.0 + tl.exp(-tl.load(b_ptr + tok * stride_b + vh, mask=live, other=0.0).to(tl.float32)))
+        beta = beta.to(tl.bfloat16).to(tl.float32)
+        sk = tl.sum(S * kn[None, :], axis=1)
+        sq = tl.sum(S * qn[None, :], axis=1)
+        delta = (v - g * sk) * beta
+        o = g * sq + delta * tl.sum(kn * qn, axis=0)
+        S = S * g + delta[:, None] * kn[None, :]
+        tl.store(out_ptr + tok * stride_ot + vh * stride_oh + rows, o.to(tl.bfloat16), mask=live & (rows < VD))
+        tl.store(d_rec_ptr + (tok * VALUE_HEADS + vh) * VD + rows, delta, mask=live & (rows < VD))
+        tl.store(g_rec_ptr + tok * VALUE_HEADS + vh, g, mask=live & (tile == 0))
+        tl.store(k_rec_ptr + (tok * KEY_HEADS + kh) * KD + cols, kn,
+                 mask=live & (tile == 0) & (vh % (VALUE_HEADS // KEY_HEADS) == 0) & (cols < KD))
+    if (vh == 0) & (tile == 0):
+        tl.store(valid_ptr + r, 1)
+
+
+_FLAT = __import__("os").environ.get("VLLM_QWEN38_GDN_VERIFY_FLAT", "0") == "1"
+# Swept at 1 live row of 4 (36 layers): 32x8 0.905 ms, 16x2 0.838, 8x1 0.763.
+_VERIFY_ROWS = 8
+_VERIFY_WARPS = 1
+
+
+def gdn_verify(
+    mixed_qkv, a, b, A_log, dt_bias, state, state_indices, num_accepted, query_start_loc,
+    num_seqs, valid, out, records, scale, max_seqs, key_heads=16, value_heads=48, flat=None,
+    rows=None, num_warps=None,
+):
+    """Spec-verify recurrence. ``out`` [T, VH, 128] bf16 gets the pre-norm output;
+    records get (g [T,VH], delta [T,VH,128], k^ [T,KH,128]) fp32; ``state`` is not
+    written. Marks ``valid[r] = 1`` for every live request row."""
+    g_rec, d_rec, k_rec = records
+    flat = _FLAT if flat is None else flat
+    rows = rows or (_VERIFY_ROWS if flat else ROWS)
+    num_warps = num_warps or (_VERIFY_WARPS if flat else 8)
+    grid = (max_seqs * value_heads * (VALUE_DIM // rows),)
+    kernel = _gdn_verify_flat_kernel if flat else _gdn_verify_kernel
+    kernel[grid](
+        mixed_qkv, a, b, A_log, dt_bias,
+        state, state_indices, num_accepted, query_start_loc, num_seqs, valid,
+        out, g_rec, d_rec, k_rec,
+        mixed_qkv.stride(0), a.stride(0), b.stride(0), out.stride(0), out.stride(1),
+        state.stride(0), state_indices.stride(0), state_indices.stride(1),
+        scale,
+        KEY_HEADS=key_heads, VALUE_HEADS=value_heads, COLUMNS=state_indices.shape[1],
+        ROWS=rows, KD=KEY_DIM, VD=VALUE_DIM,
+        num_warps=num_warps,
+    )
+
+
+@triton.jit
+def _gdn_commit_kernel(
+    state_tab, g_tab, d_tab, k_tab, idx_tab,
+    nacc_ptr, qsl_ptr, nseq_ptr, valid_ptr,
+    nsamp_ptr, req_map_ptr, ncomp_ptr, cap_ptr, cap_stride, num_rows,
+    slot_stride, idx_stride_r, idx_stride_c, mamba_block,
+    FORCE_FULL: tl.constexpr, HAS_CAPTURE: tl.constexpr, NUM_CAPTURES: tl.constexpr,
+    KEY_HEADS: tl.constexpr, VALUE_HEADS: tl.constexpr, COLUMNS: tl.constexpr,
+    ROWS: tl.constexpr, KD: tl.constexpr, VD: tl.constexpr, ROW_MODE: tl.constexpr,
+):
+    layer = tl.program_id(1)
+    pid = tl.program_id(0)
+    tiles: tl.constexpr = VD // ROWS
+    r = pid // (VALUE_HEADS * tiles)
+    vh = (pid // tiles) % VALUE_HEADS
+    tile = pid % tiles
+    # Only real batch rows: graph padding rows carry stale state indices that can alias a
+    # live request's blocks, and have no num_sampled/req_map entry.
+    if (r >= tl.load(nseq_ptr)) | (r >= num_rows):
+        return
+    if tl.load(valid_ptr + r) == 0:
+        return
+    req = tl.load(req_map_ptr + r)
+    if req < 0:
+        return
+    start = tl.load(qsl_ptr + r)
+    T = tl.load(qsl_ptr + r + 1) - start
+    n = tl.minimum(tl.maximum(tl.load(nsamp_ptr + r).to(tl.int32), 1), T)
+    full = FORCE_FULL
+    after = tl.load(ncomp_ptr + req).to(tl.int64)
+    full = full | ((after // mamba_block) != ((after - n) // mamba_block))
+    if HAS_CAPTURE:
+        caps = tl.load(cap_ptr + r * cap_stride + tl.arange(0, NUM_CAPTURES))
+        full = full | (tl.max(caps, axis=0) > 0)
+    # ROW_MODE 1: only rows that need every column; 2: only the others (0: all rows)
+    if ROW_MODE == 1:
+        if not full:
+            return
+    if ROW_MODE == 2:
+        if full:
+            return
+    last = tl.where(full, T, n)
+    state_ptr = tl.load(state_tab + layer).to(tl.pointer_type(tl.float32))
+    # Per-layer staged block ids: vLLM spreads the GDN layers over several KV-cache
+    # groups that alias the same memory with different block tables.
+    idx_ptr = tl.load(idx_tab + layer).to(tl.pointer_type(tl.int32))
+    g_rec = tl.load(g_tab + layer).to(tl.pointer_type(tl.float32))
+    d_rec = tl.load(d_tab + layer).to(tl.pointer_type(tl.float32))
+    k_rec = tl.load(k_tab + layer).to(tl.pointer_type(tl.float32))
+    kh = vh // (VALUE_HEADS // KEY_HEADS)
+    col = tl.load(nacc_ptr + r) - 1
+    src = tl.load(idx_ptr + r * idx_stride_r + col * idx_stride_c).to(tl.int64)
+    rows = tile * ROWS + tl.arange(0, ROWS)
+    cols = tl.arange(0, KD)
+    offs = vh * (VD * KD) + rows[:, None] * KD + cols[None, :]
+    S = tl.load(state_ptr + src * slot_stride + offs)
+    for t in tl.static_range(COLUMNS):
+        if t < last:
+            tok = start + t
+            g = tl.load(g_rec + tok * VALUE_HEADS + vh)
+            delta = tl.load(d_rec + (tok * VALUE_HEADS + vh) * VD + rows)
+            kn = tl.load(k_rec + (tok * KEY_HEADS + kh) * KD + cols)
+            S = S * g
+            S = S + delta[:, None] * kn[None, :]
+            if full | (t == n - 1):
+                dst = tl.load(idx_ptr + r * idx_stride_r + t * idx_stride_c).to(tl.int64)
+                tl.store(state_ptr + dst * slot_stride + offs, S)
+
+
+def gdn_commit(
+    tables, state_indices, num_accepted, query_start_loc, num_seqs, valid,
+    num_sampled, req_map, num_computed, capture_tokens, mamba_block, slot_stride, max_seqs,
+    force_full=False, key_heads=16, value_heads=48, row_mode=0,
+):
+    """Replay each deferred row's accepted tokens in every layer of ``tables``
+    (int64 device tensors of per-layer state/g/delta/k^/staged-index base addresses;
+    ``state_indices`` supplies only the index strides and column count).
+
+    ``state_indices``/``num_accepted``/``query_start_loc``/``num_seqs`` are the
+    verify-time staging buffers; ``num_sampled``/``req_map`` are per batch row;
+    ``num_computed`` is the post-step count per request slot; ``capture_tokens`` is
+    the [rows, slots] boundary-capture tensor or None. ``row_mode`` 1 commits only the
+    rows that need every column (block crossing / capture / forced), 2 only the others."""
+    state_tab, g_tab, d_tab, k_tab, idx_tab = tables
+    grid = (max_seqs * value_heads * (VALUE_DIM // ROWS), state_tab.numel())
+    cap = capture_tokens if capture_tokens is not None else num_sampled
+    ncap = capture_tokens.shape[1] if capture_tokens is not None else 1
+    assert ncap & (ncap - 1) == 0, "capture slots must be a power of two"
+    _gdn_commit_kernel[grid](
+        state_tab, g_tab, d_tab, k_tab, idx_tab,
+        num_accepted, query_start_loc, num_seqs, valid,
+        num_sampled, req_map, num_computed, cap, cap.stride(0) if capture_tokens is not None else 0,
+        min(num_sampled.numel(), req_map.numel()),
+        slot_stride, state_indices.stride(0), state_indices.stride(1), mamba_block,
+        FORCE_FULL=bool(force_full), HAS_CAPTURE=capture_tokens is not None, NUM_CAPTURES=ncap,
+        KEY_HEADS=key_heads, VALUE_HEADS=value_heads, COLUMNS=state_indices.shape[1],
+        ROWS=ROWS, KD=KEY_DIM, VD=VALUE_DIM, ROW_MODE=row_mode,
+        num_warps=8,
+    )
+
+
+@triton.jit
+def _gdn_gated_rmsnorm_kernel(
+    output, z, norm_weight, num_tokens, eps, token_capacity,
+    stride_output_token: tl.int64, stride_output_head: tl.int64,
+    stride_z_token: tl.int64, stride_z_head: tl.int64,
+    VALUE_HEADS: tl.constexpr, VALUE_HEAD_DIM: tl.constexpr,
+    SIGMOID_GATE: tl.constexpr, NORM_WEIGHT_FP32: tl.constexpr,
+):
+    """Copy of b12x gdn_decode._gated_rmsnorm_kernel (non-KDA paths). b12x freezes its own
+    Triton kernels after preparation, so a new specialization of its copy cannot JIT."""
+    token_value_head = tl.program_id(0)
+    token = token_value_head // VALUE_HEADS
+    value_head = token_value_head % VALUE_HEADS
+    cols = tl.arange(0, VALUE_HEAD_DIM)
+    mask = cols < VALUE_HEAD_DIM
+    token_i64 = token.to(tl.int64)
+    output_offsets = token_i64 * stride_output_token + value_head.to(tl.int64) * stride_output_head + cols.to(tl.int64)
+    live_tokens = tl.load(num_tokens).to(tl.int32)
+    if token >= tl.maximum(0, tl.minimum(live_tokens, token_capacity)):
+        tl.store(output + output_offsets, 0.0, mask=mask)
+        return
+    z_offsets = token_i64 * stride_z_token + value_head.to(tl.int64) * stride_z_head + cols.to(tl.int64)
+    values = tl.load(output + output_offsets, mask=mask, other=0.0).to(tl.float32)
+    variance = tl.sum(values * values, axis=0) / VALUE_HEAD_DIM
+    normalized = values * tl.rsqrt(variance + eps)
+    if NORM_WEIGHT_FP32:
+        normalized = normalized.to(tl.bfloat16)
+        weight = tl.load(norm_weight + cols, mask=mask, other=0.0).to(tl.float32)
+        weighted = normalized.to(tl.float32) * weight
+    else:
+        normalized = normalized.to(tl.bfloat16)
+        weight = tl.load(norm_weight + cols, mask=mask, other=0.0).to(tl.bfloat16)
+        weighted = (normalized * weight).to(tl.bfloat16).to(tl.float32)
+    gate_input = tl.load(z + z_offsets, mask=mask, other=0.0).to(tl.float32)
+    gate = tl.sigmoid(gate_input)
+    if not SIGMOID_GATE:
+        gate *= gate_input
+    tl.store(output + output_offsets, weighted * gate, mask=mask)
+
+
+def gdn_gated_rmsnorm(output, z, norm_weight, num_tokens, eps, value_heads, sigmoid_gate):
+    """In-place gated RMSNorm of ``output`` [T, VH, 128] with gate ``z`` (b12x semantics)."""
+    m = output.shape[0]
+    _gdn_gated_rmsnorm_kernel[(m * value_heads,)](
+        output, z, norm_weight, num_tokens, float(eps), m,
+        output.stride(0), output.stride(1), z.stride(0), z.stride(1),
+        VALUE_HEADS=value_heads, VALUE_HEAD_DIM=128,
+        SIGMOID_GATE=bool(sigmoid_gate), NORM_WEIGHT_FP32=norm_weight.dtype == torch.float32,
+        num_warps=4,
+    )

recipe/patches/qwen4_exp/nvidia/ops/gemma_norm.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+"""One-launch GemmaRMSNorm for the residual-free path.
+
+vLLM's GemmaRMSNorm.forward_cuda is its native path: it materializes
+``weight.float() + 1.0`` and then runs ir.ops.rms_norm as separate eager torch
+ops (to-fp32 copy, pow, mean, add, rsqrt, mul, mul, to-bf16). With compilation
+off that is ~10 kernels per call; QSA's q_norm/k_norm alone are ~20 tiny kernels
+per attention layer at decode. This kernel does the same fp32 math in one launch
+and reads strided rows directly (QSA normalizes a chunk() view of q/gate).
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _gemma_rms_norm_kernel(
+    x_ptr, w_ptr, out_ptr,
+    H, stride_t, stride_h, D, eps,
+    BLOCK_D: tl.constexpr,
+):
+    row = tl.program_id(0)
+    t = row // H
+    h = row % H
+    offs = tl.arange(0, BLOCK_D)
+    mask = offs < D
+    x = tl.load(x_ptr + t * stride_t + h * stride_h + offs, mask=mask, other=0.0).to(tl.float32)
+    var = tl.sum(x * x, axis=0) / D
+    y = x * (1.0 / tl.sqrt(var + eps))
+    w = tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32) + 1.0
+    tl.store(out_ptr + row * D + offs, (y * w).to(out_ptr.dtype.element_ty), mask=mask)
+
+
+def gemma_rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor | None:
+    """``x * rsqrt(mean(x^2) + eps) * (1 + weight)`` in fp32, cast to x.dtype.
+
+    Returns None when the layout is unsupported (caller falls back)."""
+    if x.dtype not in (torch.bfloat16, torch.float16) or x.dim() not in (2, 3) or x.stride(-1) != 1:
+        return None
+    x3 = x.unsqueeze(1) if x.dim() == 2 else x
+    t, h, d = x3.shape
+    if d > 8192 or weight.numel() != d:
+        return None
+    out = torch.empty(x.shape, dtype=x.dtype, device=x.device)
+    if t * h == 0:
+        return out
+    block = triton.next_power_of_2(d)
+    _gemma_rms_norm_kernel[(t * h,)](
+        x3, weight, out, h, x3.stride(0), x3.stride(1), d, eps,
+        BLOCK_D=block, num_warps=1 if block <= 256 else 4,
+    )
+    return out
+
+
+def install_gemma_norm() -> None:
+    from vllm.model_executor.layers.layernorm import GemmaRMSNorm
+
+    if getattr(GemmaRMSNorm, "_qwen38_fused", False):
+        return
+    stock = GemmaRMSNorm.forward_cuda
+
+    def forward_cuda(self, x, residual=None):
+        if residual is None:
+            out = gemma_rms_norm(x, self.weight, self.variance_epsilon)
+            if out is not None:
+                return out
+        return stock(self, x, residual)
+
+    GemmaRMSNorm.forward_cuda = forward_cuda
+    GemmaRMSNorm._qwen38_fused = True

recipe/patches/qwen4_exp/nvidia/ops/hc_fp8_gemv.py

@@ -0,0 +1,453 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Rowwise-FP8 GEMV for the HyperConnection mix projections at decode.
+
+Each GatedResidual mixer runs two bf16 projections per layer:
+
+  down (+ injection)  [rows, 10240] x [336, 10240]^T   (6.9 MB bf16)
+  up                  [rows,   320] x [10240, 320]^T   (6.6 MB bf16)
+
+About 100 mixers run per decode step, so these weights are ~1.3 GB of
+traffic per step. cuBLAS on GB10 (sm_121) falls back to sm80 WMMA kernels for
+them at 64-73% of the bf16 bandwidth floor. Storing an FP8 copy with one
+scale per output row halves the bytes; this kernel streams it with every SM
+busy: split-K partials for the narrow down projection, a single pass for the
+wide up projection. Rows > MAX_ROWS (prefill) keep the bf16 path.
+
+Split-K partials are reduced in a fixed order, so results are deterministic.
+"""
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+MAX_ROWS = 16
+FP8_MAX = 448.0
+# VLLM_QWEN38_GEMV_TUNE=1: configs re-tuned by gemv_dirty_bench.py (8 MB of dirty L2
+# before every call, closer to in-model conditions than DRAM-cold back-to-back calls).
+_TUNE = os.environ.get("VLLM_QWEN38_GEMV_TUNE", "0") == "1"
+
+# (BLOCK_N, BLOCK_K, SPLIT_K, num_warps) keyed by (N >= threshold); tuned on
+# GB10 with DRAM-cold weights by hc_fp8_gemv_bench.py.
+_NARROW = (16, 256, 10, 4) if _TUNE else (16, 256, 8, 4)  # N ~ 320-336, K = 10240; dirty 17.7 -> 14.9 us
+_WIDE = (64, 64, 1, 4)  # N = 10240, K = 320
+
+
+def quantize_rowwise_fp8(
+    weight: torch.Tensor, chunk_rows: int = 1024
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return (fp8 [N, K], fp32 scale [N]) with ``weight ~= fp8 * scale[:, None]``.
+
+    Rows are converted in chunks so the fp32 temporaries stay small: callers
+    build these copies during vLLM's memory-profiling run, whose peak sets the
+    KV cache size.
+    """
+    w = weight.detach()
+    n = w.shape[0]
+    w8 = torch.empty(w.shape, dtype=torch.float8_e4m3fn, device=w.device)
+    scale = torch.empty(n, dtype=torch.float32, device=w.device)
+    for r in range(0, n, chunk_rows):
+        rows = w[r : r + chunk_rows].float()
+        s = rows.abs().amax(dim=1).clamp_min(1e-12) / FP8_MAX
+        w8[r : r + chunk_rows] = (rows / s[:, None]).clamp_(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
+        scale[r : r + chunk_rows] = s
+    return w8, scale
+
+
+@triton.jit
+def _fp8_gemv_kernel(
+    x_ptr, w_ptr, s_ptr, out_ptr, part_ptr,
+    M, N, K,
+    stride_xm, stride_wn, stride_om,
+    K_PER_SPLIT,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    SPLIT: tl.constexpr,
+):
+    pid_n = tl.program_id(0)
+    pid_k = tl.program_id(1)
+    offs_m = tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    k_start = pid_k * K_PER_SPLIT
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for kk in range(0, K_PER_SPLIT, BLOCK_K):
+        offs_k = k_start + kk + tl.arange(0, BLOCK_K)
+        x = tl.load(
+            x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :],
+            mask=(offs_m[:, None] < M) & (offs_k[None, :] < K),
+            other=0.0,
+        )
+        w = tl.load(
+            w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :],
+            mask=(offs_n[:, None] < N) & (offs_k[None, :] < K),
+            other=0.0,
+        ).to(tl.bfloat16)
+        acc += tl.dot(x, tl.trans(w))
+    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    if SPLIT == 1:
+        s = tl.load(s_ptr + offs_n, mask=offs_n < N, other=0.0)
+        out = acc * s[None, :]
+        tl.store(
+            out_ptr + offs_m[:, None] * stride_om + offs_n[None, :],
+            out.to(out_ptr.dtype.element_ty),
+            mask=mask,
+        )
+    else:
+        tl.store(
+            part_ptr + (pid_k * M + offs_m[:, None]) * N + offs_n[None, :],
+            acc,
+            mask=mask,
+        )
+
+
+@triton.jit
+def _reduce_kernel(
+    part_ptr, s_ptr, out_ptr, M, N, stride_om,
+    SPLIT: tl.constexpr, BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < M * N
+    acc = tl.zeros((BLOCK,), dtype=tl.float32)
+    for s in tl.static_range(SPLIT):
+        acc += tl.load(part_ptr + s * M * N + offs, mask=mask, other=0.0)
+    n = offs % N
+    m = offs // N
+    scale = tl.load(s_ptr + n, mask=mask, other=0.0)
+    tl.store(
+        out_ptr + m * stride_om + n,
+        (acc * scale).to(out_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+
+def _config_for(n: int, k: int) -> tuple[int, int, int, int]:
+    if k >= 8192 and n <= 512:
+        return _NARROW
+    if n >= 2048 and k <= 512:
+        return _WIDE
+    # Other shapes (MTP projections): enough N tiles x K splits to fill ~2 waves
+    # of GB10's 48 SMs while keeping each split at least two K tiles long.
+    block_n = 64 if n >= 8192 else 32
+    block_k = 256 if k >= 2048 else 128
+    tiles = triton.cdiv(n, block_n)
+    split = 1
+    while tiles * split < 96 and k // (2 * split) >= 2 * block_k:
+        split *= 2
+    return block_n, block_k, split, 4
+
+
+def fp8_gemv(
+    x: torch.Tensor,
+    w8: torch.Tensor,
+    scale: torch.Tensor,
+    out_dtype: torch.dtype = torch.bfloat16,
+    config: tuple[int, int, int, int] | None = None,
+) -> torch.Tensor:
+    """Return ``x @ (w8 * scale[:, None]).T``; ``x`` [M <= 16, K] bf16 (row-strided ok)."""
+    m, k = x.shape
+    n = w8.shape[0]
+    assert m <= MAX_ROWS and x.dtype == torch.bfloat16 and x.stride(1) == 1
+    block_n, block_k, split_k, num_warps = config or _config_for(n, k)
+    k_per_split = triton.cdiv(triton.cdiv(k, split_k), block_k) * block_k
+    split_k = triton.cdiv(k, k_per_split)
+    out = torch.empty((m, n), dtype=out_dtype, device=x.device)
+    part = (
+        torch.empty((split_k, m, n), dtype=torch.float32, device=x.device)
+        if split_k > 1
+        else out
+    )
+    grid = (triton.cdiv(n, block_n), split_k)
+    _fp8_gemv_kernel[grid](
+        x, w8, scale, out, part,
+        m, n, k,
+        x.stride(0), w8.stride(0), out.stride(0),
+        k_per_split,
+        BLOCK_M=MAX_ROWS, BLOCK_N=block_n, BLOCK_K=block_k,
+        SPLIT=split_k,
+        num_warps=num_warps,
+    )
+    if split_k > 1:
+        block = 1024
+        _reduce_kernel[(triton.cdiv(m * n, block),)](
+            part, scale, out, m, n, out.stride(0),
+            SPLIT=split_k, BLOCK=block, num_warps=4,
+        )
+    return out
+
+
+@triton.jit
+def _reduce_silu_kernel(
+    part_ptr, s_ptr, out_ptr, y_ptr, M, N, stride_om, stride_ym, SILU_N,
+    HC: tl.constexpr, SPLIT: tl.constexpr, BLOCK: tl.constexpr,
+):
+    """_reduce_kernel that also writes hc_silu of the first SILU_N columns."""
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < M * N
+    acc = tl.zeros((BLOCK,), dtype=tl.float32)
+    for s in tl.static_range(SPLIT):
+        acc += tl.load(part_ptr + s * M * N + offs, mask=mask, other=0.0)
+    n = offs % N
+    m = offs // N
+    scale = tl.load(s_ptr + n, mask=mask, other=0.0)
+    v = (acc * scale).to(out_ptr.dtype.element_ty)
+    tl.store(out_ptr + m * stride_om + n, v, mask=mask)
+    xs = v.to(tl.float32) / HC
+    tl.store(
+        y_ptr + m * stride_ym + n,
+        (xs * tl.sigmoid(xs)).to(y_ptr.dtype.element_ty),
+        mask=mask & (n < SILU_N),
+    )
+
+
+@triton.jit
+def _fp8_gemv_down_silu_fused_kernel(
+    x_ptr, w_ptr, s_ptr, out_ptr, y_ptr, part_ptr, cnt_ptr,
+    M, N, K, stride_xm, stride_wn, stride_om, stride_ym, SILU_N, K_PER_SPLIT,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    SPLIT: tl.constexpr, HC: tl.constexpr,
+):
+    """Split-K FP8 GEMV whose last-arriving program per N tile does the reduce, the
+    row scale and hc_silu (one launch instead of GEMV + _reduce_silu_kernel). That
+    program also resets the tile counter, so CUDA-graph replays stay valid."""
+    pid_n = tl.program_id(0)
+    pid_k = tl.program_id(1)
+    offs_m = tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    k_start = pid_k * K_PER_SPLIT
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for kk in range(0, K_PER_SPLIT, BLOCK_K):
+        offs_k = k_start + kk + tl.arange(0, BLOCK_K)
+        x = tl.load(x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :],
+                    mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), other=0.0)
+        w = tl.load(w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :],
+                    mask=(offs_n[:, None] < N) & (offs_k[None, :] < K), other=0.0).to(tl.bfloat16)
+        acc += tl.dot(x, tl.trans(w))
+    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    tl.store(part_ptr + (pid_k * BLOCK_M + offs_m[:, None]) * N + offs_n[None, :], acc, mask=mask)
+    # release our partial, then count arrivals for this N tile
+    arrived = tl.atomic_add(cnt_ptr + pid_n, 1, sem="acq_rel")
+    if arrived == SPLIT - 1:
+        tot = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        for s in tl.static_range(SPLIT):
+            tot += tl.load(part_ptr + (s * BLOCK_M + offs_m[:, None]) * N + offs_n[None, :],
+                           mask=mask, other=0.0, cache_modifier=".cg")
+        sc = tl.load(s_ptr + offs_n, mask=offs_n < N, other=0.0)
+        v = (tot * sc[None, :]).to(out_ptr.dtype.element_ty)
+        tl.store(out_ptr + offs_m[:, None] * stride_om + offs_n[None, :], v, mask=mask)
+        xs = v.to(tl.float32) / HC
+        tl.store(y_ptr + offs_m[:, None] * stride_ym + offs_n[None, :],
+                 (xs * tl.sigmoid(xs)).to(y_ptr.dtype.element_ty),
+                 mask=mask & (offs_n[None, :] < SILU_N))
+        tl.atomic_xchg(cnt_ptr + pid_n, 0)
+
+
+_FUSED_DOWN = __import__("os").environ.get("VLLM_QWEN38_HC_DOWN_FUSED", "0") == "1"
+# gate-mix tiling "BLOCK_H,BLOCK_K,num_warps" (default 16,64,4)
+_MIX_CFG = tuple(int(v) for v in __import__("os").environ.get("VLLM_QWEN38_HC_MIX_CFG", "16,64,4").split(","))
+_fused_scratch: dict = {}
+
+
+def _down_scratch(device, split, n, tiles):
+    """Persistent partials + zeroed counters, created outside graph capture."""
+    key = (str(device), split, n, tiles)
+    buf = _fused_scratch.get(key)
+    if buf is None:
+        if torch.cuda.is_current_stream_capturing():
+            raise RuntimeError("HC fused down scratch must exist before graph capture")
+        buf = (torch.empty((split, MAX_ROWS, n), dtype=torch.float32, device=device),
+               torch.zeros((tiles,), dtype=torch.int32, device=device))
+        _fused_scratch[key] = buf
+    return buf
+
+
+def fp8_gemv_down_silu_fused(x, w8, scale, silu_n, hc, config=None):
+    """Single-launch fp8_gemv_down_silu (same rounding points). Calls that share a
+    shape share one scratch, which is safe while they run on one stream in order."""
+    m, k = x.shape
+    n = w8.shape[0]
+    assert m <= MAX_ROWS and x.stride(1) == 1
+    block_n, block_k, split_k, num_warps = config or _config_for(n, k)
+    k_per_split = triton.cdiv(triton.cdiv(k, split_k), block_k) * block_k
+    split_k = triton.cdiv(k, k_per_split)
+    tiles = triton.cdiv(n, block_n)
+    part, cnt = _down_scratch(x.device, split_k, n, tiles)
+    out = torch.empty((m, n), dtype=torch.bfloat16, device=x.device)
+    y = torch.empty((m, silu_n), dtype=torch.bfloat16, device=x.device)
+    _fp8_gemv_down_silu_fused_kernel[(tiles, split_k)](
+        x, w8, scale, out, y, part, cnt,
+        m, n, k, x.stride(0), w8.stride(0), out.stride(0), y.stride(0), silu_n, k_per_split,
+        BLOCK_M=MAX_ROWS, BLOCK_N=block_n, BLOCK_K=block_k, SPLIT=split_k, HC=hc,
+        num_warps=num_warps,
+    )
+    return out, y
+
+
+def fp8_gemv_down_silu(
+    x: torch.Tensor, w8: torch.Tensor, scale: torch.Tensor, silu_n: int, hc: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """HC down projection plus hc_silu of its first ``silu_n`` columns.
+
+    Returns (``x @ W.T`` [M, N] bf16, ``(v/hc) * sigmoid(v/hc)`` [M, silu_n] bf16) with
+    the stock rounding points (bf16 projection, fp32 SiLU), saving the separate
+    scaled-SiLU launch.
+    """
+    m, k = x.shape
+    n = w8.shape[0]
+    block_n, block_k, split_k, num_warps = _config_for(n, k)
+    k_per_split = triton.cdiv(triton.cdiv(k, split_k), block_k) * block_k
+    split_k = triton.cdiv(k, k_per_split)
+    out = torch.empty((m, n), dtype=torch.bfloat16, device=x.device)
+    y = torch.empty((m, silu_n), dtype=torch.bfloat16, device=x.device)
+    if split_k == 1:
+        out = fp8_gemv(x, w8, scale)
+        xs = out[:, :silu_n].float() / hc
+        y.copy_(xs * torch.sigmoid(xs))
+        return out, y
+    part = torch.empty((split_k, m, n), dtype=torch.float32, device=x.device)
+    _fp8_gemv_kernel[(triton.cdiv(n, block_n), split_k)](
+        x, w8, scale, out, part,
+        m, n, k,
+        x.stride(0), w8.stride(0), out.stride(0),
+        k_per_split,
+        BLOCK_M=MAX_ROWS, BLOCK_N=block_n, BLOCK_K=block_k,
+        SPLIT=split_k,
+        num_warps=num_warps,
+    )
+    block = 1024
+    _reduce_silu_kernel[(triton.cdiv(m * n, block),)](
+        part, scale, out, y, m, n, out.stride(0), y.stride(0), silu_n,
+        HC=hc, SPLIT=split_k, BLOCK=block, num_warps=4,
+    )
+    return out, y
+
+
+@triton.jit
+def _fp8_gemv_gate_mix_kernel(
+    x_ptr, w_ptr, s_ptr, n_ptr, out_ptr,
+    M, K, HC_DIM,
+    stride_xm, stride_wn, stride_nm, stride_om,
+    BLOCK_M: tl.constexpr, BLOCK_H: tl.constexpr, BLOCK_K: tl.constexpr,
+    HC: tl.constexpr,
+):
+    """HC up projection fused with hc_gate_mix: each program owns BLOCK_H hidden
+    columns and computes the gate logits of all HC streams for them."""
+    pid = tl.program_id(0)
+    offs_m = tl.arange(0, BLOCK_M)
+    offs_h = pid * BLOCK_H + tl.arange(0, BLOCK_H)
+    m_mask = offs_m < M
+    h_mask = offs_h < HC_DIM
+    mix = tl.zeros((BLOCK_M, BLOCK_H), dtype=tl.float32)
+    for s in tl.static_range(HC):
+        rows = s * HC_DIM + offs_h
+        acc = tl.zeros((BLOCK_M, BLOCK_H), dtype=tl.float32)
+        for k0 in range(0, K, BLOCK_K):
+            offs_k = k0 + tl.arange(0, BLOCK_K)
+            x = tl.load(
+                x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :],
+                mask=m_mask[:, None] & (offs_k[None, :] < K),
+                other=0.0,
+            )
+            w = tl.load(
+                w_ptr + rows[:, None] * stride_wn + offs_k[None, :],
+                mask=h_mask[:, None] & (offs_k[None, :] < K),
+                other=0.0,
+            ).to(tl.bfloat16)
+            acc += tl.dot(x, tl.trans(w))
+        sc = tl.load(s_ptr + rows, mask=h_mask, other=0.0)
+        g = (acc * sc[None, :]).to(tl.bfloat16).to(tl.float32)
+        xn = tl.load(
+            n_ptr + offs_m[:, None] * stride_nm + rows[None, :],
+            mask=m_mask[:, None] & h_mask[None, :],
+            other=0.0,
+        ).to(tl.float32)
+        mix += tl.sigmoid(g) * xn
+    mix = mix / HC
+    tl.store(
+        out_ptr + offs_m[:, None] * stride_om + offs_h[None, :],
+        mix.to(out_ptr.dtype.element_ty),
+        mask=m_mask[:, None] & h_mask[None, :],
+    )
+
+
+def fp8_gemv_gate_mix(
+    x: torch.Tensor,
+    w8: torch.Tensor,
+    scale: torch.Tensor,
+    normalized: torch.Tensor,
+    hc: int,
+    block_h: int = 16,
+    block_k: int = 64,
+    num_warps: int = 4,
+) -> torch.Tensor:
+    """``hc_gate_mix(normalized, x @ W.T, hc)`` without materializing the gate logits.
+
+    x [M <= 16, K] bf16; W [hc * D, K] rowwise FP8; normalized [M, hc * D] bf16.
+    Returns [M, D] bf16: mean over streams of sigmoid(bf16 gate) * normalized.
+    """
+    m, k = x.shape
+    n = w8.shape[0]
+    assert m <= MAX_ROWS and x.stride(1) == 1 and normalized.stride(1) == 1 and n % hc == 0
+    d = n // hc
+    out = torch.empty((m, d), dtype=normalized.dtype, device=x.device)
+    _fp8_gemv_gate_mix_kernel[(triton.cdiv(d, block_h),)](
+        x, w8, scale, normalized, out,
+        m, k, d,
+        x.stride(0), w8.stride(0), normalized.stride(0), out.stride(0),
+        BLOCK_M=MAX_ROWS, BLOCK_H=block_h, BLOCK_K=block_k, HC=hc,
+        num_warps=num_warps,
+    )
+    return out
+
+
+def install_fp8_decode_linears(model: torch.nn.Module) -> list[str]:
+    """Route decode rows (<= MAX_ROWS) of ``model``'s plain bf16 Linear layers
+    through ``fp8_gemv``; prefill rows keep the stock bf16 path.
+
+    Skips HyperConnection mixers (own path, VLLM_QWEN38_HC_FP8), MoE router and
+    shared-expert gates (their logits stay bf16) and anything biased/quantized.
+    b12x finishes checkpoint transfers after load_weights returns, so each FP8
+    copy is built on the first call outside graph capture (the profile run),
+    when the bf16 weight holds its final values. Returns the wrapped names.
+    """
+    from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+    names = []
+    for name, mod in model.named_modules():
+        if (
+            not isinstance(mod, LinearBase)
+            or "hyper_connection" in name
+            or name.endswith(".gate")
+            or name.endswith("shared_expert_gate")
+            or not isinstance(getattr(mod, "quant_method", None), UnquantizedLinearMethod)
+            or getattr(mod, "bias", None) is not None
+        ):
+            continue
+        weight = getattr(mod, "weight", None)
+        if (
+            weight is None
+            or weight.is_meta
+            or weight.dtype != torch.bfloat16
+            or weight.dim() != 2
+            or weight.shape[1] % 32
+        ):
+            continue
+        stock = mod.forward
+        cache: list = []
+
+        def forward(x, _mod=mod, _cache=cache, _stock=stock):
+            if not _cache:
+                if torch.cuda.is_current_stream_capturing():
+                    return _stock(x)
+                _cache.extend(quantize_rowwise_fp8(_mod.weight))
+            if x.dim() >= 2 and x.dtype == torch.bfloat16 and x.numel() <= MAX_ROWS * x.shape[-1]:
+                w8, scale = _cache
+                x2 = x.reshape(-1, x.shape[-1])
+                if x2.stride(-1) != 1:
+                    x2 = x2.contiguous()
+                out = fp8_gemv(x2, w8, scale).view(*x.shape[:-1], w8.shape[0])
+                return (out, None) if getattr(_mod, "return_bias", True) else out
+            return _stock(x)
+
+        mod.forward = forward
+        names.append(name)
+    return names

recipe/patches/qwen4_exp/nvidia/ops/mxfp8_gemv.py

@@ -0,0 +1,249 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Split-K MXFP8 GEMV over b12x's packed dense weights, for decode rows.
+
+b12x's dense MXFP8 kernel launches one CTA per 64 output columns at decode, so
+narrow projections leave most of GB10's 48 SMs idle (decode profile, M=5):
+
+  out_proj / o_proj    N=2560  K=6144   40 CTAs   101 us  (floor 58 us)
+  shared expert up     N=1280  K=2560   20 CTAs    36 us  (floor 12 us)
+  shared expert down   N=2560  K=640    40 CTAs    16 us  (floor  6 us)
+
+b12x keeps the ModelOpt layout next to its MMA-swizzled scales: values
+[N, Kp] e4m3 row-major (K padded to 128) and scale_rows [N, Kp/32] E8M0. This
+kernel reads those directly (no extra memory), dequantizes a tile to bf16
+(exact: e4m3 * 2^e fits bf16), and splits K so every SM streams weights.
+Partials are reduced in a fixed order, so results are deterministic.
+"""
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+MAX_ROWS = 16
+# VLLM_QWEN38_GEMV_TUNE=1: configs re-tuned by gemv_dirty_bench.py (8 MB of dirty L2
+# before every call, closer to in-model conditions than DRAM-cold back-to-back calls).
+_TUNE = os.environ.get("VLLM_QWEN38_GEMV_TUNE", "0") == "1"
+
+
+@triton.jit
+def _mxfp8_gemv_kernel(
+    x_ptr, w_ptr, s_ptr, out_ptr, part_ptr,
+    M, N, K,
+    stride_xm, stride_wn, stride_sn, stride_om,
+    K_PER_SPLIT,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    SPLIT: tl.constexpr,
+):
+    pid_n = tl.program_id(0)
+    pid_k = tl.program_id(1)
+    offs_m = tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_g = tl.arange(0, BLOCK_K // 32)
+    k_start = pid_k * K_PER_SPLIT
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    n_mask = offs_n[:, None] < N
+    for kk in range(0, K_PER_SPLIT, BLOCK_K):
+        k0 = k_start + kk
+        offs_k = k0 + tl.arange(0, BLOCK_K)
+        k_mask = offs_k[None, :] < K
+        x = tl.load(
+            x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :],
+            mask=(offs_m[:, None] < M) & k_mask,
+            other=0.0,
+        )
+        w = tl.load(
+            w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :],
+            mask=n_mask & k_mask,
+            other=0.0,
+        ).to(tl.float32)
+        s = tl.load(
+            s_ptr + offs_n[:, None] * stride_sn + (k0 // 32 + offs_g)[None, :],
+            mask=n_mask & ((k0 + offs_g * 32)[None, :] < K),
+            other=127,
+        )
+        # E8M0 -> 2^(s-127) by placing the byte in the fp32 exponent field.
+        scale = (s.to(tl.int32) << 23).to(tl.float32, bitcast=True)
+        w = tl.reshape(w, (BLOCK_N, BLOCK_K // 32, 32)) * scale[:, :, None]
+        w = tl.reshape(w, (BLOCK_N, BLOCK_K)).to(tl.bfloat16)
+        acc += tl.dot(x, tl.trans(w))
+    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    if SPLIT == 1:
+        tl.store(
+            out_ptr + offs_m[:, None] * stride_om + offs_n[None, :],
+            acc.to(out_ptr.dtype.element_ty),
+            mask=mask,
+        )
+    else:
+        tl.store(
+            part_ptr + (pid_k * M + offs_m[:, None]) * N + offs_n[None, :],
+            acc,
+            mask=mask,
+        )
+
+
+@triton.jit
+def _reduce_kernel(part_ptr, out_ptr, M, N, stride_om, SPLIT: tl.constexpr, BLOCK: tl.constexpr):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < M * N
+    acc = tl.zeros((BLOCK,), dtype=tl.float32)
+    for s in tl.static_range(SPLIT):
+        acc += tl.load(part_ptr + s * M * N + offs, mask=mask, other=0.0)
+    tl.store(out_ptr + (offs // N) * stride_om + offs % N, acc.to(out_ptr.dtype.element_ty), mask=mask)
+
+
+# (BLOCK_N, BLOCK_K, SPLIT_K, num_warps) per (N, K), tuned on GB10 with
+# DRAM-cold weights by mxfp8_gemv_bench.py at M=5 (MTP-4 verify). Cold-cache us,
+# b12x in-model -> this kernel: out_proj 100.9 -> 67.9, shared up 35.6 -> 16.5,
+# shared down 15.6 -> 6.8, attn qkv 167 -> 147, GDN qkvz 203.5 -> 181.
+# Shapes not listed use _default_config.
+_CONFIGS: dict[tuple[int, int], tuple[int, int, int, int]] = {
+    (2560, 6144): (32, 256, 1, 4),
+    (1280, 2560): (64, 128, 8, 4) if _TUNE else (16, 128, 4, 4),  # dirty 17.6 -> 14.8 us
+    (2560, 640): (32, 128, 1, 4),
+    (13312, 2560): (64, 256, 1, 8),
+    (16384, 2560): (32, 256, 1, 4) if _TUNE else (64, 256, 1, 8),  # dirty 186.6 -> 178.2 us
+}
+
+
+def _default_config(n: int, k: int) -> tuple[int, int, int, int]:
+    block_n = 32
+    tiles = triton.cdiv(n, block_n)
+    split = max(1, min(16, (4 * 48) // tiles))
+    while split > 1 and k // split < 256:
+        split //= 2
+    return block_n, 256 if k >= 2048 else 128, split, 4
+
+
+def mxfp8_gemv(
+    x: torch.Tensor,
+    values: torch.Tensor,
+    scale_rows: torch.Tensor,
+    out_features: int,
+    in_features: int,
+    out_dtype: torch.dtype = torch.bfloat16,
+    config: tuple[int, int, int, int] | None = None,
+) -> torch.Tensor:
+    """``x @ dequant(values, scale_rows)[:out_features, :in_features].T`` for M <= 16."""
+    m = x.shape[0]
+    n, k = int(out_features), int(in_features)
+    if scale_rows.dtype != torch.uint8:  # b12x stores E8M0 as float8_e8m0fnu
+        scale_rows = scale_rows.view(torch.uint8)
+    assert m <= MAX_ROWS and x.dtype == torch.bfloat16 and x.stride(1) == 1
+    block_n, block_k, split_k, num_warps = config or _CONFIGS.get((n, k)) or _default_config(n, k)
+    k_per_split = triton.cdiv(triton.cdiv(k, split_k), block_k) * block_k
+    split_k = triton.cdiv(k, k_per_split)
+    out = torch.empty((m, n), dtype=out_dtype, device=x.device)
+    part = (
+        torch.empty((split_k, m, n), dtype=torch.float32, device=x.device)
+        if split_k > 1
+        else out
+    )
+    grid = (triton.cdiv(n, block_n), split_k)
+    _mxfp8_gemv_kernel[grid](
+        x, values, scale_rows, out, part,
+        m, n, k,
+        x.stride(0), values.stride(0), scale_rows.stride(0), out.stride(0),
+        k_per_split,
+        BLOCK_M=MAX_ROWS, BLOCK_N=block_n, BLOCK_K=block_k,
+        SPLIT=split_k,
+        num_warps=num_warps,
+    )
+    if split_k > 1:
+        block = 1024
+        _reduce_kernel[(triton.cdiv(m * n, block),)](
+            part, out, m, n, out.stride(0), SPLIT=split_k, BLOCK=block, num_warps=4
+        )
+    return out
+
+
+@triton.jit
+def _mxfp8_swiglu_gate_kernel(
+    gu_ptr, x_ptr, wg_ptr, w_ptr, s_ptr, out_ptr,
+    M, N, K, H,
+    stride_gum, stride_xm, stride_wn, stride_sn, stride_om,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+):
+    pid_n = tl.program_id(0)
+    offs_m = tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_g = tl.arange(0, BLOCK_K // 32)
+    m_mask = offs_m[:, None] < M
+    n_mask = offs_n[:, None] < N
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k0 in range(0, K, BLOCK_K):
+        offs_k = k0 + tl.arange(0, BLOCK_K)
+        mk = m_mask & (offs_k[None, :] < K)
+        g = tl.load(gu_ptr + offs_m[:, None] * stride_gum + offs_k[None, :], mask=mk, other=0.0)
+        u = tl.load(gu_ptr + offs_m[:, None] * stride_gum + K + offs_k[None, :], mask=mk, other=0.0)
+        # vLLM act_and_mul: bf16(silu(g) in fp32), then bf16(silu * u).
+        g = g.to(tl.float32)
+        silu = (g / (1.0 + tl.exp(-g))).to(tl.bfloat16).to(tl.float32)
+        a = (silu * u.to(tl.float32)).to(tl.bfloat16)
+        w = tl.load(
+            w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :],
+            mask=n_mask & (offs_k[None, :] < K),
+            other=0.0,
+        ).to(tl.float32)
+        s = tl.load(
+            s_ptr + offs_n[:, None] * stride_sn + (k0 // 32 + offs_g)[None, :],
+            mask=n_mask & ((k0 + offs_g * 32)[None, :] < K),
+            other=127,
+        )
+        scale = (s.to(tl.int32) << 23).to(tl.float32, bitcast=True)
+        w = tl.reshape(w, (BLOCK_N, BLOCK_K // 32, 32)) * scale[:, :, None]
+        w = tl.reshape(w, (BLOCK_N, BLOCK_K)).to(tl.bfloat16)
+        acc += tl.dot(a, tl.trans(w))
+    # Shared-expert gate: sigmoid(x . w_gate) per row (L2-resident x, 5 KB w_gate).
+    gacc = tl.zeros((BLOCK_M,), dtype=tl.float32)
+    for h0 in range(0, H, BLOCK_H):
+        offs_h = h0 + tl.arange(0, BLOCK_H)
+        xv = tl.load(
+            x_ptr + offs_m[:, None] * stride_xm + offs_h[None, :],
+            mask=m_mask & (offs_h[None, :] < H),
+            other=0.0,
+        ).to(tl.float32)
+        wv = tl.load(wg_ptr + offs_h, mask=offs_h < H, other=0.0).to(tl.float32)
+        gacc += tl.sum(xv * wv[None, :], axis=1)
+    # Stock rounding: bf16 gate logit, bf16 sigmoid, bf16 down output, bf16 product.
+    gl = gacc.to(tl.bfloat16).to(tl.float32)
+    sig = (1.0 / (1.0 + tl.exp(-gl))).to(tl.bfloat16).to(tl.float32)
+    out = acc.to(tl.bfloat16).to(tl.float32) * sig[:, None]
+    tl.store(
+        out_ptr + offs_m[:, None] * stride_om + offs_n[None, :],
+        out.to(out_ptr.dtype.element_ty),
+        mask=m_mask & (offs_n[None, :] < N),
+    )
+
+
+def mxfp8_gemv_swiglu_gate(
+    gate_up: torch.Tensor,
+    x: torch.Tensor,
+    w_gate: torch.Tensor,
+    values: torch.Tensor,
+    scale_rows: torch.Tensor,
+    out_features: int,
+    in_features: int,
+    config: tuple[int, int, int] = (32, 64, 4),  # shared_fuse_bench.py: 32.1 -> 14.2 us cold, M=5
+) -> torch.Tensor:
+    """Fused shared-expert tail for decode rows:
+    ``sigmoid(x @ w_gate) * (silu(g) * u) @ dequant(values, scale_rows)[:N, :K].T``
+    with ``gate_up = [g | u]`` [M, 2K] and ``x`` [M, H] the expert input."""
+    m, k = gate_up.shape[0], int(in_features)
+    n = int(out_features)
+    if scale_rows.dtype != torch.uint8:
+        scale_rows = scale_rows.view(torch.uint8)
+    assert m <= MAX_ROWS and gate_up.shape[1] == 2 * k and gate_up.stride(1) == 1
+    assert x.shape[0] == m and x.stride(1) == 1 and w_gate.numel() == x.shape[1]
+    block_n, block_k, num_warps = config
+    out = torch.empty((m, n), dtype=torch.bfloat16, device=x.device)
+    _mxfp8_swiglu_gate_kernel[(triton.cdiv(n, block_n),)](
+        gate_up, x, w_gate, values, scale_rows, out,
+        m, n, k, x.shape[1],
+        gate_up.stride(0), x.stride(0), values.stride(0), scale_rows.stride(0), out.stride(0),
+        BLOCK_M=MAX_ROWS, BLOCK_N=block_n, BLOCK_K=block_k, BLOCK_H=512,
+        num_warps=num_warps,
+    )
+    return out

recipe/patches/qwen4_exp/nvidia/ops/nvfp4_gemv.py

@@ -0,0 +1,187 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Split-K NVFP4 GEMV for decode rows (MTP draft vocabulary subset).
+
+The draft head is evaluated 4x per MTP-4 step over a 32K-64K token subset.
+Gathered from the target's MXFP8 head it streams 1.03 bytes per weight; this
+row-major NVFP4 copy streams 0.56 (e2m1 nibbles, one e4m3 scale per 16 weights,
+one fp32 scale per row), matching the precision of the draft's own NVFP4 head.
+
+Layout (all row-major, K % 32 == 0):
+  packed  uint8 [N, K/2]   low nibble = even k, high nibble = odd k
+  scales  uint8 [N, K/16]  e4m3 bit patterns
+  gscale  fp32  [N]        weight ~= e2m1 * e4m3 * gscale
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+MAX_ROWS = 16
+E2M1_MAX = 6.0
+E4M3_MAX = 448.0
+# Midpoints between the e2m1 magnitudes 0, .5, 1, 1.5, 2, 3, 4, 6.
+_E2M1_BOUNDS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0)
+
+
+def quantize_nvfp4_rows(
+    weight: torch.Tensor, chunk_rows: int = 2048
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Quantize ``weight`` [N, K] to (packed, scales, gscale); see module doc."""
+    n, k = weight.shape
+    if k % 32:
+        raise ValueError("NVFP4 rows need K % 32 == 0")
+    dev = weight.device
+    packed = torch.empty((n, k // 2), dtype=torch.uint8, device=dev)
+    scales = torch.empty((n, k // 16), dtype=torch.uint8, device=dev)
+    gscale = torch.empty(n, dtype=torch.float32, device=dev)
+    bounds = torch.tensor(_E2M1_BOUNDS, dtype=torch.float32, device=dev)
+    for r in range(0, n, chunk_rows):
+        rows = weight[r : r + chunk_rows].float()
+        g = (rows.abs().amax(dim=1) / (E2M1_MAX * E4M3_MAX)).clamp_min(1e-30)
+        blocks = rows.view(rows.shape[0], k // 16, 16)
+        s = (blocks.abs().amax(dim=2) / E2M1_MAX / g[:, None]).clamp_(max=E4M3_MAX)
+        s8 = s.to(torch.float8_e4m3fn)
+        denom = s8.float() * g[:, None]
+        q = torch.where(denom[..., None] > 0, blocks / denom[..., None], 0.0)
+        idx = torch.bucketize(q.abs(), bounds).to(torch.uint8)
+        nib = idx | ((q < 0).to(torch.uint8) << 3)
+        nib = nib.view(rows.shape[0], k)
+        packed[r : r + chunk_rows] = nib[:, 0::2] | (nib[:, 1::2] << 4)
+        scales[r : r + chunk_rows] = s8.view(torch.uint8)
+        gscale[r : r + chunk_rows] = g
+    return packed, scales, gscale
+
+
+def dequantize_nvfp4_rows(packed, scales, gscale) -> torch.Tensor:
+    """Reference fp32 dequantization (tests)."""
+    lut = torch.tensor(
+        [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6],
+        dtype=torch.float32, device=packed.device,
+    )
+    n = packed.shape[0]
+    nib = torch.stack((packed & 0xF, packed >> 4), dim=2).view(n, -1).long()
+    v = lut[nib].view(n, -1, 16)
+    s = scales.view(torch.float8_e4m3fn).float()
+    return (v * s[..., None]).view(n, -1) * gscale[:, None]
+
+
+@triton.jit
+def _nvfp4_gemv_kernel(
+    x_ptr, w_ptr, s_ptr, g_ptr, out_ptr, part_ptr,
+    M, N, K,
+    stride_xm, stride_wn, stride_sn, stride_om,
+    K_PER_SPLIT,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    SPLIT: tl.constexpr,
+):
+    pid_n = tl.program_id(0)
+    pid_k = tl.program_id(1)
+    offs_m = tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_b = tl.arange(0, BLOCK_K // 2)
+    offs_g = tl.arange(0, BLOCK_K // 16)
+    n_mask = offs_n[:, None] < N
+    k_start = pid_k * K_PER_SPLIT
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for kk in range(0, K_PER_SPLIT, BLOCK_K):
+        k0 = k_start + kk
+        offs_k = k0 + tl.arange(0, BLOCK_K)
+        x = tl.load(
+            x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :],
+            mask=(offs_m[:, None] < M) & (offs_k[None, :] < K),
+            other=0.0,
+        )
+        b = tl.load(
+            w_ptr + offs_n[:, None] * stride_wn + (k0 // 2 + offs_b)[None, :],
+            mask=n_mask & ((k0 + 2 * offs_b)[None, :] < K),
+            other=0,
+        ).to(tl.uint32)
+        nib = tl.reshape(tl.join(b & 0xF, b >> 4), (BLOCK_N, BLOCK_K))
+        # e2m1 -> fp32 bits: e = bits 2:1, m = bit 0, sign = bit 3.
+        e = (nib >> 1) & 3
+        m = nib & 1
+        bits = tl.where(e > 0, ((e + 126) << 23) | (m << 22), m * 0x3F000000)
+        bits = bits | ((nib & 8) << 28)
+        v = bits.to(tl.float32, bitcast=True)
+        s = tl.load(
+            s_ptr + offs_n[:, None] * stride_sn + (k0 // 16 + offs_g)[None, :],
+            mask=n_mask & ((k0 + offs_g * 16)[None, :] < K),
+            other=0,
+        )
+        s = s.to(tl.float8e4nv, bitcast=True).to(tl.float32)
+        v = tl.reshape(v, (BLOCK_N, BLOCK_K // 16, 16)) * s[:, :, None]
+        w = tl.reshape(v, (BLOCK_N, BLOCK_K)).to(tl.bfloat16)  # exact: <= 6 mantissa bits
+        acc += tl.dot(x, tl.trans(w))
+    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    if SPLIT == 1:
+        g = tl.load(g_ptr + offs_n, mask=offs_n < N, other=0.0)
+        tl.store(
+            out_ptr + offs_m[:, None] * stride_om + offs_n[None, :],
+            (acc * g[None, :]).to(out_ptr.dtype.element_ty),
+            mask=mask,
+        )
+    else:
+        tl.store(
+            part_ptr + (pid_k * M + offs_m[:, None]) * N + offs_n[None, :],
+            acc,
+            mask=mask,
+        )
+
+
+@triton.jit
+def _reduce_kernel(
+    part_ptr, g_ptr, out_ptr, M, N, stride_om, SPLIT: tl.constexpr, BLOCK: tl.constexpr
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < M * N
+    acc = tl.zeros((BLOCK,), dtype=tl.float32)
+    for s in tl.static_range(SPLIT):
+        acc += tl.load(part_ptr + s * M * N + offs, mask=mask, other=0.0)
+    n = offs % N
+    g = tl.load(g_ptr + n, mask=mask, other=0.0)
+    tl.store(out_ptr + (offs // N) * stride_om + n, (acc * g).to(out_ptr.dtype.element_ty), mask=mask)
+
+
+# (BLOCK_N, BLOCK_K, SPLIT_K, num_warps); tuned by nvfp4_gemv_bench.py on GB10.
+_DEFAULT = (64, 256, 1, 4)
+_CONFIGS: dict[tuple[int, int], tuple[int, int, int, int]] = {}
+
+
+def nvfp4_gemv(
+    x: torch.Tensor,
+    packed: torch.Tensor,
+    scales: torch.Tensor,
+    gscale: torch.Tensor,
+    out_dtype: torch.dtype = torch.bfloat16,
+    config: tuple[int, int, int, int] | None = None,
+) -> torch.Tensor:
+    """``x @ dequant(packed, scales, gscale).T`` for ``x`` [M <= 16, K] bf16."""
+    m, k = x.shape
+    n = packed.shape[0]
+    assert m <= MAX_ROWS and x.dtype == torch.bfloat16 and x.stride(1) == 1
+    assert packed.shape[1] * 2 == k and scales.shape[1] * 16 == k
+    block_n, block_k, split_k, num_warps = config or _CONFIGS.get((n, k), _DEFAULT)
+    k_per_split = triton.cdiv(triton.cdiv(k, split_k), block_k) * block_k
+    split_k = triton.cdiv(k, k_per_split)
+    out = torch.empty((m, n), dtype=out_dtype, device=x.device)
+    part = (
+        torch.empty((split_k, m, n), dtype=torch.float32, device=x.device)
+        if split_k > 1
+        else out
+    )
+    grid = (triton.cdiv(n, block_n), split_k)
+    _nvfp4_gemv_kernel[grid](
+        x, packed, scales, gscale, out, part,
+        m, n, k,
+        x.stride(0), packed.stride(0), scales.stride(0), out.stride(0),
+        k_per_split,
+        BLOCK_M=MAX_ROWS, BLOCK_N=block_n, BLOCK_K=block_k,
+        SPLIT=split_k,
+        num_warps=num_warps,
+    )
+    if split_k > 1:
+        block = 1024
+        _reduce_kernel[(triton.cdiv(m * n, block),)](
+            part, gscale, out, m, n, out.stride(0), SPLIT=split_k, BLOCK=block, num_warps=4
+        )
+    return out

recipe/patches/qwen4_exp/nvidia/ops/small_gemv.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Split-K GEMV for small bf16 projections at decode (ported from the DeepSeek-V4 overlay).
+
+At decode (1-16 rows) cuBLAS on GB10 (sm_121) falls back to sm80 WMMA kernels
+for these [rows, 4096] x [N, 4096]^T projections and launches too few blocks
+to use the memory bandwidth, e.g. 4 blocks for the N=64 indexer weights_proj.
+Isolated with DRAM-cold weights the gain is ~1 ms per decode step, mostly the
+N=64 projection where cuBLAS is ~5x off; the larger ones are near bandwidth:
+
+  compressor wkv+wgate      N=2048 (c4 layers) / 1024 (c128 layers), fp32 out
+  indexer compressor        N=512,  fp32 out
+  MoE router gate           N=216,  fp32 out
+  indexer weights_proj      N=64,   bf16 out
+
+This kernel splits both N and K so every SM streams weights, accumulates in
+fp32 with tensor-core dots, and reduces split-K partials with fp32 atomics.
+Atomic reduction order is not fixed, so results can differ from torch.mm in
+the last fp32 bits. Rows > MAX_ROWS (prefill) fall back to torch.mm.
+
+Opt-in with VLLM_QWEN38_SMALL_GEMV=1.
+"""
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+ENABLED = os.environ.get("VLLM_QWEN38_SMALL_GEMV", "0") == "1"
+# VLLM_QWEN38_GEMV_TUNE=1: configs re-tuned by gemv_dirty_bench.py (8 MB of dirty L2
+# before every call, closer to in-model conditions than DRAM-cold back-to-back calls).
+_TUNE = os.environ.get("VLLM_QWEN38_GEMV_TUNE", "0") == "1"
+MAX_ROWS = 16
+
+# (BLOCK_N, BLOCK_K, SPLIT_K, num_warps) per N, tuned on GB10 by
+# ds4_small_gemv_bench.py with DRAM-cold weights (as in the real model) at the
+# DSpark verify width (6 rows). N not listed uses the nearest smaller entry.
+# Cold-cache time per call vs torch.mm: N=64 27->5 us, 216 12->11, 512 24->20,
+# 1024 44->37.5, 2048 82->73 (bandwidth floor 61 us). Only N=64 is a big win;
+# the rest are already near the bandwidth floor.
+# Qwen3.8 router (K=2560, N=512), cold: torch.mm 13.9 us -> 12.5 us (floor 9.6).
+_CONFIGS: dict[int, tuple[int, int, int, int]] = {
+    64: (32, 128, 32, 4),
+    216: (64, 128, 32, 8),
+    512: (32, 256, 10, 4) if _TUNE else (16, 256, 1, 4),  # dirty 11.6 -> 9.5 us
+    1024: (64, 256, 1, 8),
+    2048: (64, 256, 1, 8),
+}
+
+
+@triton.jit
+def _small_gemv_kernel(
+    x_ptr, w_ptr, out_ptr,
+    M, N, K,
+    stride_xm, stride_wn, stride_om,
+    K_PER_SPLIT,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    USE_ATOMIC: tl.constexpr,
+):
+    pid_n = tl.program_id(0)
+    pid_k = tl.program_id(1)
+    offs_m = tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    k_start = pid_k * K_PER_SPLIT
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for kk in range(0, K_PER_SPLIT, BLOCK_K):
+        offs_k = k_start + kk + tl.arange(0, BLOCK_K)
+        x = tl.load(
+            x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :],
+            mask=(offs_m[:, None] < M) & (offs_k[None, :] < K),
+            other=0.0,
+        )
+        w = tl.load(
+            w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :],
+            mask=(offs_n[:, None] < N) & (offs_k[None, :] < K),
+            other=0.0,
+        )
+        acc += tl.dot(x, tl.trans(w))
+    out_offs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :]
+    out_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    if USE_ATOMIC:
+        tl.atomic_add(out_offs, acc, mask=out_mask, sem="relaxed")
+    else:
+        tl.store(out_offs, acc, mask=out_mask)
+
+
+def _config_for(n: int) -> tuple[int, int, int, int]:
+    keys = [k for k in sorted(_CONFIGS) if k <= n]
+    return _CONFIGS[keys[-1] if keys else min(_CONFIGS)]
+
+
+def small_mm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    out_dtype: torch.dtype,
+    config: tuple[int, int, int, int] | None = None,
+) -> torch.Tensor:
+    """Return ``x @ weight.T`` as ``out_dtype``; ``weight`` is [N, K]."""
+    m, k = x.shape
+    n = weight.shape[0]
+    if (
+        m > MAX_ROWS
+        or m == 0
+        or x.dtype != torch.bfloat16
+        or weight.dtype != torch.bfloat16
+        or x.stride(1) != 1
+        or weight.stride(1) != 1
+    ):
+        return torch.mm(x, weight.t(), out_dtype=out_dtype)
+    block_n, block_k, split_k, num_warps = config or _config_for(n)
+    k_per_split = triton.cdiv(triton.cdiv(k, split_k), block_k) * block_k
+    split_k = triton.cdiv(k, k_per_split)
+    use_atomic = split_k > 1
+    out = (torch.zeros if use_atomic else torch.empty)(
+        (m, n), dtype=torch.float32, device=x.device
+    )
+    grid = (triton.cdiv(n, block_n), split_k)
+    _small_gemv_kernel[grid](
+        x, weight, out,
+        m, n, k,
+        x.stride(0), weight.stride(0), out.stride(0),
+        k_per_split,
+        BLOCK_M=MAX_ROWS, BLOCK_N=block_n, BLOCK_K=block_k,
+        USE_ATOMIC=use_atomic,
+        num_warps=num_warps,
+    )
+    return out if out_dtype == torch.float32 else out.to(out_dtype)

recipe/patches/vllm_core/v1/core/boundary_checkpoint.py

@@ -0,0 +1,438 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Exact request-boundary lookup and ownership for recurrent prefix caches."""
+
+import hashlib
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Literal
+
+from vllm.v1.core.kv_cache_utils import BlockHash, generate_block_hash_extra_keys
+from vllm.v1.request import Request
+
+if TYPE_CHECKING:
+    from vllm.v1.core.block_pool import BlockPool
+
+MAX_BOUNDARY_STOP_TOKENS = 128
+
+# Modalities whose placeholder spans are fully identified by their content hash.
+CHECKPOINT_MODALITIES = frozenset({"image", "video"})
+# Placeholder positions map into [2**30, 2**31): above every vocabulary and
+# inside LMCache's token wire domain. Each position carries 30 independent
+# bits of the item's full content hash.
+_CONTENT_TOKEN_BASE = 1 << 30
+# A checkpoint may not end within an item's first positions: with fewer of
+# them its key would hold under 150 bits of that item's identity.
+MIN_CHECKPOINT_ITEM_TOKENS = 5
+
+
+def content_token_ids(
+    request: Request, start: int, end: int, salt: bytes = b""
+) -> tuple[int, ...]:
+    """Token IDs of ``[start, end)`` with each multimodal span named by content.
+
+    Placeholder tokens are identical for every image, so a checkpoint keyed by
+    token IDs alone could restore one image's state for another. Positions in
+    a multimodal span are replaced by pseudo-IDs derived from the item's
+    content hash and the position inside the span. Text requests are unchanged.
+
+    Args:
+        request: Request whose ``all_token_ids`` and ``mm_features`` are read.
+        start: First position, inclusive.
+        end: Last position, exclusive.
+        salt: Extra identity (for example encoder precision) mixed into the IDs.
+    """
+    tokens = request.all_token_ids[start:end]
+    features = request.mm_features
+    if not features:
+        return tuple(tokens)
+    tokens = list(tokens)
+    for feature in features:
+        offset = feature.mm_position.offset
+        length = feature.mm_position.length
+        lo, hi = max(offset, start), min(offset + length, end)
+        if lo >= hi:
+            continue
+        key = hashlib.sha256(
+            b"vllm-mm-content-tokens-v2\0" + salt + b"\0" + feature.identifier.encode()
+        ).digest()
+        # Counter-mode BLAKE2b keyed by the full digest: 16 values per call,
+        # independent across positions, so a span is not one 30-bit seed.
+        first, last = lo - offset, hi - offset
+        for block in range(first // 16, (last - 1) // 16 + 1):
+            words = hashlib.blake2b(
+                block.to_bytes(8, "little"), key=key, digest_size=64
+            ).digest()
+            for index in range(max(first, block * 16), min(last, block * 16 + 16)):
+                word = words[(index % 16) * 4 : (index % 16) * 4 + 4]
+                value = int.from_bytes(word, "little") & (_CONTENT_TOKEN_BASE - 1)
+                tokens[offset + index - start] = _CONTENT_TOKEN_BASE + value
+    return tuple(tokens)
+
+
+def checkpoint_end_allowed(request: Request, num_tokens: int) -> bool:
+    """Whether a checkpoint may end after ``num_tokens`` of this request.
+
+    It must not end within the first MIN_CHECKPOINT_ITEM_TOKENS positions of
+    a multimodal item: its key would identify that item by too few bits.
+    """
+    for feature in request.mm_features or ():
+        offset = feature.mm_position.offset
+        depth = min(MIN_CHECKPOINT_ITEM_TOKENS, feature.mm_position.length)
+        if offset < num_tokens < offset + depth:
+            return False
+    return True
+
+
+# Optional checkpoints have fixed slots; absent slots contain no KV blocks.
+PROMPT_CHECKPOINT_SLOT = 0
+RESPONSE_CHECKPOINT_SLOT = 1
+INSTRUCTION_CHECKPOINT_SLOT = 2
+PREFILL_TAIL_CHECKPOINT_SLOT = 3
+NUM_BOUNDARY_CHECKPOINT_SLOTS = 4
+
+BoundaryCheckpointKind = Literal["instruction", "prompt", "response", "prefill_tail"]
+
+
+import os as _os
+
+# VLLM_QWEN38_TAIL_CKPT="GRID,SUFFIX": place the prefill-tail checkpoint on a GRID-token
+# grid at least SUFFIX tokens before the prompt end, instead of one to two prefill
+# chunks before it. A follow-up turn re-renders only the previous generation prompt
+# (a few tokens) differently, so it can resume almost at the old prompt's end and the
+# reusable prefix no longer depends on the chunk size. Costs one short extra prefill
+# step per new prompt. Hits stay hash-validated, so a misplaced checkpoint only misses.
+_TAIL_CKPT = tuple(
+    int(v) for v in _os.environ.get("VLLM_QWEN38_TAIL_CKPT", "").split(",") if v
+)
+
+
+def get_prefill_tail_checkpoint_position(
+    prompt_tokens: int, chunk_budget: int, instruction_boundary: int | None
+) -> int | None:
+    """Leave one to two prefill chunks after a bounded suffix-recovery point."""
+    if chunk_budget <= 0:
+        return None
+    if len(_TAIL_CKPT) == 2:
+        grid, suffix = _TAIL_CKPT
+        position = (prompt_tokens - suffix) // grid * grid
+    else:
+        position = (prompt_tokens // chunk_budget - 1) * chunk_budget
+    return position if position > (instruction_boundary or 0) else None
+
+
+def boundary_checkpoint_slots(request: Request) -> tuple[int, ...]:
+    """Allocate only the endpoints that this request can publish."""
+    slots: tuple[int, ...] = (PROMPT_CHECKPOINT_SLOT, RESPONSE_CHECKPOINT_SLOT)
+    if request.recurrent_instruction_boundary is not None:
+        slots += (INSTRUCTION_CHECKPOINT_SLOT,)
+    if request.recurrent_prefill_tail_boundary is not None:
+        slots += (PREFILL_TAIL_CHECKPOINT_SLOT,)
+    return slots
+
+
+@dataclass(frozen=True)
+class BoundaryCheckpoint:
+    checkpoint_id: int
+    num_tokens: int
+    block_ids: tuple[tuple[int, ...], ...]
+    auxiliary_block_ids: tuple[int, ...] = ()
+    draft_prefix_len: int = 0
+    kind: BoundaryCheckpointKind = "prompt"
+
+    @property
+    def dependencies(self) -> frozenset[int]:
+        return frozenset(
+            block_id
+            for group in (*self.block_ids, self.auxiliary_block_ids)
+            for block_id in group
+            if block_id != 0
+        )
+
+
+@dataclass
+class _RadixNode:
+    tokens: tuple[int, ...] = ()
+    checkpoint_id: int | None = None
+    children: dict[int, "_RadixNode"] = field(default_factory=dict)
+
+    def insert(self, tokens: tuple[int, ...], checkpoint_id: int) -> None:
+        node = self
+        while tokens:
+            child = node.children.get(tokens[0])
+            if child is None:
+                node.children[tokens[0]] = _RadixNode(tokens, checkpoint_id)
+                return
+            common = 0
+            for a, b in zip(tokens, child.tokens):
+                if a != b:
+                    break
+                common += 1
+            if common < len(child.tokens):
+                split = _RadixNode(child.tokens[:common])
+                node.children[tokens[0]] = split
+                child.tokens = child.tokens[common:]
+                split.children[child.tokens[0]] = child
+                child = split
+            node = child
+            tokens = tokens[common:]
+        node.checkpoint_id = checkpoint_id
+
+    def find(self, tokens: tuple[int, ...]) -> int | None:
+        node = self
+        found = node.checkpoint_id
+        offset = 0
+        while offset < len(tokens):
+            child = node.children.get(tokens[offset])
+            if child is None or tokens[offset : offset + len(child.tokens)] != (
+                child.tokens
+            ):
+                break
+            offset += len(child.tokens)
+            node = child
+            if node.checkpoint_id is not None:
+                found = node.checkpoint_id
+        return found
+
+    def remove(self, tokens: tuple[int, ...]) -> None:
+        node = self
+        path: list[tuple[_RadixNode, int]] = []
+        offset = 0
+        while offset < len(tokens):
+            key = tokens[offset]
+            path.append((node, key))
+            node = node.children[key]
+            offset += len(node.tokens)
+        node.checkpoint_id = None
+        for parent, key in reversed(path):
+            child = parent.children[key]
+            if child.checkpoint_id is not None:
+                break
+            if not child.children:
+                del parent.children[key]
+            elif len(child.children) == 1:
+                grandchild = next(iter(child.children.values()))
+                grandchild.tokens = child.tokens + grandchild.tokens
+                parent.children[key] = grandchild
+
+
+# A full hash authenticates preceding blocks; the first block uses the same
+# discriminators as the normal prefix cache, including cache salt and LoRA.
+_RootKey = tuple[int, BlockHash | None, tuple[Any, ...] | None]
+
+
+@dataclass
+class _PendingCheckpoint:
+    checkpoint: BoundaryCheckpoint
+    root_key: _RootKey
+    tokens: tuple[int, ...]
+    num_ranks: int
+    completed_ranks: set[int] = field(default_factory=set)
+    invalidated: bool = False
+
+
+class BoundaryCheckpointCache:
+    """Cache immutable bundles in the ordinary, evictable KV block pool.
+
+    Staged bundles pin every dependency until all workers acknowledge their
+    copies. Published bundles release those pins and participate in the pool's
+    LRU. Evicting any dependency invalidates the entire bundle. Readers acquire
+    their own pins and must copy mutable tails before running a continuation.
+    """
+
+    def __init__(self, block_pool: "BlockPool") -> None:
+        if block_pool.boundary_checkpoints is not None:
+            raise ValueError("Block pool already has a boundary checkpoint cache")
+        block_pool.boundary_checkpoints = self
+        self.block_pool = block_pool
+        self.hash_block_size = block_pool.hash_block_size
+        self._roots: dict[_RootKey, _RadixNode] = {}
+        self._entries: dict[int, _PendingCheckpoint] = {}
+        self._pending: dict[int, _PendingCheckpoint] = {}
+        self._by_block: dict[int, set[int]] = {}
+        self._pending_by_block: dict[int, set[int]] = {}
+        self._next_id = 1
+
+    def __len__(self) -> int:
+        return len(self._entries)
+
+    def next_id(self) -> int:
+        checkpoint_id = self._next_id
+        self._next_id += 1
+        return checkpoint_id
+
+    def _root_key(self, request: Request, start: int) -> _RootKey:
+        if start:
+            return start, request.block_hashes[start // self.hash_block_size - 1], None
+        extra_keys, _ = generate_block_hash_extra_keys(
+            request, 0, self.hash_block_size, 0
+        )
+        return 0, None, extra_keys
+
+    @staticmethod
+    def supports_request(request: Request) -> bool:
+        return (
+            request.prompt_token_ids is not None
+            and request.prompt_embeds is None
+            and all(
+                feature.modality in CHECKPOINT_MODALITIES and feature.identifier
+                for feature in request.mm_features or ()
+            )
+            and not request.resumable
+            and request.sampling_params is not None
+            and len(request.sampling_params.stop_token_ids or ())
+            + int(request.sampling_params.eos_token_id is not None)
+            <= MAX_BOUNDARY_STOP_TOKENS
+        )
+
+    def stage(
+        self,
+        request: Request,
+        checkpoint: BoundaryCheckpoint,
+        num_ranks: int,
+    ) -> None:
+        """Pin a complete bundle while its asynchronous copies finish."""
+        if not self.supports_request(request):
+            raise ValueError("Boundary checkpoints require a text generation request")
+        if not 0 < checkpoint.num_tokens <= request.num_tokens:
+            raise ValueError("Checkpoint must cover an existing, nonempty prefix")
+        if not checkpoint_end_allowed(request, checkpoint.num_tokens):
+            raise ValueError("Checkpoint cannot end at the start of a multimodal item")
+        if not 0 <= checkpoint.draft_prefix_len <= checkpoint.num_tokens:
+            raise ValueError("Draft prefix cannot extend beyond the target prefix")
+        if num_ranks < 1:
+            raise ValueError("Checkpoint must be acknowledged by at least one rank")
+        cid = checkpoint.checkpoint_id
+        if cid in self._entries or cid in self._pending:
+            raise ValueError("Checkpoint ID is already in use")
+        blocks = [self.block_pool.blocks[i] for i in checkpoint.dependencies]
+        if not blocks or any(block.ref_cnt == 0 for block in blocks):
+            raise ValueError("Checkpoint dependencies must still be owned by producer")
+        start = ((checkpoint.num_tokens - 1) // self.hash_block_size) * (
+            self.hash_block_size
+        )
+        pending = _PendingCheckpoint(
+            checkpoint,
+            self._root_key(request, start),
+            content_token_ids(request, start, checkpoint.num_tokens),
+            num_ranks,
+        )
+        self.block_pool.touch(blocks)
+        self._pending[cid] = pending
+        for block_id in checkpoint.dependencies:
+            self._pending_by_block.setdefault(block_id, set()).add(cid)
+
+    def acknowledge(self, checkpoint_id: int, rank: int) -> bool:
+        """Publish only after every rank's copy completion has been observed."""
+        pending = self._pending.get(checkpoint_id)
+        if pending is None:
+            return False
+        if not 0 <= rank < pending.num_ranks:
+            raise ValueError("Invalid checkpoint rank")
+        pending.completed_ranks.add(rank)
+        if len(pending.completed_ranks) != pending.num_ranks:
+            return False
+        if pending.invalidated:
+            self.discard(checkpoint_id)
+            return False
+        del self._pending[checkpoint_id]
+        self._forget_pending(pending.checkpoint)
+        root = self._roots.setdefault(pending.root_key, _RadixNode())
+        previous_id = root.find(pending.tokens)
+        if previous_id is not None:
+            previous = self._entries[previous_id]
+            if previous.tokens == pending.tokens:
+                # Deduplicate exact endpoints without disturbing active readers.
+                self.invalidate(previous_id)
+                root = self._roots.setdefault(pending.root_key, _RadixNode())
+        root.insert(pending.tokens, checkpoint_id)
+        self._entries[checkpoint_id] = pending
+        for block_id in pending.checkpoint.dependencies:
+            self._by_block.setdefault(block_id, set()).add(checkpoint_id)
+        self._release(pending.checkpoint)
+        return True
+
+    def is_pending(self, checkpoint_id: int) -> bool:
+        """Return whether a bundle still owns pins for unfinished worker copies."""
+        return checkpoint_id in self._pending
+
+    def discard(self, checkpoint_id: int) -> None:
+        """Release an unpublished bundle after its GPU writers have completed."""
+        pending = self._pending.pop(checkpoint_id, None)
+        if pending is not None:
+            self._forget_pending(pending.checkpoint)
+            self._release(pending.checkpoint)
+
+    def _forget_pending(self, checkpoint: BoundaryCheckpoint) -> None:
+        for block_id in checkpoint.dependencies:
+            ids = self._pending_by_block[block_id]
+            ids.remove(checkpoint.checkpoint_id)
+            if not ids:
+                del self._pending_by_block[block_id]
+
+    def find(self, request: Request, max_length: int) -> BoundaryCheckpoint | None:
+        if not self.supports_request(request) or request.skip_reading_prefix_cache:
+            return None
+        max_length = min(max_length, request.num_tokens)
+        if max_length <= 0:
+            return None
+        start = ((max_length - 1) // self.hash_block_size) * self.hash_block_size
+        for offset in range(start, -1, -self.hash_block_size):
+            root = self._roots.get(self._root_key(request, offset))
+            if root is None:
+                continue
+            tokens = content_token_ids(
+                request, offset, min(offset + self.hash_block_size, max_length)
+            )
+            checkpoint_id = root.find(tokens)
+            if checkpoint_id is not None:
+                return self._entries[checkpoint_id].checkpoint
+        return None
+
+    def acquire(self, checkpoint_id: int) -> BoundaryCheckpoint | None:
+        entry = self._entries.get(checkpoint_id)
+        if entry is None:
+            return None
+        checkpoint = entry.checkpoint
+        self.block_pool.touch(
+            [self.block_pool.blocks[i] for i in checkpoint.dependencies]
+        )
+        return checkpoint
+
+    def release(self, checkpoint: BoundaryCheckpoint) -> None:
+        self._release(checkpoint)
+
+    def _release(self, checkpoint: BoundaryCheckpoint) -> None:
+        self.block_pool.free_blocks(
+            self.block_pool.blocks[i] for i in checkpoint.dependencies
+        )
+
+    def contains_block(self, block_id: int) -> bool:
+        return block_id in self._by_block
+
+    def invalidate_block(self, block_id: int) -> None:
+        for checkpoint_id in self._pending_by_block.get(block_id, ()):
+            self._pending[checkpoint_id].invalidated = True
+        for checkpoint_id in tuple(self._by_block.get(block_id, ())):
+            self.invalidate(checkpoint_id)
+
+    def invalidate(self, checkpoint_id: int) -> None:
+        entry = self._entries.pop(checkpoint_id, None)
+        if entry is None:
+            return
+        root = self._roots[entry.root_key]
+        root.remove(entry.tokens)
+        if not root.children:
+            del self._roots[entry.root_key]
+        for block_id in entry.checkpoint.dependencies:
+            ids = self._by_block[block_id]
+            ids.remove(checkpoint_id)
+            if not ids:
+                del self._by_block[block_id]
+
+    def clear(self) -> None:
+        if self._pending:
+            raise RuntimeError("Cannot clear checkpoints with pending GPU writers")
+        self._entries.clear()
+        self._roots.clear()
+        self._by_block.clear()

recipe/patches/vllm_kernels/linear/mxfp8/b12x.py

@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+
+from vllm.model_executor.kernels.linear.b12x_blockscaled import B12xBlockscaledLinear
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    MXFP8_BLOCK_SIZE,
+    MXFP8_SCALE_DTYPE,
+    MXFP8_VALUE_DTYPE,
+)
+from vllm.model_executor.utils import replace_parameter
+from vllm.platforms import current_platform
+from vllm.utils.b12x import (
+    B12xWorkload,
+    b12x_layer_prefix,
+    get_b12x_dense_activation_mode,
+    get_b12x_projection_workspace_sizes,
+    register_b12x_layer,
+    reuse_packed_weight_storage,
+    run_b12x_blockscaled_linear,
+    set_b12x_preparation_provider,
+)
+from vllm.utils.b12x import (
+    get_b12x_blockscaled as _import_b12x_blockscaled,
+)
+from vllm.utils.torch_utils import _encode_layer_name
+
+from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
+
+import os as _os
+
+from vllm.logger import init_logger as _init_logger
+
+_logger = _init_logger(__name__)
+
+# Decode-row split-K GEMV over the packed values/scale_rows for the listed
+# "NxK" shapes, or "all" (qwen4_exp/nvidia/ops/mxfp8_gemv.py). Empty disables.
+_GEMV_SPEC = _os.environ.get("VLLM_QWEN38_MXFP8_GEMV", "")
+_GEMV_ALL = _GEMV_SPEC == "all"
+_GEMV_SHAPES = frozenset(
+    tuple(int(v) for v in item.split("x"))
+    for item in ("" if _GEMV_ALL else _GEMV_SPEC).split(",")
+    if item
+)
+_GEMV_LOGGED: set = set()
+
+
+class B12xMxfp8LinearKernel(Mxfp8LinearKernel):
+    """ModelOpt MXFP8 linear through a layer-held prepared b12x plan."""
+
+    @classmethod
+    def is_supported(cls, compute_capability=None):
+        del compute_capability
+        if not current_platform.is_cuda():
+            return False, "b12x MXFP8 kernels are only available on CUDA"
+        if not current_platform.is_device_capability_family(120):
+            return False, "b12x MXFP8 kernels require a Blackwell 12x device"
+        api = _import_b12x_blockscaled()
+        if api is None:
+            return False, "Install the B12X backend with `pip install vllm[b12x]`"
+        if not api.is_supported():
+            return False, "b12x.gemm.blockscaled is not supported"
+        return True, None
+
+    @classmethod
+    def can_implement(cls, config: Mxfp8LinearLayerConfig):
+        del config
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight = layer.weight.data
+        scales = layer.weight_scale.data
+        assert weight.dtype == MXFP8_VALUE_DTYPE and weight.ndim == 2
+        assert scales.dtype == MXFP8_SCALE_DTYPE and scales.ndim == 2
+        out_features, in_features = map(int, weight.shape)
+        assert in_features % MXFP8_BLOCK_SIZE == 0
+        layer.b12x_mxfp8_output_size = out_features
+        padded_out_features = (out_features + 7) // 8 * 8
+        if padded_out_features != out_features:
+            padded_weight = weight.new_zeros((padded_out_features, in_features))
+            padded_weight[:out_features].copy_(weight)
+            weight = padded_weight
+            # E8M0 byte 127 is a unit scale for the artificial zero rows.
+            padded_scales = scales.new_full(
+                (max(padded_out_features, scales.shape[0]), scales.shape[1]), 127
+            )
+            padded_scales[:out_features].copy_(scales[:out_features])
+            scales = padded_scales
+        api = _import_b12x_blockscaled()
+        assert api is not None
+        packed = api.pack_weight(
+            weight.detach(),
+            scales[:, : in_features // MXFP8_BLOCK_SIZE].detach(),
+        )
+        layer.b12x_mxfp8_packed_weight = reuse_packed_weight_storage(
+            getattr(layer, "b12x_mxfp8_packed_weight", None), packed
+        )
+        # A method that pre-quantizes or keeps BF16 activations sets the mode
+        # before weight processing; the configured default applies otherwise.
+        layer.b12x_activation_mode = getattr(
+            layer, "b12x_activation_mode", None
+        ) or get_b12x_dense_activation_mode("mxfp8")
+        layer.b12x_bf16_input_supported = (
+            in_features % 128 == 0 and padded_out_features % 8 == 0
+        )
+        name = b12x_layer_prefix(layer)
+        # A reload into the same packed storage keeps the holder and its
+        # prepared plan; new storage declares anew.
+        existing = getattr(layer, "b12x_linear", None)
+        if existing is None or not existing.holds(layer.b12x_mxfp8_packed_weight):
+            layer.b12x_linear = B12xBlockscaledLinear(
+                layer.b12x_mxfp8_packed_weight,
+                recipe="mxfp8",
+                activation_mode=layer.b12x_activation_mode,
+                layer_name=name,
+            )
+        # A layer marked b12x_prefill_a16 runs prefill-sized calls with BF16
+        # activations through a second plan over the same packed weight.
+        wants_prefill = (
+            getattr(layer, "b12x_prefill_a16", False)
+            and layer.b12x_activation_mode != "a16"
+        )
+        prefill = getattr(layer, "b12x_linear_prefill", None)
+        if not wants_prefill:
+            layer.b12x_linear_prefill = None
+        elif prefill is None or not prefill.holds(layer.b12x_mxfp8_packed_weight):
+            layer.b12x_linear_prefill = B12xBlockscaledLinear(
+                layer.b12x_mxfp8_packed_weight,
+                recipe="mxfp8",
+                activation_mode="a16",
+                layer_name=name,
+            )
+        layer.b12x_layer_name = _encode_layer_name(name)
+        register_b12x_layer(name, layer)
+        replace_parameter(layer, "weight", weight.new_empty((0,)))
+        replace_parameter(layer, "weight_scale", scales.new_empty((0,)))
+        if not getattr(layer, "b12x_preparation_suppressed", False):
+            set_b12x_preparation_provider(layer, self)
+
+    def get_b12x_preparation_units(
+        self, layer: torch.nn.Module, workload: B12xWorkload
+    ) -> Sequence[object]:
+        packed = layer.b12x_mxfp8_packed_weight
+        if packed.weight.values.is_meta:
+            return ()
+        linear = layer.b12x_linear
+        units = [linear.unit(workload, name=f"linear.mxfp8.{linear.layer_name}")]
+        prefill = getattr(layer, "b12x_linear_prefill", None)
+        if prefill is not None:
+            units.append(
+                prefill.unit(workload, name=f"linear.mxfp8.a16.{prefill.layer_name}")
+            )
+        return tuple(units)
+
+    def get_workspace_size(self, layer: torch.nn.Module, rows: int) -> int:
+        return get_b12x_projection_workspace_sizes(rows, layer)[0]
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        source = x.reshape(-1, x.shape[-1]).contiguous()
+        if source.dtype != torch.bfloat16:
+            raise ValueError("prepared vLLM MXFP8 path requires BF16 activations")
+        logical_out_features = layer.b12x_mxfp8_output_size
+        if (_GEMV_ALL or _GEMV_SHAPES) and source.shape[0] <= 16 and bias is None:
+            packed = layer.b12x_mxfp8_packed_weight
+            key = (int(logical_out_features), int(packed.in_features))
+            if _GEMV_ALL or key in _GEMV_SHAPES:
+                from vllm.models.qwen4_exp.nvidia.ops.mxfp8_gemv import mxfp8_gemv
+
+                if key not in _GEMV_LOGGED:
+                    _GEMV_LOGGED.add(key)
+                    _logger.info("MXFP8 decode GEMV active for N=%d K=%d", *key)
+                out = mxfp8_gemv(
+                    source,
+                    packed.weight.values,
+                    packed.weight.scale_rows[0],
+                    key[0],
+                    key[1],
+                )
+                return out.view(*x.shape[:-1], key[0])
+        out_features = int(layer.b12x_mxfp8_packed_weight.out_features)
+        padded = logical_out_features != out_features
+        output = run_b12x_blockscaled_linear(
+            source, None if padded else bias, out_features, layer.b12x_layer_name
+        )
+        if padded:
+            output = output[:, :logical_out_features]
+            if bias is not None:
+                output.add_(bias)
+        return output.view(*x.shape[:-1], logical_out_features)
+
+
+__all__ = ["B12xMxfp8LinearKernel"]

recipe/results/bench-sbs.txt

@@ -0,0 +1,4 @@
+  qwen3.8-flash-next ctx=2000 conc=1: prompt_tok=2074 | prefill     1850 tok/s agg (   1850/stream, TTFT   1.1s) | decode   51.7 tok/s agg ( 51.7/stream)
+  qwen3.8-flash-next ctx=2000 conc=4: prompt_tok=2063 | prefill     2230 tok/s agg (    558/stream, TTFT   4.1s) | decode  107.0 tok/s agg ( 26.7/stream)
+  qwen3.8-flash-next ctx=46000 conc=1: prompt_tok=45928 | prefill     2259 tok/s agg (   2259/stream, TTFT  20.3s) | decode   45.8 tok/s agg ( 45.8/stream)
+  qwen3.8-flash-next ctx=46000 conc=4: prompt_tok=45990 | prefill     4114 tok/s agg (   1029/stream, TTFT  84.0s) | decode   37.4 tok/s agg (  9.4/stream)

recipe/results/kl-neutral-a16-as-served.json

@@ -0,0 +1,15 @@
+{
+ "chunks": 87,
+ "positions": 178089,
+ "mean_kld": 0.04918030487709328,
+ "kld_se": 0.00030368987921901996,
+ "median_kld": 0.012365318834781647,
+ "p90_kld": 0.11989473551511765,
+ "p99_kld": 0.5401471257209778,
+ "p999_kld": 1.5018013715744019,
+ "same_top1_pct": 91.74120804766156,
+ "ppl_q": 4.076651335380382,
+ "ppl_base": 4.045180731278633,
+ "ln_ppl_ratio": 0.00774967060309506,
+ "rms_dp_pct": 6.86464061644188
+}

recipe/results/kl-neutral-a4-prefill.json

@@ -0,0 +1,15 @@
+{
+ "chunks": 87,
+ "positions": 178089,
+ "mean_kld": 0.06648334885537187,
+ "kld_se": 0.00039647044943181815,
+ "median_kld": 0.017018219456076622,
+ "p90_kld": 0.16397449374198914,
+ "p99_kld": 0.7069363594055176,
+ "p999_kld": 1.9706969261169434,
+ "same_top1_pct": 90.33067735795024,
+ "ppl_q": 4.095423230995411,
+ "ppl_base": 4.045180731278633,
+ "ln_ppl_ratio": 0.012343835396803282,
+ "rms_dp_pct": 7.96968497071922
+}

recipe/results/needle.txt

@@ -0,0 +1,6 @@
+prompt_tokens=28,842 | 17.2s (~1,672 tok/s prefill) | found=True | 'PELICAN-42'
+NEEDLE OK
+prompt_tokens=96,042 | 57.8s (~1,661 tok/s prefill) | found=True | 'PELICAN-42'
+NEEDLE OK
+prompt_tokens=124,842 | 65.5s (~1,905 tok/s prefill) | found=True | 'PELICAN-42'
+NEEDLE OK

recipe/results/quality-probe.jsonl

@@ -0,0 +1,364 @@
+{"task": "gsm8k", "ok": true, "pred": "18", "ref": "18", "finish": "stop", "out_tokens": 239}
+{"task": "gsm8k", "ok": true, "pred": "3", "ref": "3", "finish": "stop", "out_tokens": 236}
+{"task": "gsm8k", "ok": true, "pred": "70000", "ref": "70000", "finish": "stop", "out_tokens": 520}
+{"task": "gsm8k", "ok": true, "pred": "540", "ref": "540", "finish": "stop", "out_tokens": 250}
+{"task": "gsm8k", "ok": true, "pred": "20", "ref": "20", "finish": "stop", "out_tokens": 282}
+{"task": "gsm8k", "ok": true, "pred": "64", "ref": "64", "finish": "stop", "out_tokens": 476}
+{"task": "gsm8k", "ok": true, "pred": "260", "ref": "260", "finish": "stop", "out_tokens": 235}
+{"task": "gsm8k", "ok": true, "pred": "160", "ref": "160", "finish": "stop", "out_tokens": 439}
+{"task": "gsm8k", "ok": true, "pred": "45", "ref": "45", "finish": "stop", "out_tokens": 653}
+{"task": "gsm8k", "ok": true, "pred": "460", "ref": "460", "finish": "stop", "out_tokens": 308}
+{"task": "gsm8k", "ok": true, "pred": "366", "ref": "366", "finish": "stop", "out_tokens": 449}
+{"task": "gsm8k", "ok": true, "pred": "694", "ref": "694", "finish": "stop", "out_tokens": 491}
+{"task": "gsm8k", "ok": false, "pred": "12", "ref": "13", "finish": "stop", "out_tokens": 1383}
+{"task": "gsm8k", "ok": true, "pred": "18", "ref": "18", "finish": "stop", "out_tokens": 611}
+{"task": "gsm8k", "ok": true, "pred": "60", "ref": "60", "finish": "stop", "out_tokens": 314}
+{"task": "gsm8k", "ok": true, "pred": "125", "ref": "125", "finish": "stop", "out_tokens": 296}
+{"task": "gsm8k", "ok": true, "pred": "230", "ref": "230", "finish": "stop", "out_tokens": 213}
+{"task": "gsm8k", "ok": true, "pred": "57500", "ref": "57500", "finish": "stop", "out_tokens": 380}
+{"task": "gsm8k", "ok": true, "pred": "7", "ref": "7", "finish": "stop", "out_tokens": 196}
+{"task": "gsm8k", "ok": true, "pred": "6", "ref": "6", "finish": "stop", "out_tokens": 402}
+{"task": "gsm8k", "ok": true, "pred": "15", "ref": "15", "finish": "stop", "out_tokens": 577}
+{"task": "gsm8k", "ok": true, "pred": "14", "ref": "14", "finish": "stop", "out_tokens": 223}
+{"task": "gsm8k", "ok": true, "pred": "7", "ref": "7", "finish": "stop", "out_tokens": 204}
+{"task": "gsm8k", "ok": true, "pred": "8", "ref": "8", "finish": "stop", "out_tokens": 216}
+{"task": "gsm8k", "ok": true, "pred": "26", "ref": "26", "finish": "stop", "out_tokens": 255}
+{"task": "gsm8k", "ok": true, "pred": "2", "ref": "2", "finish": "stop", "out_tokens": 368}
+{"task": "gsm8k", "ok": true, "pred": "243", "ref": "243", "finish": "stop", "out_tokens": 375}
+{"task": "gsm8k", "ok": true, "pred": "16", "ref": "16", "finish": "stop", "out_tokens": 232}
+{"task": "gsm8k", "ok": true, "pred": "25", "ref": "25", "finish": "stop", "out_tokens": 295}
+{"task": "gsm8k", "ok": true, "pred": "104", "ref": "104", "finish": "stop", "out_tokens": 273}
+{"task": "gsm8k", "ok": true, "pred": "109", "ref": "109", "finish": "stop", "out_tokens": 290}
+{"task": "gsm8k", "ok": true, "pred": "80", "ref": "80", "finish": "stop", "out_tokens": 306}
+{"task": "gsm8k", "ok": true, "pred": "35", "ref": "35", "finish": "stop", "out_tokens": 140}
+{"task": "gsm8k", "ok": true, "pred": "70", "ref": "70", "finish": "stop", "out_tokens": 407}
+{"task": "gsm8k", "ok": true, "pred": "23", "ref": "23", "finish": "stop", "out_tokens": 216}
+{"task": "gsm8k", "ok": true, "pred": "9", "ref": "9", "finish": "stop", "out_tokens": 256}
+{"task": "gsm8k", "ok": true, "pred": "75", "ref": "75", "finish": "stop", "out_tokens": 240}
+{"task": "gsm8k", "ok": true, "pred": "2", "ref": "2", "finish": "stop", "out_tokens": 434}
+{"task": "gsm8k", "ok": true, "pred": "10", "ref": "10", "finish": "stop", "out_tokens": 347}
+{"task": "gsm8k", "ok": true, "pred": "18", "ref": "18", "finish": "stop", "out_tokens": 524}
+{"task": "gsm8k", "ok": true, "pred": "8", "ref": "8", "finish": "stop", "out_tokens": 315}
+{"task": "gsm8k", "ok": true, "pred": "200", "ref": "200", "finish": "stop", "out_tokens": 418}
+{"task": "gsm8k", "ok": true, "pred": "26", "ref": "26", "finish": "stop", "out_tokens": 218}
+{"task": "gsm8k", "ok": true, "pred": "48", "ref": "48", "finish": "stop", "out_tokens": 405}
+{"task": "gsm8k", "ok": true, "pred": "20", "ref": "20", "finish": "stop", "out_tokens": 338}
+{"task": "gsm8k", "ok": true, "pred": "104", "ref": "104", "finish": "stop", "out_tokens": 461}
+{"task": "gsm8k", "ok": true, "pred": "163", "ref": "163", "finish": "stop", "out_tokens": 415}
+{"task": "gsm8k", "ok": true, "pred": "800", "ref": "800", "finish": "stop", "out_tokens": 423}
+{"task": "gsm8k", "ok": true, "pred": "8", "ref": "8", "finish": "stop", "out_tokens": 218}
+{"task": "gsm8k", "ok": true, "pred": "30", "ref": "30", "finish": "stop", "out_tokens": 343}
+{"task": "gsm8k", "ok": true, "pred": "294", "ref": "294", "finish": "stop", "out_tokens": 245}
+{"task": "gsm8k", "ok": true, "pred": "5", "ref": "5", "finish": "stop", "out_tokens": 288}
+{"task": "gsm8k", "ok": true, "pred": "15", "ref": "15", "finish": "stop", "out_tokens": 346}
+{"task": "gsm8k", "ok": true, "pred": "40", "ref": "40", "finish": "stop", "out_tokens": 338}
+{"task": "gsm8k", "ok": true, "pred": "40", "ref": "40", "finish": "stop", "out_tokens": 322}
+{"task": "gsm8k", "ok": true, "pred": "14", "ref": "14", "finish": "stop", "out_tokens": 172}
+{"task": "gsm8k", "ok": true, "pred": "3", "ref": "3", "finish": "stop", "out_tokens": 208}
+{"task": "gsm8k", "ok": true, "pred": "83", "ref": "83", "finish": "stop", "out_tokens": 533}
+{"task": "gsm8k", "ok": true, "pred": "57", "ref": "57", "finish": "stop", "out_tokens": 901}
+{"task": "gsm8k", "ok": true, "pred": "187", "ref": "187", "finish": "stop", "out_tokens": 184}
+{"task": "gsm8k", "ok": true, "pred": "17", "ref": "17", "finish": "stop", "out_tokens": 330}
+{"task": "gsm8k", "ok": true, "pred": "1430", "ref": "1430", "finish": "stop", "out_tokens": 247}
+{"task": "gsm8k", "ok": true, "pred": "25000", "ref": "25000", "finish": "stop", "out_tokens": 396}
+{"task": "gsm8k", "ok": true, "pred": "1596", "ref": "1596", "finish": "stop", "out_tokens": 432}
+{"task": "gsm8k", "ok": true, "pred": "300", "ref": "300", "finish": "stop", "out_tokens": 333}
+{"task": "gsm8k", "ok": true, "pred": "36", "ref": "36", "finish": "stop", "out_tokens": 309}
+{"task": "gsm8k", "ok": true, "pred": "48", "ref": "48", "finish": "stop", "out_tokens": 358}
+{"task": "gsm8k", "ok": true, "pred": "595", "ref": "595", "finish": "stop", "out_tokens": 367}
+{"task": "gsm8k", "ok": true, "pred": "36", "ref": "36", "finish": "stop", "out_tokens": 305}
+{"task": "gsm8k", "ok": true, "pred": "60", "ref": "60", "finish": "stop", "out_tokens": 232}
+{"task": "gsm8k", "ok": true, "pred": "7425", "ref": "7425", "finish": "stop", "out_tokens": 285}
+{"task": "gsm8k", "ok": true, "pred": "60", "ref": "60", "finish": "stop", "out_tokens": 164}
+{"task": "gsm8k", "ok": true, "pred": "221", "ref": "221", "finish": "stop", "out_tokens": 268}
+{"task": "gsm8k", "ok": true, "pred": "255", "ref": "255", "finish": "stop", "out_tokens": 357}
+{"task": "gsm8k", "ok": true, "pred": "88", "ref": "88", "finish": "stop", "out_tokens": 555}
+{"task": "gsm8k", "ok": true, "pred": "60", "ref": "60", "finish": "stop", "out_tokens": 554}
+{"task": "gsm8k", "ok": true, "pred": "5", "ref": "5", "finish": "stop", "out_tokens": 386}
+{"task": "gsm8k", "ok": true, "pred": "100", "ref": "100", "finish": "stop", "out_tokens": 337}
+{"task": "gsm8k", "ok": true, "pred": "6", "ref": "6", "finish": "stop", "out_tokens": 401}
+{"task": "gsm8k", "ok": true, "pred": "70", "ref": "70", "finish": "stop", "out_tokens": 265}
+{"task": "gsm8k", "ok": true, "pred": "10", "ref": "10", "finish": "stop", "out_tokens": 295}
+{"task": "gsm8k", "ok": true, "pred": "17", "ref": "17", "finish": "stop", "out_tokens": 213}
+{"task": "gsm8k", "ok": true, "pred": "623", "ref": "623", "finish": "stop", "out_tokens": 281}
+{"task": "gsm8k", "ok": true, "pred": "600", "ref": "600", "finish": "stop", "out_tokens": 188}
+{"task": "gsm8k", "ok": true, "pred": "15", "ref": "15", "finish": "stop", "out_tokens": 401}
+{"task": "gsm8k", "ok": false, "pred": "48", "ref": "44", "finish": "stop", "out_tokens": 332}
+{"task": "gsm8k", "ok": true, "pred": "22", "ref": "22", "finish": "stop", "out_tokens": 262}
+{"task": "gsm8k", "ok": true, "pred": "9360", "ref": "9360", "finish": "stop", "out_tokens": 697}
+{"task": "gsm8k", "ok": true, "pred": "8000", "ref": "8000", "finish": "stop", "out_tokens": 251}
+{"task": "gsm8k", "ok": true, "pred": "24", "ref": "24", "finish": "stop", "out_tokens": 418}
+{"task": "gsm8k", "ok": true, "pred": "225", "ref": "225", "finish": "stop", "out_tokens": 355}
+{"task": "gsm8k", "ok": true, "pred": "28", "ref": "28", "finish": "stop", "out_tokens": 232}
+{"task": "gsm8k", "ok": true, "pred": "4", "ref": "4", "finish": "stop", "out_tokens": 270}
+{"task": "gsm8k", "ok": false, "pred": "36.36", "ref": "36", "finish": "stop", "out_tokens": 606}
+{"task": "gsm8k", "ok": true, "pred": "348", "ref": "348", "finish": "stop", "out_tokens": 375}
+{"task": "gsm8k", "ok": true, "pred": "40", "ref": "40", "finish": "stop", "out_tokens": 399}
+{"task": "gsm8k", "ok": true, "pred": "3", "ref": "3", "finish": "stop", "out_tokens": 196}
+{"task": "gsm8k", "ok": true, "pred": "12", "ref": "12", "finish": "stop", "out_tokens": 457}
+{"task": "gsm8k", "ok": true, "pred": "5", "ref": "5", "finish": "stop", "out_tokens": 467}
+{"task": "gsm8k", "ok": true, "pred": "58", "ref": "58", "finish": "stop", "out_tokens": 293}
+{"task": "gsm8k", "ok": true, "pred": "175", "ref": "175", "finish": "stop", "out_tokens": 539}
+{"task": "gsm8k", "ok": true, "pred": "6", "ref": "6", "finish": "stop", "out_tokens": 261}
+{"task": "gsm8k", "ok": true, "pred": "26", "ref": "26", "finish": "stop", "out_tokens": 326}
+{"task": "gsm8k", "ok": true, "pred": "140", "ref": "140", "finish": "stop", "out_tokens": 502}
+{"task": "gsm8k", "ok": true, "pred": "500", "ref": "500", "finish": "stop", "out_tokens": 320}
+{"task": "gsm8k", "ok": true, "pred": "20", "ref": "20", "finish": "stop", "out_tokens": 135}
+{"task": "gsm8k", "ok": true, "pred": "72", "ref": "72", "finish": "stop", "out_tokens": 560}
+{"task": "gsm8k", "ok": true, "pred": "3", "ref": "3", "finish": "stop", "out_tokens": 511}
+{"task": "gsm8k", "ok": true, "pred": "50", "ref": "50", "finish": "stop", "out_tokens": 435}
+{"task": "gsm8k", "ok": true, "pred": "28", "ref": "28", "finish": "stop", "out_tokens": 299}
+{"task": "gsm8k", "ok": true, "pred": "45", "ref": "45", "finish": "stop", "out_tokens": 257}
+{"task": "gsm8k", "ok": true, "pred": "16", "ref": "16", "finish": "stop", "out_tokens": 468}
+{"task": "gsm8k", "ok": true, "pred": "24", "ref": "24", "finish": "stop", "out_tokens": 182}
+{"task": "gsm8k", "ok": true, "pred": "25", "ref": "25", "finish": "stop", "out_tokens": 134}
+{"task": "gsm8k", "ok": true, "pred": "6", "ref": "6", "finish": "stop", "out_tokens": 627}
+{"task": "gsm8k", "ok": true, "pred": "90", "ref": "90", "finish": "stop", "out_tokens": 634}
+{"task": "gsm8k", "ok": true, "pred": "42", "ref": "42", "finish": "stop", "out_tokens": 357}
+{"task": "gsm8k", "ok": true, "pred": "360", "ref": "360", "finish": "stop", "out_tokens": 171}
+{"task": "gsm8k", "ok": true, "pred": "4", "ref": "4", "finish": "stop", "out_tokens": 264}
+{"task": "gsm8k", "ok": false, "pred": "99076.92", "ref": "95200", "finish": "stop", "out_tokens": 1311}
+{"task": "gsm8k", "ok": true, "pred": "240", "ref": "240", "finish": "stop", "out_tokens": 201}
+{"task": "gsm8k", "ok": true, "pred": "27", "ref": "27", "finish": "stop", "out_tokens": 257}
+{"task": "gsm8k", "ok": true, "pred": "48", "ref": "48", "finish": "stop", "out_tokens": 413}
+{"task": "gsm8k", "ok": true, "pred": "50", "ref": "50", "finish": "stop", "out_tokens": 389}
+{"task": "gsm8k", "ok": true, "pred": "10", "ref": "10", "finish": "stop", "out_tokens": 249}
+{"task": "gsm8k", "ok": true, "pred": "10", "ref": "10", "finish": "stop", "out_tokens": 285}
+{"task": "gsm8k", "ok": true, "pred": "82", "ref": "82", "finish": "stop", "out_tokens": 150}
+{"task": "gsm8k", "ok": true, "pred": "120", "ref": "120", "finish": "stop", "out_tokens": 209}
+{"task": "gsm8k", "ok": true, "pred": "880", "ref": "880", "finish": "stop", "out_tokens": 450}
+{"task": "gsm8k", "ok": true, "pred": "10000", "ref": "10000", "finish": "stop", "out_tokens": 493}
+{"task": "gsm8k", "ok": true, "pred": "30", "ref": "30", "finish": "stop", "out_tokens": 197}
+{"task": "gsm8k", "ok": false, "pred": "1040", "ref": "940", "finish": "stop", "out_tokens": 155}
+{"task": "gsm8k", "ok": true, "pred": "60", "ref": "60", "finish": "stop", "out_tokens": 240}
+{"task": "gsm8k", "ok": true, "pred": "13", "ref": "13", "finish": "stop", "out_tokens": 283}
+{"task": "gsm8k", "ok": true, "pred": "720", "ref": "720", "finish": "stop", "out_tokens": 150}
+{"task": "gsm8k", "ok": true, "pred": "40", "ref": "40", "finish": "stop", "out_tokens": 175}
+{"task": "gsm8k", "ok": true, "pred": "6", "ref": "6", "finish": "stop", "out_tokens": 169}
+{"task": "gsm8k", "ok": true, "pred": "29", "ref": "29", "finish": "stop", "out_tokens": 394}
+{"task": "gsm8k", "ok": true, "pred": "105", "ref": "105", "finish": "stop", "out_tokens": 457}
+{"task": "gsm8k", "ok": true, "pred": "70", "ref": "70", "finish": "stop", "out_tokens": 308}
+{"task": "gsm8k", "ok": true, "pred": "20", "ref": "20", "finish": "stop", "out_tokens": 310}
+{"task": "gsm8k", "ok": true, "pred": "400", "ref": "400", "finish": "stop", "out_tokens": 539}
+{"task": "gsm8k", "ok": true, "pred": "140", "ref": "140", "finish": "stop", "out_tokens": 176}
+{"task": "gsm8k", "ok": true, "pred": "16", "ref": "16", "finish": "stop", "out_tokens": 327}
+{"task": "gsm8k", "ok": true, "pred": "20", "ref": "20", "finish": "stop", "out_tokens": 352}
+{"task": "gsm8k", "ok": true, "pred": "4000", "ref": "4000", "finish": "stop", "out_tokens": 269}
+{"task": "gsm8k", "ok": true, "pred": "2125", "ref": "2125", "finish": "stop", "out_tokens": 1013}
+{"task": "gsm8k", "ok": true, "pred": "75", "ref": "75", "finish": "stop", "out_tokens": 525}
+{"task": "gsm8k", "ok": true, "pred": "30", "ref": "30", "finish": "stop", "out_tokens": 199}
+{"task": "gsm8k", "ok": true, "pred": "16", "ref": "16", "finish": "stop", "out_tokens": 313}
+{"task": "gsm8k", "ok": true, "pred": "4", "ref": "4", "finish": "stop", "out_tokens": 528}
+{"task": "gsm8k", "ok": true, "pred": "5", "ref": "5", "finish": "stop", "out_tokens": 426}
+{"task": "gsm8k", "ok": true, "pred": "4", "ref": "4", "finish": "stop", "out_tokens": 203}
+{"task": "gsm8k", "ok": true, "pred": "48", "ref": "48", "finish": "stop", "out_tokens": 747}
+{"task": "gsm8k", "ok": true, "pred": "272", "ref": "272", "finish": "stop", "out_tokens": 435}
+{"task": "gsm8k", "ok": true, "pred": "280", "ref": "280", "finish": "stop", "out_tokens": 471}
+{"task": "gsm8k", "ok": true, "pred": "1400", "ref": "1400", "finish": "stop", "out_tokens": 192}
+{"task": "gsm8k", "ok": true, "pred": "80", "ref": "80", "finish": "stop", "out_tokens": 494}
+{"task": "gsm8k", "ok": true, "pred": "34", "ref": "34", "finish": "stop", "out_tokens": 316}
+{"task": "gsm8k", "ok": true, "pred": "15", "ref": "15", "finish": "stop", "out_tokens": 182}
+{"task": "gsm8k", "ok": true, "pred": "16", "ref": "16", "finish": "stop", "out_tokens": 233}
+{"task": "gsm8k", "ok": true, "pred": "32", "ref": "32", "finish": "stop", "out_tokens": 1283}
+{"task": "gsm8k", "ok": true, "pred": "92", "ref": "92", "finish": "stop", "out_tokens": 739}
+{"task": "gsm8k", "ok": false, "pred": "54", "ref": "50", "finish": "stop", "out_tokens": 320}
+{"task": "gsm8k", "ok": true, "pred": "15", "ref": "15", "finish": "stop", "out_tokens": 339}
+{"task": "gsm8k", "ok": true, "pred": "77", "ref": "77", "finish": "stop", "out_tokens": 311}
+{"task": "gsm8k", "ok": true, "pred": "5", "ref": "5", "finish": "stop", "out_tokens": 196}
+{"task": "gsm8k", "ok": true, "pred": "16", "ref": "16", "finish": "stop", "out_tokens": 437}
+{"task": "gsm8k", "ok": true, "pred": "18", "ref": "18", "finish": "stop", "out_tokens": 171}
+{"task": "gsm8k", "ok": true, "pred": "120", "ref": "120", "finish": "stop", "out_tokens": 251}
+{"task": "gsm8k", "ok": true, "pred": "150", "ref": "150", "finish": "stop", "out_tokens": 378}
+{"task": "gsm8k", "ok": true, "pred": "1210", "ref": "1210", "finish": "stop", "out_tokens": 320}
+{"task": "gsm8k", "ok": true, "pred": "51", "ref": "51", "finish": "stop", "out_tokens": 441}
+{"task": "gsm8k", "ok": true, "pred": "18000", "ref": "18000", "finish": "stop", "out_tokens": 383}
+{"task": "gsm8k", "ok": true, "pred": "95", "ref": "95", "finish": "stop", "out_tokens": 305}
+{"task": "gsm8k", "ok": true, "pred": "15", "ref": "15", "finish": "stop", "out_tokens": 404}
+{"task": "gsm8k", "ok": true, "pred": "100", "ref": "100", "finish": "stop", "out_tokens": 222}
+{"task": "gsm8k", "ok": true, "pred": "350", "ref": "350", "finish": "stop", "out_tokens": 527}
+{"task": "gsm8k", "ok": true, "pred": "122", "ref": "122", "finish": "stop", "out_tokens": 288}
+{"task": "gsm8k", "ok": true, "pred": "130", "ref": "130", "finish": "stop", "out_tokens": 278}
+{"task": "gsm8k", "ok": true, "pred": "20", "ref": "20", "finish": "stop", "out_tokens": 300}
+{"task": "gsm8k", "ok": true, "pred": "160", "ref": "160", "finish": "stop", "out_tokens": 408}
+{"task": "gsm8k", "ok": true, "pred": "23", "ref": "23", "finish": "stop", "out_tokens": 272}
+{"task": "gsm8k", "ok": true, "pred": "2", "ref": "2", "finish": "stop", "out_tokens": 527}
+{"task": "gsm8k", "ok": false, "pred": "100", "ref": "25", "finish": "stop", "out_tokens": 746}
+{"task": "gsm8k", "ok": true, "pred": "30", "ref": "30", "finish": "stop", "out_tokens": 218}
+{"task": "gsm8k", "ok": true, "pred": "5", "ref": "5", "finish": "stop", "out_tokens": 495}
+{"task": "gsm8k", "ok": true, "pred": "106", "ref": "106", "finish": "stop", "out_tokens": 699}
+{"task": "gsm8k", "ok": true, "pred": "50", "ref": "50", "finish": "stop", "out_tokens": 405}
+{"task": "gsm8k", "ok": true, "pred": "34", "ref": "34", "finish": "stop", "out_tokens": 377}
+{"task": "gsm8k", "ok": true, "pred": "360", "ref": "360", "finish": "stop", "out_tokens": 178}
+{"task": "gsm8k", "ok": true, "pred": "5", "ref": "5", "finish": "stop", "out_tokens": 349}
+{"task": "gsm8k", "ok": true, "pred": "91", "ref": "91", "finish": "stop", "out_tokens": 278}
+{"task": "gsm8k", "ok": true, "pred": "24", "ref": "24", "finish": "stop", "out_tokens": 388}
+{"task": "gsm8k", "ok": true, "pred": "10", "ref": "10", "finish": "stop", "out_tokens": 125}
+{"task": "gsm8k", "ok": true, "pred": "12", "ref": "12", "finish": "stop", "out_tokens": 204}
+{"task": "gsm8k", "ok": true, "pred": "120", "ref": "120", "finish": "stop", "out_tokens": 209}
+{"task": "gsm8k", "ok": true, "pred": "6277", "ref": "6277", "finish": "stop", "out_tokens": 206}
+{"task": "gsm8k", "ok": true, "pred": "320", "ref": "320", "finish": "stop", "out_tokens": 388}
+{"task": "gsm8k", "ok": true, "pred": "7500", "ref": "7500", "finish": "stop", "out_tokens": 425}
+{"task": "humaneval", "id": "HumanEval/0", "ok": true, "finish": "stop", "out_tokens": 173}
+{"task": "humaneval", "id": "HumanEval/1", "ok": true, "finish": "stop", "out_tokens": 224}
+{"task": "humaneval", "id": "HumanEval/2", "ok": true, "finish": "stop", "out_tokens": 97}
+{"task": "humaneval", "id": "HumanEval/3", "ok": true, "finish": "stop", "out_tokens": 157}
+{"task": "humaneval", "id": "HumanEval/4", "ok": true, "finish": "stop", "out_tokens": 164}
+{"task": "humaneval", "id": "HumanEval/5", "ok": true, "finish": "stop", "out_tokens": 152}
+{"task": "humaneval", "id": "HumanEval/6", "ok": true, "finish": "stop", "out_tokens": 216}
+{"task": "humaneval", "id": "HumanEval/7", "ok": true, "finish": "stop", "out_tokens": 112}
+{"task": "humaneval", "id": "HumanEval/8", "ok": true, "finish": "stop", "out_tokens": 163}
+{"task": "humaneval", "id": "HumanEval/9", "ok": true, "finish": "stop", "out_tokens": 160}
+{"task": "humaneval", "id": "HumanEval/10", "ok": true, "finish": "stop", "out_tokens": 270}
+{"task": "humaneval", "id": "HumanEval/11", "ok": true, "finish": "stop", "out_tokens": 140}
+{"task": "humaneval", "id": "HumanEval/12", "ok": true, "finish": "stop", "out_tokens": 128}
+{"task": "humaneval", "id": "HumanEval/13", "ok": true, "finish": "stop", "out_tokens": 107}
+{"task": "humaneval", "id": "HumanEval/14", "ok": true, "finish": "stop", "out_tokens": 84}
+{"task": "humaneval", "id": "HumanEval/15", "ok": true, "finish": "stop", "out_tokens": 92}
+{"task": "humaneval", "id": "HumanEval/16", "ok": true, "finish": "stop", "out_tokens": 81}
+{"task": "humaneval", "id": "HumanEval/17", "ok": true, "finish": "stop", "out_tokens": 218}
+{"task": "humaneval", "id": "HumanEval/18", "ok": true, "finish": "stop", "out_tokens": 161}
+{"task": "humaneval", "id": "HumanEval/19", "ok": true, "finish": "stop", "out_tokens": 266}
+{"task": "humaneval", "id": "HumanEval/20", "ok": true, "finish": "stop", "out_tokens": 273}
+{"task": "humaneval", "id": "HumanEval/21", "ok": true, "finish": "stop", "out_tokens": 172}
+{"task": "humaneval", "id": "HumanEval/22", "ok": true, "finish": "stop", "out_tokens": 119}
+{"task": "humaneval", "id": "HumanEval/23", "ok": true, "finish": "stop", "out_tokens": 53}
+{"task": "humaneval", "id": "HumanEval/24", "ok": true, "finish": "stop", "out_tokens": 91}
+{"task": "humaneval", "id": "HumanEval/25", "ok": true, "finish": "stop", "out_tokens": 202}
+{"task": "humaneval", "id": "HumanEval/26", "ok": true, "finish": "stop", "out_tokens": 119}
+{"task": "humaneval", "id": "HumanEval/27", "ok": true, "finish": "stop", "out_tokens": 58}
+{"task": "humaneval", "id": "HumanEval/28", "ok": true, "finish": "stop", "out_tokens": 70}
+{"task": "humaneval", "id": "HumanEval/29", "ok": true, "finish": "stop", "out_tokens": 107}
+{"task": "humaneval", "id": "HumanEval/30", "ok": true, "finish": "stop", "out_tokens": 139}
+{"task": "humaneval", "id": "HumanEval/31", "ok": true, "finish": "stop", "out_tokens": 193}
+{"task": "humaneval", "id": "HumanEval/32", "ok": false, "finish": "stop", "out_tokens": 509}
+{"task": "humaneval", "id": "HumanEval/33", "ok": true, "finish": "stop", "out_tokens": 263}
+{"task": "humaneval", "id": "HumanEval/34", "ok": true, "finish": "stop", "out_tokens": 87}
+{"task": "humaneval", "id": "HumanEval/35", "ok": true, "finish": "stop", "out_tokens": 97}
+{"task": "humaneval", "id": "HumanEval/36", "ok": true, "finish": "stop", "out_tokens": 139}
+{"task": "humaneval", "id": "HumanEval/37", "ok": true, "finish": "stop", "out_tokens": 220}
+{"task": "humaneval", "id": "HumanEval/38", "ok": false, "finish": "stop", "out_tokens": 144}
+{"task": "humaneval", "id": "HumanEval/39", "ok": true, "finish": "stop", "out_tokens": 251}
+{"task": "humaneval", "id": "HumanEval/40", "ok": true, "finish": "stop", "out_tokens": 235}
+{"task": "humaneval", "id": "HumanEval/41", "ok": true, "finish": "stop", "out_tokens": 163}
+{"task": "humaneval", "id": "HumanEval/42", "ok": true, "finish": "stop", "out_tokens": 132}
+{"task": "humaneval", "id": "HumanEval/43", "ok": true, "finish": "stop", "out_tokens": 201}
+{"task": "humaneval", "id": "HumanEval/44", "ok": true, "finish": "stop", "out_tokens": 157}
+{"task": "humaneval", "id": "HumanEval/45", "ok": true, "finish": "stop", "out_tokens": 60}
+{"task": "humaneval", "id": "HumanEval/46", "ok": true, "finish": "stop", "out_tokens": 293}
+{"task": "humaneval", "id": "HumanEval/47", "ok": true, "finish": "stop", "out_tokens": 153}
+{"task": "humaneval", "id": "HumanEval/48", "ok": true, "finish": "stop", "out_tokens": 87}
+{"task": "humaneval", "id": "HumanEval/49", "ok": true, "finish": "stop", "out_tokens": 134}
+{"task": "humaneval", "id": "HumanEval/50", "ok": true, "finish": "stop", "out_tokens": 133}
+{"task": "humaneval", "id": "HumanEval/51", "ok": true, "finish": "stop", "out_tokens": 165}
+{"task": "humaneval", "id": "HumanEval/52", "ok": true, "finish": "stop", "out_tokens": 100}
+{"task": "humaneval", "id": "HumanEval/53", "ok": true, "finish": "stop", "out_tokens": 63}
+{"task": "humaneval", "id": "HumanEval/54", "ok": true, "finish": "stop", "out_tokens": 161}
+{"task": "humaneval", "id": "HumanEval/55", "ok": true, "finish": "stop", "out_tokens": 130}
+{"task": "humaneval", "id": "HumanEval/56", "ok": true, "finish": "stop", "out_tokens": 157}
+{"task": "humaneval", "id": "HumanEval/57", "ok": true, "finish": "stop", "out_tokens": 168}
+{"task": "humaneval", "id": "HumanEval/58", "ok": true, "finish": "stop", "out_tokens": 142}
+{"task": "humaneval", "id": "HumanEval/59", "ok": true, "finish": "stop", "out_tokens": 196}
+{"task": "humaneval", "id": "HumanEval/60", "ok": true, "finish": "stop", "out_tokens": 125}
+{"task": "humaneval", "id": "HumanEval/61", "ok": true, "finish": "stop", "out_tokens": 156}
+{"task": "humaneval", "id": "HumanEval/62", "ok": true, "finish": "stop", "out_tokens": 136}
+{"task": "humaneval", "id": "HumanEval/63", "ok": true, "finish": "stop", "out_tokens": 247}
+{"task": "humaneval", "id": "HumanEval/64", "ok": true, "finish": "stop", "out_tokens": 209}
+{"task": "humaneval", "id": "HumanEval/65", "ok": true, "finish": "stop", "out_tokens": 152}
+{"task": "humaneval", "id": "HumanEval/66", "ok": true, "finish": "stop", "out_tokens": 140}
+{"task": "humaneval", "id": "HumanEval/67", "ok": true, "finish": "stop", "out_tokens": 273}
+{"task": "humaneval", "id": "HumanEval/68", "ok": true, "finish": "stop", "out_tokens": 429}
+{"task": "humaneval", "id": "HumanEval/69", "ok": true, "finish": "stop", "out_tokens": 227}
+{"task": "humaneval", "id": "HumanEval/70", "ok": true, "finish": "stop", "out_tokens": 237}
+{"task": "humaneval", "id": "HumanEval/71", "ok": true, "finish": "stop", "out_tokens": 237}
+{"task": "humaneval", "id": "HumanEval/72", "ok": true, "finish": "stop", "out_tokens": 255}
+{"task": "humaneval", "id": "HumanEval/73", "ok": true, "finish": "stop", "out_tokens": 210}
+{"task": "humaneval", "id": "HumanEval/74", "ok": true, "finish": "stop", "out_tokens": 264}
+{"task": "humaneval", "id": "HumanEval/75", "ok": true, "finish": "stop", "out_tokens": 288}
+{"task": "humaneval", "id": "HumanEval/76", "ok": true, "finish": "stop", "out_tokens": 240}
+{"task": "humaneval", "id": "HumanEval/77", "ok": true, "finish": "stop", "out_tokens": 232}
+{"task": "humaneval", "id": "HumanEval/78", "ok": true, "finish": "stop", "out_tokens": 383}
+{"task": "humaneval", "id": "HumanEval/79", "ok": true, "finish": "stop", "out_tokens": 224}
+{"task": "humaneval", "id": "HumanEval/80", "ok": true, "finish": "stop", "out_tokens": 194}
+{"task": "humaneval", "id": "HumanEval/81", "ok": true, "finish": "stop", "out_tokens": 530}
+{"task": "humaneval", "id": "HumanEval/82", "ok": true, "finish": "stop", "out_tokens": 169}
+{"task": "humaneval", "id": "HumanEval/83", "ok": true, "finish": "stop", "out_tokens": 292}
+{"task": "humaneval", "id": "HumanEval/84", "ok": true, "finish": "stop", "out_tokens": 180}
+{"task": "humaneval", "id": "HumanEval/85", "ok": true, "finish": "stop", "out_tokens": 102}
+{"task": "humaneval", "id": "HumanEval/86", "ok": true, "finish": "stop", "out_tokens": 184}
+{"task": "humaneval", "id": "HumanEval/87", "ok": true, "finish": "stop", "out_tokens": 351}
+{"task": "humaneval", "id": "HumanEval/88", "ok": true, "finish": "stop", "out_tokens": 286}
+{"task": "humaneval", "id": "HumanEval/89", "ok": true, "finish": "stop", "out_tokens": 225}
+{"task": "humaneval", "id": "HumanEval/90", "ok": true, "finish": "stop", "out_tokens": 158}
+{"task": "humaneval", "id": "HumanEval/91", "ok": true, "finish": "stop", "out_tokens": 184}
+{"task": "humaneval", "id": "HumanEval/92", "ok": true, "finish": "stop", "out_tokens": 223}
+{"task": "humaneval", "id": "HumanEval/93", "ok": true, "finish": "stop", "out_tokens": 326}
+{"task": "humaneval", "id": "HumanEval/94", "ok": true, "finish": "stop", "out_tokens": 468}
+{"task": "humaneval", "id": "HumanEval/95", "ok": true, "finish": "stop", "out_tokens": 282}
+{"task": "humaneval", "id": "HumanEval/96", "ok": true, "finish": "stop", "out_tokens": 284}
+{"task": "humaneval", "id": "HumanEval/97", "ok": false, "finish": "stop", "out_tokens": 139}
+{"task": "humaneval", "id": "HumanEval/98", "ok": true, "finish": "stop", "out_tokens": 119}
+{"task": "humaneval", "id": "HumanEval/99", "ok": true, "finish": "stop", "out_tokens": 219}
+{"task": "humaneval", "id": "HumanEval/100", "ok": true, "finish": "stop", "out_tokens": 177}
+{"task": "humaneval", "id": "HumanEval/101", "ok": true, "finish": "stop", "out_tokens": 175}
+{"task": "humaneval", "id": "HumanEval/102", "ok": true, "finish": "stop", "out_tokens": 185}
+{"task": "humaneval", "id": "HumanEval/103", "ok": true, "finish": "stop", "out_tokens": 213}
+{"task": "humaneval", "id": "HumanEval/104", "ok": true, "finish": "stop", "out_tokens": 175}
+{"task": "humaneval", "id": "HumanEval/105", "ok": true, "finish": "stop", "out_tokens": 475}
+{"task": "humaneval", "id": "HumanEval/106", "ok": true, "finish": "stop", "out_tokens": 234}
+{"task": "humaneval", "id": "HumanEval/107", "ok": true, "finish": "stop", "out_tokens": 327}
+{"task": "humaneval", "id": "HumanEval/108", "ok": true, "finish": "stop", "out_tokens": 272}
+{"task": "humaneval", "id": "HumanEval/109", "ok": true, "finish": "stop", "out_tokens": 481}
+{"task": "humaneval", "id": "HumanEval/110", "ok": true, "finish": "stop", "out_tokens": 313}
+{"task": "humaneval", "id": "HumanEval/111", "ok": true, "finish": "stop", "out_tokens": 298}
+{"task": "humaneval", "id": "HumanEval/112", "ok": true, "finish": "stop", "out_tokens": 209}
+{"task": "humaneval", "id": "HumanEval/113", "ok": false, "finish": "stop", "out_tokens": 396}
+{"task": "humaneval", "id": "HumanEval/114", "ok": true, "finish": "stop", "out_tokens": 166}
+{"task": "humaneval", "id": "HumanEval/115", "ok": true, "finish": "stop", "out_tokens": 383}
+{"task": "humaneval", "id": "HumanEval/116", "ok": true, "finish": "stop", "out_tokens": 198}
+{"task": "humaneval", "id": "HumanEval/117", "ok": true, "finish": "stop", "out_tokens": 278}
+{"task": "humaneval", "id": "HumanEval/118", "ok": true, "finish": "stop", "out_tokens": 298}
+{"task": "humaneval", "id": "HumanEval/119", "ok": true, "finish": "stop", "out_tokens": 304}
+{"task": "humaneval", "id": "HumanEval/120", "ok": true, "finish": "stop", "out_tokens": 268}
+{"task": "humaneval", "id": "HumanEval/121", "ok": true, "finish": "stop", "out_tokens": 136}
+{"task": "humaneval", "id": "HumanEval/122", "ok": true, "finish": "stop", "out_tokens": 186}
+{"task": "humaneval", "id": "HumanEval/123", "ok": true, "finish": "stop", "out_tokens": 328}
+{"task": "humaneval", "id": "HumanEval/124", "ok": true, "finish": "stop", "out_tokens": 545}
+{"task": "humaneval", "id": "HumanEval/125", "ok": true, "finish": "stop", "out_tokens": 229}
+{"task": "humaneval", "id": "HumanEval/126", "ok": true, "finish": "stop", "out_tokens": 345}
+{"task": "humaneval", "id": "HumanEval/127", "ok": true, "finish": "stop", "out_tokens": 459}
+{"task": "humaneval", "id": "HumanEval/128", "ok": true, "finish": "stop", "out_tokens": 216}
+{"task": "humaneval", "id": "HumanEval/129", "ok": true, "finish": "stop", "out_tokens": 943}
+{"task": "humaneval", "id": "HumanEval/130", "ok": true, "finish": "stop", "out_tokens": 717}
+{"task": "humaneval", "id": "HumanEval/131", "ok": true, "finish": "stop", "out_tokens": 152}
+{"task": "humaneval", "id": "HumanEval/132", "ok": true, "finish": "stop", "out_tokens": 708}
+{"task": "humaneval", "id": "HumanEval/133", "ok": true, "finish": "stop", "out_tokens": 204}
+{"task": "humaneval", "id": "HumanEval/134", "ok": true, "finish": "stop", "out_tokens": 282}
+{"task": "humaneval", "id": "HumanEval/135", "ok": true, "finish": "stop", "out_tokens": 149}
+{"task": "humaneval", "id": "HumanEval/136", "ok": true, "finish": "stop", "out_tokens": 213}
+{"task": "humaneval", "id": "HumanEval/137", "ok": true, "finish": "stop", "out_tokens": 232}
+{"task": "humaneval", "id": "HumanEval/138", "ok": true, "finish": "stop", "out_tokens": 166}
+{"task": "humaneval", "id": "HumanEval/139", "ok": true, "finish": "stop", "out_tokens": 140}
+{"task": "humaneval", "id": "HumanEval/140", "ok": false, "finish": "stop", "out_tokens": 175}
+{"task": "humaneval", "id": "HumanEval/141", "ok": true, "finish": "stop", "out_tokens": 432}
+{"task": "humaneval", "id": "HumanEval/142", "ok": true, "finish": "stop", "out_tokens": 240}
+{"task": "humaneval", "id": "HumanEval/143", "ok": true, "finish": "stop", "out_tokens": 281}
+{"task": "humaneval", "id": "HumanEval/144", "ok": true, "finish": "stop", "out_tokens": 253}
+{"task": "humaneval", "id": "HumanEval/145", "ok": false, "finish": "stop", "out_tokens": 235}
+{"task": "humaneval", "id": "HumanEval/146", "ok": true, "finish": "stop", "out_tokens": 266}
+{"task": "humaneval", "id": "HumanEval/147", "ok": true, "finish": "stop", "out_tokens": 1157}
+{"task": "humaneval", "id": "HumanEval/148", "ok": true, "finish": "stop", "out_tokens": 331}
+{"task": "humaneval", "id": "HumanEval/149", "ok": true, "finish": "stop", "out_tokens": 239}
+{"task": "humaneval", "id": "HumanEval/150", "ok": true, "finish": "stop", "out_tokens": 199}
+{"task": "humaneval", "id": "HumanEval/151", "ok": true, "finish": "stop", "out_tokens": 190}
+{"task": "humaneval", "id": "HumanEval/152", "ok": true, "finish": "stop", "out_tokens": 257}
+{"task": "humaneval", "id": "HumanEval/153", "ok": true, "finish": "stop", "out_tokens": 391}
+{"task": "humaneval", "id": "HumanEval/154", "ok": true, "finish": "stop", "out_tokens": 182}
+{"task": "humaneval", "id": "HumanEval/155", "ok": true, "finish": "stop", "out_tokens": 138}
+{"task": "humaneval", "id": "HumanEval/156", "ok": true, "finish": "stop", "out_tokens": 285}
+{"task": "humaneval", "id": "HumanEval/157", "ok": true, "finish": "stop", "out_tokens": 152}
+{"task": "humaneval", "id": "HumanEval/158", "ok": true, "finish": "stop", "out_tokens": 194}
+{"task": "humaneval", "id": "HumanEval/159", "ok": true, "finish": "stop", "out_tokens": 341}
+{"task": "humaneval", "id": "HumanEval/160", "ok": true, "finish": "stop", "out_tokens": 273}
+{"task": "humaneval", "id": "HumanEval/161", "ok": true, "finish": "stop", "out_tokens": 183}
+{"task": "humaneval", "id": "HumanEval/162", "ok": true, "finish": "stop", "out_tokens": 119}
+{"task": "humaneval", "id": "HumanEval/163", "ok": false, "finish": "stop", "out_tokens": 170}

recipe/results/quality-probe.txt

@@ -0,0 +1,3 @@
+qwen3.8-flash-next GSM8K first 200: 193/200 = 96.5%  (truncated 0, errors 0, 757s)
+qwen3.8-flash-next HumanEval 164: 157/164 = 95.7%  (truncated 0, errors 0, 317s)
+wrote /workspace/hf-release/results-raw/quality-probe.jsonl | thinking False

recipe/results/toolcall.txt

@@ -0,0 +1,4 @@
+finish_reason: tool_calls | tool_calls: [{"id": "chatcmpl-tool-ad003c660f5c49ce", "type": "function", "function": {"name": "web_search", "arguments": "{\"query\": \"2026 Tour de France winner\"}"}}]
+reasoning field present: True | content: 
+no-tool case -> finish: stop | tool_calls: False | reasoning: True | content:   4
+TOOLCALL TEST OK

recipe/serve/seccomp.json

@@ -0,0 +1,1229 @@
+{
+ "defaultAction": "SCMP_ACT_ERRNO",
+ "defaultErrnoRet": 1,
+ "archMap": [
+  {
+   "architecture": "SCMP_ARCH_X86_64",
+   "subArchitectures": [
+    "SCMP_ARCH_X86",
+    "SCMP_ARCH_X32"
+   ]
+  },
+  {
+   "architecture": "SCMP_ARCH_AARCH64",
+   "subArchitectures": [
+    "SCMP_ARCH_ARM"
+   ]
+  },
+  {
+   "architecture": "SCMP_ARCH_MIPS64",
+   "subArchitectures": [
+    "SCMP_ARCH_MIPS",
+    "SCMP_ARCH_MIPS64N32"
+   ]
+  },
+  {
+   "architecture": "SCMP_ARCH_MIPS64N32",
+   "subArchitectures": [
+    "SCMP_ARCH_MIPS",
+    "SCMP_ARCH_MIPS64"
+   ]
+  },
+  {
+   "architecture": "SCMP_ARCH_MIPSEL64",
+   "subArchitectures": [
+    "SCMP_ARCH_MIPSEL",
+    "SCMP_ARCH_MIPSEL64N32"
+   ]
+  },
+  {
+   "architecture": "SCMP_ARCH_MIPSEL64N32",
+   "subArchitectures": [
+    "SCMP_ARCH_MIPSEL",
+    "SCMP_ARCH_MIPSEL64"
+   ]
+  },
+  {
+   "architecture": "SCMP_ARCH_S390X",
+   "subArchitectures": [
+    "SCMP_ARCH_S390"
+   ]
+  },
+  {
+   "architecture": "SCMP_ARCH_RISCV64",
+   "subArchitectures": null
+  },
+  {
+   "architecture": "SCMP_ARCH_LOONGARCH64",
+   "subArchitectures": null
+  }
+ ],
+ "syscalls": [
+  {
+   "names": [
+    "_llseek",
+    "_newselect",
+    "accept",
+    "accept4",
+    "access",
+    "adjtimex",
+    "alarm",
+    "bind",
+    "brk",
+    "cachestat",
+    "capget",
+    "capset",
+    "chdir",
+    "chmod",
+    "chown",
+    "chown32",
+    "clock_adjtime",
+    "clock_adjtime64",
+    "clock_getres",
+    "clock_getres_time64",
+    "clock_gettime",
+    "clock_gettime64",
+    "clock_nanosleep",
+    "clock_nanosleep_time64",
+    "close",
+    "close_range",
+    "connect",
+    "copy_file_range",
+    "creat",
+    "dup",
+    "dup2",
+    "dup3",
+    "epoll_create",
+    "epoll_create1",
+    "epoll_ctl",
+    "epoll_ctl_old",
+    "epoll_pwait",
+    "epoll_pwait2",
+    "epoll_wait",
+    "epoll_wait_old",
+    "eventfd",
+    "eventfd2",
+    "execve",
+    "execveat",
+    "exit",
+    "exit_group",
+    "faccessat",
+    "faccessat2",
+    "fadvise64",
+    "fadvise64_64",
+    "fallocate",
+    "fanotify_mark",
+    "fchdir",
+    "fchmod",
+    "fchmodat",
+    "fchmodat2",
+    "fchown",
+    "fchown32",
+    "fchownat",
+    "fcntl",
+    "fcntl64",
+    "fdatasync",
+    "fgetxattr",
+    "flistxattr",
+    "flock",
+    "fork",
+    "fremovexattr",
+    "fsetxattr",
+    "fstat",
+    "fstat64",
+    "fstatat64",
+    "fstatfs",
+    "fstatfs64",
+    "fsync",
+    "ftruncate",
+    "ftruncate64",
+    "futex",
+    "futex_requeue",
+    "futex_time64",
+    "futex_wait",
+    "futex_waitv",
+    "futex_wake",
+    "futimesat",
+    "get_robust_list",
+    "get_thread_area",
+    "getcpu",
+    "getcwd",
+    "getdents",
+    "getdents64",
+    "getegid",
+    "getegid32",
+    "geteuid",
+    "geteuid32",
+    "getgid",
+    "getgid32",
+    "getgroups",
+    "getgroups32",
+    "getitimer",
+    "getpeername",
+    "getpgid",
+    "getpgrp",
+    "getpid",
+    "getppid",
+    "getpriority",
+    "getrandom",
+    "getresgid",
+    "getresgid32",
+    "getresuid",
+    "getresuid32",
+    "getrlimit",
+    "getrusage",
+    "getsid",
+    "getsockname",
+    "getsockopt",
+    "gettid",
+    "gettimeofday",
+    "getuid",
+    "getuid32",
+    "getxattr",
+    "getxattrat",
+    "inotify_add_watch",
+    "inotify_init",
+    "inotify_init1",
+    "inotify_rm_watch",
+    "io_cancel",
+    "io_destroy",
+    "io_getevents",
+    "io_pgetevents",
+    "io_pgetevents_time64",
+    "io_setup",
+    "io_submit",
+    "io_uring_enter",
+    "io_uring_register",
+    "io_uring_setup",
+    "ioctl",
+    "ioprio_get",
+    "ioprio_set",
+    "ipc",
+    "kill",
+    "landlock_add_rule",
+    "landlock_create_ruleset",
+    "landlock_restrict_self",
+    "lchown",
+    "lchown32",
+    "lgetxattr",
+    "link",
+    "linkat",
+    "listen",
+    "listmount",
+    "listxattr",
+    "listxattrat",
+    "llistxattr",
+    "lremovexattr",
+    "lseek",
+    "lsetxattr",
+    "lstat",
+    "lstat64",
+    "madvise",
+    "map_shadow_stack",
+    "membarrier",
+    "memfd_create",
+    "memfd_secret",
+    "mincore",
+    "mkdir",
+    "mkdirat",
+    "mknod",
+    "mknodat",
+    "mlock",
+    "mlock2",
+    "mlockall",
+    "mmap",
+    "mmap2",
+    "mprotect",
+    "mq_getsetattr",
+    "mq_notify",
+    "mq_open",
+    "mq_timedreceive",
+    "mq_timedreceive_time64",
+    "mq_timedsend",
+    "mq_timedsend_time64",
+    "mq_unlink",
+    "mremap",
+    "mseal",
+    "msgctl",
+    "msgget",
+    "msgrcv",
+    "msgsnd",
+    "msync",
+    "munlock",
+    "munlockall",
+    "munmap",
+    "name_to_handle_at",
+    "nanosleep",
+    "newfstatat",
+    "open",
+    "openat",
+    "openat2",
+    "pause",
+    "pidfd_open",
+    "pidfd_send_signal",
+    "pipe",
+    "pipe2",
+    "pkey_alloc",
+    "pkey_free",
+    "pkey_mprotect",
+    "poll",
+    "ppoll",
+    "ppoll_time64",
+    "prctl",
+    "pread64",
+    "preadv",
+    "preadv2",
+    "prlimit64",
+    "process_mrelease",
+    "pselect6",
+    "pselect6_time64",
+    "pwrite64",
+    "pwritev",
+    "pwritev2",
+    "read",
+    "readahead",
+    "readlink",
+    "readlinkat",
+    "readv",
+    "recv",
+    "recvfrom",
+    "recvmmsg",
+    "recvmmsg_time64",
+    "recvmsg",
+    "remap_file_pages",
+    "removexattr",
+    "removexattrat",
+    "rename",
+    "renameat",
+    "renameat2",
+    "restart_syscall",
+    "riscv_hwprobe",
+    "rmdir",
+    "rseq",
+    "rt_sigaction",
+    "rt_sigpending",
+    "rt_sigprocmask",
+    "rt_sigqueueinfo",
+    "rt_sigreturn",
+    "rt_sigsuspend",
+    "rt_sigtimedwait",
+    "rt_sigtimedwait_time64",
+    "rt_tgsigqueueinfo",
+    "sched_get_priority_max",
+    "sched_get_priority_min",
+    "sched_getaffinity",
+    "sched_getattr",
+    "sched_getparam",
+    "sched_getscheduler",
+    "sched_rr_get_interval",
+    "sched_rr_get_interval_time64",
+    "sched_setaffinity",
+    "sched_setattr",
+    "sched_setparam",
+    "sched_setscheduler",
+    "sched_yield",
+    "seccomp",
+    "select",
+    "semctl",
+    "semget",
+    "semop",
+    "semtimedop",
+    "semtimedop_time64",
+    "send",
+    "sendfile",
+    "sendfile64",
+    "sendmmsg",
+    "sendmsg",
+    "sendto",
+    "set_robust_list",
+    "set_thread_area",
+    "set_tid_address",
+    "setfsgid",
+    "setfsgid32",
+    "setfsuid",
+    "setfsuid32",
+    "setgid",
+    "setgid32",
+    "setgroups",
+    "setgroups32",
+    "setitimer",
+    "setpgid",
+    "setpriority",
+    "setregid",
+    "setregid32",
+    "setresgid",
+    "setresgid32",
+    "setresuid",
+    "setresuid32",
+    "setreuid",
+    "setreuid32",
+    "setrlimit",
+    "setsid",
+    "setsockopt",
+    "setuid",
+    "setuid32",
+    "setxattr",
+    "setxattrat",
+    "shmat",
+    "shmctl",
+    "shmdt",
+    "shmget",
+    "shutdown",
+    "sigaltstack",
+    "signalfd",
+    "signalfd4",
+    "sigprocmask",
+    "sigreturn",
+    "socketcall",
+    "socketpair",
+    "splice",
+    "stat",
+    "stat64",
+    "statfs",
+    "statfs64",
+    "statmount",
+    "statx",
+    "symlink",
+    "symlinkat",
+    "sync",
+    "sync_file_range",
+    "syncfs",
+    "sysinfo",
+    "tee",
+    "tgkill",
+    "time",
+    "timer_create",
+    "timer_delete",
+    "timer_getoverrun",
+    "timer_gettime",
+    "timer_gettime64",
+    "timer_settime",
+    "timer_settime64",
+    "timerfd_create",
+    "timerfd_gettime",
+    "timerfd_gettime64",
+    "timerfd_settime",
+    "timerfd_settime64",
+    "times",
+    "tkill",
+    "truncate",
+    "truncate64",
+    "ugetrlimit",
+    "umask",
+    "uname",
+    "unlink",
+    "unlinkat",
+    "uretprobe",
+    "utime",
+    "utimensat",
+    "utimensat_time64",
+    "utimes",
+    "vfork",
+    "vmsplice",
+    "wait4",
+    "waitid",
+    "waitpid",
+    "write",
+    "writev"
+   ],
+   "action": "SCMP_ACT_ALLOW"
+  },
+  {
+   "names": [
+    "process_vm_readv",
+    "process_vm_writev",
+    "ptrace"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "minKernel": "4.8"
+   }
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 3,
+     "op": "SCMP_CMP_LT"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 7,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 8,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 9,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 10,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 13,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 14,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 15,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 16,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 17,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 18,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 21,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 22,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 24,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 26,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 27,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 28,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 29,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 30,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 31,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 32,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 33,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 35,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 36,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 39,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 41,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 42,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 43,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 44,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "socket"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 45,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "personality"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 0,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "personality"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 8,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "personality"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 131072,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "personality"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 131080,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "personality"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 4294967295,
+     "op": "SCMP_CMP_EQ"
+    }
+   ]
+  },
+  {
+   "names": [
+    "sync_file_range2",
+    "swapcontext"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "arches": [
+     "ppc64le"
+    ]
+   }
+  },
+  {
+   "names": [
+    "arm_fadvise64_64",
+    "arm_sync_file_range",
+    "sync_file_range2",
+    "breakpoint",
+    "cacheflush",
+    "set_tls"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "arches": [
+     "arm",
+     "arm64"
+    ]
+   }
+  },
+  {
+   "names": [
+    "arch_prctl"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "arches": [
+     "amd64",
+     "x32"
+    ]
+   }
+  },
+  {
+   "names": [
+    "modify_ldt"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "arches": [
+     "amd64",
+     "x32",
+     "x86"
+    ]
+   }
+  },
+  {
+   "names": [
+    "s390_pci_mmio_read",
+    "s390_pci_mmio_write",
+    "s390_runtime_instr"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "arches": [
+     "s390",
+     "s390x"
+    ]
+   }
+  },
+  {
+   "names": [
+    "riscv_flush_icache"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "arches": [
+     "riscv64"
+    ]
+   }
+  },
+  {
+   "names": [
+    "open_by_handle_at"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "caps": [
+     "CAP_DAC_READ_SEARCH"
+    ]
+   }
+  },
+  {
+   "names": [
+    "bpf",
+    "clone",
+    "clone3",
+    "fanotify_init",
+    "fsconfig",
+    "fsmount",
+    "fsopen",
+    "fspick",
+    "lookup_dcookie",
+    "lsm_get_self_attr",
+    "lsm_list_modules",
+    "lsm_set_self_attr",
+    "mount",
+    "mount_setattr",
+    "move_mount",
+    "open_tree",
+    "perf_event_open",
+    "quotactl",
+    "quotactl_fd",
+    "setdomainname",
+    "sethostname",
+    "setns",
+    "syslog",
+    "umount",
+    "umount2",
+    "unshare"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "caps": [
+     "CAP_SYS_ADMIN"
+    ]
+   }
+  },
+  {
+   "names": [
+    "clone"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 0,
+     "value": 2114060288,
+     "op": "SCMP_CMP_MASKED_EQ"
+    }
+   ],
+   "excludes": {
+    "caps": [
+     "CAP_SYS_ADMIN"
+    ],
+    "arches": [
+     "s390",
+     "s390x"
+    ]
+   }
+  },
+  {
+   "names": [
+    "clone"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "args": [
+    {
+     "index": 1,
+     "value": 2114060288,
+     "op": "SCMP_CMP_MASKED_EQ"
+    }
+   ],
+   "comment": "s390 parameter ordering for clone is different",
+   "includes": {
+    "arches": [
+     "s390",
+     "s390x"
+    ]
+   },
+   "excludes": {
+    "caps": [
+     "CAP_SYS_ADMIN"
+    ]
+   }
+  },
+  {
+   "names": [
+    "clone3"
+   ],
+   "action": "SCMP_ACT_ERRNO",
+   "errnoRet": 38,
+   "excludes": {
+    "caps": [
+     "CAP_SYS_ADMIN"
+    ]
+   }
+  },
+  {
+   "names": [
+    "reboot"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "caps": [
+     "CAP_SYS_BOOT"
+    ]
+   }
+  },
+  {
+   "names": [
+    "chroot"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "caps": [
+     "CAP_SYS_CHROOT"
+    ]
+   }
+  },
+  {
+   "names": [
+    "delete_module",
+    "init_module",
+    "finit_module"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "caps": [
+     "CAP_SYS_MODULE"
+    ]
+   }
+  },
+  {
+   "names": [
+    "acct"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "caps": [
+     "CAP_SYS_PACCT"
+    ]
+   }
+  },
+  {
+   "names": [
+    "kcmp",
+    "pidfd_getfd",
+    "process_madvise",
+    "process_vm_readv",
+    "process_vm_writev",
+    "ptrace"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "caps": [
+     "CAP_SYS_PTRACE"
+    ]
+   }
+  },
+  {
+   "names": [
+    "iopl",
+    "ioperm"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "caps": [
+     "CAP_SYS_RAWIO"
+    ]
+   }
+  },
+  {
+   "names": [
+    "settimeofday",
+    "stime",
+    "clock_settime",
+    "clock_settime64"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "caps": [
+     "CAP_SYS_TIME"
+    ]
+   }
+  },
+  {
+   "names": [
+    "vhangup"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "caps": [
+     "CAP_SYS_TTY_CONFIG"
+    ]
+   }
+  },
+  {
+   "names": [
+    "get_mempolicy",
+    "mbind",
+    "set_mempolicy",
+    "set_mempolicy_home_node"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "caps": [
+     "CAP_SYS_NICE"
+    ]
+   }
+  },
+  {
+   "names": [
+    "syslog"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "caps": [
+     "CAP_SYSLOG"
+    ]
+   }
+  },
+  {
+   "names": [
+    "bpf"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "caps": [
+     "CAP_BPF"
+    ]
+   }
+  },
+  {
+   "names": [
+    "perf_event_open"
+   ],
+   "action": "SCMP_ACT_ALLOW",
+   "includes": {
+    "caps": [
+     "CAP_PERFMON"
+    ]
+   }
+  }
+ ]
+}

recipe/serve/serve.sh

@@ -0,0 +1,290 @@
+#!/usr/bin/env bash
+# Qwen3.8-Flash-Next (NVIDIA FP4-Local) on one DGX Spark (GB10) with vLLM + b12x.
+#
+#   bash serve.sh          # OpenAI-compatible API on :8000 (docker logs -f qwen38-fp4)
+#   bash serve.sh stop
+#
+# Serves the weights in this repo (../.. from this script) with the patches in ../patches
+# mounted over a pinned eugr/spark-vllm-b12x image. Every flag below is on in the published
+# configuration; set one to 0 (or empty) to fall back to the image's stock code path.
+# Optional: API_KEY=... requires a bearer token.
+set -euo pipefail
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+NAME=qwen38-fp4
+IMAGE="${IMAGE:-eugr/spark-vllm-b12x@sha256:03a956544e9f9b4b7ba91690469fbe7848bdbef29a7d2e0f5788368467bb3595}"
+MODEL_DIR="${MODEL_DIR:-$(cd "$SCRIPT_DIR/../.." && pwd)}"
+CACHE_DIR="${CACHE_DIR:-$HOME/.cache/qwen38-flash-next-spark}"  # compile/autotune caches (first boot autotunes)
+PORT="${PORT:-8000}"
+GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.70}"
+MAX_MODEL_LEN="${MAX_MODEL_LEN:-131072}"
+MAX_NUM_SEQS="${MAX_NUM_SEQS:-4}"
+# Prefill chunk. The request-boundary prefix cache used to place its reusable prefill
+# checkpoint one to two chunks before the prompt end, so larger chunks made every agent
+# follow-up recompute more (real-session replay, 2026-10-19: follow-up TTFT 2.26 s at
+# 2048, 3.53 s at 4096). TAIL_CKPT moves that checkpoint to just before the prompt end:
+# follow-ups recompute only the new message (0.77 s at either size), so 4096 now costs
+# nothing there and prefills cold prompts 13-15% faster (~2,140 tok/s). KV pool
+# ~233K -> ~182K tokens. 8192 OOMs at 0.70 memory.
+MAX_NUM_BATCHED_TOKENS="${MAX_NUM_BATCHED_TOKENS:-4096}"
+MTP_TOKENS="${MTP_TOKENS:-4}"          # 0 disables MTP; swept 2-6, 4 is best on real agentic turns
+EXTRA_ARGS="${EXTRA_ARGS:-}"           # appended verbatim to vllm serve
+# Docker default seccomp + io_uring_{setup,enter,register} only: b12x reads the on-disk
+# PLE tables with io_uring (B12X_DISK_BACKEND supports only io_uring or gds), which the
+# default profile blocks.
+SECCOMP="${SECCOMP:-$SCRIPT_DIR/seccomp.json}"
+# MXFP8_LM_HEAD=1  image-native b12x MXFP8 target lm_head with bf16 activations
+#                  (bf16 head was an sm80 WMMA fallback, ~5 ms/step); ~3% faster steps,
+#                  +6% synthetic decode, larger KV cache. Not bitwise lossless.
+# Patches from ../patches (mounted over the image). Defaults (2026-10-16,
+# real agentic replay, server inter-token latency): 71.2 -> 63.3 ms per MTP verify step,
+# 46.9 -> 52.2 tok/s weighted decode; logprob KL vs the unpatched launcher 7.5e-3, inside
+# the 8.6e-3 run-to-run noise; tool calls 5/5. Set any flag to 0 (or empty) to disable it.
+# 2026-10-19 defaults add MTP_BF16_KV, PLE_KEEPALIVE_US=2000, PLE_OVERLAP, GEMMA_NORM, HC_FUSE:
+# replay 62.3 -> 60.3 ms/step, draft acceptance 55 -> 63%, 52.0 -> 59.1 tok/s weighted decode;
+# 40-prompt category harness median 53.8 -> 68.1 tok/s; KL vs before 6.1e-3; tool calls 5/5.
+# Cost: the bf16 draft-layer KV pads the mamba pages, KV pool ~300K -> ~220K tokens.
+# + GDN_DEFER (verify records + one post-acceptance commit): replay 60.0-60.3 -> 58.4 ms/step,
+# 40-prompt 59.0-59.2 -> 57.4; 34.7K-token greedy output identical to stock over 256 tokens.
+# + MOE_A16 (2026-10-19 evening): routed experts with bf16 activations instead of the
+# checkpoint-default FP4 activations. KLD vs BF16 (AtomicChat llama.cpp neutral set, 61K
+# positions) 0.0677 -> 0.0482, same top-1 90.3 -> 91.8%, PPL +1.9% -> +0.8%; and faster:
+# replay 58.72 -> 57.59 ms/step (60.6 -> 63.2 tok/s), 40-prompt 57.59 -> 56.40 ms/step;
+# prefill -1..3%. 6 of 7 cold boots OK (one unexplained cudaErrorNotPermitted at init).
+# Measured noise source: b12x MoE sums top-k expert outputs with bf16 atomics, so outputs
+# vary run to run (MOE_DETERMINISTIC=1 fixes it for FP4-activation MoE at +4% decode and
+# +28% TTFT; it has no effect with MOE_A16). FP8 KV costs no measurable KLD.
+#   SMALL_GEMV=1   split-K Triton GEMV for the MoE router (sm80 fallback on GB10)
+#   FP8_LM_HEAD=1  older rowwise-fp8 lm_head patch; superseded by MXFP8_LM_HEAD (its
+#                  fp32 quantize temp shrinks the KV cache). Use one or the other.
+#   HC_FP8=1       rowwise-FP8 split-K GEMVs for the HyperConnection mix projections
+#                  at decode (sm80 bf16 fallback otherwise); prefill keeps bf16.
+#   MXFP8_GEMV="2560x6144,1280x2560,..."  split-K Triton GEMV over b12x's packed MXFP8
+#                  weights for decode rows of the listed NxK projections (b12x leaves
+#                  most SMs idle on narrow N at decode); prefill stays on b12x.
+#   DRAFT_VOCAB=/cache/draft_vocab_64k.json  MTP draft logits over a token subset:
+#                  MXFP8 rows gathered from the shared target lm_head (the draft's own
+#                  NVFP4 head is dropped), scattered to the full vocab with -inf.
+#   DRAFT_NVFP4=1  with DRAFT_VOCAB: keep the subset rows as row-major NVFP4 (requantized
+#                  from the MXFP8 head; the stock draft head is NVFP4 too), ~half the bytes.
+#   MTP_FP8=1      rowwise-FP8 decode GEMVs for the MTP draft layer's bf16 projections.
+#   MTP_BF16_KV=1  keep the MTP draft layer's KV cache in bf16 (the target stays FP8).
+#   GEMMA_NORM=1   one Triton launch per GemmaRMSNorm (QSA q/k norms etc.; stock ~10 eager ops).
+#   GEMV_TUNE=1    decode GEMV configs re-tuned with dirty L2 (router, HC down, qkvz, shared up).
+#   HC_FUSE=1      with HC_FP8: scaled SiLU and gate mix fused into the HC FP8 GEMVs (4 -> 2 launches).
+#   GDN_DEFER=1    GDN verify writes per-token records instead of 5 fp32 states per layer; one
+#                  post-acceptance commit writes the column the next step reads.
+#   PLE_OVERLAP=1  run the disk-PLE read/lookup on a side stream; the forward waits only at the
+#                  PLE layer (stream wait-value node), so decoder layer 0 overlaps the read.
+#   TARGET_FP8=1   same for the target's remaining plain bf16 projections (PLE key/value).
+#   SHARED_FUSE=1  shared-expert tail at decode in one kernel: SiLU*mul + MXFP8 down GEMV +
+#                  sigmoid(shared_expert_gate) (stock runs 5-6 small kernels per MoE layer).
+#   PLE_PATCH=1    disk-PLE reader knobs (set via EXTRA_ENV): VLLM_PLE_QUEUE_DEPTH
+#                  (io_uring depth, stock 64), B12X_PLE_BUFFERED=1 (page cache instead
+#                  of O_DIRECT), VLLM_PLE_STATS=N (log reader stats every N steps).
+#   PLE_ASYNC=1    enqueue the per-step disk-PLE row reads on the CUDA stream as a host
+#                  function, so the CPU launches the model graph while the rows are read
+#                  (stock blocks the CPU on the ID copy + reads, then launches).
+#   PLE_KEEPALIVE_US=N  with PLE_ASYNC: while PLE reads are recent (2 s), one 4 KiB read
+#                  every N us keeps the NVMe out of its idle power states, which add
+#                  ~1-2 ms wake-up to each step's read burst (fio p99 2.2 ms after 60 ms idle).
+MXFP8_LM_HEAD="${MXFP8_LM_HEAD:-1}"
+FP8_LM_HEAD="${FP8_LM_HEAD:-0}"
+SMALL_GEMV="${SMALL_GEMV:-1}"
+HC_FP8="${HC_FP8:-1}"
+PLE_PATCH="${PLE_PATCH:-0}"
+DRAFT_VOCAB="${DRAFT_VOCAB-$SCRIPT_DIR/draft_vocab_64k.json}"   # set empty to disable
+MXFP8_GEMV="${MXFP8_GEMV-2560x6144,1280x2560,16384x2560,640x2560}"   # shapes that beat b12x in-model
+MTP_FP8="${MTP_FP8:-1}"
+MTP_BF16_KV="${MTP_BF16_KV:-1}"
+GEMMA_NORM="${GEMMA_NORM:-1}"
+GEMV_TUNE="${GEMV_TUNE:-0}"
+HC_FUSE="${HC_FUSE:-1}"
+PLE_OVERLAP="${PLE_OVERLAP:-1}"
+GDN_DEFER="${GDN_DEFER:-1}"
+GDN_ASYNC="${GDN_ASYNC:-0}"            # GDN_DEFER commit on a side stream overlapping the MTP draft phase
+GDN_FLAT="${GDN_FLAT:-0}"              # GDN_DEFER straight-line verify kernel (8 rows x 1 warp per program)
+DRAFT_NVFP4="${DRAFT_NVFP4:-0}"
+TARGET_FP8="${TARGET_FP8:-1}"
+SHARED_FUSE="${SHARED_FUSE:-1}"
+PLE_ASYNC="${PLE_ASYNC:-1}"
+PLE_KEEPALIVE_US="${PLE_KEEPALIVE_US:-2000}"
+MOE_DETERMINISTIC="${MOE_DETERMINISTIC:-0}"  # b12x MoE fixed-order top-k sum instead of bf16 atomic scatter (run-to-run identical)
+MOE_A16="${MOE_A16:-1}"                # routed experts with bf16 activations (W4A16) instead of the checkpoint default NVFP4 activations (W4A4)
+A4_PREFILL="${A4_PREFILL:-0}"          # FP4 activations for prefill-size MoE calls only: +6-8% prefill, but KLD 0.0492 -> 0.0665 (W4A4 level)
+TAIL_CKPT="${TAIL_CKPT:-16,16}"             # "GRID,SUFFIX": prefill-tail prefix-cache checkpoint just before the prompt end (decouples follow-up reuse from chunk size)
+CSF_INLINE="${CSF_INLINE:-1}"          # b12x W4A16 decode MoE reads NVFP4-CSF scales inline (no csf_Pair expansion pass):
+                                       # bit-identical, -5.7% replay step time (57.5 -> 54.2 ms)
+LOGIT_DUMP="${LOGIT_DUMP:-}"           # measurement only: dump dir under /cache for prompt logits
+EXPERT_STATS="${EXPERT_STATS:-}"       # measurement only: json under /cache, distinct experts per decode call
+PATCH_DIR="${PATCH_DIR:-$(cd "$SCRIPT_DIR/../patches" && pwd)}"
+VLLM_PKG=/usr/local/lib/python3.12/dist-packages/vllm
+EXTRA_ENV="${EXTRA_ENV:-}"             # extra "-e KEY=VAL" docker flags (tuning knobs)
+B12X_CACHE="${B12X_CACHE:-b12x/compile}"  # b12x compile/autotune cache under CACHE_DIR; point experiments with new shapes at a copy
+
+if [ "${1:-start}" = "stop" ]; then
+  docker stop "$NAME" >/dev/null 2>&1 || true
+  docker rm "$NAME" >/dev/null 2>&1 || true
+  exit 0
+fi
+
+# b12x keeps compiled kernels and autotune results under B12X_COMPILE_CACHE_DIR
+# (default ~/.cache/b12x/compile inside the container, lost on every restart,
+# which re-ran the ~14 min gdn_prefill autotune).
+mkdir -p "$CACHE_DIR"/{vllm,flashinfer,triton,b12x/compile}
+patch_args=()
+declare -A mounted=()
+mount() {  # mount <patch-relative source> <absolute target>, once per target
+  [ -n "${mounted[$2]:-}" ] && return
+  mounted[$2]=1
+  patch_args+=(-v "$PATCH_DIR/$1:$2:ro")
+}
+QWEN=$VLLM_PKG/models/qwen4_exp/nvidia
+B12X=/usr/local/lib/python3.12/dist-packages/flashinfer/experimental/b12x
+if [ "$FP8_LM_HEAD" = "1" ] || [ "$SMALL_GEMV" = "1" ]; then
+  mount qwen4_exp/nvidia/model.py $QWEN/model.py
+  mount qwen4_exp/nvidia/ops/small_gemv.py $QWEN/ops/small_gemv.py
+fi
+[ "$FP8_LM_HEAD" = "1" ] && patch_args+=(-e VLLM_QWEN38_FP8_LM_HEAD=1)
+[ "$SMALL_GEMV" = "1" ] && patch_args+=(-e VLLM_QWEN38_SMALL_GEMV=1)
+[ "$MXFP8_LM_HEAD" = "1" ] && patch_args+=(-e VLLM_MXFP8_LM_HEAD=1)
+if [ "$HC_FP8" = "1" ]; then
+  mount qwen4_exp/nvidia/hyperconnection.py $QWEN/hyperconnection.py
+  mount qwen4_exp/nvidia/ops/hc_fp8_gemv.py $QWEN/ops/hc_fp8_gemv.py
+  patch_args+=(-e VLLM_QWEN38_HC_FP8=1)
+fi
+if [ -n "$MXFP8_GEMV" ]; then
+  mount vllm_kernels/linear/mxfp8/b12x.py $VLLM_PKG/model_executor/kernels/linear/mxfp8/b12x.py
+  mount qwen4_exp/nvidia/ops/mxfp8_gemv.py $QWEN/ops/mxfp8_gemv.py
+  patch_args+=(-e VLLM_QWEN38_MXFP8_GEMV="$MXFP8_GEMV")
+fi
+if [ -n "$DRAFT_VOCAB" ] && [ ! -f "$DRAFT_VOCAB" ]; then
+  echo "warning: $DRAFT_VOCAB not found; draft vocabulary subset disabled" >&2
+  DRAFT_VOCAB=""
+fi
+if [ -n "$DRAFT_VOCAB" ]; then
+  # The draft shares the target's MXFP8 lm_head and gathers the subset rows from it.
+  mount qwen4_exp/nvidia/mtp.py $QWEN/mtp.py
+  mount qwen4_exp/nvidia/ops/mxfp8_gemv.py $QWEN/ops/mxfp8_gemv.py
+  patch_args+=(-v "$DRAFT_VOCAB:/opt/draft_vocab.json:ro" -e VLLM_QWEN38_DRAFT_VOCAB=/opt/draft_vocab.json -e VLLM_MTP_NVFP4_LM_HEAD=0)
+  if [ "$DRAFT_NVFP4" = "1" ]; then
+    mount qwen4_exp/nvidia/ops/nvfp4_gemv.py $QWEN/ops/nvfp4_gemv.py
+    patch_args+=(-e VLLM_QWEN38_DRAFT_NVFP4=1)
+  fi
+fi
+if [ "$MTP_FP8" = "1" ]; then
+  mount qwen4_exp/nvidia/mtp.py $QWEN/mtp.py
+  mount qwen4_exp/nvidia/ops/hc_fp8_gemv.py $QWEN/ops/hc_fp8_gemv.py
+  patch_args+=(-e VLLM_QWEN38_MTP_FP8=1)
+fi
+if [ "$MTP_BF16_KV" = "1" ]; then
+  mount qwen4_exp/nvidia/mtp.py $QWEN/mtp.py
+  patch_args+=(-e VLLM_QWEN38_MTP_BF16_KV=1)
+fi
+if [ "$GEMV_TUNE" = "1" ]; then patch_args+=(-e VLLM_QWEN38_GEMV_TUNE=1); fi
+if [ "$HC_FUSE" = "1" ] && [ "$HC_FP8" = "1" ]; then patch_args+=(-e VLLM_QWEN38_HC_FUSE=1); fi
+if [ "$GEMMA_NORM" = "1" ]; then
+  mount qwen4_exp/nvidia/model.py $QWEN/model.py
+  mount qwen4_exp/nvidia/ops/gemma_norm.py $QWEN/ops/gemma_norm.py
+  patch_args+=(-e VLLM_QWEN38_GEMMA_NORM=1)
+fi
+if [ "$SHARED_FUSE" = "1" ]; then
+  mount qwen4_exp/nvidia/model.py $QWEN/model.py
+  mount qwen4_exp/nvidia/ops/small_gemv.py $QWEN/ops/small_gemv.py
+  mount qwen4_exp/nvidia/ops/mxfp8_gemv.py $QWEN/ops/mxfp8_gemv.py
+  patch_args+=(-e VLLM_QWEN38_SHARED_FUSE=1)
+fi
+if [ "$TARGET_FP8" = "1" ]; then
+  mount qwen4_exp/nvidia/model.py $QWEN/model.py
+  mount qwen4_exp/nvidia/ops/small_gemv.py $QWEN/ops/small_gemv.py
+  mount qwen4_exp/nvidia/ops/hc_fp8_gemv.py $QWEN/ops/hc_fp8_gemv.py
+  patch_args+=(-e VLLM_QWEN38_TARGET_FP8=1)
+fi
+if [ "$GDN_DEFER" = "1" ]; then
+  mount qwen4_exp/nvidia/model.py $QWEN/model.py
+  mount qwen4_exp/nvidia/gdn_defer.py $QWEN/gdn_defer.py
+  mount qwen4_exp/nvidia/ops/gdn_deferred.py $QWEN/ops/gdn_deferred.py
+  patch_args+=(-e VLLM_QWEN38_GDN_DEFER=1)
+  [ "$GDN_ASYNC" = "1" ] && patch_args+=(-e VLLM_QWEN38_GDN_COMMIT_ASYNC=1)
+  [ "$GDN_FLAT" = "1" ] && patch_args+=(-e VLLM_QWEN38_GDN_VERIFY_FLAT=1)
+fi
+[ "$MOE_DETERMINISTIC" = "1" ] && patch_args+=(-e B12X_DYNAMIC_DETERMINISTIC_OUTPUT=1)
+[ "$MOE_A16" = "1" ] && patch_args+=(-e VLLM_B12X_MOE_FP4_FORCE_A16=1)
+[ "$A4_PREFILL" = "1" ] && patch_args+=(-e B12X_W4A16_A4_PREFILL=1)
+if [ "$CSF_INLINE" = "1" ]; then
+  mount b12x/moe/_shared/kernels/w4a16/kernel.py $B12X/moe/_shared/kernels/w4a16/kernel.py
+  mount b12x/moe/fused_moe/_impl.py $B12X/moe/fused_moe/_impl.py
+  mount b12x/moe/fused_moe/_preparation.py $B12X/moe/fused_moe/_preparation.py
+  patch_args+=(-e B12X_W4A16_CSF_INLINE=1)
+fi
+if [ -n "$TAIL_CKPT" ] && [ "$TAIL_CKPT" != off ]; then  # TAIL_CKPT=off: stock checkpoint
+  mount vllm_core/v1/core/boundary_checkpoint.py $VLLM_PKG/v1/core/boundary_checkpoint.py
+  patch_args+=(-e VLLM_QWEN38_TAIL_CKPT="$TAIL_CKPT")
+fi
+if [ -n "$LOGIT_DUMP" ]; then
+  mount qwen4_exp/nvidia/model.py $QWEN/model.py
+  mount qwen4_exp/nvidia/logit_dump.py $QWEN/logit_dump.py
+  patch_args+=(-e VLLM_QWEN38_LOGIT_DUMP="/cache/$LOGIT_DUMP")
+fi
+if [ -n "$EXPERT_STATS" ]; then
+  mount qwen4_exp/nvidia/model.py $QWEN/model.py
+  mount qwen4_exp/nvidia/expert_stats.py $QWEN/expert_stats.py
+  patch_args+=(-e VLLM_QWEN38_EXPERT_STATS="/cache/$EXPERT_STATS")
+fi
+if [ "$PLE_OVERLAP" = "1" ]; then
+  mount qwen4_exp/nvidia/b12x_ple.py $QWEN/b12x_ple.py
+  patch_args+=(-e VLLM_QWEN38_PLE_OVERLAP=1)
+fi
+if [ "$PLE_PATCH" = "1" ]; then
+  mount qwen4_exp/nvidia/b12x_ple.py $QWEN/b12x_ple.py
+  mount b12x/loader/_row_plan.h $B12X/loader/_row_plan.h
+fi
+if [ "$PLE_ASYNC" = "1" ]; then
+  mount b12x/loader/_row_plan.h $B12X/loader/_row_plan.h
+  mount b12x/loader/_ple_reader.c $B12X/loader/_ple_reader.c
+  mount b12x/loader/_storage.c $B12X/loader/_storage.c
+  mount b12x/sequence/_shared/disk_table.py $B12X/sequence/_shared/disk_table.py
+  patch_args+=(-e B12X_PLE_ASYNC=1)
+  if [ "${PLE_KEEPALIVE_US:-0}" != "0" ]; then patch_args+=(-e B12X_PLE_KEEPALIVE_US="$PLE_KEEPALIVE_US"); fi
+fi
+spec=()
+[ "$MTP_TOKENS" != "0" ] && spec=(--speculative-config "{\"method\":\"mtp\",\"num_speculative_tokens\":${MTP_TOKENS}}")
+
+docker rm -f "$NAME" >/dev/null 2>&1 || true
+docker run -d --name "$NAME" \
+  --gpus all --ipc host --network host --shm-size=16g \
+  --ulimit memlock=-1 --ulimit stack=67108864 \
+  --security-opt seccomp="$SECCOMP" \
+  -e VLLM_NO_USAGE_STATS=1 -e VLLM_DO_NOT_TRACK=1 -e DO_NOT_TRACK=1 \
+  -e HF_HUB_DISABLE_TELEMETRY=1 \
+  -e VLLM_PLE_TABLE_MEMORY=disk \
+  -e B12X_DISK_BACKEND=io_uring \
+  -e VLLM_USE_V2_MODEL_RUNNER=1 \
+  -e B12X_POLICY_MODE=auto \
+  -e VLLM_SSM_CONV_STATE_LAYOUT=DS \
+  -e VLLM_CACHE_ROOT=/cache/vllm \
+  -e FLASHINFER_WORKSPACE_BASE=/cache/flashinfer \
+  -e TRITON_CACHE_DIR=/cache/triton \
+  -e B12X_COMPILE_CACHE_DIR="/cache/$B12X_CACHE" \
+  "${patch_args[@]}" $EXTRA_ENV \
+  -v "$MODEL_DIR":/model:ro \
+  -v "$CACHE_DIR":/cache \
+  --entrypoint vllm \
+  "$IMAGE" serve /model \
+    --served-model-name nvidia/Qwen3.8-Flash-Next-FP4-Local qwen3.8-flash-next \
+    --host 0.0.0.0 --port "$PORT" \
+    --quantization modelopt_mixed \
+    --load-format b12x \
+    --moe-backend b12x \
+    --linear-backend b12x \
+    --kv-cache-dtype fp8 \
+    --max-model-len "$MAX_MODEL_LEN" \
+    --max-num-seqs "$MAX_NUM_SEQS" \
+    --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \
+    --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
+    --enable-prefix-caching \
+    --reasoning-parser qwen3 \
+    --tool-call-parser qwen3_xml \
+    --enable-auto-tool-choice \
+    "${spec[@]}" ${API_KEY:+--api-key "$API_KEY"} $EXTRA_ARGS
+echo "started $NAME on port $PORT (docker logs -f $NAME)"

tokenizer_config.json

@@ -0,0 +1,305 @@
+{
+    "add_prefix_space": false,
+    "added_tokens_decoder": {
+        "248044": {
+            "content": "<|endoftext|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248045": {
+            "content": "<|im_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248046": {
+            "content": "<|im_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248047": {
+            "content": "<|object_ref_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248048": {
+            "content": "<|object_ref_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248049": {
+            "content": "<|box_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248050": {
+            "content": "<|box_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248051": {
+            "content": "<|quad_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248052": {
+            "content": "<|quad_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248053": {
+            "content": "<|vision_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248054": {
+            "content": "<|vision_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248055": {
+            "content": "<|vision_pad|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248056": {
+            "content": "<|image_pad|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248057": {
+            "content": "<|video_pad|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248058": {
+            "content": "<tool_call>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248059": {
+            "content": "</tool_call>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248060": {
+            "content": "<|fim_prefix|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248061": {
+            "content": "<|fim_middle|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248062": {
+            "content": "<|fim_suffix|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248063": {
+            "content": "<|fim_pad|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248064": {
+            "content": "<|repo_name|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248065": {
+            "content": "<|file_sep|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248066": {
+            "content": "<tool_response>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248067": {
+            "content": "</tool_response>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248068": {
+            "content": "<think>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248069": {
+            "content": "</think>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": false
+        },
+        "248070": {
+            "content": "<|audio_start|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248071": {
+            "content": "<|audio_end|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248072": {
+            "content": "<tts_pad>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248073": {
+            "content": "<tts_text_bos>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248074": {
+            "content": "<tts_text_eod>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248075": {
+            "content": "<tts_text_bos_single>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "248076": {
+            "content": "<|audio_pad|>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        }
+    },
+    "additional_special_tokens": [
+        "<|im_start|>",
+        "<|im_end|>",
+        "<|object_ref_start|>",
+        "<|object_ref_end|>",
+        "<|box_start|>",
+        "<|box_end|>",
+        "<|quad_start|>",
+        "<|quad_end|>",
+        "<|vision_start|>",
+        "<|vision_end|>",
+        "<|vision_pad|>",
+        "<|image_pad|>",
+        "<|video_pad|>"
+    ],
+    "bos_token": null,
+    "chat_template": "{%- set image_count = namespace(value=0) %}\n{%- set video_count = namespace(value=0) %}\n{%- macro render_content(content, do_vision_count, is_system_content=false) %}\n    {%- if content is string %}\n        {{- content }}\n    {%- elif content is iterable and content is not mapping %}\n        {%- for item in content %}\n            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}\n                {%- if is_system_content %}\n                    {{- raise_exception('System message cannot contain images.') }}\n                {%- endif %}\n                {%- if do_vision_count %}\n                    {%- set image_count.value = image_count.value + 1 %}\n                {%- endif %}\n                {%- if add_vision_id %}\n                    {{- 'Picture ' ~ image_count.value ~ ': ' }}\n                {%- endif %}\n                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}\n            {%- elif 'video' in item or item.type == 'video' %}\n                {%- if is_system_content %}\n                    {{- raise_exception('System message cannot contain videos.') }}\n                {%- endif %}\n                {%- if do_vision_count %}\n                    {%- set video_count.value = video_count.value + 1 %}\n                {%- endif %}\n                {%- if add_vision_id %}\n                    {{- 'Video ' ~ video_count.value ~ ': ' }}\n                {%- endif %}\n                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}\n            {%- elif 'text' in item %}\n                {{- item.text }}\n            {%- else %}\n                {{- raise_exception('Unexpected item type in content.') }}\n            {%- endif %}\n        {%- endfor %}\n    {%- elif content is none or content is undefined %}\n        {{- '' }}\n    {%- else %}\n        {{- raise_exception('Unexpected content type.') }}\n    {%- endif %}\n{%- endmacro %}\n{%- if not messages %}\n    {{- raise_exception('No messages provided.') }}\n{%- endif %}\n{%- set reasoning_instructions = '' %}\n{%- if enable_thinking is undefined or enable_thinking is true %}\n    {%- set resolved_reasoning_effort = reasoning_effort|default('xhigh') %}\n    {%- if resolved_reasoning_effort not in ('xhigh', 'medium', 'low') %}\n        {{- raise_exception('Unexpected reasoning effort ' ~ reasoning_effort ~ '. Supported types are xhigh (default), medium, and low.') }}\n    {%- endif %}\n    {%- if resolved_reasoning_effort == 'xhigh' %}\n        {%- set reasoning_instructions = 'Reasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer.' %}\n    {%- elif resolved_reasoning_effort == 'low' %}\n        {%- set reasoning_instructions = 'Reasoning effort is set to low. Keep your thinking brief and focused, moving directly to the conclusion without unnecessary elaboration.' %}\n    {%- endif %}\n{%- endif %}\n{%- if tools and tools is iterable and tools is not mapping %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if reasoning_instructions %}\n        {{- reasoning_instructions + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou have access to the following functions:\\n\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\" }}\n    {{- '\\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>' }}\n    {%- if messages[0].role == 'system' %}\n        {%- set content = render_content(messages[0].content, false, true)|trim %}\n        {%- if content %}\n            {{- '\\n\\n' + content }}\n        {%- endif %}\n    {%- endif %}\n    {{- '<|im_end|>\\n' }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {%- set content = render_content(messages[0].content, false, true)|trim %}\n        {%- if content %}\n            {{- '<|im_start|>system\\n' + (reasoning_instructions + '\\n\\n' if reasoning_instructions else '')  + content + '<|im_end|>\\n' }}\n        {%- elif reasoning_instructions %}\n            {{- '<|im_start|>system\\n' + reasoning_instructions + '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- elif reasoning_instructions %}\n        {{- '<|im_start|>system\\n' + reasoning_instructions + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" %}\n        {%- set content = render_content(message.content, false)|trim %}\n        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}\n            {%- set ns.multi_step_tool = false %}\n            {%- set ns.last_query_index = index %}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if ns.multi_step_tool %}\n    {{- raise_exception('No user query found in messages.') }}\n{%- endif %}\n{%- for message in messages %}\n    {%- set content = render_content(message.content, true)|trim %}\n    {%- if message.role == \"system\" %}\n        {%- if not loop.first %}\n            {{- raise_exception('System message must be at the beginning.') }}\n        {%- endif %}\n    {%- elif message.role == \"user\" %}\n        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is string %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- endif %}\n        {%- set reasoning_content = reasoning_content|trim %}\n        {%- if preserve_thinking is undefined or preserve_thinking is true or loop.index0 > ns.last_query_index %}\n            {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content + '\\n</think>\\n\\n' + content }}\n        {%- else %}\n            {{- '<|im_start|>' + message.role + '\\n' + content }}\n        {%- endif %}\n        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if tool_call.function is defined %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {%- if loop.first %}\n                    {%- if content|trim %}\n                        {{- '\\n\\n<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n                    {%- else %}\n                        {{- '<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n                    {%- endif %}\n                {%- else %}\n                    {{- '\\n<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n                {%- endif %}\n                {%- if tool_call.arguments is defined and tool_call.arguments != '' %}\n                    {%- for args_name, args_value in tool_call.arguments|items %}\n                        {{- '<parameter=' + args_name + '>\\n' }}\n                        {%- set args_value = args_value | string if args_value is string else args_value | tojson | safe %}\n                        {{- args_value }}\n                        {{- '\\n</parameter>\\n' }}\n                    {%- endfor %}\n                {%- endif %}\n                {{- '</function>\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.previtem and loop.previtem.role != \"tool\" %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- content }}\n        {{- '\\n</tool_response>' }}\n        {%- if not loop.last and loop.nextitem.role != \"tool\" %}\n            {{- '<|im_end|>\\n' }}\n        {%- elif loop.last %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- else %}\n        {{- raise_exception('Unexpected message role.') }}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- else %}\n        {{- '<think>\\n' }}\n    {%- endif %}\n{%- endif %}",
+    "clean_up_tokenization_spaces": false,
+    "eos_token": "<|im_end|>",
+    "errors": "replace",
+    "model_max_length": 262144,
+    "pad_token": "<|endoftext|>",
+    "split_special_tokens": false,
+    "tokenizer_class": "Qwen2Tokenizer",
+    "unk_token": null,
+    "add_bos_token": false,
+    "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+",
+    "extra_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|>"
+    }
+}

video_preprocessor_config.json

@@ -0,0 +1,21 @@
+{
+    "size": {
+        "longest_edge": 25165824,
+        "shortest_edge": 4096
+    },
+    "patch_size": 16,
+    "temporal_patch_size": 2,
+    "merge_size": 2,
+    "image_mean": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "image_std": [
+        0.5,
+        0.5,
+        0.5
+    ],
+    "processor_class": "Qwen3VLProcessor",
+    "video_processor_type": "Qwen3VLVideoProcessor"
+}