use newtype for logprobs count (-1 for all)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml

@@ -21,6 +21,10 @@ steps:
     timeout_in_minutes: 30
     optional: true
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 2+
+      mem: 24+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -38,6 +42,10 @@ steps:
     timeout_in_minutes: 30
     optional: true
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -55,6 +63,10 @@ steps:
     timeout_in_minutes: 30
     optional: true
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"

.buildkite/intel_jobs/basic_correctness.yaml

@@ -5,6 +5,10 @@ steps:
 - label: XPU Sleep Mode
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/engine_intel.yaml

@@ -5,6 +5,10 @@ steps:
 - label: Engine (1 GPU)
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/expert_parallelism_intel.yaml

@@ -6,6 +6,10 @@ steps:
   key: eplb-algorithm
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/kernels_intel.yaml

@@ -5,6 +5,10 @@ steps:
 - label: vLLM IR Tests
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/lora_intel.yaml

@@ -5,6 +5,10 @@ steps:
 - label: LoRA Runtime + Utils
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:
@@ -34,6 +38,10 @@ steps:
 - label: LoRA Fused/MoE Kernels
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -54,6 +62,10 @@ steps:
 - label: LoRA Punica Kernels
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -74,6 +86,10 @@ steps:
 - label: LoRA Punica FP8/XPU Ops
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -94,6 +110,10 @@ steps:
 - label: LoRA Models
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:
@@ -117,6 +137,10 @@ steps:
 - label: LoRA Multimodal
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/misc_intel.yaml

@@ -5,6 +5,10 @@ steps:
 - label: V1 Core + KV + Metrics
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -31,6 +35,10 @@ steps:
 - label: V1 Sample + Logits
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -71,6 +79,10 @@ steps:
 - label: XPU CPU Offload
   timeout_in_minutes: 60
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -95,6 +107,10 @@ steps:
   key: regression
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -126,6 +142,10 @@ steps:
   timeout_in_minutes: 30
   num_devices: 2
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -157,6 +177,10 @@ steps:
   key: async-engine-inputs-utils-worker
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/model_runner_v2_intel.yaml

@@ -5,6 +5,10 @@ steps:
 - label: Model Runner V2 Core Tests (Intel)
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -30,6 +34,10 @@ steps:
 - label: Model Runner V2 Examples (Intel)
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/models_multimodal_intel.yaml

@@ -6,6 +6,10 @@ steps:
   key: multi-modal-models-standard-1-qwen2
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -27,6 +31,10 @@ steps:
   key: multi-modal-models-standard-2-qwen3-gemma
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -47,6 +55,10 @@ steps:
   key: multi-modal-models-standard-3-llava-qwen2-vl
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:
@@ -68,6 +80,10 @@ steps:
   key: multi-modal-models-standard-4-other-whisper
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -88,6 +104,10 @@ steps:
   key: multi-modal-processor
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/test-intel.yaml

@@ -19,6 +19,10 @@ steps:
       - image-build-xpu
     timeout_in_minutes: 30
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 2+
+      mem: 24+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -49,6 +53,10 @@ steps:
       - image-build-xpu
     timeout_in_minutes: 30
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -74,6 +82,10 @@ steps:
       - image-build-xpu
     timeout_in_minutes: 30
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -93,6 +105,10 @@ steps:
       - image-build-xpu
     timeout_in_minutes: 30
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"

.buildkite/scripts/install-kv-connectors.sh

@@ -4,6 +4,11 @@
 
 set -euo pipefail
 
+if python3 -c "import torch; raise SystemExit(0 if torch.version.hip is not None else 1)"; then
+    uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+    exit 0
+fi
+
 REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
 
 uv pip install --system -r "${REQUIREMENTS_FILE}"

.buildkite/test-amd.yaml

@@ -647,7 +647,7 @@ steps:
   - pytest -v -s v1/cudagraph/test_cudagraph_mode.py
 
 - label: e2e Core (1 GPU) # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 35
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
   agent_pool: mi250_1
   optional: true
@@ -1594,9 +1594,10 @@ steps:
 #----------------------------------------------------------  mi300 · kernels  ----------------------------------------------------------#
 
 - label: Kernels Attention Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
+  optional: true
   parallelism: 2
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -1627,10 +1628,11 @@ steps:
   - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py  kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py
 
 - label: Kernels MoE Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 50
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
-  parallelism: 4
+  optional: true
+  parallelism: 5
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - csrc/quantization/cutlass_w8a8/moe/
@@ -2075,19 +2077,6 @@ steps:
   - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
   - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
 
-- label: e2e Core (1 GPU) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/
-  - tests/v1/e2e/
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
-
 - label: e2e Scheduling (1 GPU) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2133,9 +2122,10 @@ steps:
   - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
 
 - label: Spec Decode Eagle # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
+  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/v1/spec_decode/
@@ -3053,7 +3043,7 @@ steps:
 #----------------------------------------------------------  mi355 · kernels  ----------------------------------------------------------#
 
 - label: Kernels (B200-MI355) # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 15
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   agent_pool: mi355_1
   working_dir: "/vllm-workspace/"
@@ -3077,11 +3067,10 @@ steps:
   - pytest -v -s tests/kernels/attention/test_attention_selector.py
 
 - label: Kernels Attention Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   agent_pool: mi355_1
   parallelism: 2
-  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - csrc/attention/
@@ -3095,10 +3084,10 @@ steps:
   - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 - label: Kernels MoE Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 50
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   agent_pool: mi355_1
-  parallelism: 4
+  parallelism: 5
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - csrc/quantization/cutlass_w8a8/moe/

.buildkite/test_areas/engine.yaml

@@ -74,6 +74,16 @@ steps:
     - tests/v1/e2e/general/
   commands:
     - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
+  mirror:
+    amd:
+      device: mi250_1
+      timeout_in_minutes: 35
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/
+      - tests/v1/e2e/general/
+      - vllm/platforms/rocm.py
 
 - label: V1 e2e (2 GPUs)
   key: v1-e2e-2-gpus

.buildkite/test_areas/kernels.yaml

@@ -74,6 +74,20 @@ steps:
   commands:
     - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
   parallelism: 2
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 55
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/attention/
+      - vllm/v1/attention
+      - vllm/model_executor/layers/attention
+      - tests/kernels/attention
+      - vllm/_aiter_ops.py
+      - vllm/envs.py
+      - vllm/platforms/rocm.py
 
 - label: Kernels Attention DiffKV Test (H100)
   key: kernels-attention-diffkv-test-h100
@@ -104,6 +118,7 @@ steps:
       source_file_dependencies:
       - csrc/quantization/
       - vllm/model_executor/layers/quantization
+      - vllm/config/
       - tests/kernels/quantization
       - tests/kernels/quantization/test_rocm_skinny_gemms.py
       - vllm/_aiter_ops.py
@@ -127,6 +142,22 @@ steps:
     - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
     - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
   parallelism: 5
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 50
+      source_file_dependencies:
+      - csrc/quantization/cutlass_w8a8/moe/
+      - csrc/moe/
+      - tests/kernels/moe
+      - vllm/model_executor/layers/fused_moe/
+      - vllm/distributed/device_communicators/
+      - vllm/envs.py
+      - vllm/config
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      depends_on:
+      - image-build-amd
 
 - label: Kernels Mamba Test
   key: kernels-mamba-test

.buildkite/test_areas/lm_eval.yaml

@@ -101,6 +101,16 @@ steps:
   num_devices: 8
   commands:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
+  mirror:
+    amd:
+      device: mi300_8
+      timeout_in_minutes: 180
+      depends_on:
+      - image-build-amd
+      commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
+      - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
 
 - label: MoE Refactor Integration Test (H100 - TEMPORARY)
   key: moe-refactor-integration-test-h100-temporary

.buildkite/test_areas/misc.yaml

@@ -105,6 +105,12 @@ steps:
     # Integration test for streaming correctness (requires special branch).
     - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
     - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 60
+      depends_on:
+      - image-build-amd
 
 - label: V1 Others (CPU)
   key: v1-others-cpu

.buildkite/test_areas/models_multimodal.yaml

@@ -68,7 +68,6 @@ steps:
     - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
   mirror:
     amd:
-      soft_fail: true
       device: mi325_1
       depends_on:
       - image-build-amd

.buildkite/test_areas/pytorch.yaml

@@ -107,6 +107,12 @@ steps:
     - tests/compile/passes
   commands:
   - pytest -s -v compile/passes --ignore compile/passes/distributed
+  mirror:
+    amd:
+      device: mi300_1
+      timeout_in_minutes: 180
+      depends_on:
+      - image-build-amd
 
 - label: PyTorch Fullgraph Smoke Test
   key: pytorch-fullgraph-smoke-test

.buildkite/test_areas/spec_decode.yaml

@@ -12,6 +12,20 @@ steps:
     - tests/v1/e2e/spec_decode/
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 45
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/spec_decode/
+      - vllm/v1/worker/gpu/spec_decode/
+      - vllm/model_executor/model_loader/
+      - vllm/v1/sample/
+      - vllm/model_executor/layers/
+      - tests/v1/e2e/spec_decode/
+      - vllm/platforms/rocm.py
 
 - label: Spec Decode Eagle Nightly B200
   key: spec-decode-eagle-nightly-b200

.claude/skills/ci-fails-buildkite/SKILL.md

@@ -0,0 +1,35 @@
+---
+name: ci-fails-buildkite
+description: Fetch and diagnose vLLM Buildkite CI failure logs. Use when investigating failing CI jobs on a PR or build, when the user pastes a buildkite.com URL, or asks to fetch/diagnose CI logs.
+---
+
+# Diagnosing vLLM Buildkite CI Failures
+
+Buildkite logs are public; no login needed.
+
+`.buildkite/scripts/ci-fetch-log.sh` saves each log as `ci-<build>-<job-name>.log`, stripped of timestamps and ANSI codes. Existing files are kept; set `CI_FETCH_LOG_FORCE=1` to refetch.
+
+## Fetching logs
+
+```bash
+# All failed jobs in a PR's latest build (current branch's PR if omitted):
+.buildkite/scripts/ci-fetch-log.sh --pr <PR>
+
+# All failed jobs in a build (--soft also includes soft-failed jobs;
+# --all fetches every finished job):
+.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>"
+
+# One job — `gh pr checks` URLs (#<job_uuid>) and web UI URLs (?sid=) both
+# work; pass "-" as a second argument to stream to stdout:
+.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>#<job_uuid>"
+```
+
+To clean an already-downloaded log with `.buildkite/scripts/ci-clean-log.sh`:
+
+```bash
+./ci-clean-log.sh ci.log
+```
+
+## Reference
+
+See [docs/contributing/ci/failures.md](../../../docs/contributing/ci/failures.md) for the full guide: filing CI failure issues, investigating/bisecting, reproducing flaky tests, and daily triage.

.github/CODEOWNERS

@@ -2,15 +2,14 @@
 # for more info about CODEOWNERS file
 
 # This lists cover the "core" components of vLLM that require careful review
-/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng @vadiklyutiy
+/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng
 /vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi
 /vllm/lora @jeejeelee
 /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
 /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye
 /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye
 /vllm/model_executor/layers/mamba @tdoublep @tomeras91
-/vllm/model_executor/layers/mamba/gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
-/vllm/model_executor/layers/rotary_embedding.py @vadiklyutiy
+/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
 /vllm/model_executor/model_loader @22quinn
 /vllm/model_executor/layers/batch_invariant.py @yewentao256 
 /vllm/ir @ProExpertProg

.gitignore

@@ -199,7 +199,9 @@ cython_debug/
 .vscode/
 
 # Claude
-.claude/
+.claude/*
+!.claude/skills/
+!.claude/skills/**
 
 # Codex
 .codex/

AGENTS.md

@@ -114,17 +114,6 @@ Follow these rules for all code changes in this repository:
 - Keep comments and docstrings minimal and concise.
 - Assume the reader is familiar with vLLM.
 
-### Diagnosing CI failures
-
-Buildkite logs are public; no login needed. Details: [docs/contributing/ci/failures.md](docs/contributing/ci/failures.md).
-
-```bash
-# All failed-job logs for a PR's latest build (current branch's PR if omitted):
-.buildkite/scripts/ci-fetch-log.sh --pr <PR>
-# Any Buildkite build or job URL also works:
-.buildkite/scripts/ci-fetch-log.sh "<buildkite_url>"
-```
-
 ### Commit messages
 
 Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:

CMakeLists.txt

@@ -319,82 +319,35 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
 endif()
 
 #
-# _C extension
+# Legacy _C extension (ROCm only — CUDA ops migrated to _C_stable_libtorch)
 #
 
-set(VLLM_EXT_SRC
-  "csrc/quantization/activation_kernels.cu"
-  "csrc/torch_bindings.cpp")
-
-if(VLLM_GPU_LANG STREQUAL "CUDA")
-  SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
-
-  # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
-  set(CUTLASS_REVISION "v4.4.2")
-
-  # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
-  if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
-    set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
-  endif()
-
-  if(VLLM_CUTLASS_SRC_DIR)
-    if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
-      get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
-    endif()
-    message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
-    FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
-  else()
-    FetchContent_Declare(
-        cutlass
-        GIT_REPOSITORY https://github.com/nvidia/cutlass.git
-        # Please keep this in sync with CUTLASS_REVISION line above.
-        GIT_TAG ${CUTLASS_REVISION}
-        GIT_PROGRESS TRUE
-
-        # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
-        # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
-        # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
-        GIT_SHALLOW TRUE
-    )
-  endif()
-  FetchContent_MakeAvailable(cutlass)
-
-  set_gencode_flags_for_srcs(
-    SRCS "${VLLM_EXT_SRC}"
-    CUDA_ARCHS "${CUDA_ARCHS}")
-
-# if CUDA endif
-endif()
-
-if (VLLM_GPU_LANG STREQUAL "HIP")
-  # Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
-  # TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
-  list(APPEND VLLM_EXT_SRC
+if(VLLM_GPU_LANG STREQUAL "HIP")
+  set(VLLM_EXT_SRC
+    "csrc/torch_bindings.cpp"
     "csrc/custom_quickreduce.cu"
     "csrc/cuda_view.cu"
-    "csrc/libtorch_stable/cuda_utils_kernels.cu"
-  )
-# if ROCM endif
-endif()
+    "csrc/libtorch_stable/cuda_utils_kernels.cu")
 
-message(STATUS "Enabling C extension.")
-define_extension_target(
-  _C
-  DESTINATION vllm
-  LANGUAGE ${VLLM_GPU_LANG}
-  SOURCES ${VLLM_EXT_SRC}
-  COMPILE_FLAGS ${VLLM_GPU_FLAGS}
-  ARCHITECTURES ${VLLM_GPU_ARCHES}
-  INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
-  INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
-  USE_SABI 3
-  WITH_SOABI)
+  message(STATUS "Enabling C extension.")
+  define_extension_target(
+    _C
+    DESTINATION vllm
+    LANGUAGE ${VLLM_GPU_LANG}
+    SOURCES ${VLLM_EXT_SRC}
+    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+    ARCHITECTURES ${VLLM_GPU_ARCHES}
+    INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
+    INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
+    USE_SABI 3
+    WITH_SOABI)
 
-# If CUTLASS is compiled on NVCC >= 12.5, it by default uses
-# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
-# driver API. This causes problems when linking with earlier versions of CUDA.
-# Setting this variable sidesteps the issue by calling the driver directly.
-target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+  # If CUTLASS is compiled on NVCC >= 12.5, it by default uses
+  # cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
+  # driver API. This causes problems when linking with earlier versions of CUDA.
+  # Setting this variable sidesteps the issue by calling the driver directly.
+  target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+endif() # _C HIP endif
 
 if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
   #
@@ -403,6 +356,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
   set(VLLM_STABLE_EXT_SRC
     "csrc/libtorch_stable/torch_bindings.cpp"
     "csrc/libtorch_stable/activation_kernels.cu"
+    "csrc/libtorch_stable/quantization/activation_kernels.cu"
     "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
     "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
     "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
@@ -429,6 +383,38 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
 
   if(VLLM_GPU_LANG STREQUAL "CUDA")
+    SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
+
+    # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
+    set(CUTLASS_REVISION "v4.4.2")
+
+    # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
+    if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
+      set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
+    endif()
+
+    if(VLLM_CUTLASS_SRC_DIR)
+      if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
+        get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
+      endif()
+      message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
+      FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
+    else()
+      FetchContent_Declare(
+          cutlass
+          GIT_REPOSITORY https://github.com/nvidia/cutlass.git
+          # Please keep this in sync with CUTLASS_REVISION line above.
+          GIT_TAG ${CUTLASS_REVISION}
+          GIT_PROGRESS TRUE
+
+          # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
+          # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
+          # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
+          GIT_SHALLOW TRUE
+      )
+    endif()
+    FetchContent_MakeAvailable(cutlass)
+
     list(APPEND VLLM_STABLE_EXT_SRC
       "csrc/libtorch_stable/cuda_view.cu"
       "csrc/libtorch_stable/cuda_utils_kernels.cu"
@@ -929,7 +915,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       SRCS "${FP4_SM120_SRCS}"
       CUDA_ARCHS "${FP4_SM120_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
-    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
     message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}")
@@ -962,7 +947,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       SRCS "${FP4_SM100_SRCS}"
       CUDA_ARCHS "${FP4_SM100_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
-    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
     message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}")

cmake/external_projects/qutlass.cmake

@@ -60,6 +60,7 @@ endif()
 
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
   set(QUTLASS_SOURCES
+    csrc/qutlass_registration.cpp
     ${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp
     ${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu
     ${qutlass_SOURCE_DIR}/qutlass/csrc/gemm_ada.cu
@@ -78,8 +79,19 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
 
   if(CUTLASS_INCLUDE_DIR AND EXISTS "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h")
     list(APPEND QUTLASS_INCLUDES "${CUTLASS_INCLUDE_DIR}")
+    if(CUTLASS_TOOLS_UTIL_INCLUDE_DIR AND
+       EXISTS "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}/cutlass/util/packed_stride.hpp")
+      list(APPEND QUTLASS_INCLUDES "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}")
+    else()
+      get_filename_component(_qutlass_cutlass_root "${CUTLASS_INCLUDE_DIR}" DIRECTORY)
+      if(EXISTS "${_qutlass_cutlass_root}/tools/util/include/cutlass/util/packed_stride.hpp")
+        list(APPEND QUTLASS_INCLUDES "${_qutlass_cutlass_root}/tools/util/include")
+      endif()
+    endif()
   elseif(EXISTS "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include/cutlass/cutlass.h")
-    list(APPEND QUTLASS_INCLUDES "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include")
+    list(APPEND QUTLASS_INCLUDES
+      "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include"
+      "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/tools/util/include")
     message(STATUS "[QUTLASS] Using QuTLASS vendored CUTLASS headers (no vLLM CUTLASS detected).")
   else()
     message(FATAL_ERROR "[QUTLASS] CUTLASS headers not found. "
@@ -91,12 +103,23 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
     CUDA_ARCHS "${QUTLASS_ARCHS}"
   )
 
-  target_sources(_C PRIVATE ${QUTLASS_SOURCES})
-  target_include_directories(_C PRIVATE ${QUTLASS_INCLUDES})
-  target_compile_definitions(_C PRIVATE
+  # QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
+  # Keep it as its own extension (registers torch.ops._qutlass_C).
+  define_extension_target(
+    _qutlass_C
+    DESTINATION vllm
+    LANGUAGE ${VLLM_GPU_LANG}
+    SOURCES ${QUTLASS_SOURCES}
+    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+    ARCHITECTURES ${VLLM_GPU_ARCHES}
+    INCLUDE_DIRECTORIES ${QUTLASS_INCLUDES}
+    USE_SABI 3
+    WITH_SOABI)
+
+  target_compile_definitions(_qutlass_C PRIVATE
     QUTLASS_DISABLE_PYBIND=1
     TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
-  )
+    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
 
   set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
     $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -111,4 +134,5 @@ else()
       "[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in "
       "CUDA_ARCHS='${CUDA_ARCHS}'.")
   endif()
+  add_custom_target(_qutlass_C)
 endif()

csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu

@@ -268,9 +268,14 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba
   using TileShapeD = typename MlaSm100Type::TileShapeD;
   arguments.problem_shape =
       cute::make_tuple(TileShapeH{}, static_cast<int>(max_seq_len), TileShapeD{}, static_cast<int>(num_batches));
-  // Assumes device 0 when getting sm_count.
-  arguments.hw_info.sm_count =
-      sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count;
+  if (sm_count <= 0) {
+    int current_device = 0;
+    cudaGetDevice(&current_device);
+    arguments.hw_info.sm_count =
+        cutlass::KernelHardwareInfo::query_device_multiprocessor_count(current_device);
+  } else {
+    arguments.hw_info.sm_count = sm_count;
+  }
   arguments.split_kv = static_cast<int>(num_kv_splits);
   MlaSm100Type::Fmha::set_split_kv(arguments);
 

csrc/libtorch_stable/moe/moe_align_sum_kernels.cu

@@ -9,7 +9,7 @@
 #include <torch/headeronly/core/ScalarType.h>
 
 #include "../../cuda_compat.h"
-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "libtorch_stable/dispatch_utils.h"
 #include "libtorch_stable/torch_utils.h"
 

csrc/libtorch_stable/ops.h

@@ -2,9 +2,25 @@
 
 #include <torch/csrc/stable/library.h>
 #include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/util/Exception.h>
 
 #include <optional>
 #include <string>
+#include <vector>
+
+#include <torch/csrc/stable/ops.h>
+
+inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) {
+  // Ensure tensor is on CUDA
+  STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device");
+
+  // Get the raw data pointer
+  void* data_ptr = tensor.mutable_data_ptr();
+
+  /// Create a new tensor from the raw data pointer
+  return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(),
+                                  tensor.device(), tensor.scalar_type());
+}
 
 void per_token_group_quant_fp8(const torch::stable::Tensor& input,
                                torch::stable::Tensor& output_q,
@@ -371,6 +387,18 @@ void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void silu_and_mul_clamp(torch::stable::Tensor& out,
                         torch::stable::Tensor& input, double limit,
                         double alpha = 1.0, double beta = 0.0);
+
+void silu_and_mul_quant(torch::stable::Tensor& out,
+                        torch::stable::Tensor& input,
+                        torch::stable::Tensor& scale);
+
+void persistent_masked_m_silu_mul_quant(
+    const torch::stable::Tensor& input,              // (E, T, 2*H)
+    const torch::stable::Tensor& tokens_per_expert,  // (E)
+    torch::stable::Tensor& y_q,                      // (E, T, H) [OUT]
+    torch::stable::Tensor& y_s,  // (E, T, H//group_size) [OUT]
+    bool use_ue8m0);
+
 void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_tanh_and_mul(torch::stable::Tensor& out,

csrc/libtorch_stable/quantization/activation_kernels.cu

@@ -1,16 +1,12 @@
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "libtorch_stable/torch_utils.h"
 
 #include <cmath>
-#include "core/math.hpp"
-#include "../cuda_compat.h"
-#include "dispatch_utils.h"
 
+#include "libtorch_stable/core/math.hpp"
+#include "cuda_compat.h"
+#include "libtorch_stable/dispatch_utils.h"
 #include "quantization/w8a8/fp8/common.cuh"
 
-#include <c10/util/Float8_e4m3fn.h>
-
 #ifndef USE_ROCM
   #include <cuda_bf16.h>
   #include <cuda_fp16.h>
@@ -33,7 +29,6 @@ typedef __hip_fp8x4_e4m3_fnuz __nv_fp8x4_e4m3;
   #endif
 #endif
 
-#include "core/registration.h"
 namespace vllm {
 
 template <typename T>
@@ -564,41 +559,47 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
 }  // namespace vllm
 
 // Launch activation, gating, and quantize kernel.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL)                               \
-  int d = input.size(-1) / 2;                                               \
-  int64_t num_tokens = input.numel() / input.size(-1);                      \
-  dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4);     \
-  dim3 block(std::min(d, 512));                                             \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));         \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();             \
-  VLLM_DISPATCH_FLOATING_TYPES(                                             \
-      input.scalar_type(), "act_and_mul_kernel", [&] {                      \
-        VLLM_DISPATCH_FP8_TYPES(                                            \
-            out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] {  \
-              vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>,    \
-                                             fp8_t>                         \
-                  <<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(),       \
-                                               input.data_ptr<scalar_t>(),  \
-                                               scale.data_ptr<float>(), d); \
-            });                                                             \
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL)                             \
+  int d = input.size(-1) / 2;                                             \
+  int64_t num_tokens = input.numel() / input.size(-1);                    \
+  dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4);   \
+  dim3 block(std::min(d, 512));                                           \
+  const torch::stable::accelerator::DeviceGuard device_guard(             \
+      input.get_device_index());                                          \
+  const cudaStream_t stream =                                             \
+      get_current_cuda_stream(input.get_device_index());                  \
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                    \
+      input.scalar_type(), "act_and_mul_kernel", [&] {                    \
+        VLLM_STABLE_DISPATCH_FP8_TYPES(                                   \
+            out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] { \
+              vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>,  \
+                                             fp8_t>                       \
+                  <<<grid, block, 0, stream>>>(                           \
+                      out.mutable_data_ptr<fp8_t>(),                      \
+                      input.const_data_ptr<scalar_t>(),                   \
+                      scale.const_data_ptr<float>(), d);                  \
+            });                                                           \
       });
 
-void silu_and_mul_quant(torch::Tensor& out,    // [..., d]
-                        torch::Tensor& input,  // [..., 2 * d]
-                        torch::Tensor& scale) {
-  TORCH_CHECK(out.dtype() == torch::kFloat8_e4m3fn ||
-              out.dtype() == torch::kFloat8_e4m3fnuz);
-  TORCH_CHECK(input.dtype() == torch::kFloat16 ||
-              input.dtype() == torch::kBFloat16);
-  TORCH_CHECK(input.size(-1) % 2 == 0);
+void silu_and_mul_quant(torch::stable::Tensor& out,    // [..., d]
+                        torch::stable::Tensor& input,  // [..., 2 * d]
+                        torch::stable::Tensor& scale) {
+  STD_TORCH_CHECK(
+      out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
+      out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
+  STD_TORCH_CHECK(
+      input.scalar_type() == torch::headeronly::ScalarType::Half ||
+          input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "Input must be FP16 or BF16");
+  STD_TORCH_CHECK(input.size(-1) % 2 == 0);
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel);
 }
 
 void persistent_masked_m_silu_mul_quant(
-    const at::Tensor& input,              // (E, T, 2*H)
-    const at::Tensor& tokens_per_expert,  // (E)
-    at::Tensor& y_q,                      // (E, T, H) [OUT]
-    at::Tensor& y_s,                      // (E, T, H//group_size) [OUT]
+    const torch::stable::Tensor& input,              // (E, T, 2*H)
+    const torch::stable::Tensor& tokens_per_expert,  // (E)
+    torch::stable::Tensor& y_q,                      // (E, T, H) [OUT]
+    torch::stable::Tensor& y_s,  // (E, T, H//group_size) [OUT]
     bool cast_scale_ue8m0) {
 #ifndef USE_ROCM
 
@@ -606,14 +607,18 @@ void persistent_masked_m_silu_mul_quant(
   // fixed GROUP_SIZE of 128.
   static constexpr int GROUP_SIZE = 128;
 
-  TORCH_CHECK(input.dtype() == torch::kBFloat16);
-  TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn ||
-              y_q.dtype() == torch::kFloat8_e4m3fnuz);
-  TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
+  STD_TORCH_CHECK(input.scalar_type() ==
+                  torch::headeronly::ScalarType::BFloat16);
+  STD_TORCH_CHECK(
+      y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
+      y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
+  STD_TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
 
   bool const is_packed_ue8m0 =
-      (y_s.dtype() == torch::kInt32 && cast_scale_ue8m0);
-  TORCH_CHECK(y_s.dtype() == torch::kFloat32 || is_packed_ue8m0);
+      (y_s.scalar_type() == torch::headeronly::ScalarType::Int &&
+       cast_scale_ue8m0);
+  STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Float ||
+                  is_packed_ue8m0);
 
   using Idx_t = int64_t;
 
@@ -631,7 +636,7 @@ void persistent_masked_m_silu_mul_quant(
 
   int const NUM_GROUPS = H / GROUP_SIZE;
 
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
 
   // TODO: Get this from cuda_arch ?
   static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32;
@@ -643,18 +648,21 @@ void persistent_masked_m_silu_mul_quant(
     static constexpr int max_shared_mem_bytes =                                \
         GROUP_SIZE * 2 * STAGES * NUM_WARPS * 2;                               \
     dim3 grid(sms), block(THREAD_COUNT);                                       \
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(input));          \
-    VLLM_DISPATCH_FP8_TYPES(                                                   \
+    const torch::stable::accelerator::DeviceGuard device_guard(                \
+        input.get_device_index());                                             \
+    VLLM_STABLE_DISPATCH_FP8_TYPES(                                            \
         y_q.scalar_type(), "silu_mul_fp8_quant_deep_gemm_kernel", [&] {        \
           vllm::silu_mul_fp8_quant_deep_gemm_kernel<                           \
               BLOCK_COUNT, max_shared_mem_bytes, fp8_t, scale_t, THREAD_COUNT, \
               Idx_t, CEIL_UE8M0, GROUP_SIZE, STAGES>                           \
               <<<grid, block, max_shared_mem_bytes + (E + 1) * 16, stream>>>(  \
-                  reinterpret_cast<__nv_bfloat16*>(input.data_ptr()),          \
-                  (fp8_t*)y_q.data_ptr(),                                      \
-                  reinterpret_cast<scale_t*>(y_s.data_ptr()),                  \
-                  reinterpret_cast<int32_t*>(tokens_per_expert.data_ptr()), E, \
-                  T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e,       \
+                  reinterpret_cast<const __nv_bfloat16*>(                      \
+                      input.const_data_ptr()),                                 \
+                  y_q.mutable_data_ptr<fp8_t>(),                               \
+                  reinterpret_cast<scale_t*>(y_s.mutable_data_ptr()),          \
+                  reinterpret_cast<const int32_t*>(                            \
+                      tokens_per_expert.const_data_ptr()),                     \
+                  E, T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e,    \
                   stride_yq_t, stride_yq_h, STRIDE_YS_E, STRIDE_YS_T,          \
                   STRIDE_YS_G, STRIDE_YS_P, stride_counts_e);                  \
         });
@@ -679,7 +687,7 @@ void persistent_masked_m_silu_mul_quant(
   Idx_t stride_ys_g = y_s.stride(2);
   Idx_t stride_ys_p = 0;
   if (!cast_scale_ue8m0) {
-    TORCH_CHECK(!is_packed_ue8m0);
+    STD_TORCH_CHECK(!is_packed_ue8m0);
     LAUNCH_ON_H(float, stride_ys_e, stride_ys_t, stride_ys_g, stride_ys_p,
                 false);
     return;
@@ -692,8 +700,8 @@ void persistent_masked_m_silu_mul_quant(
     return;
   }
 
-  TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
-  TORCH_CHECK(y_s.dtype() == torch::kInt32);
+  STD_TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
+  STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Int);
 
   // Int32 packed ue8m0 scales tensor.
   // Let E, T, G be the number to experts, number of tokens and number of groups

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu

@@ -31,7 +31,7 @@
 
 #include "cutlass/util/packed_stride.hpp"
 
-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "core/batch_invariant.hpp"
 
 using namespace cute;

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu

@@ -31,7 +31,7 @@
 
 #include "cutlass/util/packed_stride.hpp"
 
-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "core/batch_invariant.hpp"
 
 using namespace cute;

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/cutlass_gemm_caller.cuh

@@ -19,7 +19,7 @@
 #include "cutlass/gemm/collective/collective_builder.hpp"
 #include "cutlass/util/packed_stride.hpp"
 
-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "libtorch_stable/cutlass_extensions/common.hpp"
 // clang-format on
 

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh

@@ -14,7 +14,7 @@
 #include "cutlass/epilogue/collective/collective_builder.hpp"
 #include "cutlass/gemm/collective/collective_builder.hpp"
 
-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "libtorch_stable/cutlass_extensions/common.hpp"
 // clang-format on
 

csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh

@@ -22,7 +22,7 @@
 #include "cutlass/epilogue/threadblock/fusion/visitors.hpp"
 #include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"
 
-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "libtorch_stable/cutlass_extensions/common.hpp"
 // clang-format on
 

csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu

@@ -301,8 +301,9 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
 
   const int sf_k_local = local_group_id % kGroupsPerBlockX;
   const int row_local = local_group_id / kGroupsPerBlockX;
-  const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
-  const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
+  // Rows on grid.x: mn scales with tokens and can exceed the 65535 grid.y cap.
+  const int sf_k_idx = blockIdx.y * kGroupsPerBlockX + sf_k_local;
+  const int mn_idx = blockIdx.x * kRowsPerBlock + row_local;
 
 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
   asm volatile("griddepcontrol.wait;");
@@ -496,14 +497,15 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
                   " is not a multiple of 4.");
   const int kx = GetGroupsPerBlockX(padded_groups_per_row);
   const int ry = 16 / kx;
-  const int64_t blocks_x = padded_groups_per_row / kx;
-  const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry;
+  const int64_t row_blocks = (tma_aligned_mn + ry - 1) / ry;
+  const int64_t sf_k_blocks = padded_groups_per_row / kx;
   const int num_threads = (kx * ry) * THREADS_PER_GROUP;
-  // CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs.
-  STD_TORCH_CHECK(blocks_x <= static_cast<int64_t>(INT32_MAX) &&
-                      blocks_y <= static_cast<int64_t>(INT32_MAX),
+  // CUDA caps grid.x at 2^31 - 1 and grid.y at 2^16 - 1 (65535).
+  constexpr int64_t kMaxGridDimYZ = 65535;
+  STD_TORCH_CHECK(row_blocks <= static_cast<int64_t>(INT32_MAX) &&
+                      sf_k_blocks <= kMaxGridDimYZ,
                   "per_token_group_quant_8bit_packed grid too large: (",
-                  blocks_x, ", ", blocks_y, ").");
+                  row_blocks, ", ", sf_k_blocks, ").");
 
   auto dst_type = output_q.scalar_type();
 
@@ -513,8 +515,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
   #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY)                         \
     do {                                                                       \
       cudaLaunchConfig_t config = {};                                          \
-      config.gridDim = dim3(static_cast<unsigned int>(blocks_x),               \
-                            static_cast<unsigned int>(blocks_y));              \
+      config.gridDim = dim3(static_cast<unsigned int>(row_blocks),             \
+                            static_cast<unsigned int>(sf_k_blocks));           \
       config.blockDim = dim3(num_threads);                                     \
       config.dynamicSmemBytes = 0;                                             \
       config.stream = stream;                                                  \
@@ -539,8 +541,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
 #else
   #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY)                         \
     do {                                                                       \
-      dim3 grid(static_cast<unsigned int>(blocks_x),                           \
-                static_cast<unsigned int>(blocks_y));                          \
+      dim3 grid(static_cast<unsigned int>(row_blocks),                         \
+                static_cast<unsigned int>(sf_k_blocks));                       \
       dim3 block(num_threads);                                                 \
       per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
                                                         RY>                    \

csrc/libtorch_stable/torch_bindings.cpp

@@ -34,6 +34,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   // TODO: Remove this once ROCm upgrade to torch 2.11.
   ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
 
+  // Note about marlin kernel 'workspace' arguments:
+  // Technically these should be mutable since they are modified by the kernel.
+  // But since they are set back to zero once the kernel is finished we can
+  // hand wave and say that they have no net effect.
+  //
+  // The reason to mark 'workspace' as immutable is so that they don't interfere
+  // with using ScalarType arguments in the ops. If they are marked as mutable,
+  // pytorch throws an assert in
+  // 'torch._higher_order_ops._register_effectful_op' that prevents these
+  // kernels from being torch.compile'd.
+  // See the following document for more info on custom types and ops that use
+  // custom types:
+  // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
+
   // Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
   ops.def(
       "machete_supported_schedules("
@@ -480,6 +494,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "Tensor workspace, int k, int max_seq_len) -> ()");
 
   // Activation ops
+  ops.def(
+      "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
+      "y_q, Tensor! y_s, bool use_ue8m0) -> ()");
+  ops.def("weak_ref_tensor(Tensor input) -> Tensor");
+
   // Activation function used in SwiGLU.
   ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
 
@@ -492,6 +511,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
       "float alpha=1.0, float beta=0.0) -> ()");
 
+  // SwiGLU activation with FP8 quantization.
+  ops.def(
+      "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
+
   // Activation function used in GeGLU with `none` approximation.
   ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
 
@@ -690,6 +713,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
 
   // Activation kernels (shared CUDA/ROCm)
+  ops.impl("persistent_masked_m_silu_mul_quant",
+           TORCH_BOX(&persistent_masked_m_silu_mul_quant));
+  ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor));
+  ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant));
   ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
   ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
   ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul));

csrc/ops.h

@@ -9,28 +9,6 @@
 
 #include <vector>
 
-torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
-  // Ensure tensor is on CUDA
-  if (!tensor.is_cuda()) {
-    throw std::runtime_error("Tensor must be on CUDA device");
-  }
-
-  // Get the raw data pointer
-  void* data_ptr = tensor.data_ptr();
-
-  // Get tensor sizes and strides
-  std::vector<int64_t> sizes = tensor.sizes().vec();
-  std::vector<int64_t> strides = tensor.strides().vec();
-
-  // Get tensor options (dtype, device)
-  auto options = tensor.options();
-
-  // Create a new tensor from the raw data pointer
-  auto new_tensor = torch::from_blob(data_ptr, sizes, strides, options);
-
-  return new_tensor;
-}
-
 // rms_norm and fused_add_rms_norm declarations also exist in
 // csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
 // because the CPU build still uses these torch::Tensor declarations.
@@ -53,16 +31,6 @@ void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
 void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit,
                         double alpha = 1.0, double beta = 0.0);
 
-void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
-                        torch::Tensor& scale);
-
-void persistent_masked_m_silu_mul_quant(
-    const at::Tensor& input,   // (E, T, 2*H)
-    const at::Tensor& counts,  // (E)
-    at::Tensor& y_q,           // (E, T, H) [OUT]
-    at::Tensor& y_s,           // (E, T, H//group_size) [OUT]
-    bool use_ue8m0);
-
 void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
 
 void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);

csrc/qutlass_registration.cpp

@@ -0,0 +1,5 @@
+#include "core/registration.h"
+
+// QuTLASS registers torch.ops._qutlass_C via TORCH_LIBRARY in bindings.cpp.
+// This stub lets Python import vllm._qutlass_C to trigger op registration.
+REGISTER_EXTENSION(_qutlass_C)

csrc/torch_bindings.cpp

@@ -20,17 +20,6 @@
 
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   // vLLM custom ops
-  //
-
-  ops.def(
-      "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
-      "y_q, Tensor! y_s,"
-      "bool use_ue8m0) -> ()");
-  ops.impl("persistent_masked_m_silu_mul_quant", torch::kCUDA,
-           &persistent_masked_m_silu_mul_quant);
-
-  ops.def("weak_ref_tensor(Tensor input) -> Tensor");
-  ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
 
 #ifdef USE_ROCM
   // TODO: Remove this once we upgrade to torch 2.11.
@@ -39,35 +28,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
   ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
            &get_cuda_view_from_cpu_tensor);
-#endif
-
-  // Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
-  ops.def(
-      "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
-  ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
-
-  // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
-  // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
-  // kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
-  // full-cache bf16/fp8 variants).
-
-  // Quantization ops
-#ifndef USE_ROCM
-
-  // Note about marlin kernel 'workspace' arguments:
-  // Technically these should be mutable since they are modified by the kernel.
-  // But since they are set back to zero once the kernel is finished we can
-  // hand wave and say that they have no net effect.
-  //
-  // The reason to mark 'workspace' as immutable is so that they don't interfere
-  // with using ScalarType arguments in the ops. If they are marked as mutable,
-  // pytorch throws an assert in
-  // 'torch._higher_order_ops._register_effectful_op' that prevents these
-  // kernels from being torch.compile'd.
-  // See the following document for more info on custom types and ops that use
-  // custom types:
-  // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
-
 #endif
 }
 

docs/contributing/model/basic.md

@@ -133,10 +133,10 @@ The model should inherit protocol `IsAttentionFree` and also implement class met
 For the mamba layers themselves, please use the [`MambaMixer`](../../../vllm/model_executor/layers/mamba/mamba_mixer.py) (for Mamba-1) or [`MambaMixer2`](../../../vllm/model_executor/layers/mamba/mamba_mixer2.py) (for Mamba-2) classes.
 The model should also be added to the `MODELS_CONFIG_MAP` dictionary in [vllm/model_executor/models/config.py](../../../vllm/model_executor/models/config.py) to ensure that the runtime defaults are optimized.
 
-For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`BambaForCausalLM`](../../../vllm/model_executor/models/bamba.py) (for an example of a model that uses Mamba-2 and attention together).
+For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`NemotronHForCausalLM`](../../../vllm/model_executor/models/nemotron_h.py) (for an example of a model that uses Mamba-2 and attention together).
 These models should follow the same instructions as case (1), but they should inherit protocol `IsHybrid` (instead of `IsAttentionFree`) and it is *not* necessary to add them to the `MODELS_CONFIG_MAP` (their runtime defaults will be inferred from the protocol).
 
-For case (3), we recommend looking at the implementation of [`MiniMaxText01ForCausalLM`](../../../vllm/model_executor/models/minimax_text_01.py) or [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which use custom "mamba-like" layers `MiniMaxText01LinearAttention` and `ShortConv` respectively.
+For case (3), we recommend looking at the implementation of [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which uses a custom "mamba-like" layer `ShortConv`.
 Please follow the same guidelines as case (2) for implementing these models.
 We use "mamba-like" to refer to layers that possess a state that is updated in-place, rather than being appended-to (like KV cache for attention).
 For implementing new custom mamba-like layers, one should inherit from `MambaBase` and implement the methods `get_state_dtype`, `get_state_shape` to calculate the data types and state shapes at runtime, as well as `mamba_type` and `get_attn_backend`.
@@ -144,5 +144,5 @@ It is also necessary to implement the "attention meta-data" class which handles
 Please see [`LinearAttentionMetadata`](../../../vllm/v1/attention/backends/linear_attn.py) or [`ShortConvAttentionMetadata`](../../../vllm/v1/attention/backends/short_conv_attn.py) for examples of this.
 It is also worth noting that we should update `MambaAttentionBackendEnum` in [`registry.py`](../../../vllm/v1/attention/backends/registry.py) when adding a new mamba backend.
 Finally, if one wants to support torch compile and CUDA graphs, it necessary to wrap the call to the mamba-like layer inside a custom op and register it.
-Please see the calls to `direct_register_custom_op` in [vllm/model_executor/models/minimax_text_01.py](../../../vllm/model_executor/models/minimax_text_01.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this.
+Please see the calls to `direct_register_custom_op` in [vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py](../../../vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this.
 The new custom op should then be added to the list `_attention_ops` in [vllm/config/compilation.py](../../../vllm/config/compilation.py) to ensure that piecewise CUDA graphs works as intended.

docs/design/attention_backends.md

@@ -170,8 +170,8 @@ Priority is **1 = highest** (tried first).
 | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
 | `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
-| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
+| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 7.x-9.x |
+| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
 | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
 | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |

docs/features/kv_offloading_usage.md

@@ -74,6 +74,7 @@ vllm serve <model> \
 | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
 | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
 | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
+| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
 | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
 
 ## Secondary Tiers

docs/features/tool_calling.md

@@ -321,15 +321,6 @@ For Qwen2.5, the chat template in tokenizer_config.json has already included sup
 
 Flags: `--tool-call-parser hermes`
 
-### MiniMax Models (`minimax_m1`)
-
-Supported models:
-
-* `MiniMaxAi/MiniMax-M1-40k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja))
-* `MiniMaxAi/MiniMax-M1-80k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja))
-
-Flags: `--tool-call-parser minimax --chat-template examples/tool_chat_template_minimax_m1.jinja`
-
 ### DeepSeek-V3 Models (`deepseek_v3`)
 
 Supported models:

docs/models/pooling_models/token_embed.md

@@ -61,7 +61,7 @@ Models of any architecture can be converted into embedding models using `--conve
 | `ColModernVBertForRetrieval` | ColModernVBERT | T / I | `ModernVBERT/colmodernvbert-merged` | | |
 | `ColPaliForRetrieval` | ColPali | T / I | `vidore/colpali-v1.3-hf` | | |
 | `ColQwen3` | Qwen3-VL | T / I | `TomoroAI/tomoro-colqwen3-embed-4b`, `TomoroAI/tomoro-colqwen3-embed-8b` | | |
-| `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3` | | |
+| `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3`, `vultr/VultronRetrieverPrime-Qwen3.5-8B` | | |
 | `OpsColQwen3Model` | Qwen3-VL | T / I | `OpenSearch-AI/Ops-Colqwen3-4B`, `OpenSearch-AI/Ops-Colqwen3-8B` | | |
 | `Qwen3VLNemotronEmbedModel` | Qwen3-VL | T / I | `nvidia/nemotron-colembed-vl-4b-v2`, `nvidia/nemotron-colembed-vl-8b-v2` | ✅︎ | ✅︎ |
 | `*ForConditionalGeneration`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | \* | N/A | \* | \* |

docs/models/supported_models.md

@@ -441,7 +441,6 @@ th {
 | `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
 | `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
 | `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
-| `MiniMaxForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01-hf`, etc. | | |
 | `MiniMaxM2ForCausalLM` | MiniMax-M2, MiniMax-M2.1 | `MiniMaxAI/MiniMax-M2`, etc. | ✅︎ | ✅︎ |
 | `MistralForCausalLM` | Ministral-3, Mistral, Mistral-Instruct | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-7B-v0.1`, `mistralai/Mistral-7B-Instruct-v0.1`, etc. | ✅︎ | ✅︎ |
 | `MistralLarge3ForCausalLM` | Mistral-Large-3-675B-Base-2512, Mistral-Large-3-675B-Instruct-2512 | `mistralai/Mistral-Large-3-675B-Base-2512`, `mistralai/Mistral-Large-3-675B-Instruct-2512`, etc. | ✅︎ | ✅︎ |
@@ -487,8 +486,6 @@ th {
 | `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
 | `TeleChat3ForCausalLM` | TeleChat3 | `Tele-AI/TeleChat3-36B-Thinking`, `Tele-AI/TeleChat3-Coder-36B-Thinking`, etc. | ✅︎ | ✅︎ |
 | `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ |
-| `MiniMaxM1ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-M1-40k`, `MiniMaxAI/MiniMax-M1-80k`, etc. | | |
-| `MiniMaxText01ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01`, etc. | | |
 | `Zamba2ForCausalLM` | Zamba2 | `Zyphra/Zamba2-7B-instruct`, `Zyphra/Zamba2-2.7B-instruct`, `Zyphra/Zamba2-1.2B-instruct`, etc. | | |
 
 !!! note
@@ -595,6 +592,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
 | `MiniCPMO` | MiniCPM-O | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>E+</sup> | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ |
 | `MiniCPMV` | MiniCPM-V | T + I<sup>E+</sup> + V<sup>E+</sup> | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, etc. | ✅︎ | |
+| `MiniMaxM3SparseForConditionalGeneration` | MiniMax-M3 | T + I<sup>+</sup> + V<sup>+</sup> | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | |
 | `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
 | `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
 | `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |

docs/usage/v1_guide.md

@@ -128,7 +128,7 @@ Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaFor
 Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., `BambaForCausalLM`,
 `Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`, `Plamo2ForCausalLM`).
 
-Hybrid models with mechanisms different to Mamba are also supported (e.g, `MiniMaxText01ForCausalLM`, `MiniMaxM1ForCausalLM`, `Lfm2ForCausalLM`).
+Hybrid models with mechanisms different to Mamba are also supported (e.g, `Lfm2ForCausalLM`).
 
 Please note that prefix caching is not yet supported for any of the above models.
 

examples/generate/multimodal/vision_language_offline.py

@@ -1481,39 +1481,6 @@ def run_minicpmv(questions: list[str], modality: str) -> ModelRequestData:
     return run_minicpmv_base(questions, modality, "openbmb/MiniCPM-V-2_6")
 
 
-def run_minimax_vl_01(questions: list[str], modality: str) -> ModelRequestData:
-    assert modality == "image"
-
-    model_name = "MiniMaxAI/MiniMax-VL-01"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        max_num_seqs=2,
-        limit_mm_per_prompt={modality: 1},
-        trust_remote_code=True,
-        tensor_parallel_size=8,
-    )
-
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    messages = [
-        [
-            {
-                "role": "user",
-                "content": [{"type": "image"}, {"type": "text", "text": question}],
-            }
-        ]
-        for question in questions
-    ]
-    prompts = tokenizer.apply_chat_template(
-        messages, add_generation_prompt=True, tokenize=False
-    )
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompts=prompts,
-    )
-
-
 # Mistral-3 HF-format
 def run_mistral3(questions: list[str], modality: str) -> ModelRequestData:
     assert modality == "image"
@@ -2485,7 +2452,6 @@ model_example_map = {
     "mantis": run_mantis,
     "minicpmo": run_minicpmo,
     "minicpmv": run_minicpmv,
-    "minimax_vl_01": run_minimax_vl_01,
     "mistral3": run_mistral3,
     "molmo": run_molmo,
     "molmo2": run_molmo2,

examples/pooling/score/colqwen3_5_rerank_online.py

@@ -7,11 +7,27 @@ ColQwen3.5 is a multi-modal ColBERT-style model based on Qwen3.5.
 It produces per-token embeddings and uses MaxSim scoring for retrieval
 and reranking. Supports both text and image inputs.
 
+Works for any ColQwen3.5 checkpoint, e.g. `athrael-soju/colqwen3.5-4.5B-v3`
+or `vultr/VultronRetrieverPrime-Qwen3.5-8B`.
+
 Start the server with:
-    vllm serve athrael-soju/colqwen3.5-4.5B --max-model-len 4096
+    vllm serve athrael-soju/colqwen3.5-4.5B-v3 --max-model-len 4096 \
+        --mm-processor-kwargs '{"min_pixels": 65536, "max_pixels": 1835008}'
 
 Then run this script:
     python colqwen3_5_rerank_online.py
+
+Parity note (matching the native colpali ColQwen3_5Processor pipeline):
+  - Visual-token budget: ColQwen3_5Processor uses max_num_visual_tokens=1792,
+    i.e. max_pixels = 1792 * (patch_size*merge_size)^2 = 1792 * 32^2 = 1835008
+    (with min_pixels = shortest_edge = 65536). Pass these via --mm-processor-kwargs
+    as above; the default budget gives fewer visual tokens and lower retrieval ndcg.
+  - When you build prompts yourself (token_embed), reproduce the processor exactly:
+      image (document): wrap in the instruction template
+        "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
+        "Describe the image.<|im_end|><|endoftext|>"
+      query: append the augmentation suffix  <text> + "<|endoftext|>" * 10
+    Omitting these reproduces a silent ~2.5 ndcg@10 drop vs the native pipeline.
 """
 
 import requests

examples/tool_chat_template_minimax_m1.jinja

@@ -1,91 +0,0 @@
-{{ '<begin_of_document>' -}}
-{%- if custom_tools is defined %}
-    {%- set tools = custom_tools %}
-{%- endif %}
-{%- if not tools is defined %}
-    {%- set tools = none %}
-{%- endif %}
-
-{#- Extract system message #}
-{% set ns = namespace(system_prompt='') -%}
-{%- if messages[0]['role'] == 'system' %}
-    {%- if messages[0]['content'] is string %}
-        {%- set ns.system_prompt = messages[0]['content']|trim %}
-    {%- else %}
-        {%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %}
-    {%- endif %}
-    {%- set messages = messages[1:] %}
-{%- else %}
-    {%- if tools is not none %}
-        {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
-    {%- else %}
-        {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
-    {%- endif %}
-{%- endif %}
-
-{#- System message #}
-{%- if ns.system_prompt != '' %}
-{{ '<beginning_of_sentence>system ai_setting=assistant\n' + ns.system_prompt + '<end_of_sentence>\n' -}}
-{%- endif %}
-
-{#- Tools configuration #}
-{%- if tools is not none %}
-{{ '<beginning_of_sentence>system tool_setting=tools\nYou are provided with these tools:\n<tools>\n' -}}
-{%- for tool in tools %}
-{{ tool | tojson ~ '\n' -}}
-{%- endfor %}
-{{ '</tools>\n\nIf you need to call tools, please respond with <tool_calls></tool_calls> XML tags, and provide tool-name and json-object of arguments, following the format below:\n<tool_calls>\n{"name": <tool-name>, "arguments": <args-json-object>}\n...\n</tool_calls><end_of_sentence>\n' -}}
-{%- endif %}
-
-{#- Process messages #}
-{%- for message in messages %}
-    {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
-        {%- if message['role'] == 'user' %}
-{{ '<beginning_of_sentence>user name=user\n' -}}
-{%- if message['content'] is string %}
-{{ message['content']|trim -}}
-{%- else %}
-{%- for content in message['content'] %}
-{%- if content['type'] == 'text' %}
-{{ content['text']|trim -}}
-{%- endif %}
-{%- endfor %}
-{%- endif %}
-{{ '<end_of_sentence>\n' -}}
-        {%- elif message['role'] == 'assistant' %}
-{{ '<beginning_of_sentence>ai name=assistant\n' -}}
-{%- if message['content'] is string %}
-{{ message['content']|trim -}}
-{%- else %}
-{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %}
-{{ content['text']|trim -}}
-{%- endfor %}
-{%- endif %}
-{{ '<end_of_sentence>\n' -}}
-        {%- endif %}
-    {%- elif 'tool_calls' in message %}
-{{ '<beginning_of_sentence>ai name=assistant\n<tool_calls>\n' -}}
-{%- for tool_call in message.tool_calls %}
-{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}}
-{%- endfor %}
-{{ '</tool_calls><end_of_sentence>\n' -}}
-    {%- elif message.role == "tool" or message.role == "ipython" %}
-{{ '<beginning_of_sentence>tool name=tools\n' -}}
-{%- if message.content is string %}
-{{ 'tool result: ' + message.content + '\n\n' -}}
-{%- else %}
-{%- for content in message['content'] %}
-{%- if content['type'] == 'text' %}
-{{ 'tool result: ' + content['text'] + '\n\n' -}}
-{%- elif content.get('name') %}
-{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}}
-{%- endif %}
-{%- endfor %}
-{%- endif %}
-{{ '<end_of_sentence>\n' -}}
-    {%- endif %}
-{%- endfor %}
-
-{%- if add_generation_prompt %}
-{{ '<beginning_of_sentence>ai name=assistant\n' -}}
-{%- endif %}

requirements/test/cuda.in

@@ -40,7 +40,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
 transformers==5.5.3
 tokenizers==0.22.2
-schemathesis>=3.39.15 # Required for openai schema test.
+schemathesis>=4.0.0 # Required for openai schema test.
 # quantization
 bitsandbytes==0.49.2
 buildkite-test-collector==0.1.9

requirements/test/nightly-torch.txt

@@ -31,7 +31,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
 transformers==5.5.3
 tokenizers==0.22.2
-schemathesis>=3.39.15 # Required for openai schema test.
+schemathesis>=4.0.0 # Required for openai schema test.
 # quantization
 bitsandbytes>=0.49.2
 buildkite-test-collector==0.1.9

requirements/test/rocm.in

@@ -39,7 +39,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
 transformers==5.5.3
 tokenizers==0.22.2
-schemathesis>=3.39.15 # Required for openai schema test
+schemathesis>=4.0.0 # Required for openai schema test
 # quantization
 bitsandbytes==0.49.2
 buildkite-test-collector==0.1.9

rust/src/chat/src/renderer/hf/format.rs

@@ -386,7 +386,6 @@ mod tests {
             tool_chat_template_llama3.2_pythonic.jinja         => String
             tool_chat_template_llama4_json.jinja               => OpenAi
             tool_chat_template_llama4_pythonic.jinja           => OpenAi
-            tool_chat_template_minimax_m1.jinja                => OpenAi
             tool_chat_template_mistral.jinja                   => String
             tool_chat_template_mistral3.jinja                  => OpenAi
             tool_chat_template_mistral_parallel.jinja          => String

rust/src/chat/tests/chat.rs

@@ -16,7 +16,8 @@ use vllm_engine_core_client::protocol::logprobs::{
     Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob,
 };
 use vllm_engine_core_client::protocol::{
-    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, StopReason,
+    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest, LogprobsCount,
+    StopReason,
 };
 use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
 use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
@@ -1387,8 +1388,8 @@ async fn chat_stream_and_collect_preserve_prompt_and_sample_logprobs() {
     .await;
 
     let mut request = sample_request("chat-logprobs");
-    request.sampling_params.logprobs = Some(1);
-    request.sampling_params.prompt_logprobs = Some(1);
+    request.sampling_params.logprobs = Some(LogprobsCount::Top(1));
+    request.sampling_params.prompt_logprobs = Some(LogprobsCount::Top(1));
 
     let mut stream = chat.chat(request.clone()).await.unwrap();
     match next_semantic(&mut stream).await.unwrap().unwrap() {

rust/src/chat/tests/templates/vllm_examples/tool_chat_template_minimax_m1.jinja

@@ -1,91 +0,0 @@
-{{ '<begin_of_document>' -}}
-{%- if custom_tools is defined %}
-    {%- set tools = custom_tools %}
-{%- endif %}
-{%- if not tools is defined %}
-    {%- set tools = none %}
-{%- endif %}
-
-{#- Extract system message #}
-{% set ns = namespace(system_prompt='') -%}
-{%- if messages[0]['role'] == 'system' %}
-    {%- if messages[0]['content'] is string %}
-        {%- set ns.system_prompt = messages[0]['content']|trim %}
-    {%- else %}
-        {%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %}
-    {%- endif %}
-    {%- set messages = messages[1:] %}
-{%- else %}
-    {%- if tools is not none %}
-        {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
-    {%- else %}
-        {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
-    {%- endif %}
-{%- endif %}
-
-{#- System message #}
-{%- if ns.system_prompt != '' %}
-{{ '<beginning_of_sentence>system ai_setting=assistant\n' + ns.system_prompt + '<end_of_sentence>\n' -}}
-{%- endif %}
-
-{#- Tools configuration #}
-{%- if tools is not none %}
-{{ '<beginning_of_sentence>system tool_setting=tools\nYou are provided with these tools:\n<tools>\n' -}}
-{%- for tool in tools %}
-{{ tool | tojson ~ '\n' -}}
-{%- endfor %}
-{{ '</tools>\n\nIf you need to call tools, please respond with <tool_calls></tool_calls> XML tags, and provide tool-name and json-object of arguments, following the format below:\n<tool_calls>\n{"name": <tool-name>, "arguments": <args-json-object>}\n...\n</tool_calls><end_of_sentence>\n' -}}
-{%- endif %}
-
-{#- Process messages #}
-{%- for message in messages %}
-    {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
-        {%- if message['role'] == 'user' %}
-{{ '<beginning_of_sentence>user name=user\n' -}}
-{%- if message['content'] is string %}
-{{ message['content']|trim -}}
-{%- else %}
-{%- for content in message['content'] %}
-{%- if content['type'] == 'text' %}
-{{ content['text']|trim -}}
-{%- endif %}
-{%- endfor %}
-{%- endif %}
-{{ '<end_of_sentence>\n' -}}
-        {%- elif message['role'] == 'assistant' %}
-{{ '<beginning_of_sentence>ai name=assistant\n' -}}
-{%- if message['content'] is string %}
-{{ message['content']|trim -}}
-{%- else %}
-{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %}
-{{ content['text']|trim -}}
-{%- endfor %}
-{%- endif %}
-{{ '<end_of_sentence>\n' -}}
-        {%- endif %}
-    {%- elif 'tool_calls' in message %}
-{{ '<beginning_of_sentence>ai name=assistant\n<tool_calls>\n' -}}
-{%- for tool_call in message.tool_calls %}
-{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}}
-{%- endfor %}
-{{ '</tool_calls><end_of_sentence>\n' -}}
-    {%- elif message.role == "tool" or message.role == "ipython" %}
-{{ '<beginning_of_sentence>tool name=tools\n' -}}
-{%- if message.content is string %}
-{{ 'tool result: ' + message.content + '\n\n' -}}
-{%- else %}
-{%- for content in message['content'] %}
-{%- if content['type'] == 'text' %}
-{{ 'tool result: ' + content['text'] + '\n\n' -}}
-{%- elif content.get('name') %}
-{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}}
-{%- endif %}
-{%- endfor %}
-{%- endif %}
-{{ '<end_of_sentence>\n' -}}
-    {%- endif %}
-{%- endfor %}
-
-{%- if add_generation_prompt %}
-{{ '<beginning_of_sentence>ai name=assistant\n' -}}
-{%- endif %}

rust/src/cmd/src/cli.rs

@@ -20,6 +20,7 @@ use serde_with::{DefaultOnNull, OneOrMany, serde_as};
 use thiserror_ext::AsReport as _;
 use uuid::Uuid;
 use vllm_engine_core_client::TransportMode;
+use vllm_engine_core_client::protocol::LogprobsCount;
 use vllm_managed_engine::ManagedEngineConfig;
 use vllm_managed_engine::cli::{ManagedEngineArgs, repartition_managed_engine_args};
 use vllm_server::{
@@ -136,9 +137,9 @@ pub struct SharedRuntimeArgs {
     pub max_model_len: Option<u32>,
     /// Maximum number of log probabilities to return when `logprobs` is
     /// specified in sampling parameters. `-1` means no cap.
-    #[arg(long, value_parser = clap::value_parser!(i32).range(-1..), allow_negative_numbers = true)]
+    #[arg(long, allow_negative_numbers = true)]
     #[serde(default)]
-    pub max_logprobs: Option<i32>,
+    pub max_logprobs: Option<LogprobsCount>,
     /// TCP port for the gRPC Generate service. When not set, no gRPC server is
     /// started.
     #[arg(long)]
@@ -529,7 +530,7 @@ impl ServeArgs {
         self.managed_engine.clone().into_config(
             self.runtime.model.clone(),
             self.runtime.max_model_len,
-            self.runtime.max_logprobs,
+            self.runtime.max_logprobs.map(managed_max_logprobs_to_i32),
             self.runtime.language_model_only,
             self.runtime.disable_log_stats,
             self.runtime.shutdown_timeout,
@@ -555,5 +556,9 @@ fn frontend_ipc_addresses() -> (String, String) {
     )
 }
 
+fn managed_max_logprobs_to_i32(count: LogprobsCount) -> i32 {
+    i32::try_from(count).expect("max_logprobs is parsed through i32")
+}
+
 #[cfg(test)]
 mod tests;

rust/src/cmd/src/cli/tests.rs

@@ -1,5 +1,6 @@
 use expect_test::expect;
 use vllm_engine_core_client::TransportMode;
+use vllm_engine_core_client::protocol::LogprobsCount;
 use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection};
 
 use super::{Cli, Command};
@@ -165,10 +166,10 @@ fn serve_args_forward_max_logprobs_to_frontend_and_managed_engine() {
     let Command::Serve(args) = cli.command else {
         panic!("expected serve args");
     };
-    assert_eq!(args.runtime.max_logprobs, Some(-1));
+    assert_eq!(args.runtime.max_logprobs, Some(LogprobsCount::All));
 
     let frontend_config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
-    assert_eq!(frontend_config.max_logprobs, Some(-1));
+    assert_eq!(frontend_config.max_logprobs, Some(LogprobsCount::All));
 
     let engine_config = args.to_managed_engine_config(5555);
     assert_eq!(engine_config.python_args, vec!["--max-logprobs", "-1"]);
@@ -529,7 +530,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() {
     assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32);
     assert!(args.runtime.language_model_only);
     assert_eq!(args.runtime.max_model_len, Some(8192));
-    assert_eq!(args.runtime.max_logprobs, Some(-1));
+    assert_eq!(args.runtime.max_logprobs, Some(LogprobsCount::All));
     assert_eq!(args.runtime.shutdown_timeout, 3);
 }
 

rust/src/engine-core-client/examples/external_engine_logprobs.rs

@@ -6,7 +6,7 @@ use futures::StreamExt as _;
 use tokio::time::timeout;
 use tracing_subscriber::EnvFilter;
 use vllm_engine_core_client::protocol::{
-    EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams,
+    EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams, LogprobsCount,
 };
 use vllm_engine_core_client::{
     EngineCoreClient, EngineCoreClientConfig, EngineCoreStreamOutput, TransportMode,
@@ -33,10 +33,10 @@ struct Args {
     output_timeout_secs: u64,
     #[arg(long, default_value_t = 1)]
     max_tokens: u32,
-    #[arg(long, default_value_t = 2)]
-    logprobs: i32,
-    #[arg(long, default_value_t = 1)]
-    prompt_logprobs: i32,
+    #[arg(long, default_value_t = LogprobsCount::Top(2), allow_negative_numbers = true)]
+    logprobs: LogprobsCount,
+    #[arg(long, default_value_t = LogprobsCount::Top(1), allow_negative_numbers = true)]
+    prompt_logprobs: LogprobsCount,
     #[arg(long, default_value_t = 96)]
     prompt_repeats: usize,
 }
@@ -64,8 +64,8 @@ fn build_request(
     request_id: String,
     prompt_token_ids: Vec<u32>,
     max_tokens: u32,
-    logprobs: i32,
-    prompt_logprobs: i32,
+    logprobs: LogprobsCount,
+    prompt_logprobs: LogprobsCount,
     client_index: u32,
 ) -> EngineCoreRequest {
     EngineCoreRequest {

rust/src/engine-core-client/src/protocol/logprobs_count.rs

@@ -0,0 +1,135 @@
+use std::fmt;
+use std::str::FromStr;
+
+use serde::{Deserialize, Deserializer, Serialize, Serializer};
+
+/// Number of log probabilities requested for a token position.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum LogprobsCount {
+    /// Return the full model vocabulary.
+    All,
+    /// Return the top-N tokens by probability.
+    Top(u32),
+}
+
+impl LogprobsCount {
+    /// Expands the count to the actual number of logprobs to return, given the vocabulary size.
+    pub fn expanded(self, vocab_size: usize) -> usize {
+        match self {
+            Self::All => vocab_size,
+            Self::Top(count) => count as usize,
+        }
+    }
+}
+
+impl TryFrom<i32> for LogprobsCount {
+    type Error = String;
+
+    fn try_from(value: i32) -> Result<Self, Self::Error> {
+        match value {
+            -1 => Ok(Self::All),
+            value if value < -1 => Err(format!("must be non-negative or -1, got {value}")),
+            value => Ok(Self::Top(value as u32)),
+        }
+    }
+}
+
+impl TryFrom<LogprobsCount> for i32 {
+    type Error = String;
+
+    fn try_from(value: LogprobsCount) -> Result<Self, Self::Error> {
+        match value {
+            LogprobsCount::All => Ok(-1),
+            LogprobsCount::Top(count) => {
+                i32::try_from(count).map_err(|_| format!("must fit within i32, got {count}"))
+            }
+        }
+    }
+}
+
+impl FromStr for LogprobsCount {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let value = s
+            .parse::<i32>()
+            .map_err(|e| format!("must be an i32 integer, got {s:?}: {e}"))?;
+        Self::try_from(value)
+    }
+}
+
+impl fmt::Display for LogprobsCount {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::All => (-1).fmt(f),
+            Self::Top(count) => count.fmt(f),
+        }
+    }
+}
+
+impl Serialize for LogprobsCount {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        let value: i32 = (*self).try_into().map_err(serde::ser::Error::custom)?;
+        value.serialize(serializer)
+    }
+}
+
+impl<'de> Deserialize<'de> for LogprobsCount {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let value = i32::deserialize(deserializer)?;
+        Self::try_from(value).map_err(serde::de::Error::custom)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use rmpv::Value;
+
+    use super::*;
+    use crate::protocol::{decode_msgpack, encode_msgpack};
+
+    #[test]
+    fn logprobs_count_serializes_as_wire_integer() {
+        assert_eq!(serde_json::to_value(LogprobsCount::All).unwrap(), -1);
+        assert_eq!(serde_json::to_value(LogprobsCount::Top(3)).unwrap(), 3);
+    }
+
+    #[test]
+    fn logprobs_count_deserializes_wire_integer() {
+        assert_eq!(
+            serde_json::from_value::<LogprobsCount>(serde_json::json!(-1)).unwrap(),
+            LogprobsCount::All
+        );
+        assert_eq!(
+            serde_json::from_value::<LogprobsCount>(serde_json::json!(3)).unwrap(),
+            LogprobsCount::Top(3)
+        );
+        assert!(serde_json::from_value::<LogprobsCount>(serde_json::json!(-2)).is_err());
+        assert!(
+            serde_json::from_value::<LogprobsCount>(serde_json::json!(i64::from(i32::MAX) + 1))
+                .is_err()
+        );
+    }
+
+    #[test]
+    fn logprobs_count_decodes_msgpack_signed_and_unsigned() {
+        let mut encoded = Vec::new();
+        rmpv::encode::write_value(&mut encoded, &Value::from(-1)).unwrap();
+        assert_eq!(
+            decode_msgpack::<LogprobsCount>(&encoded).unwrap(),
+            LogprobsCount::All
+        );
+
+        let encoded = encode_msgpack(&LogprobsCount::Top(7)).unwrap();
+        assert_eq!(
+            decode_msgpack::<LogprobsCount>(&encoded).unwrap(),
+            LogprobsCount::Top(7)
+        );
+    }
+}

rust/src/engine-core-client/src/protocol/mod.rs

@@ -56,6 +56,7 @@ mod classified_outputs;
 pub mod dtype;
 pub mod handshake;
 pub mod logprobs;
+mod logprobs_count;
 pub mod lora;
 pub mod multimodal;
 pub mod stats;
@@ -66,6 +67,7 @@ pub use classified_outputs::{
 };
 pub use dtype::ModelDtype;
 pub use logprobs::decode_engine_core_outputs;
+pub use logprobs_count::LogprobsCount;
 
 /// Request types are encoded as single-byte protocol constants so they can be
 /// sent over the ZMQ socket without an extra encoding step.
@@ -277,14 +279,20 @@ pub struct EngineCoreSamplingParams {
     pub max_tokens: u32,
     /// Minimum number of tokens to generate before EOS or stop-token handling.
     pub min_tokens: u32,
+    /// Maximum number of reasoning ("thinking") tokens to emit before the
+    /// reasoning section is force-closed. `None` means unlimited; the
+    /// user-facing `-1` sentinel is normalized to `None` by the frontend before
+    /// reaching this DTO, so only non-negative values are sent. Enforced
+    /// engine-side (and only when a reasoning parser is configured).
+    pub thinking_token_budget: Option<u64>,
     /// Number of log probabilities to return per generated token.
     ///
-    /// `None` disables sample logprobs. `-1` requests the full vocabulary.
-    pub logprobs: Option<i32>,
+    /// `None` disables sample logprobs.
+    pub logprobs: Option<LogprobsCount>,
     /// Number of log probabilities to return per prompt token.
     ///
-    /// `None` disables prompt logprobs. `-1` requests the full vocabulary.
-    pub prompt_logprobs: Option<i32>,
+    /// `None` disables prompt logprobs.
+    pub prompt_logprobs: Option<LogprobsCount>,
     /// Minimum probability threshold for token sampling.
     pub min_p: f32,
     /// Frequency penalty applied by the sampler.
@@ -345,6 +353,7 @@ impl EngineCoreSamplingParams {
             seed: None,
             max_tokens: 65536,
             min_tokens: 0,
+            thinking_token_budget: None,
             logprobs: None,
             prompt_logprobs: None,
             min_p: 0.0,

rust/src/engine-core-client/src/tests/client.rs

@@ -150,6 +150,7 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest {
             top_k: 8,
             max_tokens: 32,
             min_tokens: 1,
+            thinking_token_budget: Some(256),
             stop_token_ids: vec![151643],
             eos_token_id: Some(151645),
             all_stop_token_ids: BTreeSet::from([151643, 151645]),
@@ -2502,6 +2503,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
             seed: None,
             max_tokens: 16,
             min_tokens: 0,
+            thinking_token_budget: None,
             logprobs: None,
             prompt_logprobs: None,
             min_p: 0.0,

rust/src/engine-core-client/src/tests/python_compat.py

@@ -39,6 +39,7 @@ class EngineCoreSamplingParams(msgspec.Struct, dict=True, omit_defaults=True):
     seed: int | None = None
     max_tokens: int = 16
     min_tokens: int = 0
+    thinking_token_budget: int | None = None
     min_p: float = 0.0
     frequency_penalty: float = 0.0
     presence_penalty: float = 0.0
@@ -122,6 +123,7 @@ request = EngineCoreRequest(
         seed=None,
         max_tokens=32,
         min_tokens=1,
+        thinking_token_budget=256,
         min_p=0.0,
         frequency_penalty=0.0,
         presence_penalty=0.0,

rust/src/server/src/config.rs

@@ -2,12 +2,13 @@ use std::collections::HashMap;
 use std::fmt;
 use std::time::Duration;
 
-use anyhow::{Result, bail};
+use anyhow::Result;
 use axum::http::{HeaderName, HeaderValue, Method};
 use educe::Educe;
 use serde::Serialize;
 use serde_json::Value;
 use vllm_chat::{ChatTemplateContentFormatOption, ParserSelection, RendererSelection};
+use vllm_engine_core_client::protocol::LogprobsCount;
 use vllm_engine_core_client::{CoordinatorMode as EngineCoreCoordinatorMode, TransportMode};
 
 /// How the HTTP server obtains its listening socket.
@@ -133,7 +134,7 @@ pub struct Config {
     pub chat_template_content_format: ChatTemplateContentFormatOption,
     /// Optional maximum number of top log probabilities accepted by the
     /// frontend. `None` delegates to the text layer default.
-    pub max_logprobs: Option<i32>,
+    pub max_logprobs: Option<LogprobsCount>,
     /// HTTP/API-server behavior switches.
     pub api_server_options: ApiServerOptions,
     /// CORS settings applied to every HTTP response.
@@ -158,15 +159,6 @@ impl Config {
     pub fn validate(&self) -> Result<()> {
         vllm_chat::validate_parser_overrides(&self.tool_call_parser, &self.reasoning_parser)?;
         self.cors.validate()?;
-        if let Some(max_logprobs) = self.max_logprobs
-            && max_logprobs < -1
-        {
-            bail!(
-                "max_logprobs must be non-negative or -1, got {}",
-                max_logprobs
-            );
-        }
-
         Ok(())
     }
 

rust/src/server/src/error.rs

@@ -103,6 +103,7 @@ fn is_request_validation_error(error: &vllm_text::Error) -> bool {
             | vllm_text::Error::EmptyPromptTokenIds { .. }
             | vllm_text::Error::Logprobs(_)
             | vllm_text::Error::OutOfVocab(_)
+            | vllm_text::Error::InvalidThinkingTokenBudget
             // An empty tokenized prompt detected later, at request prepare
             // time, surfaces through the transparent Llm wrapper.
             | vllm_text::Error::Llm(vllm_llm::Error::EmptyPromptTokenIds { .. })
@@ -127,6 +128,18 @@ mod tests {
         assert!(response.error.message.contains("9000"));
     }
 
+    #[test]
+    fn invalid_thinking_token_budget_maps_to_invalid_request() {
+        let api_error = text_submit_error(
+            "failed to submit completion request",
+            vllm_text::Error::InvalidThinkingTokenBudget,
+        );
+        assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
+        let response = api_error.to_error_response();
+        assert_eq!(response.error.error_type, "invalid_request_error");
+        assert!(response.error.message.contains("thinking_token_budget"));
+    }
+
     #[test]
     fn chat_wrapped_prompt_too_long_maps_to_invalid_request() {
         let error = vllm_chat::Error::Text(vllm_text::Error::PromptTooLong {

rust/src/server/src/grpc/convert.rs

@@ -3,7 +3,7 @@
 
 use tonic::Status;
 use uuid::Uuid;
-use vllm_engine_core_client::protocol::{StopReason, StructuredOutputsParams};
+use vllm_engine_core_client::protocol::{LogprobsCount, StopReason, StructuredOutputsParams};
 use vllm_text::{
     DecodedLogprobs, DecodedPromptLogprobs, FinishReason, Finished, Prompt, SamplingParams,
     TextDecodeOptions, TextRequest,
@@ -202,18 +202,22 @@ fn build_sampling_params(
 /// Map the proto `CandidateTokens` selector to a `(logprobs_count,
 /// logprob_token_ids)` pair.
 ///
-/// - `top_n(k)` → `(k, None)` — return top-k candidates by probability
-/// - `all` → `(-1, None)` — return the full vocabulary
+/// - `top_n(k)` → `(Top(k), None)` — return top-k candidates by probability
+/// - `all` → `(All, None)` — return the full vocabulary
 /// - `token_ids(n)` → `(1, Some(vec of n token ids))` — return logprobs for specific tokens (the
 ///   count `n` is stored in the proto as the number of token IDs that follow, but the actual IDs
 ///   are carried via `logprob_token_ids` on `SamplingParams`)
-/// - absent → `(1, None)` — just the sampled/scored token
-fn candidate_logprob_spec(candidates: Option<&pb::CandidateTokens>) -> (i32, Option<Vec<u32>>) {
+/// - absent → `(Top(1), None)` — just the sampled/scored token
+fn candidate_logprob_spec(
+    candidates: Option<&pb::CandidateTokens>,
+) -> (LogprobsCount, Option<Vec<u32>>) {
     match candidates.and_then(|c| c.select.as_ref()) {
-        Some(pb::candidate_tokens::Select::TopN(n)) => (*n as i32, None),
-        Some(pb::candidate_tokens::Select::All(true)) => (-1, None),
-        Some(pb::candidate_tokens::Select::TokenIds(ids)) => (1, Some(ids.ids.clone())),
-        _ => (1, None),
+        Some(pb::candidate_tokens::Select::TopN(n)) => (LogprobsCount::Top(*n), None),
+        Some(pb::candidate_tokens::Select::All(true)) => (LogprobsCount::All, None),
+        Some(pb::candidate_tokens::Select::TokenIds(ids)) => {
+            (LogprobsCount::Top(1), Some(ids.ids.clone()))
+        }
+        _ => (LogprobsCount::Top(1), None),
     }
 }
 

rust/src/server/src/routes/inference/generate/convert.rs

@@ -87,6 +87,7 @@ pub(super) fn prepare_generate_request(
 #[cfg(test)]
 mod tests {
     use serde_json::json;
+    use vllm_engine_core_client::protocol::LogprobsCount;
     use vllm_text::Prompt;
 
     use super::prepare_generate_request;
@@ -132,10 +133,13 @@ mod tests {
             Prompt::TokenIds(vec![11, 22, 33])
         );
         assert_eq!(prepared.text_request.sampling_params.max_tokens, Some(7));
-        assert_eq!(prepared.text_request.sampling_params.logprobs, Some(2));
+        assert_eq!(
+            prepared.text_request.sampling_params.logprobs,
+            Some(LogprobsCount::Top(2))
+        );
         assert_eq!(
             prepared.text_request.sampling_params.prompt_logprobs,
-            Some(1)
+            Some(LogprobsCount::Top(1))
         );
         assert!(prepared.text_request.sampling_params.ignore_eos);
         assert_eq!(prepared.text_request.priority, -3);
@@ -150,6 +154,33 @@ mod tests {
         );
     }
 
+    #[test]
+    fn prepare_generate_request_forwards_thinking_token_budget() {
+        let request: GenerateRequest = serde_json::from_value(json!({
+            "model": "Qwen/Qwen1.5-0.5B-Chat",
+            "token_ids": [11, 22, 33],
+            "sampling_params": {
+                "thinking_token_budget": 64
+            }
+        }))
+        .expect("parse request");
+
+        let prepared = prepare_generate_request(
+            request,
+            &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
+            ResolvedRequestContext::default(),
+        )
+        .expect("prepare");
+
+        // The raw inference route shares `vllm_text::SamplingParams`, so the
+        // field is carried through to lowering exactly like the OpenAI routes
+        // (normalization/validation then happens in `lower_sampling_params`).
+        assert_eq!(
+            prepared.text_request.sampling_params.thinking_token_budget,
+            Some(64)
+        );
+    }
+
     #[test]
     fn prepare_generate_request_gates_continuous_usage_on_include_usage() {
         let request: GenerateRequest = serde_json::from_value(json!({

rust/src/server/src/routes/inference/generate/validate.rs

@@ -34,16 +34,6 @@ pub(super) fn validate_request_compat(
         );
     }
 
-    if let Some(prompt_logprobs) = request.sampling_params.prompt_logprobs
-        && prompt_logprobs < 0
-        && prompt_logprobs != -1
-    {
-        bail_invalid_request!(
-            param = "sampling_params",
-            "`prompt_logprobs` must be a non-negative value or -1."
-        );
-    }
-
     Ok(())
 }
 

rust/src/server/src/routes/openai/chat_completions/convert.rs

@@ -4,6 +4,7 @@ use vllm_chat::{
     ChatMessage as VllmChatMessage, ChatOptions, ChatRequest, ChatTool, ChatToolChoice,
     GenerationPromptMode, SamplingParams,
 };
+use vllm_engine_core_client::protocol::LogprobsCount;
 
 use super::types::ChatCompletionRequest;
 use super::validate;
@@ -94,7 +95,7 @@ pub(super) fn prepare_chat_request(
 
     // Auto-enable prompt logprobs for non-streaming echo, matching Python vLLM's
     // behavior.
-    let top_logprobs = request.top_logprobs.unwrap_or(0);
+    let top_logprobs = request.top_logprobs.unwrap_or(LogprobsCount::Top(0));
     let prompt_logprobs = request
         .prompt_logprobs
         .or((request.echo && !request.stream).then_some(top_logprobs));
@@ -115,6 +116,7 @@ pub(super) fn prepare_chat_request(
             seed: request.seed,
             max_tokens: request.max_completion_tokens,
             min_tokens: request.min_tokens,
+            thinking_token_budget: request.thinking_token_budget,
             logprobs: request.logprobs.then_some(top_logprobs),
             prompt_logprobs,
             min_p: request.min_p,
@@ -377,6 +379,7 @@ mod tests {
         ChatTool as VllmChatTool, ChatToolChoice, GenerationPromptMode,
         SamplingParams as VllmSamplingParams,
     };
+    use vllm_engine_core_client::protocol::LogprobsCount;
     use vllm_text::output::TextDecodeOptions;
 
     use super::prepare_chat_request;
@@ -613,6 +616,31 @@ mod tests {
         assert_eq!(prepared.chat_request.sampling_params, expected);
     }
 
+    #[test]
+    fn prepare_chat_request_passes_through_thinking_token_budget() {
+        let prepare = |budget: Option<i64>| {
+            prepare_chat_request(
+                ChatCompletionRequest {
+                    thinking_token_budget: budget,
+                    ..base_request()
+                },
+                &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
+                ResolvedRequestContext::default(),
+            )
+            .expect("request is valid")
+            .chat_request
+            .sampling_params
+            .thinking_token_budget
+        };
+
+        // The convert layer forwards the raw value verbatim (including the `-1`
+        // "unlimited" sentinel); normalization/validation happens during
+        // lowering (see `vllm_text::lower`).
+        assert_eq!(prepare(Some(64)), Some(64));
+        assert_eq!(prepare(Some(-1)), Some(-1));
+        assert_eq!(prepare(None), None);
+    }
+
     #[test]
     fn prepare_chat_request_accepts_developer_messages() {
         let request = ChatCompletionRequest {
@@ -941,7 +969,7 @@ mod tests {
         let request = ChatCompletionRequest {
             stream: false,
             logprobs: true,
-            prompt_logprobs: Some(2),
+            prompt_logprobs: Some(LogprobsCount::Top(2)),
             ..base_request()
         };
 
@@ -954,10 +982,13 @@ mod tests {
 
         assert!(prepared.options.requested_logprobs);
         assert!(prepared.options.include_prompt_logprobs);
-        assert_eq!(prepared.chat_request.sampling_params.logprobs, Some(0));
+        assert_eq!(
+            prepared.chat_request.sampling_params.logprobs,
+            Some(LogprobsCount::Top(0))
+        );
         assert_eq!(
             prepared.chat_request.sampling_params.prompt_logprobs,
-            Some(2)
+            Some(LogprobsCount::Top(2))
         );
     }
 
@@ -965,7 +996,7 @@ mod tests {
     fn prepare_chat_request_keeps_prompt_logprobs_independent_from_echo() {
         let request = ChatCompletionRequest {
             logprobs: true,
-            top_logprobs: Some(3),
+            top_logprobs: Some(LogprobsCount::Top(3)),
             echo: true,
             ..base_request()
         };
@@ -977,7 +1008,10 @@ mod tests {
         )
         .expect("request is valid");
 
-        assert_eq!(prepared.chat_request.sampling_params.logprobs, Some(3));
+        assert_eq!(
+            prepared.chat_request.sampling_params.logprobs,
+            Some(LogprobsCount::Top(3))
+        );
         assert_eq!(prepared.chat_request.sampling_params.prompt_logprobs, None);
         assert!(!prepared.options.include_prompt_logprobs);
     }

rust/src/server/src/routes/openai/chat_completions/types.rs

@@ -6,6 +6,7 @@ use serde_json::Value;
 use serde_with::SerializeDisplay;
 use validator::Validate;
 use vllm_chat::ReasoningEffort;
+use vllm_engine_core_client::protocol::LogprobsCount;
 
 use crate::routes::openai::utils::structured_outputs::ResponseFormat;
 use crate::routes::openai::utils::types::{
@@ -44,10 +45,8 @@ pub struct ChatCompletionRequest {
     #[serde(default)]
     pub logprobs: bool,
 
-    /// An integer specifying the number of most likely tokens to return
-    /// -1 means return all
-    #[validate(range(min = -1))]
-    pub top_logprobs: Option<i32>,
+    /// Number of most likely tokens to return. `-1` means return full vocab.
+    pub top_logprobs: Option<LogprobsCount>,
 
     /// Deprecated: Replaced by max_completion_tokens
     #[deprecated(note = "Use max_completion_tokens instead")]
@@ -155,8 +154,8 @@ pub struct ChatCompletionRequest {
     /// Truncate prompt tokens to this length
     pub truncate_prompt_tokens: Option<i64>,
 
-    /// Number of prompt logprobs to return
-    pub prompt_logprobs: Option<i32>,
+    /// Number of prompt logprobs to return. `-1` means return full vocab.
+    pub prompt_logprobs: Option<LogprobsCount>,
 
     /// Restrict output to these token IDs only
     pub allowed_token_ids: Option<Vec<u32>>,
@@ -165,8 +164,10 @@ pub struct ChatCompletionRequest {
     pub bad_words: Option<Vec<String>>,
 
     // -------- Extra vLLM Parameters --------
-    /// Token budget for reasoning/thinking
-    pub thinking_token_budget: Option<u32>,
+    /// Token budget for reasoning/thinking. Accepts a non-negative integer, or
+    /// `-1` for unlimited (mirroring the Python frontend, which normalizes `-1`
+    /// to "no budget").
+    pub thinking_token_budget: Option<i64>,
 
     /// Whether to include reasoning content in the response
     #[serde(default = "default_true")]

rust/src/server/src/routes/openai/chat_completions/validate.rs

@@ -1,6 +1,7 @@
 use super::types::ChatCompletionRequest;
 use crate::error::{ApiError, bail_invalid_request};
 use crate::routes::openai::utils::types::{ChatMessage, Tool, ToolChoice, ToolChoiceValue};
+use vllm_engine_core_client::protocol::LogprobsCount;
 
 /// Enforce the minimal compatibility contract for the Rust OpenAI server.
 pub(super) fn validate_request_compat(
@@ -30,14 +31,12 @@ pub(super) fn validate_request_compat(
     }
 
     if let Some(prompt_logprobs) = request.prompt_logprobs {
-        if prompt_logprobs < 0 && prompt_logprobs != -1 {
-            bail_invalid_request!(
-                param = "prompt_logprobs",
-                "prompt_logprobs must be a non-negative value or -1."
-            );
-        }
-
-        if request.stream && (prompt_logprobs > 0 || prompt_logprobs == -1) {
+        if request.stream
+            && matches!(
+                prompt_logprobs,
+                LogprobsCount::All | LogprobsCount::Top(1..)
+            )
+        {
             bail_invalid_request!(
                 param = "prompt_logprobs",
                 "prompt_logprobs are not available when stream=true."
@@ -108,11 +107,6 @@ pub(super) fn validate_request_compat(
         "truncate_prompt_tokens",
         "truncate_prompt_tokens is not supported.",
     )?;
-    reject_non_default(
-        request.thinking_token_budget.as_ref(),
-        "thinking_token_budget",
-        "thinking_token_budget is not supported.",
-    )?;
     reject_non_default(
         request.media_io_kwargs.as_ref(),
         "media_io_kwargs",
@@ -159,6 +153,7 @@ mod tests {
 
     use serde_json::json;
     use vllm_chat::ReasoningEffort;
+    use vllm_engine_core_client::protocol::LogprobsCount;
 
     use super::validate_request_compat;
     use crate::routes::openai::chat_completions::types::ChatCompletionRequest;
@@ -304,7 +299,7 @@ mod tests {
     #[test]
     fn validate_request_compat_rejects_top_logprobs_without_logprobs() {
         let request = ChatCompletionRequest {
-            top_logprobs: Some(0),
+            top_logprobs: Some(LogprobsCount::Top(0)),
             ..base_request()
         };
         assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
@@ -313,26 +308,26 @@ mod tests {
     #[test]
     fn validate_request_compat_rejects_streaming_prompt_logprobs_requests() {
         let request = ChatCompletionRequest {
-            prompt_logprobs: Some(1),
+            prompt_logprobs: Some(LogprobsCount::Top(1)),
             ..base_request()
         };
         assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
 
         let request = ChatCompletionRequest {
-            prompt_logprobs: Some(-1),
+            prompt_logprobs: Some(LogprobsCount::All),
             ..base_request()
         };
         assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
     }
 
     #[test]
-    fn validate_request_compat_rejects_invalid_prompt_logprobs_value() {
-        let request = ChatCompletionRequest {
-            stream: false,
-            prompt_logprobs: Some(-2),
-            ..base_request()
-        };
-        assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
+    fn chat_request_deserialization_rejects_invalid_prompt_logprobs_value() {
+        let result = serde_json::from_value::<ChatCompletionRequest>(json!({
+            "model": "Qwen/Qwen1.5-0.5B-Chat",
+            "messages": [{"role": "user", "content": "hello"}],
+            "prompt_logprobs": -2
+        }));
+        assert!(result.is_err());
     }
 
     #[test]

rust/src/server/src/routes/openai/completions/convert.rs

@@ -1,3 +1,4 @@
+use vllm_engine_core_client::protocol::LogprobsCount;
 use vllm_text::{SamplingParams, TextDecodeOptions, TextRequest};
 
 use super::types::CompletionRequest;
@@ -61,15 +62,7 @@ pub(super) fn prepare_completion_request(
         .map(|request| request.lora_name.clone())
         .unwrap_or_else(|| lora_resolution.model_names.first().cloned().unwrap_or_default());
 
-    let logprobs = match request.logprobs {
-        Some(logprobs) => Some(i32::try_from(logprobs).map_err(|_| {
-            ApiError::invalid_request(
-                "`logprobs` must fit within a signed 32-bit integer.".to_string(),
-                Some("logprobs"),
-            )
-        })?),
-        None => None,
-    };
+    let logprobs = request.logprobs.map(LogprobsCount::Top);
     let prompt_only = request.echo && request.max_tokens == Some(0);
     let prompt_logprobs =
         request.prompt_logprobs.or(if request.echo && (!request.stream || prompt_only) {
@@ -108,6 +101,7 @@ pub(super) fn prepare_completion_request(
             seed: request.seed,
             max_tokens,
             min_tokens: request.min_tokens,
+            thinking_token_budget: request.thinking_token_budget,
             logprobs,
             prompt_logprobs,
             min_p: request.min_p,
@@ -162,6 +156,7 @@ pub(super) fn prepare_completion_request(
 mod tests {
     use axum::http::HeaderMap;
     use serde_json::json;
+    use vllm_engine_core_client::protocol::LogprobsCount;
     use vllm_text::Prompt;
 
     use super::prepare_completion_request;
@@ -246,7 +241,10 @@ mod tests {
             Prompt::TokenIds(vec![11, 22, 33])
         );
         assert_eq!(prepared.text_request.sampling_params.max_tokens, Some(7));
-        assert_eq!(prepared.text_request.sampling_params.logprobs, Some(2));
+        assert_eq!(
+            prepared.text_request.sampling_params.logprobs,
+            Some(LogprobsCount::Top(2))
+        );
         assert_eq!(prepared.text_request.sampling_params.top_p, Some(0.9));
         assert_eq!(prepared.text_request.sampling_params.top_k, Some(42));
         assert_eq!(prepared.text_request.sampling_params.min_p, Some(0.1));
@@ -266,6 +264,34 @@ mod tests {
         assert!(!prepared.text_request.decode_options.skip_special_tokens);
     }
 
+    #[test]
+    fn prepare_completion_request_passes_through_thinking_token_budget() {
+        let prepare = |budget: serde_json::Value| {
+            let request: CompletionRequest = serde_json::from_value(json!({
+                "model": "Qwen/Qwen1.5-0.5B-Chat",
+                "prompt": "hello",
+                "thinking_token_budget": budget,
+            }))
+            .expect("parse request");
+            prepare_completion_request(
+                request,
+                &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
+                ResolvedRequestContext::default(),
+            )
+            .expect("prepare")
+            .text_request
+            .sampling_params
+            .thinking_token_budget
+        };
+
+        // The convert layer forwards the raw value verbatim (including the `-1`
+        // "unlimited" sentinel); normalization/validation happens during
+        // lowering (see `vllm_text::lower`).
+        assert_eq!(prepare(json!(64)), Some(64));
+        assert_eq!(prepare(json!(-1)), Some(-1));
+        assert_eq!(prepare(json!(null)), None);
+    }
+
     #[test]
     fn prepare_completion_request_maps_stream_usage_and_token_format_options() {
         let request: CompletionRequest = serde_json::from_value(json!({
@@ -381,10 +407,13 @@ mod tests {
         .expect("prepare");
 
         assert!(prepared.options.prompt_only);
-        assert_eq!(prepared.text_request.sampling_params.logprobs, Some(3));
+        assert_eq!(
+            prepared.text_request.sampling_params.logprobs,
+            Some(LogprobsCount::Top(3))
+        );
         assert_eq!(
             prepared.text_request.sampling_params.prompt_logprobs,
-            Some(3)
+            Some(LogprobsCount::Top(3))
         );
     }
 
@@ -406,10 +435,13 @@ mod tests {
         )
         .expect("prepare");
 
-        assert_eq!(prepared.text_request.sampling_params.logprobs, Some(3));
+        assert_eq!(
+            prepared.text_request.sampling_params.logprobs,
+            Some(LogprobsCount::Top(3))
+        );
         assert_eq!(
             prepared.text_request.sampling_params.prompt_logprobs,
-            Some(3)
+            Some(LogprobsCount::Top(3))
         );
     }
 
@@ -450,10 +482,13 @@ mod tests {
             ResolvedRequestContext::default(),
         )
         .expect("prepare");
-        assert_eq!(prepared.text_request.sampling_params.logprobs, Some(1));
+        assert_eq!(
+            prepared.text_request.sampling_params.logprobs,
+            Some(LogprobsCount::Top(1))
+        );
         assert_eq!(
             prepared.text_request.sampling_params.prompt_logprobs,
-            Some(2)
+            Some(LogprobsCount::Top(2))
         );
     }
 

rust/src/server/src/routes/openai/completions/types.rs

@@ -3,6 +3,7 @@ use std::collections::HashMap;
 use serde::{Deserialize, Serialize};
 use serde_json::{Map, Value};
 use validator::Validate;
+use vllm_engine_core_client::protocol::LogprobsCount;
 use vllm_text::Prompt;
 
 use crate::routes::openai::utils::types::{
@@ -131,8 +132,8 @@ pub struct CompletionRequest {
     /// Restrict output to these token IDs only
     pub allowed_token_ids: Option<Vec<u32>>,
 
-    /// Number of prompt logprobs to return
-    pub prompt_logprobs: Option<i32>,
+    /// Number of prompt logprobs to return. `-1` means return full vocab.
+    pub prompt_logprobs: Option<LogprobsCount>,
 
     // -------- Extra vLLM Parameters --------
     /// Whether to add special tokens (e.g. BOS) to the prompt
@@ -146,6 +147,11 @@ pub struct CompletionRequest {
     /// Additional kwargs for structured outputs
     pub structured_outputs: Option<Value>,
 
+    /// Token budget for reasoning/thinking. Accepts a non-negative integer, or
+    /// `-1` for unlimited (mirroring the Python frontend, which normalizes `-1`
+    /// to "no budget").
+    pub thinking_token_budget: Option<i64>,
+
     /// Request scheduling priority (lower means earlier; default 0)
     pub priority: Option<i32>,
 

rust/src/server/src/routes/openai/completions/validate.rs

@@ -1,3 +1,4 @@
+use vllm_engine_core_client::protocol::LogprobsCount;
 use vllm_text::Prompt;
 
 use super::types::CompletionRequest;
@@ -44,29 +45,18 @@ pub(super) fn validate_request_compat(
         bail_invalid_request!(param = "suffix", "suffix is not supported.");
     }
 
-    if let Some(logprobs) = request.logprobs
-        && logprobs > i32::MAX as u32
-    {
-        bail_invalid_request!(
-            param = "logprobs",
-            "`logprobs` must fit within a signed 32-bit integer."
-        );
-    }
-
     if let Some(prompt_logprobs) = request.prompt_logprobs {
-        if request.stream && (prompt_logprobs > 0 || prompt_logprobs == -1) {
+        if request.stream
+            && matches!(
+                prompt_logprobs,
+                LogprobsCount::All | LogprobsCount::Top(1..)
+            )
+        {
             bail_invalid_request!(
                 param = "prompt_logprobs",
                 "`prompt_logprobs` are not available when `stream=true`."
             );
         }
-
-        if prompt_logprobs < 0 && prompt_logprobs != -1 {
-            bail_invalid_request!(
-                param = "prompt_logprobs",
-                "`prompt_logprobs` must be a non-negative value or -1."
-            );
-        }
     }
 
     if request.use_beam_search {
@@ -101,6 +91,7 @@ pub(super) fn validate_request_compat(
 #[cfg(test)]
 mod tests {
     use serde_json::json;
+    use vllm_engine_core_client::protocol::LogprobsCount;
 
     use super::validate_request_compat;
     use crate::routes::openai::completions::types::CompletionRequest;
@@ -150,7 +141,7 @@ mod tests {
     #[test]
     fn validate_request_compat_rejects_streaming_prompt_logprobs() {
         let request = CompletionRequest {
-            prompt_logprobs: Some(1),
+            prompt_logprobs: Some(LogprobsCount::Top(1)),
             ..base_request()
         };
         assert!(
@@ -162,7 +153,7 @@ mod tests {
     fn validate_request_compat_accepts_non_stream_prompt_logprobs() {
         let request = CompletionRequest {
             stream: false,
-            prompt_logprobs: Some(-1),
+            prompt_logprobs: Some(LogprobsCount::All),
             ..base_request()
         };
         assert!(

rust/src/text/src/backend/mod.rs

@@ -2,6 +2,7 @@ pub mod hf;
 
 use std::sync::Arc;
 
+use vllm_engine_core_client::protocol::LogprobsCount;
 use vllm_tokenizer::DynTokenizer;
 
 use crate::error::Result;
@@ -26,9 +27,7 @@ pub struct SamplingLimits {
     /// Runtime context window size reported by the engine startup handshake.
     pub max_model_len: u32,
     /// Maximum number of top log probabilities accepted by this frontend.
-    ///
-    /// `-1` means allowing requests up to the model vocabulary size.
-    pub max_logprobs: i32,
+    pub max_logprobs: LogprobsCount,
 
     /// Model vocabulary size from the model config, used to bound generated
     /// token IDs and logits-domain sampling controls.
@@ -41,7 +40,7 @@ pub struct SamplingLimits {
 impl SamplingLimits {
     /// Original Python definition:
     /// <https://github.com/vllm-project/vllm/blob/b5adb027ad03c29b46181752ba3b1cb84eff1dd4/vllm/config/model.py#L216-L220>
-    pub const DEFAULT_MAX_LOGPROBS: i32 = 20;
+    pub const DEFAULT_MAX_LOGPROBS: LogprobsCount = LogprobsCount::Top(20);
     /// Original Python definition:
     /// <https://github.com/vllm-project/vllm/blob/b5adb027ad03c29b46181752ba3b1cb84eff1dd4/vllm/sampling_params.py#L30-L32>
     pub const MAX_LOGPROB_TOKEN_IDS: usize = 128;

rust/src/text/src/error.rs

@@ -20,6 +20,8 @@ pub enum Error {
     Logprobs(#[from] LogprobsError),
     #[error(transparent)]
     OutOfVocab(#[from] OutOfVocabError),
+    #[error("`thinking_token_budget` must be a non-negative integer or -1 for unlimited.")]
+    InvalidThinkingTokenBudget,
     #[error("text request stream `{request_id}` closed before terminal output")]
     StreamClosedBeforeTerminalOutput { request_id: String },
     #[error(transparent)]

rust/src/text/src/lib.rs

@@ -19,6 +19,7 @@ pub use output::{
 pub use request::{Prompt, SamplingParams, TextRequest};
 use trait_set::trait_set;
 use vllm_engine_core_client::EngineCoreClient;
+use vllm_engine_core_client::protocol::LogprobsCount;
 pub use vllm_llm::FinishReason;
 use vllm_llm::{GenerateOutputStream, Llm};
 use vllm_tokenizer::DynTokenizer;
@@ -48,7 +49,7 @@ pub struct TextLlm {
     /// Runtime context window size reported by the engine startup handshake.
     max_model_len: u32,
     /// Maximum number of top log probabilities accepted by this text facade.
-    max_logprobs: i32,
+    max_logprobs: LogprobsCount,
 }
 
 impl TextLlm {
@@ -68,7 +69,7 @@ impl TextLlm {
     }
 
     /// Override the maximum accepted logprobs count.
-    pub fn with_max_logprobs(mut self, max_logprobs: Option<i32>) -> Self {
+    pub fn with_max_logprobs(mut self, max_logprobs: Option<LogprobsCount>) -> Self {
         if let Some(max_logprobs) = max_logprobs {
             self.max_logprobs = max_logprobs;
         }

rust/src/text/src/lower.rs

@@ -87,6 +87,7 @@ pub fn lower_sampling_params(
         seed,
         max_tokens,
         min_tokens,
+        thinking_token_budget,
         logprobs,
         prompt_logprobs,
         min_p,
@@ -128,6 +129,7 @@ pub fn lower_sampling_params(
         prompt_len,
     )?;
     let min_tokens = min_tokens.unwrap_or(0);
+    let thinking_token_budget = normalize_thinking_token_budget(thinking_token_budget)?;
     let frequency_penalty = frequency_penalty.unwrap_or(0.0);
     let presence_penalty = presence_penalty.unwrap_or(0.0);
 
@@ -149,6 +151,7 @@ pub fn lower_sampling_params(
         seed,
         max_tokens,
         min_tokens,
+        thinking_token_budget,
         logprobs,
         prompt_logprobs,
         min_p,
@@ -170,6 +173,21 @@ pub fn lower_sampling_params(
     Ok(params)
 }
 
+/// Normalize the user-facing `thinking_token_budget` into the engine value.
+///
+/// Mirrors Python's `validate_thinking_token_budget`
+/// (<https://github.com/vllm-project/vllm/blob/ecf9d83520eb217401b47d8a5451a27c5231b8c2/vllm/sampling_params.py#L35-L55>):
+/// `None` and the `-1` "unlimited" sentinel both map to `None`; any other
+/// negative value is rejected; non-negative values pass through unchanged. Like
+/// Python's `int`, no upper bound is imposed.
+fn normalize_thinking_token_budget(value: Option<i64>) -> Result<Option<u64>> {
+    match value {
+        None | Some(-1) => Ok(None),
+        Some(budget) if budget >= 0 => Ok(Some(budget as u64)),
+        Some(_) => Err(Error::InvalidThinkingTokenBudget),
+    }
+}
+
 /// Convert bad-word strings into token-ID sequences, following the Python vLLM
 /// logic in `SamplingParams.update_from_tokenizer()`.
 ///
@@ -251,6 +269,7 @@ mod tests {
     use std::collections::{BTreeSet, HashMap};
 
     use serial_test::file_serial;
+    use vllm_engine_core_client::protocol::LogprobsCount;
 
     use super::*;
     use crate::backend::hf::HfTextBackend;
@@ -366,6 +385,36 @@ mod tests {
         )
     }
 
+    #[test]
+    fn lower_sampling_params_normalizes_thinking_token_budget() {
+        let lower = |budget: Option<i64>| {
+            lower_sampling_params_with_limits(
+                SamplingParams {
+                    thinking_token_budget: budget,
+                    ..SamplingParams::default()
+                },
+                sample_sampling_limits(),
+            )
+        };
+
+        // Non-negative budgets (including 0) pass through unchanged.
+        assert_eq!(lower(Some(256)).unwrap().thinking_token_budget, Some(256));
+        assert_eq!(lower(Some(0)).unwrap().thinking_token_budget, Some(0));
+        // `None` and the `-1` "unlimited" sentinel both disable the budget.
+        assert_eq!(lower(None).unwrap().thinking_token_budget, None);
+        assert_eq!(lower(Some(-1)).unwrap().thinking_token_budget, None);
+        // No upper bound is imposed, matching Python's `int`.
+        assert_eq!(
+            lower(Some(i64::from(u32::MAX) + 1)).unwrap().thinking_token_budget,
+            Some(u64::from(u32::MAX) + 1)
+        );
+        // Other negatives are rejected.
+        assert!(matches!(
+            lower(Some(-2)),
+            Err(Error::InvalidThinkingTokenBudget)
+        ));
+    }
+
     #[test]
     fn lower_text_request_applies_python_style_eos_hints() {
         let prepared = lower_text_request(
@@ -386,6 +435,7 @@ mod tests {
                 seed: None,
                 max_tokens: 999997,
                 min_tokens: 0,
+                thinking_token_budget: None,
                 logprobs: None,
                 prompt_logprobs: None,
                 min_p: 0.0,
@@ -437,6 +487,7 @@ mod tests {
                 seed: None,
                 max_tokens: 999997,
                 min_tokens: 0,
+                thinking_token_budget: None,
                 logprobs: None,
                 prompt_logprobs: None,
                 min_p: 0.0,
@@ -567,6 +618,7 @@ mod tests {
                 seed: None,
                 max_tokens: 40957,
                 min_tokens: 0,
+                thinking_token_budget: None,
                 logprobs: None,
                 prompt_logprobs: None,
                 min_p: 0.0,
@@ -628,6 +680,7 @@ mod tests {
                 seed: None,
                 max_tokens: 999997,
                 min_tokens: 0,
+                thinking_token_budget: None,
                 logprobs: None,
                 prompt_logprobs: None,
                 min_p: 0.0,
@@ -697,6 +750,7 @@ mod tests {
                 seed: None,
                 max_tokens: 32,
                 min_tokens: 2,
+                thinking_token_budget: None,
                 logprobs: None,
                 prompt_logprobs: None,
                 min_p: 0.1,
@@ -721,8 +775,8 @@ mod tests {
     #[test]
     fn lower_sampling_params_passes_logprobs_fields_through() {
         let sampling_params = SamplingParams {
-            logprobs: Some(3),
-            prompt_logprobs: Some(-1),
+            logprobs: Some(LogprobsCount::Top(3)),
+            prompt_logprobs: Some(LogprobsCount::All),
             ..Default::default()
         };
 
@@ -739,7 +793,7 @@ mod tests {
                 default_max_tokens: None,
             },
             SamplingLimits {
-                max_logprobs: -1,
+                max_logprobs: LogprobsCount::All,
                 ..sample_sampling_limits()
             },
             3,
@@ -747,15 +801,15 @@ mod tests {
         )
         .unwrap();
 
-        assert_eq!(params.logprobs, Some(3));
-        assert_eq!(params.prompt_logprobs, Some(-1));
+        assert_eq!(params.logprobs, Some(LogprobsCount::Top(3)));
+        assert_eq!(params.prompt_logprobs, Some(LogprobsCount::All));
     }
 
     #[test]
     fn lower_sampling_params_rejects_full_vocab_logprobs_over_default_cap() {
         let error = lower_sampling_params_with_limits(
             SamplingParams {
-                logprobs: Some(-1),
+                logprobs: Some(LogprobsCount::All),
                 ..Default::default()
             },
             sample_sampling_limits(),
@@ -776,24 +830,24 @@ mod tests {
     fn lower_sampling_params_expands_full_vocab_logprobs_from_model_vocab() {
         let params = lower_sampling_params_with_limits(
             SamplingParams {
-                logprobs: Some(-1),
+                logprobs: Some(LogprobsCount::All),
                 ..Default::default()
             },
             SamplingLimits {
-                max_logprobs: 1500,
+                max_logprobs: LogprobsCount::Top(1500),
                 ..sample_sampling_limits()
             },
         )
         .unwrap();
 
-        assert_eq!(params.logprobs, Some(-1));
+        assert_eq!(params.logprobs, Some(LogprobsCount::All));
     }
 
     #[test]
     fn lower_sampling_params_rejects_invalid_logprob_token_ids() {
         let error = lower_sampling_params_with_limits(
             SamplingParams {
-                logprobs: Some(1),
+                logprobs: Some(LogprobsCount::Top(1)),
                 logprob_token_ids: Some(vec![1000]),
                 ..Default::default()
             },
@@ -929,6 +983,7 @@ mod tests {
                 seed: None,
                 max_tokens: 128,
                 min_tokens: 0,
+                thinking_token_budget: None,
                 logprobs: None,
                 prompt_logprobs: None,
                 min_p: 0.1,

rust/src/text/src/lower/logprobs.rs

@@ -1,15 +1,14 @@
 //! Python-compatible validation for logprobs sampling params.
 //!
-//! `-1` is expanded only for bounds checks. The original request values are
+//! `All` is expanded only for bounds checks. The original request values are
 //! passed through to engine-core.
 
 use crate::backend::SamplingLimits;
 use thiserror::Error;
+use vllm_engine_core_client::protocol::LogprobsCount;
 
 #[derive(Debug, Error)]
 pub enum LogprobsError {
-    #[error("{parameter} must be non-negative or -1, got {value}")]
-    InvalidCount { parameter: &'static str, value: i32 },
     #[error(
         "requested {parameter} of {requested}, which is greater than max allowed: {max_allowed}"
     )]
@@ -30,19 +29,21 @@ pub enum LogprobsError {
         "when both logprobs and logprob_token_ids are set, logprobs must equal \
          len(logprob_token_ids). Got logprobs={logprobs}, len(logprob_token_ids)={num_token_ids}."
     )]
-    TokenIdsMismatch { logprobs: i32, num_token_ids: usize },
+    TokenIdsMismatch {
+        logprobs: LogprobsCount,
+        num_token_ids: usize,
+    },
 }
 
 /// Validate logprobs count sampling parameters.
 pub(super) fn validate_logprobs(
-    logprobs: Option<i32>,
-    prompt_logprobs: Option<i32>,
+    logprobs: Option<LogprobsCount>,
+    prompt_logprobs: Option<LogprobsCount>,
     logprob_token_ids: Option<&[u32]>,
     sampling_limits: SamplingLimits,
 ) -> Result<(), LogprobsError> {
     let vocab_size = sampling_limits.model_vocab_size;
-    let max_logprobs =
-        normalize_logprobs_count(sampling_limits.max_logprobs, vocab_size, "max_logprobs")?;
+    let max_logprobs = sampling_limits.max_logprobs.expanded(vocab_size);
 
     validate_logprobs_count(logprobs, max_logprobs, vocab_size, "logprobs")?;
     validate_logprobs_count(prompt_logprobs, max_logprobs, vocab_size, "prompt_logprobs")?;
@@ -50,7 +51,7 @@ pub(super) fn validate_logprobs(
 }
 
 fn validate_logprobs_count(
-    requested: Option<i32>,
+    requested: Option<LogprobsCount>,
     max_logprobs: usize,
     vocab_size: usize,
     parameter: &'static str,
@@ -59,7 +60,7 @@ fn validate_logprobs_count(
         return Ok(());
     };
 
-    let requested = normalize_logprobs_count(requested, vocab_size, parameter)?;
+    let requested = requested.expanded(vocab_size);
     if requested > max_logprobs {
         return Err(LogprobsError::TooManyCount {
             parameter,
@@ -72,7 +73,7 @@ fn validate_logprobs_count(
 }
 
 pub(super) fn validate_logprob_token_ids(
-    logprobs: Option<i32>,
+    logprobs: Option<LogprobsCount>,
     logprob_token_ids: Option<&[u32]>,
 ) -> Result<(), LogprobsError> {
     let Some(logprob_token_ids) = logprob_token_ids else {
@@ -88,7 +89,7 @@ pub(super) fn validate_logprob_token_ids(
     }
 
     if let Some(logprobs) = logprobs
-        && logprobs != n as i32
+        && logprobs != LogprobsCount::Top(n as u32)
     {
         return Err(LogprobsError::TokenIdsMismatch {
             logprobs,
@@ -98,15 +99,3 @@ pub(super) fn validate_logprob_token_ids(
 
     Ok(())
 }
-
-fn normalize_logprobs_count(
-    value: i32,
-    vocab_size: usize,
-    parameter: &'static str,
-) -> Result<usize, LogprobsError> {
-    match value {
-        -1 => Ok(vocab_size),
-        value if value < 0 => Err(LogprobsError::InvalidCount { parameter, value }),
-        value => Ok(value as usize),
-    }
-}

rust/src/text/src/output/decoded.rs

@@ -309,7 +309,7 @@ fn matches_stop_string(stops: &[String], output: &str, new_bytes: usize) -> Opti
         .find_map(|(ss_idx, (ss, len, start_off))| {
             output[start_off..]
                 .windows(len)
-                .rposition(|w| w == ss)
+                .position(|w| w == ss)
                 .map(|pos| (ss_idx, start_off + pos))
         })
 }
@@ -562,6 +562,13 @@ mod tests {
         assert_eq!(result, Some((0, 4)));
     }
 
+    #[test]
+    fn stop_string_matches_leftmost_with_multiple_new_bytes() {
+        let stops = vec!["\n".to_string()];
+        let result = matches_stop_string(&stops, "Answer\n\n", 2);
+        assert_eq!(result, Some((0, 6)));
+    }
+
     #[test]
     fn stop_string_matches_at_beginning() {
         let stops = vec!["say".to_string()];

rust/src/text/src/request.rs

@@ -3,9 +3,9 @@ use std::collections::HashMap;
 use enum_as_inner::EnumAsInner;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
-use vllm_engine_core_client::protocol::StructuredOutputsParams;
 use vllm_engine_core_client::protocol::lora::LoraRequest;
 use vllm_engine_core_client::protocol::multimodal::MmFeatures;
+use vllm_engine_core_client::protocol::{LogprobsCount, StructuredOutputsParams};
 
 use crate::error::{Error, Result};
 use crate::output::TextDecodeOptions;
@@ -56,14 +56,20 @@ pub struct SamplingParams {
     pub max_tokens: Option<u32>,
     /// Minimum number of tokens to generate before EOS or stop-token handling.
     pub min_tokens: Option<u32>,
+    /// Maximum number of reasoning ("thinking") tokens to emit before the
+    /// reasoning section is force-closed. `None` or the user-facing `-1`
+    /// "unlimited" sentinel both disable the budget. The raw value is carried
+    /// here; `-1` is normalized to `None` (and other negatives rejected) during
+    /// lowering (see `lower_sampling_params`).
+    pub thinking_token_budget: Option<i64>,
     /// Number of log probabilities to return per generated token.
     ///
-    /// `None` disables sample logprobs. `-1` requests the full vocabulary.
-    pub logprobs: Option<i32>,
+    /// `None` disables sample logprobs.
+    pub logprobs: Option<LogprobsCount>,
     /// Number of log probabilities to return per prompt token.
     ///
-    /// `None` disables prompt logprobs. `-1` requests the full vocabulary.
-    pub prompt_logprobs: Option<i32>,
+    /// `None` disables prompt logprobs.
+    pub prompt_logprobs: Option<LogprobsCount>,
     /// Minimum probability threshold for token sampling. `None` means no
     /// explicit user override.
     pub min_p: Option<f32>,
@@ -116,6 +122,7 @@ impl Default for SamplingParams {
             seed: None,
             max_tokens: None,
             min_tokens: None,
+            thinking_token_budget: None,
             logprobs: None,
             prompt_logprobs: None,
             min_p: None,

setup.py

@@ -769,6 +769,7 @@ class precompiled_wheel_utils:
                             "vllm/_C.abi3.so",
                             "vllm/_C_stable_libtorch.abi3.so",
                             "vllm/_moe_C_stable_libtorch.abi3.so",
+                            "vllm/_qutlass_C.abi3.so",
                             "vllm/_flashmla_C.abi3.so",
                             "vllm/_flashmla_extension_C.abi3.so",
                             "vllm/_sparse_flashmla_C.abi3.so",
@@ -1135,6 +1136,7 @@ if _is_cuda():
         # DeepGEMM requires CUDA 12.3+ (SM90/SM100)
         # Optional since it won't build on unsupported architectures
         ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", optional=True))
+        ext_modules.append(CMakeExtension(name="vllm._qutlass_C", optional=True))
     # fmha_sm100 is a Python/CuTe-DSL package installed into vllm.third_party.
     ext_modules.append(CMakeExtension(name="vllm.fmha_sm100", optional=True))
 
@@ -1149,7 +1151,8 @@ if _is_cpu():
         ext_modules.append(CMakeExtension(name="vllm._C"))
 
 if _build_custom_ops():
-    ext_modules.append(CMakeExtension(name="vllm._C"))
+    if _is_hip():
+        ext_modules.append(CMakeExtension(name="vllm._C"))
     if _is_cuda() or _is_hip():
         ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch"))
         ext_modules.append(CMakeExtension(name="vllm._moe_C_stable_libtorch"))

tests/compile/passes/test_double_aiter_rms_quant_fusion.py

@@ -22,7 +22,7 @@ import torch
 
 import vllm.config
 from tests.compile.backend import TestBackend
-from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
 from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
 from vllm.config import (
@@ -83,9 +83,8 @@ class _ViewDoubleQuantModel(torch.nn.Module):
     [_NoViewDoubleQuantModel, _ViewDoubleQuantModel],
     ids=["no_view", "with_view"],
 )
-@pytest.mark.skipif(
-    not is_aiter_found_and_supported(),
-    reason="Only test on ROCm with AITER installed and supported",
+@pytest.mark.skip(
+    reason="Skipping for now because pytorch compiler removes one the two quant ops"
 )
 def test_double_aiter_rms_fp8_group_quant_fusion(
     model_cls: type[torch.nn.Module],

tests/distributed/test_pipeline_parallel.py

@@ -175,7 +175,7 @@ MULTIMODAL_MODELS = {
     "facebook/chameleon-7b": PPTestSettings.fast(),
     "adept/fuyu-8b": PPTestSettings.fast(),
     "zai-org/glm-4v-9b": PPTestSettings.fast(),
-    "OpenGVLab/InternVL2-1B": PPTestSettings.fast(),
+    "OpenGVLab/InternVL3-1B": PPTestSettings.fast(),
     "llava-hf/llava-1.5-7b-hf": PPTestSettings.fast(),
     "llava-hf/llava-v1.6-mistral-7b-hf": PPTestSettings.fast(),
     "llava-hf/LLaVA-NeXT-Video-7B-hf": PPTestSettings.fast(),
@@ -203,7 +203,7 @@ TEST_MODELS = [
     "intfloat/e5-mistral-7b-instruct",
     "BAAI/bge-multilingual-gemma2",
     # [MULTIMODAL GENERATION]
-    "OpenGVLab/InternVL2-1B",
+    "OpenGVLab/InternVL3-1B",
     "microsoft/Phi-3.5-vision-instruct",
     "fixie-ai/ultravox-v0_5-llama-3_2-1b",
     # [LANGUAGE GENERATION - HYBRID ARCH]

tests/engine/test_arg_utils.py

@@ -649,3 +649,196 @@ def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch):
     args = EngineArgs(model="s3://bucket/model", tokenizer="s3://bucket/tokenizer")
     assert args.model == "s3://bucket/model"
     assert args.tokenizer == "s3://bucket/tokenizer"
+
+
+class TestDeviceIds:
+    def test_device_ids_with_cvd_out_of_range(self, monkeypatch):
+        """--device-ids index beyond the CVD set raises ValueError."""
+        from vllm.platforms import current_platform
+
+        key = current_platform.device_control_env_var
+        monkeypatch.setenv(key, "4,5")
+        args = EngineArgs(model="m", device_ids=[0, 2])
+        with pytest.raises(ValueError, match="out of range"):
+            args._resolve_device_ids()
+
+    def test_device_ids_with_cvd_resolve_to_physical_ids(self, monkeypatch):
+        """--device-ids are CVD-local indices resolved to physical ids."""
+        from vllm.platforms import current_platform
+
+        key = current_platform.device_control_env_var
+        monkeypatch.setenv(key, "4,5")
+        args = EngineArgs(model="m", device_ids=[0, 1])
+        assert args._resolve_device_ids() == [4, 5]
+
+    def test_device_ids_with_uuid_cvd_resolve_to_physical_ids(self, monkeypatch):
+        """--device-ids support UUID CVD values resolved by the platform."""
+        from vllm.platforms import current_platform
+
+        key = current_platform.device_control_env_var
+        monkeypatch.setenv(key, "GPU-abcd1234,GPU-ef567890")
+        monkeypatch.setattr(
+            type(current_platform),
+            "device_control_id_to_physical_device_id",
+            classmethod(
+                lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
+            ),
+        )
+
+        args = EngineArgs(model="m", device_ids=[0, 1])
+        assert args._resolve_device_ids() == [4, 5]
+
+    def test_device_ids_with_uuid_args_resolve_to_physical_ids(self, monkeypatch):
+        """UUID --device-ids are resolved to physical IDs immediately."""
+        from vllm.platforms import current_platform
+
+        monkeypatch.setattr(
+            type(current_platform),
+            "device_control_id_to_physical_device_id",
+            classmethod(lambda cls, device_id: {"GPU-abcd1234": 4}[device_id]),
+        )
+
+        args = EngineArgs(model="m", device_ids=["GPU-abcd1234"])
+        assert args._resolve_device_ids() == [4]
+
+    def test_device_ids_reject_mixed_integer_and_uuid_args(self):
+        """--device-ids must not mix CVD indices and UUIDs."""
+        args = EngineArgs(model="m", device_ids=[0, "GPU-abcd1234"])
+        with pytest.raises(ValueError, match="must not mix"):
+            args._resolve_device_ids()
+
+    def test_no_device_ids(self):
+        """No --device-ids returns None."""
+        args = EngineArgs(model="m")
+        assert args._resolve_device_ids() is None
+
+    def test_cli_parsing(self):
+        """--device-ids parses comma-separated string from CLI."""
+        parser = FlexibleArgumentParser()
+        EngineArgs.add_cli_args(parser)
+        parsed = parser.parse_args(["--model", "m", "--device-ids", "0,2,4"])
+        assert parsed.device_ids == [0, 2, 4]
+
+    def test_cli_parsing_uuid(self):
+        """--device-ids parses comma-separated UUID strings from CLI."""
+        parser = FlexibleArgumentParser()
+        EngineArgs.add_cli_args(parser)
+        parsed = parser.parse_args(
+            ["--model", "m", "--device-ids", "GPU-abcd1234,GPU-ef567890"]
+        )
+        assert parsed.device_ids == ["GPU-abcd1234", "GPU-ef567890"]
+
+    def test_assigned_physical_gpu_ids_are_physical_with_cvd(self, monkeypatch):
+        """assigned_physical_gpu_ids are already physical and not composed with CVD."""
+        import vllm.platforms.interface as platform_interface
+        from vllm.platforms import current_platform
+
+        monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [4, 5])
+        monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
+
+        assert current_platform.device_id_to_physical_device_id(0) == 4
+        assert current_platform.device_id_to_physical_device_id(1) == 5
+        assert current_platform.logical_device_id_to_visible_device_id(0) == 0
+        assert current_platform.logical_device_id_to_visible_device_id(1) == 1
+
+    def test_assigned_physical_gpu_ids_map_to_visible_uuid_cvd(self, monkeypatch):
+        """Physical IDs map back to visible ordinals when CVD uses UUIDs."""
+        import vllm.platforms.interface as platform_interface
+        from vllm.platforms import current_platform
+
+        monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [5])
+        monkeypatch.setenv(
+            current_platform.device_control_env_var,
+            "GPU-abcd1234,GPU-ef567890",
+        )
+        monkeypatch.setattr(
+            type(current_platform),
+            "device_control_id_to_physical_device_id",
+            classmethod(
+                lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
+            ),
+        )
+
+        assert current_platform.logical_device_id_to_visible_device_id(0) == 1
+
+    def test_device_ids_reject_duplicates(self):
+        """--device-ids must not contain duplicate entries."""
+        args = EngineArgs(model="m", device_ids=[2, 2])
+        with pytest.raises(ValueError, match="duplicates"):
+            args._resolve_device_ids()
+
+    def test_cli_parsing_strips_whitespace(self):
+        """--device-ids tolerates whitespace around commas."""
+        parser = FlexibleArgumentParser()
+        EngineArgs.add_cli_args(parser)
+        parsed = parser.parse_args(["--model", "m", "--device-ids", "0, 2, 4"])
+        assert parsed.device_ids == [0, 2, 4]
+
+    def test_visible_ordinal_to_physical_ignores_assigned_ids(self, monkeypatch):
+        """visible_device_id_to_physical_device_id maps torch device ordinals,
+        independent of the logical-to-physical mapping.
+
+        Regression test: CustomAllreduce passes device.index (a visible
+        ordinal) and must not index into assigned_physical_gpu_ids, which
+        raised IndexError for non-identity --device-ids like [2, 3].
+        """
+        import vllm.platforms.interface as platform_interface
+        from vllm.platforms import current_platform
+
+        monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [2, 3])
+        monkeypatch.delenv(current_platform.device_control_env_var, raising=False)
+
+        # CVD unset: visible ordinal == physical ID, even beyond the
+        # assigned list's length.
+        assert current_platform.visible_device_id_to_physical_device_id(2) == 2
+        assert current_platform.visible_device_id_to_physical_device_id(3) == 3
+
+        monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
+        assert current_platform.visible_device_id_to_physical_device_id(1) == 5
+        with pytest.raises(IndexError, match="out of range"):
+            current_platform.visible_device_id_to_physical_device_id(2)
+
+
+class TestDpDeviceIdSharding:
+    def test_dp_supervisor_device_ids_stay_env_relative(self):
+        """Regression test: the DP supervisor must pass env-relative indices,
+        not physical IDs, because each child re-resolves --device-ids
+        against its inherited device-control env var."""
+        import argparse
+
+        from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
+
+        args = argparse.Namespace(
+            tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=None
+        )
+        assert _build_device_ids(args, local_rank=0) == [0, 1]
+        assert _build_device_ids(args, local_rank=1) == [2, 3]
+
+    def test_dp_supervisor_shards_user_device_ids(self):
+        """User-provided --device-ids are sharded across DP children."""
+        import argparse
+
+        from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
+
+        args = argparse.Namespace(
+            tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=[4, 5, 6, 7]
+        )
+        assert _build_device_ids(args, local_rank=0) == [4, 5]
+        assert _build_device_ids(args, local_rank=1) == [6, 7]
+        with pytest.raises(ValueError, match="needs devices"):
+            _build_device_ids(args, local_rank=2)
+
+    def test_dp_rank_shards_user_assigned_gpu_ids(self):
+        """get_physical_gpu_ids_for_local_dp_rank slices the user-provided
+        --device-ids list instead of recomputing from the env var."""
+        from vllm.platforms import current_platform
+        from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank
+
+        evar = current_platform.device_control_env_var
+        assert get_physical_gpu_ids_for_local_dp_rank(
+            evar, local_dp_rank=1, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
+        ) == [6, 7]
+        with pytest.raises(ValueError, match="needs devices"):
+            get_physical_gpu_ids_for_local_dp_rank(
+                evar, local_dp_rank=2, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
+            )

tests/entrypoints/anthropic/test_anthropic_messages_conversion.py

@@ -8,6 +8,8 @@ AnthropicServingMessages._convert_anthropic_to_openai_request().
 Also covers extended-thinking edge cases such as ``redacted_thinking``
 blocks echoed back by Anthropic clients, and streaming conversion in
 ``message_stream_converter``.
+
+Also covers cache usage computation in ``_build_anthropic_usage``.
 """
 
 import json
@@ -18,7 +20,11 @@ import pytest
 from vllm.entrypoints.anthropic.protocol import (
     AnthropicMessagesRequest,
 )
-from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
+from vllm.entrypoints.anthropic.serving import (
+    AnthropicServingMessages,
+    _build_anthropic_usage,
+    _get_cached_tokens,
+)
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionResponseStreamChoice,
     ChatCompletionStreamResponse,
@@ -27,6 +33,7 @@ from vllm.entrypoints.openai.engine.protocol import (
     DeltaFunctionCall,
     DeltaMessage,
     DeltaToolCall,
+    PromptTokenUsageInfo,
     UsageInfo,
 )
 
@@ -653,6 +660,108 @@ class TestThinkingBlockConversion:
         assert asst.get("content") == "Hi!"
 
 
+# ======================================================================
+# Cache usage computation
+# ======================================================================
+
+
+class TestGetCachedTokens:
+    """Tests for _get_cached_tokens helper."""
+
+    def test_none_usage(self):
+        assert _get_cached_tokens(None) is None
+
+    def test_no_prompt_tokens_details(self):
+        usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
+        assert _get_cached_tokens(usage) is None
+
+    def test_cached_tokens_present(self):
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
+        )
+        assert _get_cached_tokens(usage) == 80
+
+    def test_cached_tokens_zero(self):
+        """Zero cached tokens should return 0, not None."""
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
+        )
+        assert _get_cached_tokens(usage) == 0
+
+    def test_cached_tokens_none_in_details(self):
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=None),
+        )
+        assert _get_cached_tokens(usage) is None
+
+
+class TestBuildAnthropicUsage:
+    """Tests for _build_anthropic_usage helper.
+
+    Anthropic defines: total_input = input_tokens + cache_read + cache_creation
+    vLLM's prompt_tokens is the total.
+    """
+
+    def test_no_cache_info(self):
+        """When cache info is unavailable, return raw prompt_tokens."""
+        result = _build_anthropic_usage(100, 10, None)
+        assert result.input_tokens == 100
+        assert result.output_tokens == 10
+        assert result.cache_read_input_tokens is None
+        assert result.cache_creation_input_tokens is None
+
+    def test_cache_hit(self):
+        """When cache is hit, input_tokens excludes cached tokens."""
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
+        )
+        result = _build_anthropic_usage(100, 10, usage)
+        assert result.input_tokens == 20  # 100 - 80
+        assert result.output_tokens == 10
+        assert result.cache_read_input_tokens == 80
+        assert result.cache_creation_input_tokens == 0
+
+    def test_zero_cached_tokens(self):
+        """Zero cached tokens should still set cache_creation to 0."""
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
+        )
+        result = _build_anthropic_usage(100, 10, usage)
+        assert result.input_tokens == 100  # 100 - 0
+        assert result.cache_read_input_tokens == 0
+        assert result.cache_creation_input_tokens == 0
+
+    def test_all_tokens_cached(self):
+        """When all tokens are cached, input_tokens should be 0."""
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=100),
+        )
+        result = _build_anthropic_usage(100, 10, usage)
+        assert result.input_tokens == 0
+        assert result.cache_read_input_tokens == 100
+        assert result.cache_creation_input_tokens == 0
+
+    def test_no_prompt_tokens_details(self):
+        """UsageInfo without prompt_tokens_details returns no cache info."""
+        usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
+        result = _build_anthropic_usage(100, 10, usage)
+        assert result.input_tokens == 100
+        assert result.cache_read_input_tokens is None
+        assert result.cache_creation_input_tokens is None
+
+
 class TestInlineSystemMessageInMessagesArray:
     """Verify that ``role: system`` messages embedded inside the ``messages``
     array are preserved in their original position.
@@ -1098,6 +1207,135 @@ class TestMessageStartIncludesTypeAndRole:
         assert message["role"] == "assistant"
 
 
+class TestStreamingCacheUsageSemantics:
+    """Locks in the documented streaming behavior of cache usage fields.
+
+    vLLM's OpenAI chat completion streaming only attaches
+    ``prompt_tokens_details`` to the terminal usage chunk. The Anthropic layer
+    mirrors that contract: cache fields are omitted on ``message_start`` (key
+    absence signals "unknown") and populated on ``message_delta`` (the final
+    cumulative count). This is intentionally consistent with vLLM's OpenAI
+    behavior, even though Anthropic's upstream API populates cache fields on
+    ``message_start``; closing that gap requires plumbing cache info into the
+    first chunk at the OpenAI layer, which is out of scope here.
+    """
+
+    @pytest.mark.asyncio
+    async def test_streaming_cache_fields_absent_then_populated(self):
+        """First chunk lacks prompt_tokens_details (vLLM contract);
+        message_start omits cache fields. The final chunk carries
+        prompt_tokens_details, so message_delta carries resolved values."""
+
+        async def sse_input():
+            yield _make_stream_chunk(
+                delta=DeltaMessage(role="assistant", content="hi"),
+                usage=UsageInfo(prompt_tokens=100, total_tokens=100),
+            )
+            yield _make_stream_chunk(finish_reason="stop")
+            yield _make_stream_chunk(
+                choices=[],
+                usage=UsageInfo(
+                    prompt_tokens=100,
+                    completion_tokens=5,
+                    total_tokens=105,
+                    prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
+                ),
+            )
+            yield "data: [DONE]"
+
+        converter = _make_stream_converter()
+        output = []
+        async for event in converter.message_stream_converter(sse_input()):
+            output.append(event)
+        events = _parse_sse_events(output)
+
+        # message_start: cache fields unknown → omitted from JSON entirely.
+        start_usage = events[0][1]["message"]["usage"]
+        assert events[0][0] == "message_start"
+        assert start_usage["input_tokens"] == 100
+        assert "cache_read_input_tokens" not in start_usage
+        assert "cache_creation_input_tokens" not in start_usage
+
+        # message_delta: authoritative usage with cache fields populated.
+        delta_usage = next(
+            data["usage"] for ev, data in events if ev == "message_delta"
+        )
+        assert delta_usage["input_tokens"] == 20  # 100 - 80
+        assert delta_usage["cache_read_input_tokens"] == 80
+        assert delta_usage["cache_creation_input_tokens"] == 0
+
+    @pytest.mark.asyncio
+    async def test_streaming_no_cache_hit(self):
+        """When the final chunk reports cached_tokens=0, message_delta carries
+        cache fields = 0 (cache miss); message_start still omits them."""
+
+        async def sse_input():
+            yield _make_stream_chunk(
+                delta=DeltaMessage(role="assistant"),
+                usage=UsageInfo(prompt_tokens=50, total_tokens=50),
+            )
+            yield _make_stream_chunk(finish_reason="stop")
+            yield _make_stream_chunk(
+                choices=[],
+                usage=UsageInfo(
+                    prompt_tokens=50,
+                    completion_tokens=5,
+                    total_tokens=55,
+                    prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
+                ),
+            )
+            yield "data: [DONE]"
+
+        converter = _make_stream_converter()
+        output = []
+        async for event in converter.message_stream_converter(sse_input()):
+            output.append(event)
+        events = _parse_sse_events(output)
+
+        start_usage = events[0][1]["message"]["usage"]
+        delta_usage = next(
+            data["usage"] for ev, data in events if ev == "message_delta"
+        )
+        assert start_usage["input_tokens"] == 50
+        assert "cache_read_input_tokens" not in start_usage
+        assert "cache_creation_input_tokens" not in start_usage
+        assert delta_usage["input_tokens"] == 50  # 50 - 0
+        assert delta_usage["cache_read_input_tokens"] == 0
+        assert delta_usage["cache_creation_input_tokens"] == 0
+
+    @pytest.mark.asyncio
+    async def test_streaming_no_prompt_tokens_details_at_all(self):
+        """If --enable-prompt-tokens-details is off, no chunk carries cache
+        info; both message_start and message_delta omit cache fields."""
+
+        async def sse_input():
+            yield _make_stream_chunk(
+                delta=DeltaMessage(role="assistant"),
+                usage=UsageInfo(prompt_tokens=30, total_tokens=30),
+            )
+            yield _make_stream_chunk(finish_reason="stop")
+            yield _make_stream_chunk(
+                choices=[],
+                usage=UsageInfo(prompt_tokens=30, completion_tokens=2, total_tokens=32),
+            )
+            yield "data: [DONE]"
+
+        converter = _make_stream_converter()
+        output = []
+        async for event in converter.message_stream_converter(sse_input()):
+            output.append(event)
+        events = _parse_sse_events(output)
+
+        start_usage = events[0][1]["message"]["usage"]
+        delta_usage = next(
+            data["usage"] for ev, data in events if ev == "message_delta"
+        )
+        assert "cache_read_input_tokens" not in start_usage
+        assert "cache_creation_input_tokens" not in start_usage
+        assert "cache_read_input_tokens" not in delta_usage
+        assert "cache_creation_input_tokens" not in delta_usage
+
+
 # ======================================================================
 # Auto-detection of system-first template requirement
 # ======================================================================

tests/entrypoints/openai/test_dp_supervisor.py

@@ -364,7 +364,7 @@ class MockVLLMServer:
         await self._serve_task
 
 
-def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]):
+def launch_mock_vllm(child_args: argparse.Namespace):
     logger.info("Launching mock vLLM on port %s", child_args.port)
     mock_vllm = MockVLLMServer(
         port=child_args.port,
@@ -375,7 +375,7 @@ def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]
 
 
 def launch_mock_vllm_with_drain(
-    child_args: argparse.Namespace, env_updates: dict[str, str]
+    child_args: argparse.Namespace,
 ):
     logger.info("Launching mock vLLM with 15s drain on port %s", child_args.port)
     mock_vllm = MockVLLMServer(

tests/entrypoints/openai/test_openai_schema.py

@@ -6,15 +6,22 @@ from typing import Final
 import pytest
 import schemathesis
 from hypothesis import HealthCheck, settings
-from schemathesis import GenerationConfig
-from schemathesis.models import Case
+from schemathesis import GenerationMode
+from schemathesis.config import (
+    ChecksConfig,
+    CoveragePhaseConfig,
+    GenerationConfig,
+    PhasesConfig,
+    PositiveDataAcceptanceConfig,
+    ProjectConfig,
+    ProjectsConfig,
+    SchemathesisConfig,
+)
 
 from vllm.platforms import current_platform
 
 from ...utils import RemoteOpenAIServer
 
-schemathesis.experimental.OPEN_API_3_1.enable()
-
 MODEL_NAME = "HuggingFaceTB/SmolVLM-256M-Instruct"
 MAXIMUM_IMAGES = 2
 _ROCM_TIMEOUT_MULTIPLIER = 3 if current_platform.is_rocm() else 1
@@ -44,21 +51,38 @@ def server():
 @pytest.fixture(scope="module")
 def get_schema(server):
     # avoid generating null (\x00) bytes in strings during test case generation
-    return schemathesis.openapi.from_uri(
+    return schemathesis.openapi.from_url(
         f"{server.url_root}/openapi.json",
-        generation_config=GenerationConfig(allow_x00=False),
+        config=SchemathesisConfig(
+            projects=ProjectsConfig(
+                default=ProjectConfig(
+                    generation=GenerationConfig(
+                        allow_x00=False,
+                        modes=[GenerationMode.POSITIVE],
+                    ),
+                    checks=ChecksConfig(
+                        positive_data_acceptance=PositiveDataAcceptanceConfig(
+                            enabled=False,
+                        ),
+                    ),
+                    phases=PhasesConfig(
+                        coverage=CoveragePhaseConfig(enabled=False),
+                    ),
+                ),
+            ),
+        ),
     )
 
 
-schema = schemathesis.from_pytest_fixture("get_schema")
+schema = schemathesis.pytest.from_fixture("get_schema")
 
 
 @schemathesis.hook
-def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
+def before_generate_case(context: schemathesis.HookContext, strategy):
     op = context.operation
     assert op is not None
 
-    def no_invalid_types(case: schemathesis.models.Case):
+    def no_invalid_types(case: schemathesis.Case):
         """
         Skips tool_calls with `"type": "custom"` which schemathesis incorrectly
         generates instead of the valid `"type": "function"`.
@@ -68,39 +92,25 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
             -d '{"messages": [{"role": "assistant", "tool_calls": [{"custom": {"input": "", "name": ""}, "id": "", "type": "custom"}]}]}' \
             http://localhost:8000/v1/chat/completions
         """  # noqa: E501
-        if hasattr(case, "body") and isinstance(case.body, dict):
-            if (
-                "messages" in case.body
-                and isinstance(case.body["messages"], list)
-                and len(case.body["messages"]) > 0
-            ):
-                for message in case.body["messages"]:
-                    if not isinstance(message, dict):
-                        continue
+        if (
+            hasattr(case, "body")
+            and isinstance(case.body, dict)
+            and "messages" in case.body
+            and isinstance(case.body["messages"], list)
+            and len(case.body["messages"]) > 0
+        ):
+            for message in case.body["messages"]:
+                if not isinstance(message, dict):
+                    continue
 
-                    tool_calls = message.get("tool_calls", [])
-                    if isinstance(tool_calls, list):
-                        for tool_call in tool_calls:
-                            if isinstance(tool_call, dict):
-                                if tool_call.get("type") != "function":
-                                    return False
-                                if "custom" in tool_call:
-                                    return False
-
-            # Sometimes structured_outputs.grammar is generated to be empty
-            # Causing a server error in EBNF grammar parsing
-            # https://github.com/vllm-project/vllm/pull/22587#issuecomment-3195253421
-            structured_outputs = case.body.get("structured_outputs", {})
-            grammar = (
-                structured_outputs.get("grammar")
-                if isinstance(structured_outputs, dict)
-                else None
-            )
-
-            if grammar == "":
-                # Allow None (will be handled as no grammar)
-                # But skip empty strings
-                return False
+                tool_calls = message.get("tool_calls", [])
+                if isinstance(tool_calls, list):
+                    for tool_call in tool_calls:
+                        if isinstance(tool_call, dict):
+                            if tool_call.get("type") != "function":
+                                return False
+                            if "custom" in tool_call:
+                                return False
 
         return True
 
@@ -108,7 +118,6 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
 
 
 @schema.parametrize()
-@schema.override(headers={"Content-Type": "application/json"})
 @settings(
     deadline=LONG_TIMEOUT_SECONDS * 1000,
     max_examples=50,
@@ -122,7 +131,7 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
     # generating large-but-valid request bodies before vLLM is called.
     suppress_health_check=[HealthCheck.filter_too_much, HealthCheck.data_too_large],
 )
-def test_openapi_stateless(case: Case):
+def test_openapi_stateless(case: schemathesis.Case):
     key = (
         case.operation.method.upper(),
         case.operation.path,
@@ -151,4 +160,8 @@ def test_openapi_stateless(case: Case):
     }.get(key, DEFAULT_TIMEOUT_SECONDS)
 
     # No need to verify SSL certificate for localhost
-    case.call_and_validate(verify=False, timeout=timeout)
+    case.call_and_validate(
+        verify=False,
+        timeout=timeout,
+        headers={"Content-Type": "application/json"},
+    )

tests/entrypoints/pooling/classify/test_online_vision.py

@@ -25,7 +25,7 @@ def server():
         "--runner",
         "pooling",
         "--max-model-len",
-        "5000",
+        "16384",
         "--enforce-eager",
         "--limit-mm-per-prompt",
         json.dumps({"video": MAXIMUM_VIDEOS}),
@@ -143,4 +143,4 @@ def test_chat_video_url_request(server: RemoteOpenAIServer, model_name: str):
     assert output.model == model_name
     assert len(output.data) == 1
     assert len(output.data[0].probs) == 2
-    assert output.usage.prompt_tokens == 4807
+    assert output.usage.prompt_tokens == 8993

tests/entrypoints/serve/render/test_derender.py

@@ -8,6 +8,7 @@ import pytest
 import pytest_asyncio
 
 from tests.utils import RemoteLaunchRenderServer
+from vllm.tokenizers import get_tokenizer
 
 MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
 
@@ -486,3 +487,438 @@ async def test_derender_completion_kv_transfer_params_passthrough(client):
     )
     assert response.status_code == 200
     assert response.json()["kv_transfer_params"] == kv
+
+
+# ---------------------------------------------------------------------------
+# E2E: render -> derender roundtrip with parser (reasoning + tool calls)
+# ---------------------------------------------------------------------------
+
+PARSER_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+
+_E2E_TOOLS = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_weather",
+            "description": "Get weather for a city",
+            "parameters": {
+                "type": "object",
+                "properties": {"city": {"type": "string"}},
+            },
+        },
+    }
+]
+
+
+@pytest.fixture(scope="module")
+def parser_server():
+    args = [
+        "--enable-auto-tool-choice",
+        "--tool-call-parser",
+        "hermes",
+        "--reasoning-parser",
+        "deepseek_r1",
+    ]
+    with RemoteLaunchRenderServer(PARSER_MODEL, args) as remote_server:
+        yield remote_server
+
+
+@pytest_asyncio.fixture
+async def parser_client(parser_server):
+    async with httpx.AsyncClient(
+        base_url=parser_server.url_for(""), timeout=60.0
+    ) as http_client:
+        yield http_client
+
+
+@pytest.fixture(scope="module")
+def parser_tokenizer():
+    return get_tokenizer(PARSER_MODEL)
+
+
+def _encode(tokenizer, text: str) -> list[int]:
+    return tokenizer.encode(text, add_special_tokens=False)
+
+
+def _decoded(tokenizer, token_ids: list[int]) -> str:
+    return tokenizer.decode(token_ids, skip_special_tokens=True)
+
+
+def _require_markers_survive(tokenizer, text: str, *markers: str) -> list[int]:
+    """Encode text and skip the test if any marker is lost in roundtrip."""
+    ids = _encode(tokenizer, text)
+    decoded = tokenizer.decode(ids, skip_special_tokens=False)
+    for m in markers:
+        if m not in decoded:
+            pytest.skip(f"Marker {m!r} lost in encode->decode roundtrip")
+    return ids
+
+
+async def _e2e_render_chat(
+    client: httpx.AsyncClient,
+    model: str,
+    messages: list[dict],
+) -> dict:
+    resp = await client.post(
+        "/v1/chat/completions/render",
+        json={"model": model, "messages": messages},
+    )
+    assert resp.status_code == 200, resp.text
+    return resp.json()
+
+
+def _e2e_generate_response(
+    token_ids: list[int],
+    request_id: str = "chatcmpl-e2e-test",
+) -> dict:
+    return {
+        "request_id": request_id,
+        "choices": [
+            {
+                "index": 0,
+                "token_ids": token_ids,
+                "finish_reason": "stop",
+            }
+        ],
+    }
+
+
+@pytest.mark.asyncio
+async def test_e2e_plain_roundtrip(parser_client, parser_tokenizer):
+    """Plain text without reasoning markers roundtrips correctly."""
+    messages = [{"role": "user", "content": "What is 2+2?"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    answer = "The answer is four."
+    output_ids = _encode(parser_tokenizer, answer)
+    expected = _decoded(parser_tokenizer, output_ids)
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    content = resp.json()["choices"][0]["message"]["content"]
+    assert content == expected
+
+
+@pytest.mark.asyncio
+async def test_e2e_token_identity(parser_client, parser_tokenizer):
+    """encode(derender(token_ids)) == token_ids (RL invariant)."""
+    messages = [{"role": "user", "content": "Hi"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    answer = "Hello! How can I help?"
+    output_ids = _encode(parser_tokenizer, answer)
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+        },
+    )
+    assert resp.status_code == 200
+    content = resp.json()["choices"][0]["message"]["content"]
+    re_encoded = _encode(parser_tokenizer, content)
+    assert output_ids == re_encoded
+
+
+@pytest.mark.asyncio
+async def test_e2e_non_ascii_roundtrip(parser_client, parser_tokenizer):
+    """CJK + emoji roundtrip without U+FFFD."""
+    messages = [{"role": "user", "content": "Reply in Chinese"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    answer = "你好世界 😀"
+    output_ids = _encode(parser_tokenizer, answer)
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+        },
+    )
+    assert resp.status_code == 200
+    content = resp.json()["choices"][0]["message"]["content"]
+    assert "�" not in content
+
+
+@pytest.mark.asyncio
+async def test_e2e_parsed_reasoning(parser_client, parser_tokenizer):
+    """<think>...</think> splits into reasoning + content."""
+    messages = [{"role": "user", "content": "What is 2+3?"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    reasoning_text = "The user wants 2 plus 3. That is 5."
+    answer_text = "The answer is 5."
+    output_text = f"<think>{reasoning_text}</think>{answer_text}"
+    output_ids = _require_markers_survive(parser_tokenizer, output_text, "</think>")
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {
+                "model": PARSER_MODEL,
+                "messages": messages,
+                "include_reasoning": True,
+            },
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    msg = resp.json()["choices"][0]["message"]
+    assert msg["reasoning"] is not None
+    assert reasoning_text in msg["reasoning"]
+    assert answer_text in msg["content"]
+    assert "<think>" not in msg["content"]
+
+
+@pytest.mark.asyncio
+async def test_e2e_parsed_tool_call(parser_client, parser_tokenizer):
+    """<tool_call> extracted into tool_calls field."""
+    messages = [{"role": "user", "content": "Weather in Paris?"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    output_text = (
+        "<think>Let me check the weather.</think>"
+        '<tool_call>\n{"name": "get_weather", '
+        '"arguments": {"city": "Paris"}}\n</tool_call>'
+    )
+    output_ids = _require_markers_survive(
+        parser_tokenizer,
+        output_text,
+        "</think>",
+        "<tool_call>",
+        "</tool_call>",
+    )
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {
+                "model": PARSER_MODEL,
+                "messages": messages,
+                "tools": _E2E_TOOLS,
+                "tool_choice": "auto",
+            },
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    choice = resp.json()["choices"][0]
+    assert choice["message"]["tool_calls"]
+    assert choice["message"]["tool_calls"][0]["function"]["name"] == "get_weather"
+
+
+@pytest.mark.asyncio
+async def test_e2e_parsed_reasoning_and_tool_call(parser_client, parser_tokenizer):
+    """Reasoning + tool call in the same output."""
+    messages = [{"role": "user", "content": "Weather in Paris?"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    reasoning_text = "I should look up the weather."
+    tool_text = (
+        '<tool_call>\n{"name": "get_weather", '
+        '"arguments": {"city": "Paris"}}\n</tool_call>'
+    )
+    output_text = f"<think>{reasoning_text}</think>{tool_text}"
+    output_ids = _require_markers_survive(
+        parser_tokenizer, output_text, "</think>", "<tool_call>"
+    )
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {
+                "model": PARSER_MODEL,
+                "messages": messages,
+                "tools": _E2E_TOOLS,
+                "tool_choice": "auto",
+                "include_reasoning": True,
+            },
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    choice = resp.json()["choices"][0]
+    assert choice["message"]["reasoning"] is not None
+    assert reasoning_text in choice["message"]["reasoning"]
+    assert choice["message"]["tool_calls"]
+
+
+@pytest.mark.asyncio
+async def test_e2e_no_chat_request_fallback(parser_client, parser_tokenizer):
+    """Without chat_request, derender falls back to plain detokenization."""
+    messages = [{"role": "user", "content": "Hello"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    answer = "Hi there!"
+    output_ids = _encode(parser_tokenizer, answer)
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+        },
+    )
+    assert resp.status_code == 200
+    content = resp.json()["choices"][0]["message"]["content"]
+    assert "Hi" in content
+
+
+# ---------------------------------------------------------------------------
+# E2E: HarmonyParser + GPT-OSS
+# ---------------------------------------------------------------------------
+
+HARMONY_MODEL = "openai/gpt-oss-20b"
+
+
+def _ensure_harmony_vocab():
+    """Pre-cache the o200k_base BPE file needed by openai-harmony.
+
+    The Rust tiktoken-rs backend downloads from Azure Blob Storage, which
+    may be unreachable in some environments.  When the cache is cold we
+    fetch the file ourselves and place it in ``/tmp/tiktoken-rs-cache/``
+    using the SHA-1(URL) filename that tiktoken-rs expects.
+    """
+    import hashlib
+    import urllib.request
+    from pathlib import Path
+
+    url = "https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken"
+    cache_dir = Path("/tmp/tiktoken-rs-cache")
+    cache_key = hashlib.sha1(url.encode()).hexdigest()
+    cache_file = cache_dir / cache_key
+    if not cache_file.exists():
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        urllib.request.urlretrieve(url, cache_file)
+
+
+@pytest.fixture(scope="module")
+def harmony_server():
+    _ensure_harmony_vocab()
+    args = [
+        "--trust-remote-code",
+        "--enable-auto-tool-choice",
+        "--tool-call-parser",
+        "openai",
+        "--reasoning-parser",
+        "openai_gptoss",
+    ]
+    with RemoteLaunchRenderServer(HARMONY_MODEL, args) as remote_server:
+        yield remote_server
+
+
+@pytest_asyncio.fixture
+async def harmony_client(harmony_server):
+    async with httpx.AsyncClient(
+        base_url=harmony_server.url_for(""), timeout=60.0
+    ) as http_client:
+        yield http_client
+
+
+@pytest.fixture(scope="module")
+def harmony_tokenizer():
+    return get_tokenizer(HARMONY_MODEL, trust_remote_code=True)
+
+
+def _harmony_extract_assistant_ids(
+    tokenizer, assistant_msg: dict, user_content: str = "test"
+) -> list[int]:
+    """Extract assistant token IDs via apply_chat_template diff."""
+    prompt = [{"role": "user", "content": user_content}]
+    full = prompt + [assistant_msg]
+    text_prompt = tokenizer.apply_chat_template(
+        prompt, add_generation_prompt=True, tokenize=False
+    )
+    text_full = tokenizer.apply_chat_template(
+        full, add_generation_prompt=False, tokenize=False
+    )
+    prompt_ids = tokenizer.encode(text_prompt)
+    full_ids = tokenizer.encode(text_full)
+    assistant_ids = list(full_ids[len(prompt_ids) :])
+    if not assistant_ids:
+        pytest.skip("Could not extract assistant tokens for Harmony")
+    return assistant_ids
+
+
+@pytest.mark.asyncio
+async def test_e2e_harmony_plain_roundtrip(harmony_client, harmony_tokenizer):
+    """GPT-OSS content-only roundtrip."""
+    messages = [{"role": "user", "content": "What is 2+2?"}]
+    gen_req = await _e2e_render_chat(harmony_client, HARMONY_MODEL, messages)
+
+    assistant_msg = {"role": "assistant", "content": "Four."}
+    output_ids = _harmony_extract_assistant_ids(harmony_tokenizer, assistant_msg)
+
+    resp = await harmony_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": HARMONY_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {
+                "model": HARMONY_MODEL,
+                "messages": messages,
+            },
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    content = resp.json()["choices"][0]["message"]["content"]
+    assert content is not None and len(content) > 0
+    assert "Four" in content
+
+
+@pytest.mark.asyncio
+async def test_e2e_harmony_reasoning(harmony_client, harmony_tokenizer):
+    """GPT-OSS reasoning: analysis channel extracted."""
+    messages = [{"role": "user", "content": "Add 2 and 3."}]
+    gen_req = await _e2e_render_chat(harmony_client, HARMONY_MODEL, messages)
+
+    reasoning_text = "The user wants 2 plus 3."
+    answer_text = "The answer is 5."
+    assistant_msg = {
+        "role": "assistant",
+        "thinking": reasoning_text,
+        "content": answer_text,
+    }
+    output_ids = _harmony_extract_assistant_ids(harmony_tokenizer, assistant_msg)
+
+    decoded = harmony_tokenizer.decode(output_ids)
+    if reasoning_text not in decoded:
+        pytest.skip("Harmony template did not render thinking")
+
+    resp = await harmony_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": HARMONY_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {
+                "model": HARMONY_MODEL,
+                "messages": messages,
+                "include_reasoning": True,
+            },
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    msg = resp.json()["choices"][0]["message"]
+    assert msg["reasoning"] is not None
+    assert reasoning_text in msg["reasoning"]
+    assert answer_text in (msg["content"] or "")

tests/entrypoints/serve/tokenize/test_serving_tokenization.py

@@ -15,7 +15,7 @@ from vllm.entrypoints.serve.tokenize.protocol import (
     TokenizeChatRequest,
     TokenizeCompletionRequest,
 )
-from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
+from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
 from vllm.v1.engine.async_llm import AsyncLLM
 
 MODEL_NAME = "openai-community/gpt2"
@@ -58,7 +58,7 @@ class MockModelConfig:
         return self.diff_sampling_param or {}
 
 
-def _build_serving_tokenization(engine: AsyncLLM) -> OpenAIServingTokenization:
+def _build_serving_tokenization(engine: AsyncLLM) -> ServingTokenization:
     models = OpenAIServingModels(
         engine_client=engine,
         base_model_paths=BASE_MODEL_PATHS,
@@ -71,8 +71,7 @@ def _build_serving_tokenization(engine: AsyncLLM) -> OpenAIServingTokenization:
         chat_template=None,
         chat_template_content_format="auto",
     )
-    return OpenAIServingTokenization(
-        engine,
+    return ServingTokenization(
         models,
         openai_serving_render=serving_render,
         request_logger=None,

tests/evals/gsm8k/test_gsm8k_correctness.py

@@ -78,7 +78,16 @@ def test_gsm8k_correctness(config_filename):
             "Skipping DeepSeek-V3.2 and DeepSeek-R1 on ROCm platforms "
             "due to agent pool disk space issues and pod evictions."
         )
+    if current_platform.is_rocm() and (
+        "Qwen3.5-35B-A3B-MXFP4-AITER-TP2" in config_filename.name
+    ):
+        from vllm.platforms.rocm import on_gfx950
 
+        if not on_gfx950():
+            pytest.skip(
+                "Skipping Qwen3.5-35B-A3B-MXFP4-AITER-TP2 on non-GFX950 platforms. "
+                "The quantization scheme is not supported on non-GFX950 platforms."
+            )
     # Parse server arguments from config (use shlex to handle quoted strings)
     server_args_str = eval_config.get("server_args", "")
     server_args = shlex.split(server_args_str) if server_args_str else []

tests/kernels/attention/test_attention_selector.py

@@ -15,16 +15,14 @@ from vllm.config import (
 from vllm.platforms import current_platform
 from vllm.platforms.cpu import CpuPlatform
 
-# CudaPlatform and RocmPlatform import their respective compiled C extensions
-# at module level, raising ModuleNotFoundError on incompatible builds.
-try:
+if current_platform.is_cuda():
     from vllm.platforms.cuda import CudaPlatform
-except (ImportError, ModuleNotFoundError):
+else:
     CudaPlatform = None
 
-try:
+if current_platform.is_rocm():
     from vllm.platforms.rocm import RocmPlatform
-except (ImportError, ModuleNotFoundError):
+else:
     RocmPlatform = None
 
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -434,9 +432,15 @@ def test_per_head_quant_scales_backend_selection(
     [
         ("FLASH_ATTN", True, True),  # FlashAttn supports non-causal
         ("FLASH_ATTN", False, True),  # FlashAttn also works with causal
-        ("FLASHINFER", True, False),  # FlashInfer does not support non-causal
-        ("FLASHINFER", False, True),  # FlashInfer works with causal
-    ],
+    ]
+    + (
+        [
+            ("FLASHINFER", True, True),  # FlashInfer supports non-causal
+            ("FLASHINFER", False, True),  # FlashInfer works with causal
+        ]
+        if CudaPlatform is not None
+        else []
+    ),
 )
 def test_non_causal_backend_selection(
     backend_name: str, use_non_causal: bool, should_succeed: bool
@@ -459,11 +463,12 @@ def test_non_causal_backend_selection(
         attention_config=attention_config, cache_config=cache_config
     )
 
-    if CudaPlatform is None:
-        pytest.skip("CudaPlatform not available")
+    platform = CudaPlatform or RocmPlatform
+    if platform is None:
+        pytest.skip("CudaPlatform and RocmPlatform are not available")
     with (
         set_current_vllm_config(vllm_config),
-        patch("vllm.platforms.current_platform", CudaPlatform()),
+        patch("vllm.platforms.current_platform", platform()),
     ):
         if should_succeed:
             backend = get_attn_backend(

tests/kernels/attention/test_prefix_prefill.py

@@ -5,10 +5,12 @@ import math
 import random
 import time
 from collections.abc import Callable
+from contextlib import nullcontext
 
 import pytest
 import torch
 import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
 
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
@@ -557,15 +559,21 @@ def test_contexted_kv_attention_alibi(
             query_len, seq_len, alibi_slopes, device, dtype
         )
 
-        # Compute attention
-        out = F.scaled_dot_product_attention(
-            q_sdpa,
-            k_sdpa,
-            v_sdpa,
-            attn_mask=alibi_mask,
-            dropout_p=0.0,
-            scale=scale,
-        )
+        # Compute attention. On ROCm we force use of the Math SDPA backend rather than
+        # the Flash or Mem-Efficient backends for increased numerical accuracy
+        if current_platform.is_rocm():
+            sdpa_context = sdpa_kernel(SDPBackend.MATH)
+        else:
+            sdpa_context = nullcontext()
+        with sdpa_context:
+            out = F.scaled_dot_product_attention(
+                q_sdpa,
+                k_sdpa,
+                v_sdpa,
+                attn_mask=alibi_mask,
+                dropout_p=0.0,
+                scale=scale,
+            )
 
         # Reshape output back to [query_len, num_heads, head_size]
         out = out.view(num_heads, query_len, head_size).permute(1, 0, 2)

tests/kernels/attention/test_rocm_triton_attn_dsv4.py

@@ -90,7 +90,9 @@ def _ref_sparse_prefill_ragged(
     return out.to(torch.bfloat16)
 
 
-def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
+def _pack_fp8_ds_mla_cache(
+    kv: torch.Tensor, block_size: int, is_extra: bool = False
+) -> torch.Tensor:
     assert kv.shape[-1] == HEAD_DIM
     num_tokens = kv.shape[0]
     num_blocks = (num_tokens + block_size - 1) // block_size
@@ -101,7 +103,9 @@ def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
     )
     cache_flat = cache.view(torch.uint8).flatten()
     kv_nope_fp8 = (
-        kv[:, :NOPE_HEAD_DIM].to(current_platform.fp8_dtype()).view(torch.uint8)
+        kv[:, :NOPE_HEAD_DIM]
+        .to(torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype())
+        .view(torch.uint8)
     )
     kv_rope_u8 = kv[:, NOPE_HEAD_DIM:].contiguous().view(torch.uint8)
 
@@ -120,7 +124,7 @@ def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
 
 
 def _read_fp8_ds_mla_cache(
-    cache: torch.Tensor, slot: int, block_size: int
+    cache: torch.Tensor, slot: int, block_size: int, is_extra: bool = False
 ) -> torch.Tensor:
     cache_flat = cache.view(torch.uint8).flatten()
     block_idx = slot // block_size
@@ -129,7 +133,9 @@ def _read_fp8_ds_mla_cache(
     token_base = block_base + pos * 576
 
     nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM]
-    nope = nope_u8.view(current_platform.fp8_dtype()).to(torch.float32)
+    nope = nope_u8.view(
+        torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype()
+    ).to(torch.float32)
     rope_u8 = cache_flat[
         token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
     ]
@@ -157,7 +163,9 @@ def _ref_sparse_decode_ragged(
         ]
         if extra_cache is not None and extra_rows is not None:
             row_kv.extend(
-                _read_fp8_ds_mla_cache(extra_cache, int(slot), block_size)
+                _read_fp8_ds_mla_cache(
+                    extra_cache, int(slot), block_size, is_extra=True
+                )
                 for slot in extra_rows[query_idx]
             )
 
@@ -326,7 +334,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None:
     main_kv = torch.randn(6, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
     extra_kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
     main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size)
-    extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
+    extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
     main_indices = torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device)
     main_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device=device)
     extra_indices = torch.tensor([1, 3, 0], dtype=torch.int32, device=device)
@@ -477,7 +485,7 @@ def test_sparse_attn_decode_split_k_kernel(
         rows = [[1, 3, 0, 5, 2, 4], [3, 0, 6]]
         extra_kv = torch.randn(7, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
         extra_rows = rows
-        extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
+        extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
         extra_indices, extra_indptr = _ragged_from_rows(rows, device)
 
     attn_sink = (

tests/kernels/attention/test_triton_unified_attention.py

@@ -18,11 +18,7 @@ HEAD_SIZES = [128, 256]
 BLOCK_SIZES = [16]
 
 DTYPES = [torch.bfloat16]
-QDTYPES = (
-    [None, torch.float8_e4m3fn]
-    if not current_platform.is_rocm()
-    else [None, torch.float8_e4m3fnuz]
-)
+QDTYPES = [None, current_platform.fp8_dtype()]
 FP8_DTYPE = current_platform.fp8_dtype()
 
 # one value large enough to test overflow in index calculation.

tests/kernels/core/test_minimax_reduce_rms.py

@@ -10,8 +10,12 @@ from torch.multiprocessing import spawn
 from tests.kernels.utils import opcheck
 from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
 from vllm.distributed import cleanup_dist_env_and_memory
-from vllm.model_executor.layers.minimax_rms_norm import MiniMaxText01RMSNormTP
+from vllm.model_executor.layers.minimax_rms_norm import (
+    MiniMaxText01RMSNormTP,
+    rms_norm_tp,
+)
 from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON
 from vllm.utils.network_utils import get_open_port
 from vllm.utils.torch_utils import set_random_seed
 
@@ -54,8 +58,19 @@ def _worker_forward_qk(
     torch.manual_seed(seed + 1000 + local_rank)
     qkv = torch.randn(num_tokens, hq + hk + hk, dtype=dtype, device="cuda")
 
-    q_ref, k_ref, v_ref = qkv.clone().split([hq, hk, hk], dim=-1)
-    ref_q, ref_k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q_ref, k_ref)
+    # Reference: eager all-reduce path. ``forward_qk`` no longer all-reduces
+    # the variance (it is the tp==1 / already-reduced building block), so the
+    # multi-rank reference must use the eager path that performs the global
+    # variance all-reduce, matching the fused kernel below.
+    ref_q, ref_k = rms_norm_tp._minimax_qk_norm_tp_eager(
+        qkv.clone(),
+        q_norm.weight,
+        k_norm.weight,
+        hq,
+        hk,
+        world_size,
+        eps,
+    )
 
     # Set up Lamport workspace.
     from vllm.distributed.parallel_state import get_tp_group
@@ -150,3 +165,44 @@ def test_minimax_reduce_rms_qk(
         nprocs=world_size,
         join=True,
     )
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda() or not HAS_TRITON,
+    reason="CUDA and Triton required",
+)
+@pytest.mark.parametrize("num_tokens", [1, 7, 128, 333, 2049])
+@pytest.mark.parametrize("hidden_dims", [(3072, 512), (768, 256), (3000, 500)])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("tp_world", [1, 4, 8])
+@pytest.mark.parametrize("eps", [1e-6])
+@pytest.mark.parametrize("seed", [42])
+def test_minimax_qk_norm_triton_fallback(
+    monkeypatch, num_tokens, hidden_dims, dtype, tp_world, eps, seed
+):
+    """Single-GPU check: Triton fallback kernels vs the pure-torch reference.
+
+    The all-reduce is a TP communication barrier, so it is monkeypatched to
+    identity here; both the Triton path and the reference see the same
+    (patched) reduction. This validates the kernel math and the folded
+    ``/ tp_world`` scaling without needing multiple ranks -- ``hidden_dims``
+    are the per-rank q/k segment widths.
+    """
+    monkeypatch.setattr(rms_norm_tp, "_all_reduce_variance", lambda v: v)
+
+    q_size, kv_size = hidden_dims
+    device = "cuda"
+    torch.manual_seed(seed)
+    qkv = torch.randn(num_tokens, q_size + 2 * kv_size, dtype=dtype, device=device)
+    q_weight = torch.randn(q_size, dtype=dtype, device=device)
+    k_weight = torch.randn(kv_size, dtype=dtype, device=device)
+
+    q_triton, k_triton = rms_norm_tp._minimax_qk_norm_tp_fallback(
+        qkv, q_weight, k_weight, q_size, kv_size, 0, tp_world, eps
+    )
+    q_ref, k_ref = rms_norm_tp._minimax_qk_norm_tp_eager(
+        qkv, q_weight, k_weight, q_size, kv_size, tp_world, eps
+    )
+
+    torch.testing.assert_close(q_triton, q_ref, atol=3e-2, rtol=3e-2)
+    torch.testing.assert_close(k_triton, k_ref, atol=3e-2, rtol=3e-2)

tests/kernels/moe/test_ocp_mx_moe.py

@@ -9,6 +9,7 @@ import pytest
 import torch
 from packaging import version
 
+from vllm._aiter_ops import is_aiter_found
 from vllm.platforms import current_platform
 from vllm.utils.flashinfer import has_flashinfer
 
@@ -31,17 +32,15 @@ HOPPER_MXFP4_BF16_AVAILABLE = (
 # ROCm platform and dependencies
 ROCM_AVAILABLE = current_platform.is_rocm()
 ROCM_TRITON_KERNELS_AVAILABLE = False
-ROCM_AITER_AVAILABLE = False
+ROCM_AITER_AVAILABLE = is_aiter_found()
 ROCM_GFX950 = False
 
 if ROCM_AVAILABLE:
-    from vllm._aiter_ops import rocm_aiter_ops
     from vllm.platforms.rocm import on_gfx950
     from vllm.utils.import_utils import has_triton_kernels
 
     ROCM_TRITON_KERNELS_AVAILABLE = has_triton_kernels()
     ROCM_GFX950 = on_gfx950()
-    ROCM_AITER_AVAILABLE = rocm_aiter_ops.is_enabled()
 
     if ROCM_AITER_AVAILABLE:
         from aiter.ops.triton.moe.quant_moe import upcast_from_mxfp
@@ -83,7 +82,7 @@ def enable_pickle(monkeypatch):
     [
         ModelCase("fxmarty/qwen_1.5-moe-a2.7b-mxfp4", tp=2),
         ModelCase("fxmarty/deepseek_r1_3_layers_mxfp4", tp=8),
-        ModelCase("fxmarty/Llama-4-Scout-17B-16E-Instruct-2-layers-mxfp4", tp=1),
+        ModelCase("mawong-amd/Llama-4-Scout-17B-16E-Instruct-2-layers-mxfp4", tp=1),
         ModelCase("fxmarty/Llama-3.1-70B-Instruct-2-layers-mxfp6", tp=1),
         ModelCase("fxmarty/Llama-3.1-70B-Instruct-2-layers-mxfp6", tp=4),
     ],
@@ -102,6 +101,7 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase):
         tensor_parallel_size=model_case.tp,
         load_format="dummy",
         compilation_config={"cudagraph_capture_sizes": [16]},
+        gpu_memory_utilization=0.8,  # mxfp6 models use more scratch space
     ) as llm:
         # Disabled as check_model is broken: https://github.com/vllm-project/vllm/pull/18465#issuecomment-3329880562
         # def check_model(model):
@@ -1267,7 +1267,7 @@ def test_rocm_mxfp4_moe_oracle(
 
     This test validates that the oracle functions work end-to-end:
     - select_mxfp4_moe_backend() selects a valid backend
-    - convert_to_mxfp4_moe_kernel_format() converts weights without error
+    - convert_gpt_oss_weight_to_mxfp4_moe_kernel_format() converts weights without error
     - make_mxfp4_moe_quant_config() builds a valid quant config
     - make_mxfp4_moe_kernel() creates a kernel that runs without error
     - The kernel output is within accuracy tolerance of reference
@@ -1287,7 +1287,7 @@ def test_rocm_mxfp4_moe_oracle(
     from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
         Mxfp4MoeBackend,
         backend_to_kernel_cls,
-        convert_to_mxfp4_moe_kernel_format,
+        convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,
         make_mxfp4_moe_kernel,
         make_mxfp4_moe_quant_config,
     )
@@ -1387,7 +1387,7 @@ def test_rocm_mxfp4_moe_oracle(
 
     # Convert weights using oracle
     w13_conv, w2_conv, w13_scale_conv, w2_scale_conv, w13_bias_conv, w2_bias_conv = (
-        convert_to_mxfp4_moe_kernel_format(
+        convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
             mxfp4_backend=backend,
             layer=layer,  # type: ignore[arg-type]
             w13_weight=w13_quant,
@@ -1423,7 +1423,7 @@ def test_rocm_mxfp4_moe_oracle(
             mxfp4_backend=backend,
             experts_cls=experts_cls,
             routing_tables=None,
-            shared_experts=None,
+            layer=None,
         )
 
         # Create inputs

tests/kernels/quantization/test_per_token_group_quant.py

@@ -345,6 +345,63 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(
         )
 
 
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="packed FP8 per-token-group quant kernel requires a CUDA-alike GPU",
+)
+def test_per_token_group_quant_fp8_packed_large_mn():
+    """Regression test for https://github.com/vllm-project/vllm/issues/45099.
+
+    Some background: gridDim.x and gridDim.y have different limits of 2^31 - 1 and
+    2^16 - 1, respectively.
+    Prior code introduced a bug where it incorrectly assumed grid.x and y both have
+    2^31 - 1 limits and mixed them up, which doesn't surface until the kernel is
+    launched with a large mn that exceeds grid.y limit (2^16 - 1).
+
+    This issue doesn't surface often because each forward pass only processes a
+    bounded token batch, not the full context.
+    Quantizing tensors with more rows than that will fail at launch with
+    "CUDA error: invalid argument".
+    This is a differential test that compares fp8 output against Triton output
+    reference when token size sits just above the gridDim.y 2^16 - 1 limit.
+    """
+
+    device = "cuda"
+    group_size = 128
+    # hidden 2048 -> 2048/128 = 16 groups per row -> kx=16, ry=1: one grid row per mn
+    # row, so any mn > 65535 overflowed grid.y before the fix.
+    num_tokens, hidden_dim = 65537, 2048
+    torch.manual_seed(42)
+    x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
+
+    out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
+        x,
+        group_size=group_size,
+        use_ue8m0=True,
+    )
+
+    with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
+        ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
+            x, group_size, use_ue8m0=True
+        )
+
+    assert torch.equal(out_q, ref_q), "Quantized output mismatch"
+
+    # Vectorized packed-scale check; the per-element loop used by the smaller
+    # tests is too slow at this size. groups_per_row is a multiple of 4 here,
+    # so there is no K padding and the packed view lines up.
+    mn = num_tokens
+    groups_per_row = hidden_dim // group_size
+    k_num_packed = (groups_per_row + 3) // 4
+    assert groups_per_row % 4 == 0
+    ref_exponents = (ref_s.reshape(mn, groups_per_row).view(torch.int32) >> 23) & 0xFF
+    exp = ref_exponents.view(mn, k_num_packed, 4)
+    expected = (
+        exp[..., 0] | (exp[..., 1] << 8) | (exp[..., 2] << 16) | (exp[..., 3] << 24)
+    )
+    assert torch.equal(out_s_packed.cpu(), expected.cpu()), "Packed scale mismatch"
+
+
 @pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
 @pytest.mark.parametrize("group_size", [64, 128])
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")

tests/kernels/quantization/test_triton_scaled_mm.py

@@ -60,8 +60,10 @@ def test_rocm_compressed_tensors_w8a8(
     vllm_runner, example_prompts, model_path, max_tokens, num_logprobs
 ):
     dtype = "bfloat16"
-
-    with vllm_runner(model_path, dtype=dtype) as vllm_model:
+    # Pin to TRITON_ATTN, see https://github.com/vllm-project/vllm/issues/46179
+    with vllm_runner(
+        model_path, dtype=dtype, attention_backend="TRITON_ATTN"
+    ) as vllm_model:
         vllm_model.generate_greedy_logprobs(example_prompts, max_tokens, num_logprobs)