forked from Karylab-cklius/vllm
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
76c973e13c | ||
|
|
53275a22d6 | ||
|
|
9729e05917 | ||
|
|
be37de73a9 | ||
|
|
e633514f50 | ||
|
|
6647a1a88b | ||
|
|
56afa45bf7 | ||
|
|
41cf8bbd60 | ||
|
|
427b2793f0 | ||
|
|
f4ab9994b5 | ||
|
|
8c9b156eca | ||
|
|
738af995a7 | ||
|
|
293c3895e2 | ||
|
|
e7fef86e50 | ||
|
|
9cb8ed5008 | ||
|
|
c5905f7760 | ||
|
|
22c6542fa7 | ||
|
|
483bda03a7 | ||
|
|
92b3a243d5 | ||
|
|
14682903c8 | ||
|
|
cbfaaeceeb | ||
|
|
cf6c0d2518 | ||
|
|
1b563d1134 | ||
|
|
cbdfa83c84 | ||
|
|
12846bbf88 | ||
|
|
f074bd6cef | ||
|
|
579d7b3705 | ||
|
|
cb7eb5c9b4 | ||
|
|
4f3c528941 | ||
|
|
0f854f78e8 | ||
|
|
a3028cebbf | ||
|
|
e327716282 | ||
|
|
de1828b58f | ||
|
|
6542ed479c | ||
|
|
485f33d796 | ||
|
|
1cad156ac5 | ||
|
|
99bd07204b |
@@ -16,7 +16,6 @@ steps:
|
||||
- tests/kernels/test_onednn.py
|
||||
- tests/kernels/test_awq_int4_to_int8.py
|
||||
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
|
||||
- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
|
||||
@@ -25,8 +24,7 @@ steps:
|
||||
pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
|
||||
pytest -x -v -s tests/kernels/test_onednn.py
|
||||
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
|
||||
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
|
||||
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py"
|
||||
|
||||
- label: CPU-Compatibility Tests
|
||||
depends_on: []
|
||||
|
||||
@@ -37,8 +37,7 @@ function cpu_tests() {
|
||||
pytest -x -v -s tests/kernels/test_onednn.py
|
||||
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
|
||||
pytest -x -v -s tests/kernels/core/test_cpu_activation.py
|
||||
pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
|
||||
pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"
|
||||
|
||||
# skip tests requiring model downloads if HF_TOKEN is not set
|
||||
# due to rate-limits
|
||||
|
||||
@@ -1,39 +0,0 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
|
||||
|
||||
uv pip install --system -r "${REQUIREMENTS_FILE}"
|
||||
|
||||
NIXL_METADATA=$(python3 - <<'PY'
|
||||
import importlib.metadata as metadata
|
||||
|
||||
import torch
|
||||
|
||||
cuda_version = torch.version.cuda
|
||||
if cuda_version is None:
|
||||
raise SystemExit("torch.version.cuda is not set")
|
||||
|
||||
print(cuda_version.split(".", 1)[0], metadata.version("nixl"))
|
||||
PY
|
||||
)
|
||||
read -r CUDA_MAJOR NIXL_VERSION <<<"${NIXL_METADATA}"
|
||||
|
||||
# nixl>=1.1.0 can install multiple CUDA wheel variants. Keep only the variant
|
||||
# matching this CI image so nixl_ep_cpp links against the available libcudart.
|
||||
uv pip uninstall --system nixl-cu12 nixl-cu13 2>/dev/null || true
|
||||
uv pip install --system --no-deps "nixl-cu${CUDA_MAJOR}==${NIXL_VERSION}"
|
||||
|
||||
python3 - <<'PY'
|
||||
import importlib.metadata as metadata
|
||||
|
||||
for package_name in ("nixl", "nixl-cu12", "nixl-cu13"):
|
||||
try:
|
||||
version = metadata.version(package_name)
|
||||
except metadata.PackageNotFoundError:
|
||||
version = "not installed"
|
||||
print(f"{package_name}: {version}")
|
||||
PY
|
||||
+13
-15
@@ -1238,11 +1238,14 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/serve
|
||||
- tests/entrypoints/rpc
|
||||
- tests/entrypoints/serve/instrumentator
|
||||
- tests/tool_use
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
- pytest -v -s entrypoints/serve/instrumentator
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
|
||||
- pytest -v -s tool_use
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -1272,14 +1275,10 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
- tests/entrypoints/generate
|
||||
- tests/tool_use
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s tool_use
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -1369,7 +1368,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/tool_parsers
|
||||
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
|
||||
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling
|
||||
|
||||
- label: OpenAI API correctness # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -2746,11 +2745,14 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/serve
|
||||
- tests/entrypoints/rpc
|
||||
- tests/entrypoints/serve/instrumentator
|
||||
- tests/tool_use
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
- pytest -v -s entrypoints/serve/instrumentator
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
|
||||
- pytest -v -s tool_use
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -2780,14 +2782,10 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
- tests/entrypoints/generate
|
||||
- tests/tool_use
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s tool_use
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
|
||||
timeout_in_minutes: 180
|
||||
|
||||
@@ -11,7 +11,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
|
||||
key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
|
||||
@@ -22,7 +22,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
@@ -34,7 +34,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
@@ -46,7 +46,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
|
||||
@@ -58,7 +58,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
|
||||
@@ -73,7 +73,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
|
||||
|
||||
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
|
||||
@@ -87,7 +87,7 @@ steps:
|
||||
- vllm/v1/worker/kv_connector_model_runner_mixin.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
|
||||
@@ -102,5 +102,5 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
|
||||
@@ -11,7 +11,7 @@ steps:
|
||||
- tests/entrypoints/
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/tool_parsers
|
||||
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
|
||||
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text
|
||||
|
||||
- label: Entrypoints Integration (LLM)
|
||||
key: entrypoints-integration-llm
|
||||
@@ -60,13 +60,9 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
- tests/entrypoints/generate
|
||||
- tests/tool_use
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
- pytest -v -s entrypoints/generate
|
||||
- pytest -v -s tool_use
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -102,11 +98,14 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/serve
|
||||
- tests/entrypoints/rpc
|
||||
- tests/entrypoints/serve/instrumentator
|
||||
- tests/tool_use
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
- pytest -v -s entrypoints/serve/instrumentator
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
|
||||
- pytest -v -s tool_use
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -154,5 +153,6 @@ steps:
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/model_executor/models/whisper.py
|
||||
commands: # LMEval
|
||||
- pytest -s entrypoints/openai/correctness/
|
||||
|
||||
@@ -86,7 +86,7 @@ steps:
|
||||
- tests/v1/metrics
|
||||
- tests/entrypoints/openai/correctness/test_lmeval.py
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# split the test to avoid interference
|
||||
- pytest -v -s -m 'not cpu_test' v1/core
|
||||
|
||||
@@ -45,19 +45,19 @@ steps:
|
||||
- vllm/entrypoints/serve/
|
||||
- vllm/v1/engine/
|
||||
- tests/utils.py
|
||||
# - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
# - tests/entrypoints/rpc/test_collective_rpc.py
|
||||
- tests/entrypoints/serve/disagg/test_serving_tokens.py
|
||||
- tests/entrypoints/serve/instrumentator/test_basic.py
|
||||
- tests/entrypoints/serve/instrumentator/test_metrics.py
|
||||
# - tests/entrypoints/serve/dev/test_sleep.py
|
||||
# - tests/entrypoints/serve/instrumentator/test_sleep.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
# - pytest -v -s entrypoints/rpc/test_collective_rpc.py
|
||||
- pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
|
||||
- pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
|
||||
- pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
|
||||
# - pytest -v -s entrypoints/serve/dev/test_sleep.py
|
||||
# - pytest -v -s entrypoints/serve/instrumentator/test_sleep.py
|
||||
|
||||
- label: Rust Frontend Core Correctness
|
||||
timeout_in_minutes: 30
|
||||
|
||||
@@ -10,7 +10,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Add label
|
||||
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
|
||||
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
|
||||
with:
|
||||
script: |
|
||||
github.rest.issues.addLabels({
|
||||
|
||||
@@ -14,7 +14,7 @@ jobs:
|
||||
steps:
|
||||
- name: Label issues based on keywords
|
||||
id: label-step
|
||||
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
|
||||
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
|
||||
with:
|
||||
script: |
|
||||
// Configuration: Add new labels and keywords here
|
||||
@@ -315,7 +315,7 @@ jobs:
|
||||
|
||||
- name: CC users for labeled issues
|
||||
if: steps.label-step.outputs.labels_added != '[]'
|
||||
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
|
||||
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
|
||||
with:
|
||||
script: |
|
||||
// Configuration: Map labels to GitHub users to CC
|
||||
@@ -392,7 +392,7 @@ jobs:
|
||||
|
||||
- name: Request missing ROCm info from issue author
|
||||
if: contains(steps.label-step.outputs.labels_added, 'rocm') && contains(toJSON(github.event.issue.labels.*.name), 'bug')
|
||||
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
|
||||
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
|
||||
with:
|
||||
script: |
|
||||
const body = (context.payload.issue.body || '').toLowerCase();
|
||||
|
||||
@@ -12,7 +12,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Update PR description
|
||||
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
|
||||
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
|
||||
with:
|
||||
script: |
|
||||
const { owner, repo } = context.repo;
|
||||
@@ -55,7 +55,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Post welcome comment for first-time contributors
|
||||
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
|
||||
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
|
||||
with:
|
||||
script: |
|
||||
const { owner, repo } = context.repo;
|
||||
|
||||
@@ -20,7 +20,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Check PR label and author merge count
|
||||
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
|
||||
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
|
||||
with:
|
||||
script: |
|
||||
const { data: pr } = await github.rest.pulls.get({
|
||||
@@ -49,7 +49,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
|
||||
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
|
||||
- uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
- run: echo "::add-matcher::.github/workflows/matchers/actionlint.json"
|
||||
|
||||
@@ -21,7 +21,7 @@ repos:
|
||||
rev: v21.1.2
|
||||
hooks:
|
||||
- id: clang-format
|
||||
exclude: 'csrc/(moe/topk_softmax_kernels.cu|libtorch_stable/quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
|
||||
exclude: 'csrc/(moe/topk_softmax_kernels.cu|quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
|
||||
types_or: [c++, cuda]
|
||||
args: [--style=file, --verbose]
|
||||
- repo: https://github.com/DavidAnson/markdownlint-cli2
|
||||
|
||||
+3
-1
@@ -315,7 +315,8 @@ set(VLLM_EXT_SRC
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
list(APPEND VLLM_EXT_SRC
|
||||
"csrc/minimax_reduce_rms_kernel.cu")
|
||||
"csrc/minimax_reduce_rms_kernel.cu"
|
||||
"csrc/minimax_m3_build_k2q_csr.cu")
|
||||
|
||||
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
|
||||
|
||||
@@ -637,6 +638,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
"csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"
|
||||
"csrc/libtorch_stable/pos_encoding_kernels.cu"
|
||||
"csrc/libtorch_stable/fused_qknorm_rope_kernel.cu"
|
||||
"csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu"
|
||||
"csrc/libtorch_stable/layernorm_kernels.cu"
|
||||
"csrc/libtorch_stable/layernorm_quant_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
|
||||
|
||||
@@ -369,18 +369,6 @@ else()
|
||||
add_compile_definitions(-DVLLM_NUMA_DISABLED)
|
||||
endif()
|
||||
|
||||
# check if the pytorch wheel ships libopenblas.so.
|
||||
set(VLLM_OPENBLAS_LIB "")
|
||||
if (NOT ENABLE_X86_ISA)
|
||||
file(GLOB _VLLM_TORCH_OPENBLAS_LIBS
|
||||
"${TORCH_INSTALL_PREFIX}/lib/libopenblas*.so*")
|
||||
# Note: we don't link openblas directly to _C extension, as it's available through libtorch.so
|
||||
if (_VLLM_TORCH_OPENBLAS_LIBS)
|
||||
list(GET _VLLM_TORCH_OPENBLAS_LIBS 0 VLLM_OPENBLAS_LIB)
|
||||
message(STATUS "CPU OpenBLAS library: ${VLLM_OPENBLAS_LIB}")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
#
|
||||
# Generate CPU attention dispatch header
|
||||
#
|
||||
@@ -399,7 +387,6 @@ endif()
|
||||
#
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/cpu/activation.cpp"
|
||||
"csrc/cpu/sgl-kernels/fla.cpp"
|
||||
"csrc/cpu/utils.cpp"
|
||||
"csrc/cpu/spec_decode_utils.cpp"
|
||||
"csrc/cpu/layernorm.cpp"
|
||||
@@ -409,13 +396,6 @@ set(VLLM_EXT_SRC
|
||||
"csrc/cpu/cpu_attn.cpp"
|
||||
"csrc/cpu/torch_bindings.cpp")
|
||||
|
||||
if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND
|
||||
VLLM_RVV_VLEN GREATER 0 AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/cpu/cpu_wna16.cpp"
|
||||
${VLLM_EXT_SRC})
|
||||
endif()
|
||||
|
||||
if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/cpu/shm.cpp"
|
||||
@@ -423,12 +403,6 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
|
||||
${VLLM_EXT_SRC})
|
||||
endif()
|
||||
|
||||
if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/cpu/shm.cpp"
|
||||
${VLLM_EXT_SRC})
|
||||
endif()
|
||||
|
||||
if(USE_ONEDNN)
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/cpu/dnnl_kernels.cpp"
|
||||
@@ -437,6 +411,7 @@ endif()
|
||||
|
||||
if (ENABLE_X86_ISA)
|
||||
set(VLLM_EXT_SRC_SGL
|
||||
"csrc/cpu/sgl-kernels/fla.cpp"
|
||||
"csrc/cpu/sgl-kernels/conv.cpp"
|
||||
"csrc/cpu/sgl-kernels/gemm.cpp"
|
||||
"csrc/cpu/sgl-kernels/gemm_int8.cpp"
|
||||
@@ -448,7 +423,6 @@ if (ENABLE_X86_ISA)
|
||||
"csrc/cpu/sgl-kernels/moe_fp8.cpp")
|
||||
|
||||
set(VLLM_EXT_SRC_AVX512
|
||||
"csrc/cpu/sgl-kernels/fla.cpp"
|
||||
"csrc/cpu/shm.cpp"
|
||||
"csrc/cpu/cpu_wna16.cpp"
|
||||
"csrc/cpu/cpu_fused_moe.cpp"
|
||||
@@ -465,7 +439,6 @@ if (ENABLE_X86_ISA)
|
||||
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp")
|
||||
|
||||
set(VLLM_EXT_SRC_AVX2
|
||||
"csrc/cpu/sgl-kernels/fla.cpp"
|
||||
"csrc/cpu/utils.cpp"
|
||||
"csrc/cpu/spec_decode_utils.cpp"
|
||||
"csrc/cpu/cpu_attn.cpp"
|
||||
@@ -539,9 +512,6 @@ else()
|
||||
USE_SABI 3
|
||||
WITH_SOABI
|
||||
)
|
||||
if (VLLM_OPENBLAS_LIB)
|
||||
target_compile_definitions(_C PRIVATE VLLM_HAS_OPENBLAS)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
message(STATUS "Enabling C extension.")
|
||||
|
||||
@@ -31,7 +31,7 @@ endif()
|
||||
|
||||
if(VLLM_FLASH_ATTN_SRC_DIR)
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn SOURCE_DIR
|
||||
vllm-flash-attn SOURCE_DIR
|
||||
${VLLM_FLASH_ATTN_SRC_DIR}
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
)
|
||||
@@ -39,7 +39,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn
|
||||
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
|
||||
GIT_TAG dd62dac706b1cf7895bd99b18c6cb7e7e117ee25
|
||||
GIT_TAG bce29425653ec0fbc579d329883030e832d15ada
|
||||
GIT_PROGRESS TRUE
|
||||
# Don't share the vllm-flash-attn build between build types
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
|
||||
@@ -94,10 +94,6 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
|
||||
: reg(RVVI(__riscv_vle16_v_f16, LMUL_256)(
|
||||
static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
|
||||
|
||||
explicit FP16Vec16(const c10::Half v)
|
||||
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _f16, LMUL_256)(
|
||||
RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
|
||||
|
||||
explicit FP16Vec16(const FP32Vec16& vec);
|
||||
|
||||
void save(void* ptr) const {
|
||||
@@ -169,9 +165,6 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
|
||||
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
|
||||
|
||||
explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {};
|
||||
explicit BF16Vec16(const c10::BFloat16 v)
|
||||
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16, LMUL_256)(
|
||||
RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
|
||||
explicit BF16Vec16(const FP32Vec16&);
|
||||
|
||||
void save(void* ptr) const {
|
||||
@@ -297,9 +290,6 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
|
||||
}
|
||||
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
|
||||
}
|
||||
explicit BF16Vec16(const c10::BFloat16 v)
|
||||
: reg_fp32(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(static_cast<float>(v),
|
||||
VEC_ELEM_NUM)) {}
|
||||
explicit BF16Vec16(const FP32Vec16&);
|
||||
void save(void* ptr) const {
|
||||
float tmp[16];
|
||||
@@ -639,19 +629,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
: reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)(
|
||||
data.reg, data.reg)) {};
|
||||
explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
|
||||
explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
|
||||
const uint64_t q_values = static_cast<uint64_t>(value);
|
||||
auto packed = RVVI(__riscv_vmv_v_x_u64, LMUL_1024)(q_values, VEC_ELEM_NUM);
|
||||
auto lane_ids = RVVI(__riscv_vid_v_u64, LMUL_1024)(VEC_ELEM_NUM);
|
||||
auto shifts =
|
||||
RVVI(__riscv_vsll_vx_u64, LMUL_1024)(lane_ids, 2, VEC_ELEM_NUM);
|
||||
auto shifted =
|
||||
RVVI(__riscv_vsrl_vv_u64, LMUL_1024)(packed, shifts, VEC_ELEM_NUM);
|
||||
auto idx64 =
|
||||
RVVI(__riscv_vand_vx_u64, LMUL_1024)(shifted, 0xF, VEC_ELEM_NUM);
|
||||
auto idx32 = RVVI(__riscv_vnsrl_wx_u32, LMUL_512)(idx64, 0, VEC_ELEM_NUM);
|
||||
reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx32, VEC_ELEM_NUM);
|
||||
}
|
||||
explicit FP32Vec16(const FP16Vec16& v);
|
||||
|
||||
#ifdef __riscv_zvfbfmin
|
||||
@@ -664,10 +641,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
|
||||
#endif
|
||||
|
||||
// FP8 stub: dead code on RISC-V (fp8 KV cache is x86-only), needed for
|
||||
// load_b_pair_vec template to compile on all platforms.
|
||||
explicit FP32Vec16(const BF16Vec32&, int) : FP32Vec16() {}
|
||||
|
||||
FP32Vec16 operator+(const FP32Vec16& b) const {
|
||||
return FP32Vec16(
|
||||
RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
|
||||
@@ -918,30 +891,6 @@ inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
|
||||
acc = acc.fma(a, b);
|
||||
}
|
||||
|
||||
template <typename VecT>
|
||||
static void interleave_save_16b(const VecT& vec0, const VecT& vec1, void* ptr) {
|
||||
alignas(64) uint16_t values0[VecT::VEC_ELEM_NUM];
|
||||
alignas(64) uint16_t values1[VecT::VEC_ELEM_NUM];
|
||||
vec0.save(values0);
|
||||
vec1.save(values1);
|
||||
|
||||
auto* packed = reinterpret_cast<uint32_t*>(ptr);
|
||||
for (int32_t i = 0; i < VecT::VEC_ELEM_NUM; ++i) {
|
||||
packed[i] = static_cast<uint32_t>(values0[i]) |
|
||||
(static_cast<uint32_t>(values1[i]) << 16);
|
||||
}
|
||||
}
|
||||
|
||||
static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1,
|
||||
void* ptr) {
|
||||
interleave_save_16b(vec0, vec1, ptr);
|
||||
}
|
||||
|
||||
static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1,
|
||||
void* ptr) {
|
||||
interleave_save_16b(vec0, vec1, ptr);
|
||||
}
|
||||
|
||||
#ifdef __riscv_zvfbfmin
|
||||
template <>
|
||||
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
|
||||
|
||||
+1
-106
@@ -89,35 +89,6 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
|
||||
}
|
||||
};
|
||||
|
||||
struct FP16Vec16 : public Vec<FP16Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
ss16x8x2_t reg;
|
||||
|
||||
explicit FP16Vec16(const void* ptr) {
|
||||
reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr);
|
||||
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
|
||||
}
|
||||
|
||||
explicit FP16Vec16(bool, const void* ptr) : FP16Vec16(ptr) {}
|
||||
|
||||
explicit FP16Vec16(const FP32Vec16&);
|
||||
|
||||
void save(void* ptr) const {
|
||||
vec_xst(reg.val[0], 0, (signed short*)ptr);
|
||||
vec_xst(reg.val[1], 16, (signed short*)ptr);
|
||||
}
|
||||
|
||||
void save(void* ptr, int elem_num) const {
|
||||
int num = std::max(0, std::min(elem_num, VEC_ELEM_NUM));
|
||||
if (num <= 8) {
|
||||
vec_xst_len(reg.val[0], (signed short*)ptr, num * 2);
|
||||
} else {
|
||||
vec_xst(reg.val[0], 0, (signed short*)ptr);
|
||||
vec_xst_len(reg.val[1], (signed short*)ptr + 8, (num - 8) * 2);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
struct BF16Vec16 : public Vec<BF16Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
|
||||
@@ -129,8 +100,6 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
|
||||
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
|
||||
}
|
||||
|
||||
explicit BF16Vec16(bool, const void* ptr) : BF16Vec16(ptr) {}
|
||||
|
||||
explicit BF16Vec16(const FP32Vec16&);
|
||||
|
||||
void save(void* ptr) const {
|
||||
@@ -410,8 +379,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
reg.val[3] = vec_xl(48, ptr);
|
||||
}
|
||||
|
||||
explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
|
||||
|
||||
explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
|
||||
|
||||
explicit FP32Vec16(const FP32Vec16& data) {
|
||||
@@ -435,7 +402,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
reg.val[3] = data.reg.val[1];
|
||||
}
|
||||
|
||||
explicit FP32Vec16(const FP16Vec16& v);
|
||||
explicit FP32Vec16(const BF16Vec16& v) {
|
||||
reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]);
|
||||
reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]);
|
||||
@@ -769,40 +735,6 @@ inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
|
||||
#endif
|
||||
}
|
||||
|
||||
inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
|
||||
alignas(16) float temp_fp32[16];
|
||||
alignas(16) c10::Half temp_fp16[16];
|
||||
|
||||
vec_xst(v.reg.val[0], 0, temp_fp32);
|
||||
vec_xst(v.reg.val[1], 16, temp_fp32);
|
||||
vec_xst(v.reg.val[2], 32, temp_fp32);
|
||||
vec_xst(v.reg.val[3], 48, temp_fp32);
|
||||
|
||||
for (int i = 0; i < 16; i++) {
|
||||
temp_fp16[i] = c10::Half(temp_fp32[i]);
|
||||
}
|
||||
|
||||
reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)temp_fp16);
|
||||
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)temp_fp16);
|
||||
}
|
||||
|
||||
inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
|
||||
alignas(16) c10::Half temp_fp16[16];
|
||||
alignas(16) float temp_fp32[16];
|
||||
|
||||
vec_xst(v.reg.val[0], 0, (signed short*)temp_fp16);
|
||||
vec_xst(v.reg.val[1], 16, (signed short*)temp_fp16);
|
||||
|
||||
for (int i = 0; i < 16; i++) {
|
||||
temp_fp32[i] = float(temp_fp16[i]);
|
||||
}
|
||||
|
||||
reg.val[0] = vec_xl(0, temp_fp32);
|
||||
reg.val[1] = vec_xl(16, temp_fp32);
|
||||
reg.val[2] = vec_xl(32, temp_fp32);
|
||||
reg.val[3] = vec_xl(48, temp_fp32);
|
||||
}
|
||||
|
||||
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
|
||||
#ifdef _ARCH_PWR10
|
||||
__vector signed short ret[4];
|
||||
@@ -862,43 +794,6 @@ inline void prefetch(const void* addr) {
|
||||
__asm__ __volatile__("dcbt 0, %0" : : "r"(addr) : "memory");
|
||||
}
|
||||
|
||||
struct INT8Vec64 {
|
||||
__vector signed char data[4];
|
||||
|
||||
INT8Vec64() = default;
|
||||
|
||||
explicit INT8Vec64(const int8_t* ptr) {
|
||||
data[0] = vec_xl(0, ptr);
|
||||
data[1] = vec_xl(16, ptr);
|
||||
data[2] = vec_xl(32, ptr);
|
||||
data[3] = vec_xl(48, ptr);
|
||||
}
|
||||
|
||||
explicit INT8Vec64(bool, const int8_t* ptr) : INT8Vec64(ptr) {}
|
||||
|
||||
void save(int8_t* ptr) const {
|
||||
vec_xst(data[0], 0, ptr);
|
||||
vec_xst(data[1], 16, ptr);
|
||||
vec_xst(data[2], 32, ptr);
|
||||
vec_xst(data[3], 48, ptr);
|
||||
}
|
||||
|
||||
void save(int8_t* ptr, int elem_num) const {
|
||||
if (elem_num <= 0) return;
|
||||
|
||||
int full_vecs = elem_num / 16;
|
||||
for (int i = 0; i < full_vecs && i < 4; i++) {
|
||||
vec_xst(data[i], i * 16, ptr);
|
||||
}
|
||||
|
||||
int remaining = elem_num % 16;
|
||||
if (remaining > 0 && full_vecs < 4) {
|
||||
vec_xst_len(data[full_vecs], ptr + full_vecs * 16, remaining);
|
||||
}
|
||||
}
|
||||
|
||||
void nt_save(int8_t* ptr) const { save(ptr); }
|
||||
};
|
||||
} // namespace vec_op
|
||||
}; // namespace vec_op
|
||||
|
||||
#endif
|
||||
|
||||
@@ -1,82 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
#include <ATen/native/CPUBlas.h>
|
||||
|
||||
// Unlike brgemm, PyTorch does not publicly expose at::native::cpublas::gemm
|
||||
// If OpenBLS is available in the PyTorch wheel, we rely on it for fast
|
||||
// bf16:bf16->fp32 GEMMs Otherwise, we fall back to PyTorch reference BLAS path.
|
||||
#if defined(VLLM_HAS_OPENBLAS)
|
||||
extern "C" void sbgemm_(char* transa, char* transb, int* m, int* n, int* k,
|
||||
float* alpha, const at::BFloat16* a, int* lda,
|
||||
const at::BFloat16* b, int* ldb, float* beta, float* c,
|
||||
int* ldc);
|
||||
|
||||
extern "C" void sgemm_(char* transa, char* transb, int* m, int* n, int* k,
|
||||
float* alpha, const float* a, int* lda, const float* b,
|
||||
int* ldb, float* beta, float* c, int* ldc);
|
||||
|
||||
inline char blas_transpose(at::native::TransposeType trans) {
|
||||
switch (trans) {
|
||||
case at::native::TransposeType::NoTranspose:
|
||||
return 'n';
|
||||
case at::native::TransposeType::Transpose:
|
||||
return 't';
|
||||
case at::native::TransposeType::ConjTranspose:
|
||||
return 'c';
|
||||
}
|
||||
return 'n';
|
||||
}
|
||||
|
||||
inline void blas_gemm(at::native::TransposeType transa,
|
||||
at::native::TransposeType transb, int64_t m, int64_t n,
|
||||
int64_t k, float alpha, const at::BFloat16* a,
|
||||
int64_t lda, const at::BFloat16* b, int64_t ldb,
|
||||
float beta, float* c, int64_t ldc) {
|
||||
char transa_ = blas_transpose(transa);
|
||||
char transb_ = blas_transpose(transb);
|
||||
int m_ = static_cast<int>(m);
|
||||
int n_ = static_cast<int>(n);
|
||||
int k_ = static_cast<int>(k);
|
||||
int lda_ = static_cast<int>(lda);
|
||||
int ldb_ = static_cast<int>(ldb);
|
||||
int ldc_ = static_cast<int>(ldc);
|
||||
sbgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
|
||||
c, &ldc_);
|
||||
}
|
||||
|
||||
inline void blas_gemm(at::native::TransposeType transa,
|
||||
at::native::TransposeType transb, int64_t m, int64_t n,
|
||||
int64_t k, float alpha, const float* a, int64_t lda,
|
||||
const float* b, int64_t ldb, float beta, float* c,
|
||||
int64_t ldc) {
|
||||
char transa_ = blas_transpose(transa);
|
||||
char transb_ = blas_transpose(transb);
|
||||
int m_ = static_cast<int>(m);
|
||||
int n_ = static_cast<int>(n);
|
||||
int k_ = static_cast<int>(k);
|
||||
int lda_ = static_cast<int>(lda);
|
||||
int ldb_ = static_cast<int>(ldb);
|
||||
int ldc_ = static_cast<int>(ldc);
|
||||
sgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
|
||||
c, &ldc_);
|
||||
}
|
||||
|
||||
inline void blas_gemm(at::native::TransposeType, at::native::TransposeType,
|
||||
int64_t, int64_t, int64_t, float, const at::Half*,
|
||||
int64_t, const at::Half*, int64_t, float, float*,
|
||||
int64_t) {
|
||||
TORCH_CHECK(false, "CPU OpenBLAS hgemm is not available.");
|
||||
}
|
||||
#else
|
||||
template <typename scalar_t>
|
||||
inline void blas_gemm(at::native::TransposeType transa,
|
||||
at::native::TransposeType transb, int64_t m, int64_t n,
|
||||
int64_t k, float alpha, const scalar_t* a, int64_t lda,
|
||||
const scalar_t* b, int64_t ldb, float beta, float* c,
|
||||
int64_t ldc) {
|
||||
auto gemm = at::native::cpublas::gemm_no_downcast_stub.DEFAULT;
|
||||
gemm(c10::CppTypeToScalarType<scalar_t>::value, transa, transb, m, n, k,
|
||||
at::Scalar(alpha), a, lda, b, ldb, at::Scalar(beta), c, ldc);
|
||||
}
|
||||
#endif
|
||||
+141
-278
@@ -301,42 +301,25 @@ void chunk_gated_delta_rule_kernel_impl(
|
||||
// attn = k_beta @ key.transpose(-1, -2)
|
||||
// attn: [B, HV, num_chunk, chunk_size, chunk_size]
|
||||
// transpose and pack for key
|
||||
if constexpr (brgemm_supported()) {
|
||||
pack_vnni<scalar_t>(
|
||||
/* dst */ k_transpose,
|
||||
/* src */ curr_k_pad,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* ld_src */ qk_head_size,
|
||||
/* ld_dst */ chunk_size);
|
||||
// k_beta @ key.transpose(-1, -2)
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ chunk_size,
|
||||
/* ldc */ chunk_size,
|
||||
/* add_C */ false,
|
||||
/* A */ curr_k_beta,
|
||||
/* B */ k_transpose,
|
||||
/* C */ curr_attn);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::Transpose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
chunk_size,
|
||||
chunk_size,
|
||||
qk_head_size,
|
||||
1.0f,
|
||||
curr_k_pad,
|
||||
qk_head_size,
|
||||
curr_k_beta,
|
||||
qk_head_size,
|
||||
0.0f,
|
||||
curr_attn,
|
||||
chunk_size);
|
||||
}
|
||||
pack_vnni<scalar_t>(
|
||||
/* dst */ k_transpose,
|
||||
/* src */ curr_k_pad,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* ld_src */ qk_head_size,
|
||||
/* ld_dst */ chunk_size);
|
||||
// k_beta @ key.transpose(-1, -2)
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ chunk_size,
|
||||
/* ldc */ chunk_size,
|
||||
/* add_C */ false,
|
||||
/* A */ curr_k_beta,
|
||||
/* B */ k_transpose,
|
||||
/* C */ curr_attn);
|
||||
// attn = attn * decay_mask
|
||||
for (int64_t m = 0; m < chunk_size; m++) {
|
||||
at::vec::map2<float>(
|
||||
@@ -430,42 +413,25 @@ void chunk_gated_delta_rule_kernel_impl(
|
||||
// k_beta_g = k_beta * g: [B, HV, num_chunk, chunk_size, EK]
|
||||
// k_cumdecay: [B, HV, num_chunk, chunk_size, EK]
|
||||
// pack for value
|
||||
if constexpr (brgemm_supported()) {
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ v_pack,
|
||||
/* src */ curr_v_beta,
|
||||
/* N */ chunk_size,
|
||||
/* K */ v_head_size,
|
||||
/* ld_src */ v_head_size,
|
||||
/* ld_dst */ v_head_size);
|
||||
// value = attn @ v_beta
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ curr_attn_reduced,
|
||||
/* B */ v_pack,
|
||||
/* C */ curr_value);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::NoTranspose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
v_head_size,
|
||||
chunk_size,
|
||||
chunk_size,
|
||||
1.0f,
|
||||
curr_v_beta,
|
||||
v_head_size,
|
||||
curr_attn_reduced,
|
||||
chunk_size,
|
||||
0.0f,
|
||||
curr_value,
|
||||
v_head_size);
|
||||
}
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ v_pack,
|
||||
/* src */ curr_v_beta,
|
||||
/* N */ chunk_size,
|
||||
/* K */ v_head_size,
|
||||
/* ld_src */ v_head_size,
|
||||
/* ld_dst */ v_head_size);
|
||||
// value = attn @ v_beta
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ curr_attn_reduced,
|
||||
/* B */ v_pack,
|
||||
/* C */ curr_value);
|
||||
// k_beta_g = k_beta * g.exp().unsqueeze(-1)
|
||||
for (int64_t j = 0; j < chunk_size; j++) {
|
||||
int64_t i = 0;
|
||||
@@ -479,42 +445,25 @@ void chunk_gated_delta_rule_kernel_impl(
|
||||
}
|
||||
}
|
||||
// pack for k_beta_g
|
||||
if constexpr (brgemm_supported()) {
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ k_beta_g_pack,
|
||||
/* src */ k_beta_g,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* ld_src */ qk_head_size,
|
||||
/* ld_dst */ qk_head_size);
|
||||
// k_cumdecay = attn @ k_beta_g
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ qk_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ qk_head_size,
|
||||
/* ldc */ qk_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ curr_attn_reduced,
|
||||
/* B */ k_beta_g_pack,
|
||||
/* C */ k_cumdecay);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::NoTranspose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
qk_head_size,
|
||||
chunk_size,
|
||||
chunk_size,
|
||||
1.0f,
|
||||
k_beta_g,
|
||||
qk_head_size,
|
||||
curr_attn_reduced,
|
||||
chunk_size,
|
||||
0.0f,
|
||||
k_cumdecay,
|
||||
qk_head_size);
|
||||
}
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ k_beta_g_pack,
|
||||
/* src */ k_beta_g,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* ld_src */ qk_head_size,
|
||||
/* ld_dst */ qk_head_size);
|
||||
// k_cumdecay = attn @ k_beta_g
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ qk_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ qk_head_size,
|
||||
/* ldc */ qk_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ curr_attn_reduced,
|
||||
/* B */ k_beta_g_pack,
|
||||
/* C */ k_cumdecay);
|
||||
for (int i = 0; i < chunk_size; i++) {
|
||||
at::vec::map<scalar_t>(
|
||||
[](fVec x) { return x; },
|
||||
@@ -602,42 +551,25 @@ void chunk_gated_delta_rule_kernel_impl(
|
||||
|
||||
// attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
|
||||
// k_transpose_i = k_i.transpose(-1, -2)
|
||||
if constexpr (brgemm_supported()) {
|
||||
pack_vnni<scalar_t>(
|
||||
/* dst */ k_transpose_i,
|
||||
/* src */ k_i,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* ld_src */ qk_head_size,
|
||||
/* ld_dst */ chunk_size);
|
||||
// attn_i = q_i @ k_transpose_i
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ chunk_size,
|
||||
/* ldc */ chunk_size,
|
||||
/* add_C */ false,
|
||||
/* A */ q_i,
|
||||
/* B */ k_transpose_i,
|
||||
/* C */ attn_i);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::Transpose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
chunk_size,
|
||||
chunk_size,
|
||||
qk_head_size,
|
||||
1.0f,
|
||||
k_i,
|
||||
qk_head_size,
|
||||
q_i,
|
||||
qk_head_size,
|
||||
0.0f,
|
||||
attn_i,
|
||||
chunk_size);
|
||||
}
|
||||
pack_vnni<scalar_t>(
|
||||
/* dst */ k_transpose_i,
|
||||
/* src */ k_i,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* ld_src */ qk_head_size,
|
||||
/* ld_dst */ chunk_size);
|
||||
// attn_i = q_i @ k_transpose_i
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ chunk_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ chunk_size,
|
||||
/* ldc */ chunk_size,
|
||||
/* add_C */ false,
|
||||
/* A */ q_i,
|
||||
/* B */ k_transpose_i,
|
||||
/* C */ attn_i);
|
||||
// attn_i = attn_i * decay_mask_i
|
||||
for (int64_t m = 0; m < chunk_size; m++) {
|
||||
auto attn_i_m = attn_i + m * chunk_size;
|
||||
@@ -677,45 +609,28 @@ void chunk_gated_delta_rule_kernel_impl(
|
||||
}
|
||||
|
||||
// pack for curr_last_recurrent_state
|
||||
if constexpr (brgemm_supported()) {
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ curr_last_recurrent_state_pack_reduced,
|
||||
/* src */ curr_last_recurrent_state_reduced,
|
||||
/* N */ qk_head_size,
|
||||
/* K */ v_head_size,
|
||||
/* ld_src */ v_head_size,
|
||||
/* ld_dst */ v_head_size);
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ curr_last_recurrent_state_pack_reduced,
|
||||
/* src */ curr_last_recurrent_state_reduced,
|
||||
/* N */ qk_head_size,
|
||||
/* K */ v_head_size,
|
||||
/* ld_src */ v_head_size,
|
||||
/* ld_dst */ v_head_size);
|
||||
|
||||
// v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
|
||||
// k_cumdecay_i: [chunk_size, EK]
|
||||
// curr_last_recurrent_state: [EK, EV]
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ k_cumdecay_i_reduced,
|
||||
/* B */ curr_last_recurrent_state_pack_reduced,
|
||||
/* C */ v_prime);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::NoTranspose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
v_head_size,
|
||||
chunk_size,
|
||||
qk_head_size,
|
||||
1.0f,
|
||||
curr_last_recurrent_state_reduced,
|
||||
v_head_size,
|
||||
k_cumdecay_i_reduced,
|
||||
qk_head_size,
|
||||
0.0f,
|
||||
v_prime,
|
||||
v_head_size);
|
||||
}
|
||||
// v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
|
||||
// k_cumdecay_i: [chunk_size, EK]
|
||||
// curr_last_recurrent_state: [EK, EV]
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ k_cumdecay_i_reduced,
|
||||
/* B */ curr_last_recurrent_state_pack_reduced,
|
||||
/* C */ v_prime);
|
||||
|
||||
// v_new = v_prime = v_i - v_prime
|
||||
// v_i: [chunk_size, EV]
|
||||
@@ -748,75 +663,41 @@ void chunk_gated_delta_rule_kernel_impl(
|
||||
}
|
||||
// attn_inter = qg @ curr_last_recurrent_state: [chunk_size, EV]
|
||||
// curr_last_recurrent_state: [EK, EV]
|
||||
if constexpr (brgemm_supported()) {
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ qg,
|
||||
/* B */ curr_last_recurrent_state_pack_reduced,
|
||||
/* C */ attn_inter);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::NoTranspose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
v_head_size,
|
||||
chunk_size,
|
||||
qk_head_size,
|
||||
1.0f,
|
||||
curr_last_recurrent_state_reduced,
|
||||
v_head_size,
|
||||
qg,
|
||||
qk_head_size,
|
||||
0.0f,
|
||||
attn_inter,
|
||||
v_head_size);
|
||||
}
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ qk_head_size,
|
||||
/* lda */ qk_head_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ qg,
|
||||
/* B */ curr_last_recurrent_state_pack_reduced,
|
||||
/* C */ attn_inter);
|
||||
|
||||
// core_attn_out[:, :, i] = attn_inter + attn_i @ v_new
|
||||
// pack for v_prime
|
||||
if constexpr (brgemm_supported()) {
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ v_prime_pack_reduced,
|
||||
/* src */ v_prime_reduced,
|
||||
/* N */ chunk_size,
|
||||
/* K */ v_head_size,
|
||||
/* ld_src */ v_head_size,
|
||||
/* ld_dst */ v_head_size);
|
||||
// attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
|
||||
// attn_i: [chunk_size, chunk_size]
|
||||
// v_new: [chunk_size, EV]
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ true,
|
||||
/* A */ attn_i_reduced,
|
||||
/* B */ v_prime_pack_reduced,
|
||||
/* C */ attn_inter);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::NoTranspose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
v_head_size,
|
||||
chunk_size,
|
||||
chunk_size,
|
||||
1.0f,
|
||||
v_prime_reduced,
|
||||
v_head_size,
|
||||
attn_i_reduced,
|
||||
chunk_size,
|
||||
1.0f,
|
||||
attn_inter,
|
||||
v_head_size);
|
||||
}
|
||||
pack_vnni2<scalar_t>(
|
||||
/* dst */ v_prime_pack_reduced,
|
||||
/* src */ v_prime_reduced,
|
||||
/* N */ chunk_size,
|
||||
/* K */ v_head_size,
|
||||
/* ld_src */ v_head_size,
|
||||
/* ld_dst */ v_head_size);
|
||||
// attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
|
||||
// attn_i: [chunk_size, chunk_size]
|
||||
// v_new: [chunk_size, EV]
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ chunk_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ true,
|
||||
/* A */ attn_i_reduced,
|
||||
/* B */ v_prime_pack_reduced,
|
||||
/* C */ attn_inter);
|
||||
|
||||
// core_attn_out[:, :, i] = attn_inter
|
||||
for (int64_t m = 0; m < chunk_size; m++) {
|
||||
@@ -881,34 +762,17 @@ void chunk_gated_delta_rule_kernel_impl(
|
||||
/* ld_dst */ chunk_size);
|
||||
// kgv = kg.transpose(-1, -2) @ v_new
|
||||
// v_new: [chunk_size, EV]
|
||||
if constexpr (brgemm_supported()) {
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ qk_head_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ kg_transpose,
|
||||
/* B */ v_prime_pack_reduced,
|
||||
/* C */ kgv);
|
||||
} else {
|
||||
blas_gemm(
|
||||
at::native::TransposeType::NoTranspose,
|
||||
at::native::TransposeType::NoTranspose,
|
||||
v_head_size,
|
||||
qk_head_size,
|
||||
chunk_size,
|
||||
1.0f,
|
||||
v_prime_reduced,
|
||||
v_head_size,
|
||||
kg_transpose,
|
||||
chunk_size,
|
||||
0.0f,
|
||||
kgv,
|
||||
v_head_size);
|
||||
}
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ qk_head_size,
|
||||
/* N */ v_head_size,
|
||||
/* K */ chunk_size,
|
||||
/* lda */ chunk_size,
|
||||
/* ldb */ v_head_size,
|
||||
/* ldc */ v_head_size,
|
||||
/* add_C */ false,
|
||||
/* A */ kg_transpose,
|
||||
/* B */ v_prime_pack_reduced,
|
||||
/* C */ kgv);
|
||||
// last_recurrent_state = 1) + 2)
|
||||
for (int64_t m = 0; m < qk_head_size; m++) {
|
||||
at::vec::map2<float>(
|
||||
@@ -1057,8 +921,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
|
||||
float k_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[k_scale_offset] : 1.0f;
|
||||
int64_t v_offset = si * v_strideS + bi * v_strideB + ni * v_strideH;
|
||||
int64_t o_offset = ((bi * seq_len + si) * v_num_heads + ni) * v_head_dim;
|
||||
// See: https://github.com/sgl-project/sglang/pull/26634
|
||||
float beta_val = 1 / (1 + std::exp(-b_ptr[bi * v_num_heads + ni]));
|
||||
float beta_val = 1 / (1 + std::exp(-b_ptr[ni]));
|
||||
fVec beta_vec = fVec(beta_val);
|
||||
int64_t dvi = 0;
|
||||
for (; dvi <= v_head_dim - VecSize; dvi += VecSize) {
|
||||
|
||||
@@ -4,12 +4,9 @@
|
||||
// clang-format off
|
||||
|
||||
#pragma once
|
||||
#include "common.h"
|
||||
#include "blas_gemm.h"
|
||||
#include <ATen/native/CPUBlas.h>
|
||||
|
||||
#if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
|
||||
#define CPU_CAPABILITY_AVX512
|
||||
#endif
|
||||
#include "common.h"
|
||||
|
||||
// amx-bf16
|
||||
#define TILE_M 16
|
||||
@@ -24,39 +21,31 @@ constexpr int block_size_n() {
|
||||
return 2 * TILE_N;
|
||||
}
|
||||
|
||||
constexpr bool brgemm_supported() {
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
return true;
|
||||
#else
|
||||
return false;
|
||||
#endif
|
||||
}
|
||||
|
||||
// define threshold using brgemm (intel AMX)
|
||||
template <typename T>
|
||||
inline bool can_use_brgemm(int M);
|
||||
template <>
|
||||
inline bool can_use_brgemm<at::BFloat16>(int M) {
|
||||
return brgemm_supported() && M > 4;
|
||||
return M > 4;
|
||||
}
|
||||
template <>
|
||||
inline bool can_use_brgemm<at::Half>(int M) {
|
||||
return brgemm_supported();
|
||||
return true;
|
||||
}
|
||||
// this requires PyTorch 2.7 or above
|
||||
template <>
|
||||
inline bool can_use_brgemm<int8_t>(int M) {
|
||||
return brgemm_supported() && M > 4;
|
||||
return M > 4;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline bool can_use_brgemm<uint8_t>(int M) {
|
||||
return brgemm_supported() && M > 4;
|
||||
return M > 4;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) {
|
||||
return brgemm_supported() && M > 4;
|
||||
return M > 4;
|
||||
}
|
||||
|
||||
// work around compiler internal error
|
||||
|
||||
@@ -11,9 +11,7 @@
|
||||
|
||||
#include <ATen/cpu/vec/functional.h>
|
||||
#include <ATen/cpu/vec/vec.h>
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
#include <immintrin.h>
|
||||
#endif
|
||||
namespace {
|
||||
|
||||
using namespace at::vec;
|
||||
|
||||
+8
-10
@@ -5,7 +5,7 @@
|
||||
#include <sys/stat.h>
|
||||
#include <unistd.h>
|
||||
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
#ifdef __aarch64__
|
||||
#include <atomic>
|
||||
#endif
|
||||
|
||||
@@ -38,7 +38,7 @@ struct KernelVecType<c10::Half> {
|
||||
};
|
||||
|
||||
struct ThreadSHMContext {
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
#ifdef __aarch64__
|
||||
// memory model is weaker on AArch64, so we use atomic variables for
|
||||
// consumer (load-acquire) and producer (store-release) to make sure
|
||||
// that a stamp cannot be ready before the corresponding data is ready.
|
||||
@@ -75,7 +75,7 @@ struct ThreadSHMContext {
|
||||
TORCH_CHECK(group_size <= MAX_SHM_RANK_NUM);
|
||||
TORCH_CHECK((size_t)this % 64 == 0);
|
||||
TORCH_CHECK((size_t)thread_shm_ptr % 64 == 0);
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
#ifdef __aarch64__
|
||||
_curr_thread_stamp[0].store(1, std::memory_order_relaxed);
|
||||
_curr_thread_stamp[1].store(1, std::memory_order_relaxed);
|
||||
_ready_thread_stamp[0].store(0, std::memory_order_relaxed);
|
||||
@@ -124,7 +124,7 @@ struct ThreadSHMContext {
|
||||
}
|
||||
|
||||
char get_curr_stamp(int idx) const {
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
#ifdef __aarch64__
|
||||
return _curr_thread_stamp[idx].load(std::memory_order_acquire);
|
||||
#else
|
||||
return _curr_thread_stamp[idx];
|
||||
@@ -132,7 +132,7 @@ struct ThreadSHMContext {
|
||||
}
|
||||
|
||||
char get_ready_stamp(int idx) const {
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
#ifdef __aarch64__
|
||||
return _ready_thread_stamp[idx].load(std::memory_order_acquire);
|
||||
#else
|
||||
return _ready_thread_stamp[idx];
|
||||
@@ -140,7 +140,7 @@ struct ThreadSHMContext {
|
||||
}
|
||||
|
||||
void next_stamp() {
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
#ifdef __aarch64__
|
||||
_curr_thread_stamp[local_stamp_buffer_idx].fetch_add(
|
||||
1, std::memory_order_release);
|
||||
#else
|
||||
@@ -150,7 +150,7 @@ struct ThreadSHMContext {
|
||||
}
|
||||
|
||||
void commit_ready_stamp() {
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
#ifdef __aarch64__
|
||||
_ready_thread_stamp[local_stamp_buffer_idx].store(
|
||||
_curr_thread_stamp[local_stamp_buffer_idx].load(
|
||||
std::memory_order_relaxed),
|
||||
@@ -186,10 +186,8 @@ struct ThreadSHMContext {
|
||||
break;
|
||||
}
|
||||
++_spinning_count;
|
||||
#if defined(__aarch64__)
|
||||
#ifdef __aarch64__
|
||||
__asm__ __volatile__("yield");
|
||||
#elif defined(__powerpc64__)
|
||||
__asm__ __volatile__("or 1,1,1");
|
||||
#else
|
||||
_mm_pause();
|
||||
#endif // __aarch64__
|
||||
|
||||
+21
-22
@@ -378,8 +378,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
#endif
|
||||
|
||||
// SHM CCL
|
||||
#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
|
||||
defined(__powerpc64__)
|
||||
#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__))
|
||||
ops.def(
|
||||
"init_shm_manager(str name, int group_size, int rank, int thread_num) -> "
|
||||
"int",
|
||||
@@ -448,25 +447,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"bool is_vnni) -> Tensor");
|
||||
ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
|
||||
|
||||
// Adapted from sglang: casual_conv1d kernels
|
||||
ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
|
||||
ops.impl("causal_conv1d_weight_pack", torch::kCPU,
|
||||
&causal_conv1d_weight_pack);
|
||||
ops.def(
|
||||
"causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
|
||||
"conv_states, Tensor? query_start_loc,"
|
||||
"Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
|
||||
"int pad_slot_id, bool is_vnni) -> "
|
||||
"Tensor");
|
||||
ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
|
||||
ops.def(
|
||||
"causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
|
||||
"weight, Tensor? bias, bool silu_activation,"
|
||||
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
|
||||
"bool is_vnni) -> Tensor");
|
||||
ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
|
||||
#endif
|
||||
|
||||
// Adapted from sglang: GDN kernels
|
||||
ops.def(
|
||||
"chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "
|
||||
@@ -490,6 +470,25 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"-> (Tensor, Tensor)");
|
||||
ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
|
||||
|
||||
// Adapted from sglang: casual_conv1d kernels
|
||||
ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
|
||||
ops.impl("causal_conv1d_weight_pack", torch::kCPU,
|
||||
&causal_conv1d_weight_pack);
|
||||
ops.def(
|
||||
"causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
|
||||
"conv_states, Tensor? query_start_loc,"
|
||||
"Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
|
||||
"int pad_slot_id, bool is_vnni) -> "
|
||||
"Tensor");
|
||||
ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
|
||||
ops.def(
|
||||
"causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
|
||||
"weight, Tensor? bias, bool silu_activation,"
|
||||
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
|
||||
"bool is_vnni) -> Tensor");
|
||||
ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
|
||||
#endif
|
||||
|
||||
// CPU attention kernels
|
||||
ops.def(
|
||||
"get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
|
||||
@@ -519,7 +518,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
|
||||
|
||||
// WNA16
|
||||
#if defined(__AVX512F__) || defined(__riscv_v)
|
||||
#if defined(__AVX512F__)
|
||||
ops.def(
|
||||
"cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
|
||||
"Tensor scales, Tensor? zeros, Tensor? g_idx, Tensor? bias, SymInt "
|
||||
|
||||
@@ -10,11 +10,20 @@
|
||||
|
||||
namespace vllm {
|
||||
|
||||
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
|
||||
// `alpha` and `beta` are applied to opposite operands:
|
||||
// - alpha lives INSIDE the activation (the activated half): the gated
|
||||
// activation computes act_half * sigmoid(alpha * act_half).
|
||||
// - beta is added to the OTHER (non-activated) half before the multiply.
|
||||
// So the result is always ACT(act_half, alpha) * (other_half + beta).
|
||||
// Which half is which depends on `act_first` (see below). Defaults
|
||||
// alpha=1.0, beta=0.0 reproduce the plain SwiGLU/GeGLU behavior.
|
||||
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float),
|
||||
bool act_first, bool HAS_CLAMP>
|
||||
__device__ __forceinline__ scalar_t compute(const scalar_t& x,
|
||||
const scalar_t& y,
|
||||
const float limit) {
|
||||
const float limit,
|
||||
const float alpha,
|
||||
const float beta) {
|
||||
if constexpr (act_first) {
|
||||
scalar_t gate = x;
|
||||
scalar_t up = y;
|
||||
@@ -22,7 +31,9 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
|
||||
gate = (scalar_t)fminf((float)gate, limit);
|
||||
up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
|
||||
}
|
||||
return ACT_FN(gate) * up;
|
||||
// act_first: gate is the activated half -> alpha applies to gate;
|
||||
// beta is added to up (the non-activated half).
|
||||
return ACT_FN(gate, alpha) * (scalar_t)((float)up + beta);
|
||||
} else {
|
||||
scalar_t gate = x;
|
||||
scalar_t up = y;
|
||||
@@ -30,55 +41,66 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
|
||||
gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
|
||||
up = (scalar_t)fminf((float)up, limit);
|
||||
}
|
||||
return gate * ACT_FN(up);
|
||||
// !act_first: up is the activated half -> alpha applies to up;
|
||||
// beta is added to gate (the non-activated half).
|
||||
return (scalar_t)((float)gate + beta) * ACT_FN(up, alpha);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
|
||||
template <typename packed_t,
|
||||
packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
|
||||
bool act_first, bool HAS_CLAMP>
|
||||
__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
|
||||
const packed_t& y,
|
||||
const float limit) {
|
||||
const float limit,
|
||||
const float alpha,
|
||||
const float beta) {
|
||||
if constexpr (act_first) {
|
||||
packed_t gate = x;
|
||||
packed_t up = y;
|
||||
float2 u = cast_to_float2(up);
|
||||
if constexpr (HAS_CLAMP) {
|
||||
float2 g = cast_to_float2(gate);
|
||||
float2 u = cast_to_float2(up);
|
||||
g.x = fminf(g.x, limit);
|
||||
g.y = fminf(g.y, limit);
|
||||
u.x = fmaxf(fminf(u.x, limit), -limit);
|
||||
u.y = fmaxf(fminf(u.y, limit), -limit);
|
||||
gate = cast_to_packed<packed_t>(g);
|
||||
up = cast_to_packed<packed_t>(u);
|
||||
}
|
||||
return packed_mul(PACKED_ACT_FN(gate), up);
|
||||
// act_first: gate is the activated half -> alpha applies to gate;
|
||||
// beta is added to up (the non-activated half).
|
||||
u.x += beta;
|
||||
u.y += beta;
|
||||
return packed_mul(PACKED_ACT_FN(gate, alpha), cast_to_packed<packed_t>(u));
|
||||
} else {
|
||||
packed_t gate = x;
|
||||
packed_t up = y;
|
||||
float2 g = cast_to_float2(gate);
|
||||
if constexpr (HAS_CLAMP) {
|
||||
float2 g = cast_to_float2(gate);
|
||||
float2 u = cast_to_float2(up);
|
||||
g.x = fmaxf(fminf(g.x, limit), -limit);
|
||||
g.y = fmaxf(fminf(g.y, limit), -limit);
|
||||
u.x = fminf(u.x, limit);
|
||||
u.y = fminf(u.y, limit);
|
||||
gate = cast_to_packed<packed_t>(g);
|
||||
up = cast_to_packed<packed_t>(u);
|
||||
}
|
||||
return packed_mul(gate, PACKED_ACT_FN(up));
|
||||
// !act_first: up is the activated half -> alpha applies to up;
|
||||
// beta is added to gate (the non-activated half).
|
||||
g.x += beta;
|
||||
g.y += beta;
|
||||
return packed_mul(cast_to_packed<packed_t>(g), PACKED_ACT_FN(up, alpha));
|
||||
}
|
||||
}
|
||||
|
||||
// Activation and gating kernel template.
|
||||
template <typename scalar_t, typename packed_t,
|
||||
scalar_t (*ACT_FN)(const scalar_t&),
|
||||
packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
|
||||
bool use_vec, bool HAS_CLAMP, bool use_256b = false>
|
||||
scalar_t (*ACT_FN)(const scalar_t&, const float),
|
||||
packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
|
||||
bool act_first, bool use_vec, bool HAS_CLAMP, bool use_256b = false>
|
||||
__global__ void act_and_mul_kernel(
|
||||
scalar_t* __restrict__ out, // [..., d]
|
||||
const scalar_t* __restrict__ input, // [..., 2, d]
|
||||
const int d, const float limit) {
|
||||
const int d, const float limit, const float alpha, const float beta) {
|
||||
const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
|
||||
const scalar_t* y_ptr = x_ptr + d;
|
||||
scalar_t* out_ptr = out + blockIdx.x * d;
|
||||
@@ -105,7 +127,7 @@ __global__ void act_and_mul_kernel(
|
||||
for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
|
||||
x.elts[j] =
|
||||
packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
|
||||
x.elts[j], y.elts[j], limit);
|
||||
x.elts[j], y.elts[j], limit, alpha, beta);
|
||||
}
|
||||
if constexpr (use_256b) {
|
||||
st256(x, &out_vec[i]);
|
||||
@@ -118,29 +140,34 @@ __global__ void act_and_mul_kernel(
|
||||
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
|
||||
const scalar_t x = VLLM_LDG(&x_ptr[idx]);
|
||||
const scalar_t y = VLLM_LDG(&y_ptr[idx]);
|
||||
out_ptr[idx] =
|
||||
compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
|
||||
out_ptr[idx] = compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(
|
||||
x, y, limit, alpha, beta);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Gated activations take an `alpha` argument that scales the sigmoid input
|
||||
// (`x * sigmoid(alpha * x)`). alpha defaults to 1.0 at all call sites, which
|
||||
// is exactly SiLU; only the clamp path (silu_and_mul_with_clamp) passes a
|
||||
// non-default alpha. Activations that do not use alpha simply ignore it.
|
||||
template <typename T>
|
||||
__device__ __forceinline__ T silu_kernel(const T& x) {
|
||||
// x * sigmoid(x)
|
||||
return (T)(((float)x) / (1.0f + expf((float)-x)));
|
||||
__device__ __forceinline__ T silu_kernel(const T& x, const float alpha) {
|
||||
// x * sigmoid(alpha * x)
|
||||
return (T)(((float)x) / (1.0f + expf((float)-x * alpha)));
|
||||
}
|
||||
|
||||
template <typename packed_t>
|
||||
__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val) {
|
||||
// x * sigmoid(x)
|
||||
__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val,
|
||||
const float alpha) {
|
||||
// x * sigmoid(alpha * x)
|
||||
float2 fval = cast_to_float2(val);
|
||||
fval.x = fval.x / (1.0f + expf(-fval.x));
|
||||
fval.y = fval.y / (1.0f + expf(-fval.y));
|
||||
fval.x = fval.x / (1.0f + expf(-fval.x * alpha));
|
||||
fval.y = fval.y / (1.0f + expf(-fval.y * alpha));
|
||||
return cast_to_packed<packed_t>(fval);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ __forceinline__ T gelu_kernel(const T& x) {
|
||||
__device__ __forceinline__ T gelu_kernel(const T& x, const float /*alpha*/) {
|
||||
// Equivalent to PyTorch GELU with 'none' approximation.
|
||||
// Refer to:
|
||||
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
|
||||
@@ -150,7 +177,8 @@ __device__ __forceinline__ T gelu_kernel(const T& x) {
|
||||
}
|
||||
|
||||
template <typename packed_t>
|
||||
__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
|
||||
__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val,
|
||||
const float /*alpha*/) {
|
||||
// Equivalent to PyTorch GELU with 'none' approximation.
|
||||
// Refer to:
|
||||
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
|
||||
@@ -162,7 +190,8 @@ __device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
|
||||
__device__ __forceinline__ T gelu_tanh_kernel(const T& x,
|
||||
const float /*alpha*/) {
|
||||
// Equivalent to PyTorch GELU with 'tanh' approximation.
|
||||
// Refer to:
|
||||
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
|
||||
@@ -176,7 +205,7 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
|
||||
|
||||
template <typename packed_t>
|
||||
__device__ __forceinline__ packed_t
|
||||
packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
packed_gelu_tanh_kernel(const packed_t& val, const float /*alpha*/) {
|
||||
// Equivalent to PyTorch GELU with 'tanh' approximation.
|
||||
// Refer to:
|
||||
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
|
||||
@@ -202,7 +231,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
// clamped (max only) and up input is clamped (both sides) before the
|
||||
// activation function is applied.
|
||||
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST, \
|
||||
HAS_CLAMP, LIMIT) \
|
||||
HAS_CLAMP, LIMIT, ALPHA, BETA) \
|
||||
auto dtype = input.scalar_type(); \
|
||||
int d = input.size(-1) / 2; \
|
||||
int64_t num_tokens = input.numel() / input.size(-1); \
|
||||
@@ -230,7 +259,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>( \
|
||||
out.mutable_data_ptr<scalar_t>(), \
|
||||
input.const_data_ptr<scalar_t>(), d, LIMIT); \
|
||||
input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA); \
|
||||
}); \
|
||||
} else { \
|
||||
VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
|
||||
@@ -240,7 +269,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>( \
|
||||
out.mutable_data_ptr<scalar_t>(), \
|
||||
input.const_data_ptr<scalar_t>(), d, LIMIT); \
|
||||
input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA); \
|
||||
}); \
|
||||
} \
|
||||
} else { \
|
||||
@@ -252,7 +281,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>( \
|
||||
out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(), \
|
||||
d, LIMIT); \
|
||||
d, LIMIT, ALPHA, BETA); \
|
||||
}); \
|
||||
}
|
||||
|
||||
@@ -260,14 +289,18 @@ void silu_and_mul(torch::stable::Tensor& out, // [..., d]
|
||||
torch::stable::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
true, false, 0.0f);
|
||||
true, false, 0.0f, 1.0f, 0.0f);
|
||||
}
|
||||
|
||||
void silu_and_mul_clamp(torch::stable::Tensor& out, // [..., d]
|
||||
torch::stable::Tensor& input, // [..., 2 * d]
|
||||
double limit) {
|
||||
double limit, double alpha, double beta) {
|
||||
// out = (gate.clamp(max=limit) * sigmoid(alpha * gate.clamp(max=limit)))
|
||||
// * (up.clamp(+-limit) + beta)
|
||||
// alpha=1.0, beta=0.0 reduce this to silu(gate) * up.
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
true, true, (float)limit);
|
||||
true, true, (float)limit, (float)alpha,
|
||||
(float)beta);
|
||||
}
|
||||
|
||||
void mul_and_silu(torch::stable::Tensor& out, // [..., d]
|
||||
@@ -276,21 +309,22 @@ void mul_and_silu(torch::stable::Tensor& out, // [..., d]
|
||||
// The difference between mul_and_silu and silu_and_mul is that mul_and_silu
|
||||
// applies the silu to the latter half of the input.
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
false, false, 0.0f);
|
||||
false, false, 0.0f, 1.0f, 0.0f);
|
||||
}
|
||||
|
||||
void gelu_and_mul(torch::stable::Tensor& out, // [..., d]
|
||||
torch::stable::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
|
||||
true, false, 0.0f);
|
||||
true, false, 0.0f, 1.0f, 0.0f);
|
||||
}
|
||||
|
||||
void gelu_tanh_and_mul(torch::stable::Tensor& out, // [..., d]
|
||||
torch::stable::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(
|
||||
vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
|
||||
vllm::packed_gelu_tanh_kernel, true, false,
|
||||
0.0f, 1.0f, 0.0f);
|
||||
}
|
||||
|
||||
namespace vllm {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
|
||||
#include "broadcast_load_epilogue_c2x.hpp"
|
||||
#include "cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp"
|
||||
|
||||
/*
|
||||
This file defines custom epilogues for fusing channel scales, token scales,
|
||||
|
||||
@@ -0,0 +1,615 @@
|
||||
/*
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
*
|
||||
* Horizontally-fused MiniMax-M3 attention pre-processing kernel.
|
||||
*
|
||||
* Replaces the per-token Python sequence in
|
||||
* ``MiniMaxM3SparseAttention.forward`` / ``MiniMaxM3Attention.forward``:
|
||||
*
|
||||
* q = q_norm(q); k = k_norm(k); q, k = rotary_emb(pos, q, k)
|
||||
* index_q = index_q_norm(index_q); index_k = index_k_norm(index_k)
|
||||
* index_q, index_k = rotary_emb(pos, index_q, index_k)
|
||||
* _insert_kv(k, v, index_k)
|
||||
*
|
||||
* All branches share head_dim=128 and the *same* partial-NeoX RoPE table
|
||||
* (``rotary_dim`` rotated, the trailing dims pass through). The four norms
|
||||
* are Gemma-style RMSNorm (``x * rsqrt(mean(x^2)+eps) * (1 + weight)``) with
|
||||
* independent weights.
|
||||
*
|
||||
* Everything lives in a single fused ``qkv`` tensor. The sparse layer's
|
||||
* fused projection (MinimaxM3QKVParallelLinearWithIndexer) emits, per token::
|
||||
*
|
||||
* [ q | k | v | index_q | index_k ] (the "5 results")
|
||||
*
|
||||
* while the dense layer emits just ``[ q | k | v ]``. The kernel reads the
|
||||
* index branch straight out of that packed row -- no separate index tensors.
|
||||
*
|
||||
* One kernel, one grid; each warp owns one (token, head-slot) pair. Slot
|
||||
* enumeration per token:
|
||||
* [0, nq) Q heads -> norm(q_w) + RoPE, write
|
||||
* qkv [nq, nq+nkv) K heads -> norm(k_w) + RoPE, write
|
||||
* qkv
|
||||
* (+ insert into key cache)
|
||||
* [nq+nkv, nq+2*nkv) V heads -> insert into value cache
|
||||
* IQ heads (niq) -> norm(iq_w) + RoPE, write iq
|
||||
* IK (1) -> norm(ik_w) + RoPE
|
||||
* (+ insert into index cache)
|
||||
*
|
||||
* The IQ/IK warps address the index_q/index_k sub-blocks *inside* qkv at the
|
||||
* fixed physical offsets (nq+2*nkv)*128 and (nq+2*nkv+niq)*128.
|
||||
*
|
||||
* Dense vs sparse is a compile-time choice via the ``kIsSparse``/``kInsertKV``
|
||||
* template bools (3 instantiations: dense <false,false>, sparse-profiling
|
||||
* <true,false>, sparse-serving <true,true>), so the index slots, the V slots
|
||||
* and the cache inserts fold away entirely on paths that don't use them. The
|
||||
* dense layer passes no caches/index: norm+RoPE happens in place and the
|
||||
* generic ``Attention`` layer owns the cache write.
|
||||
*
|
||||
* Q/K and (sparse) index_q/index_k are all rewritten in place inside the fused
|
||||
* ``qkv`` tensor. Caches (bf16) are scatter-written by slot.
|
||||
*/
|
||||
|
||||
#include <cmath>
|
||||
#include <cuda_runtime.h>
|
||||
#include <type_traits>
|
||||
|
||||
#include "torch_utils.h"
|
||||
|
||||
#include "../cuda_compat.h"
|
||||
#include "../type_convert.cuh"
|
||||
#include "dispatch_utils.h"
|
||||
|
||||
#ifndef FINAL_MASK
|
||||
#ifdef USE_ROCM
|
||||
#define FINAL_MASK 0xffffffffffffffffULL
|
||||
#else
|
||||
#define FINAL_MASK 0xffffffffu
|
||||
#endif
|
||||
#endif
|
||||
|
||||
namespace vllm {
|
||||
namespace minimax_m3_fused_ops {
|
||||
|
||||
namespace {
|
||||
inline int getSMVersion() {
|
||||
auto* props = get_device_prop();
|
||||
return props->major * 10 + props->minor;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Constants (hard-coded for MiniMax-M3-preview).
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
constexpr int kHeadDim = 128;
|
||||
constexpr int kNumLanes = 32;
|
||||
constexpr int kElemsPerLane = kHeadDim / kNumLanes; // 4
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Helpers
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
__device__ __forceinline__ float warpReduceSum(float val) {
|
||||
#pragma unroll
|
||||
for (int mask = 16; mask > 0; mask >>= 1) {
|
||||
val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
// Gemma RMSNorm over the full head (no-op when ``weight == nullptr``) followed
|
||||
// by partial NeoX RoPE on the leading ``rotary_dim`` dims, all in fp32. Each
|
||||
// lane owns ``kElemsPerLane`` contiguous dims [laneId*4, laneId*4+4).
|
||||
template <typename scalar_t>
|
||||
__device__ __forceinline__ void normAndRope(
|
||||
float (&elems)[kElemsPerLane], int const laneId, float const eps,
|
||||
scalar_t const* __restrict__ weight, // [kHeadDim] or nullptr (no norm)
|
||||
bool const do_rope, int const rotary_dim,
|
||||
scalar_t const* __restrict__ cos_ptr, // cos_sin_cache + pos*rotary_dim
|
||||
bool const apply_norm) {
|
||||
// ── Gemma RMSNorm: x * rsqrt(mean(x^2)+eps) * (1 + w) ──────────────────
|
||||
if (apply_norm) {
|
||||
float sumsq = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) sumsq += elems[i] * elems[i];
|
||||
sumsq = warpReduceSum(sumsq);
|
||||
float const rms_rcp = rsqrtf(sumsq / static_cast<float>(kHeadDim) + eps);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
int const dim = laneId * kElemsPerLane + i;
|
||||
float const w = 1.0f + static_cast<float>(weight[dim]);
|
||||
elems[i] = elems[i] * rms_rcp * w;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Partial NeoX RoPE on dims [0, rotary_dim) ──────────────────────────
|
||||
// half = rotary_dim/2. Pair (i, i+half) for i in [0, half). Lane L owns
|
||||
// dims [4L, 4L+4); since half is a multiple of 4, a lane lies wholly in the
|
||||
// first half (own=x[i]) or second half (own=x[i+half]); its partner lives
|
||||
// ``half/4`` lanes away (XOR with that distance).
|
||||
if (do_rope) {
|
||||
int const half = rotary_dim / 2;
|
||||
int const dim0 = laneId * kElemsPerLane;
|
||||
bool const in_rope = dim0 < rotary_dim;
|
||||
int const lane_xor = half / kElemsPerLane; // partner-lane distance
|
||||
|
||||
float partner[kElemsPerLane];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
partner[i] = __shfl_xor_sync(FINAL_MASK, elems[i], lane_xor, 32);
|
||||
}
|
||||
if (in_rope) {
|
||||
bool const first_half = dim0 < half;
|
||||
int const i_base = first_half ? dim0 : (dim0 - half); // cos/sin index
|
||||
scalar_t const* sin_ptr = cos_ptr + half;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
float const c = static_cast<float>(cos_ptr[i_base + i]);
|
||||
float const s = static_cast<float>(sin_ptr[i_base + i]);
|
||||
if (first_half) {
|
||||
elems[i] = elems[i] * c - partner[i] * s;
|
||||
} else {
|
||||
elems[i] = elems[i] * c + partner[i] * s;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Load 4 contiguous bf16 -> 4 fp32 registers.
|
||||
template <typename scalar_t>
|
||||
__device__ __forceinline__ void loadElems(scalar_t const* __restrict__ src,
|
||||
float (&elems)[kElemsPerLane]) {
|
||||
using Converter = vllm::_typeConvert<scalar_t>;
|
||||
uint2 v = *reinterpret_cast<uint2 const*>(src);
|
||||
auto const* p =
|
||||
reinterpret_cast<typename Converter::packed_hip_type const*>(&v);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane / 2; i++) {
|
||||
float2 f2 = Converter::convert(p[i]);
|
||||
elems[2 * i] = f2.x;
|
||||
elems[2 * i + 1] = f2.y;
|
||||
}
|
||||
}
|
||||
|
||||
// Store 4 fp32 registers -> 4 contiguous bf16.
|
||||
template <typename scalar_t>
|
||||
__device__ __forceinline__ void storeElems(
|
||||
scalar_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
|
||||
using Converter = vllm::_typeConvert<scalar_t>;
|
||||
uint2 v;
|
||||
auto* p = reinterpret_cast<typename Converter::packed_hip_type*>(&v);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane / 2; i++) {
|
||||
p[i] = Converter::convert(make_float2(elems[2 * i], elems[2 * i + 1]));
|
||||
}
|
||||
*reinterpret_cast<uint2*>(dst) = v;
|
||||
}
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Kernel
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Grid: 1D, ceil(num_tokens * slots_per_token / warps_per_block).
|
||||
// Each warp = one (token, slot).
|
||||
//
|
||||
// `kIsSparse` and `kInsertKV` are compile-time template bools, so all the
|
||||
// branch decisions that distinguish the dense layer from the sparse layer
|
||||
// (index slots, KV/index inserts, V slots) fold away per instantiation.
|
||||
// Three instantiations are built: dense <false,false>, sparse-profiling
|
||||
// <true,false> and sparse-serving <true,true>. Slots per token:
|
||||
// Q : nq (always — norm+RoPE)
|
||||
// K : nkv (always — norm+RoPE; +K-cache insert)
|
||||
// V : nkv only if kInsertKV (V-cache insert; no warps in dense)
|
||||
// IQ: niq only if kIsSparse (norm+RoPE)
|
||||
// IK: 1 only if kIsSparse (norm+RoPE; +index-cache insert)
|
||||
template <typename scalar_t, bool kIsSparse, bool kInsertKV>
|
||||
__global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
|
||||
scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse)
|
||||
scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr
|
||||
scalar_t* __restrict__ index_q_out, // [N, niq*128] contiguous, or nullptr
|
||||
scalar_t const* __restrict__ q_norm_w,
|
||||
scalar_t const* __restrict__ k_norm_w,
|
||||
scalar_t const* __restrict__ iq_norm_w,
|
||||
scalar_t const* __restrict__ ik_norm_w,
|
||||
scalar_t const* __restrict__ cos_sin_cache, // [max_pos, rotary_dim]
|
||||
int64_t const* __restrict__ positions, // [N] i64
|
||||
int64_t const* __restrict__ slot_mapping, // [N] i64 or nullptr
|
||||
scalar_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr
|
||||
scalar_t* __restrict__ index_cache, // [nb*bs, 128] or nullptr
|
||||
float const eps, int const rotary_dim, int const num_tokens, int const nq,
|
||||
int const nkv, int const niq, int const block_size,
|
||||
// kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
|
||||
// The head_dim (last) dim is always innermost-contiguous (stride 1), so the
|
||||
// NHD/HND layout choice is fully captured by these four strides: NHD keeps
|
||||
// s_token < s_head, HND swaps them. dim_base addresses head_dim directly.
|
||||
int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
|
||||
int64_t const kv_s_head) {
|
||||
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
|
||||
// _typeConvert<BFloat16> is unavailable on pre-Ampere; the M3 kernel only
|
||||
// runs with bf16/fp16 inputs in practice. Discard the bf16 body there.
|
||||
if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
|
||||
return;
|
||||
} else {
|
||||
#endif
|
||||
int const warpsPerBlock = blockDim.x / 32;
|
||||
int const laneId = threadIdx.x % 32;
|
||||
int const globalWarpIdx = blockIdx.x * warpsPerBlock + (threadIdx.x / 32);
|
||||
|
||||
// Slot layout (compile-time gated: dense has neither V nor index slots).
|
||||
int const v_slots = kInsertKV ? nkv : 0;
|
||||
int const idx_slots = kIsSparse ? niq + 1 : 0;
|
||||
int const slots_per_token = nq + nkv + v_slots + idx_slots;
|
||||
|
||||
int const tokenIdx = globalWarpIdx / slots_per_token;
|
||||
int const slot = globalWarpIdx % slots_per_token;
|
||||
if (tokenIdx >= num_tokens) return;
|
||||
|
||||
// Slot boundaries.
|
||||
int const k_begin = nq;
|
||||
int const v_begin = nq + nkv; // valid only when kInsertKV
|
||||
int const iq_begin = nq + nkv + v_slots; // index block start
|
||||
int const ik_slot = iq_begin + niq; // valid only when kIsSparse
|
||||
|
||||
bool const isQ = slot < k_begin;
|
||||
bool const isK = slot >= k_begin && slot < v_begin;
|
||||
bool isV = false;
|
||||
if constexpr (kInsertKV) isV = slot >= v_begin && slot < v_begin + nkv;
|
||||
bool isIQ = false, isIK = false;
|
||||
if constexpr (kIsSparse) {
|
||||
isIQ = slot >= iq_begin && slot < ik_slot;
|
||||
isIK = slot == ik_slot;
|
||||
}
|
||||
|
||||
int const dim_base = laneId * kElemsPerLane;
|
||||
// Physical row width of qkv: the dense layer packs [q|k|v]; the sparse
|
||||
// layer additionally packs [index_q (niq heads) | index_k (1 head)].
|
||||
int const qkv_row = (nq + 2 * nkv + (kIsSparse ? (niq + 1) : 0)) * kHeadDim;
|
||||
|
||||
// ── Resolve source pointer + per-branch parameters. ────────────────────
|
||||
scalar_t* row_ptr = nullptr; // in-place output location
|
||||
scalar_t const* norm_w = nullptr; // nullptr -> skip norm (V)
|
||||
bool do_rope = true;
|
||||
int head = 0; // kv head index for inserts
|
||||
|
||||
if (isQ) {
|
||||
row_ptr =
|
||||
qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
|
||||
norm_w = q_norm_w;
|
||||
} else if (isK) {
|
||||
head = slot - k_begin;
|
||||
row_ptr =
|
||||
qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
|
||||
norm_w = k_norm_w;
|
||||
} else if (isV) {
|
||||
// qkv V section starts at slot index (nq + nkv): slot * kHeadDim is the
|
||||
// correct in-tensor offset.
|
||||
head = slot - v_begin;
|
||||
row_ptr =
|
||||
qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
|
||||
norm_w = nullptr; // V: no norm, no rope
|
||||
do_rope = false;
|
||||
} else if (isIQ) {
|
||||
// index_q sub-block lives at physical offset (nq+2*nkv)*128 in qkv.
|
||||
int const ih = slot - iq_begin;
|
||||
row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
|
||||
(nq + 2 * nkv + ih) * kHeadDim;
|
||||
norm_w = iq_norm_w;
|
||||
} else { // isIK -- single shared index key at (nq+2*nkv+niq)*128.
|
||||
row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
|
||||
(nq + 2 * nkv + niq) * kHeadDim;
|
||||
norm_w = ik_norm_w;
|
||||
}
|
||||
|
||||
// Store destination. Q and index_q are gathered into dedicated contiguous
|
||||
// output buffers (when provided) so the downstream SM100 sparse kernel's
|
||||
// flat TMA descriptor can address them as [tokens*heads, head_dim]; this
|
||||
// folds the de-interleaving into the store the kernel already does, instead
|
||||
// of a separate q.contiguous() copy. Everything else stays in place.
|
||||
scalar_t* store_ptr = row_ptr;
|
||||
if (isQ && q_out != nullptr) {
|
||||
store_ptr = q_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
|
||||
slot * kHeadDim;
|
||||
} else if (isIQ && index_q_out != nullptr) {
|
||||
store_ptr = index_q_out +
|
||||
static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
|
||||
(slot - iq_begin) * kHeadDim;
|
||||
}
|
||||
|
||||
// PDL: wait for the predecessor kernel (the qkv-projection GEMM that
|
||||
// produces ``qkv``) to finish before touching any global memory. No-op
|
||||
// when PDL is not enabled on the launch. The CUDA runtime wrapper emits
|
||||
// the griddepcontrol.wait PTX with the required memory clobber internally.
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaGridDependencySynchronize();
|
||||
#endif
|
||||
|
||||
// ── Load -> norm+rope (fp32) -> store back in place. ───────────────────
|
||||
float elems[kElemsPerLane];
|
||||
loadElems<scalar_t>(row_ptr + dim_base, elems);
|
||||
|
||||
if (!isV) {
|
||||
int64_t const pos = positions[tokenIdx];
|
||||
scalar_t const* cos_ptr = cos_sin_cache + pos * rotary_dim;
|
||||
normAndRope<scalar_t>(elems, laneId, eps, norm_w, do_rope, rotary_dim,
|
||||
cos_ptr, /*apply_norm=*/norm_w != nullptr);
|
||||
storeElems<scalar_t>(store_ptr + dim_base, elems);
|
||||
}
|
||||
|
||||
// ── Cache inserts (sparse serving only). ───────────────────────────────
|
||||
if constexpr (kInsertKV) {
|
||||
// Guard (not early-return) so every thread reaches the PDL trigger below.
|
||||
int64_t const sm = (isK || isV || isIK) ? slot_mapping[tokenIdx] : -1;
|
||||
if (sm >= 0) { // skip padded / unscheduled tokens
|
||||
if (isIK) {
|
||||
scalar_t* dst = index_cache + sm * kHeadDim + dim_base;
|
||||
storeElems<scalar_t>(dst, elems);
|
||||
} else if (isK || isV) {
|
||||
// kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim].
|
||||
// Paging is logical (block = sm/block_size, token = sm%block_size);
|
||||
// the physical NHD/HND layout is honoured via the passed strides.
|
||||
int64_t const b = sm / block_size;
|
||||
int64_t const t = sm % block_size;
|
||||
int const kv = isK ? 0 : 1;
|
||||
int64_t const off =
|
||||
b * kv_s_block + kv * kv_s_kv + t * kv_s_token + head * kv_s_head;
|
||||
storeElems<scalar_t>(kv_cache + off + dim_base, elems);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// PDL: signal that this kernel is done so a dependent successor may launch
|
||||
// early. No-op when PDL is not enabled on the launch.
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Launch wrapper
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
template <typename scalar_t>
|
||||
void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
|
||||
scalar_t const* q_norm_w, scalar_t const* k_norm_w,
|
||||
scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
|
||||
scalar_t const* cos_sin_cache,
|
||||
int64_t const* positions, int64_t const* slot_mapping,
|
||||
scalar_t* kv_cache, scalar_t* index_cache,
|
||||
float const eps, int const rotary_dim,
|
||||
int const num_tokens, int const nq, int const nkv,
|
||||
int const niq, int const block_size,
|
||||
int64_t const kv_s_block, int64_t const kv_s_kv,
|
||||
int64_t const kv_s_token, int64_t const kv_s_head,
|
||||
bool const has_index, bool const insert_kv,
|
||||
cudaStream_t stream) {
|
||||
// Slot count must match the kernel's compile-time gating.
|
||||
int const v_slots = insert_kv ? nkv : 0;
|
||||
int const idx_slots = has_index ? niq + 1 : 0;
|
||||
int const slots_per_token = nq + nkv + v_slots + idx_slots;
|
||||
|
||||
constexpr int kBlockSize = 256;
|
||||
constexpr int kWarpsPerBlock = kBlockSize / 32;
|
||||
int64_t const total_warps =
|
||||
static_cast<int64_t>(num_tokens) * slots_per_token;
|
||||
int const grid =
|
||||
static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
|
||||
if (grid == 0) return;
|
||||
|
||||
#ifndef USE_ROCM
|
||||
// PDL: enable programmatic stream serialization whenever the hardware
|
||||
// supports it (SM90+). On pre-Hopper GPUs the attribute is unavailable, so
|
||||
// leave numAttrs = 0 and launch as a regular kernel via cudaLaunchKernelEx.
|
||||
static int const sm_version = getSMVersion();
|
||||
cudaLaunchConfig_t config;
|
||||
config.gridDim = dim3(grid);
|
||||
config.blockDim = dim3(kBlockSize);
|
||||
config.dynamicSmemBytes = 0;
|
||||
config.stream = stream;
|
||||
cudaLaunchAttribute attrs[1];
|
||||
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
attrs[0].val.programmaticStreamSerializationAllowed = 1;
|
||||
config.attrs = attrs;
|
||||
config.numAttrs = (sm_version >= 90) ? 1 : 0;
|
||||
|
||||
#define LAUNCH(IS_SPARSE, INSERT) \
|
||||
cudaLaunchKernelEx( \
|
||||
&config, \
|
||||
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT>, \
|
||||
qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, \
|
||||
cos_sin_cache, positions, slot_mapping, kv_cache, index_cache, eps, \
|
||||
rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv, \
|
||||
kv_s_token, kv_s_head)
|
||||
#else
|
||||
// ROCm: standard kernel launch syntax (no PDL/stream serialization).
|
||||
// clang-format off
|
||||
#define LAUNCH(IS_SPARSE, INSERT) \
|
||||
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT> \
|
||||
<<<grid, kBlockSize, 0, stream>>>( \
|
||||
qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, \
|
||||
ik_norm_w, cos_sin_cache, positions, slot_mapping, kv_cache, \
|
||||
index_cache, eps, rotary_dim, num_tokens, nq, nkv, niq, \
|
||||
block_size, kv_s_block, kv_s_kv, kv_s_token, kv_s_head)
|
||||
// clang-format on
|
||||
#endif
|
||||
|
||||
if (has_index) {
|
||||
if (insert_kv) {
|
||||
LAUNCH(true, true); // sparse serving
|
||||
} else {
|
||||
LAUNCH(true, false); // sparse profiling
|
||||
}
|
||||
} else {
|
||||
// Dense layer: never has an index branch and never inserts here (the
|
||||
// generic Attention layer owns the KV insert).
|
||||
LAUNCH(false, false);
|
||||
}
|
||||
#undef LAUNCH
|
||||
}
|
||||
|
||||
} // namespace minimax_m3_fused_ops
|
||||
} // namespace vllm
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Torch op wrapper
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
void fused_minimax_m3_qknorm_rope_kv_insert(
|
||||
torch::stable::Tensor& qkv, // [N, qkv_row] (packs index if sparse)
|
||||
torch::stable::Tensor const& q_norm_weight, // [128]
|
||||
torch::stable::Tensor const& k_norm_weight, // [128]
|
||||
torch::stable::Tensor const& cos_sin_cache, // [max_pos, rotary_dim]
|
||||
torch::stable::Tensor const& positions, // [N] i64
|
||||
int64_t num_heads, int64_t num_kv_heads, int64_t rotary_dim, double eps,
|
||||
std::optional<torch::stable::Tensor> index_q_norm_weight, // [128]
|
||||
std::optional<torch::stable::Tensor> index_k_norm_weight, // [128]
|
||||
int64_t num_index_heads, // niq; 0 => dense
|
||||
std::optional<torch::stable::Tensor> slot_mapping, // [N] i64
|
||||
std::optional<torch::stable::Tensor> kv_cache, // [nb,2,bs,nkv,128]
|
||||
std::optional<torch::stable::Tensor> index_cache, // [nb,bs,128]
|
||||
int64_t block_size,
|
||||
std::optional<torch::stable::Tensor> q_out, // [N, nq*128] contiguous
|
||||
std::optional<torch::stable::Tensor>
|
||||
index_q_out) { // [N, niq*128] contiguous
|
||||
STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
|
||||
"qkv must be contiguous CUDA");
|
||||
STD_TORCH_CHECK(
|
||||
positions.is_cuda() &&
|
||||
positions.scalar_type() == torch::headeronly::ScalarType::Long,
|
||||
"positions must be int64 CUDA");
|
||||
STD_TORCH_CHECK(cos_sin_cache.is_cuda() && cos_sin_cache.is_contiguous(),
|
||||
"cos_sin_cache must be contiguous CUDA");
|
||||
STD_TORCH_CHECK(cos_sin_cache.scalar_type() == qkv.scalar_type(),
|
||||
"cos_sin_cache dtype must match qkv");
|
||||
STD_TORCH_CHECK(
|
||||
cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == rotary_dim,
|
||||
"cos_sin_cache shape [max_pos, rotary_dim]");
|
||||
|
||||
STD_TORCH_CHECK(q_norm_weight.scalar_type() == qkv.scalar_type() &&
|
||||
k_norm_weight.scalar_type() == qkv.scalar_type(),
|
||||
"q/k norm weight dtype must match qkv");
|
||||
STD_TORCH_CHECK(
|
||||
q_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim &&
|
||||
k_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim,
|
||||
"q/k norm weight must have 128 elements");
|
||||
STD_TORCH_CHECK(rotary_dim > 0 && rotary_dim % 8 == 0 &&
|
||||
rotary_dim <= vllm::minimax_m3_fused_ops::kHeadDim,
|
||||
"rotary_dim must be a positive multiple of 8 and <= 128");
|
||||
|
||||
int const num_tokens = static_cast<int>(qkv.size(0));
|
||||
int const nq = static_cast<int>(num_heads);
|
||||
int const nkv = static_cast<int>(num_kv_heads);
|
||||
int const niq = static_cast<int>(num_index_heads);
|
||||
|
||||
// The sparse layer packs the index branch ([index_q (niq heads) | index_k
|
||||
// (1 head)]) right after [q|k|v] in the same row; the dense layer does not.
|
||||
bool const has_index = niq > 0;
|
||||
bool const insert_kv = kv_cache.has_value();
|
||||
int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim;
|
||||
int const expected_row =
|
||||
(nq + 2 * nkv + (has_index ? niq + 1 : 0)) * kHeadDim;
|
||||
STD_TORCH_CHECK(qkv.size(1) == expected_row,
|
||||
"qkv last dim must be (num_heads + 2*num_kv_heads"
|
||||
" + num_index_heads + 1) * 128 for sparse, "
|
||||
"(num_heads + 2*num_kv_heads) * 128 for dense");
|
||||
|
||||
// Only the sparse layer inserts here (dense lets the generic Attention layer
|
||||
// own the KV write); there is no dense+insert kernel instantiation.
|
||||
STD_TORCH_CHECK(
|
||||
!insert_kv || has_index,
|
||||
"insert mode (kv_cache) requires the index branch (sparse layer)");
|
||||
if (has_index) {
|
||||
STD_TORCH_CHECK(
|
||||
index_q_norm_weight.has_value() && index_k_norm_weight.has_value(),
|
||||
"index branch requires both index norm weights");
|
||||
STD_TORCH_CHECK(index_q_norm_weight->scalar_type() == qkv.scalar_type() &&
|
||||
index_k_norm_weight->scalar_type() == qkv.scalar_type(),
|
||||
"index norm weights dtype must match qkv");
|
||||
STD_TORCH_CHECK(index_q_norm_weight->numel() == kHeadDim &&
|
||||
index_k_norm_weight->numel() == kHeadDim,
|
||||
"index norm weights must have 128 elements");
|
||||
}
|
||||
// kv_cache strides (logical shape [nb, 2, bs, nkv, head_dim]). Read straight
|
||||
// off the tensor so the kernel honours whatever physical layout the attention
|
||||
// backend allocated (NHD: stride order (0,1,2,3,4); HND: (0,1,3,2,4)). No new
|
||||
// op argument is needed -- the strides ride along with the tensor itself.
|
||||
int64_t kv_s_block = 0, kv_s_kv = 0, kv_s_token = 0, kv_s_head = 0;
|
||||
if (insert_kv) {
|
||||
STD_TORCH_CHECK(
|
||||
slot_mapping.has_value() &&
|
||||
slot_mapping->scalar_type() == torch::headeronly::ScalarType::Long,
|
||||
"insert mode requires int64 slot_mapping");
|
||||
STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
|
||||
"kv_cache dtype must match qkv (bf16 cache only)");
|
||||
STD_TORCH_CHECK(index_cache.has_value() &&
|
||||
index_cache->scalar_type() == qkv.scalar_type(),
|
||||
"insert mode requires matching index_cache");
|
||||
STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1,
|
||||
"kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous "
|
||||
"head_dim (stride(4)==1)");
|
||||
kv_s_block = kv_cache->stride(0);
|
||||
kv_s_kv = kv_cache->stride(1);
|
||||
kv_s_token = kv_cache->stride(2);
|
||||
kv_s_head = kv_cache->stride(3);
|
||||
}
|
||||
// Optional contiguous gather targets: when given, the normed/roped q (and
|
||||
// index_q) are written here instead of in place, so callers avoid a separate
|
||||
// .contiguous() copy. index_q_out only makes sense on the sparse path.
|
||||
if (q_out.has_value()) {
|
||||
STD_TORCH_CHECK(
|
||||
q_out->is_cuda() && q_out->is_contiguous() &&
|
||||
q_out->scalar_type() == qkv.scalar_type(),
|
||||
"q_out must be a contiguous CUDA tensor matching qkv dtype");
|
||||
STD_TORCH_CHECK(
|
||||
q_out->numel() == static_cast<int64_t>(num_tokens) * nq * kHeadDim,
|
||||
"q_out must have num_tokens * num_heads * 128 elements");
|
||||
}
|
||||
if (index_q_out.has_value()) {
|
||||
STD_TORCH_CHECK(
|
||||
has_index,
|
||||
"index_q_out requires the index branch (num_index_heads > 0)");
|
||||
STD_TORCH_CHECK(
|
||||
index_q_out->is_cuda() && index_q_out->is_contiguous() &&
|
||||
index_q_out->scalar_type() == qkv.scalar_type(),
|
||||
"index_q_out must be a contiguous CUDA tensor matching qkv dtype");
|
||||
STD_TORCH_CHECK(index_q_out->numel() ==
|
||||
static_cast<int64_t>(num_tokens) * niq * kHeadDim,
|
||||
"index_q_out must have num_tokens * num_index_heads * 128 "
|
||||
"elements");
|
||||
}
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
qkv.get_device_index());
|
||||
auto stream = get_current_cuda_stream(qkv.get_device_index());
|
||||
|
||||
VLLM_STABLE_DISPATCH_HALF_TYPES(
|
||||
qkv.scalar_type(), "fused_minimax_m3_qknorm_rope_kv_insert", [&] {
|
||||
using st = scalar_t;
|
||||
vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st>(
|
||||
reinterpret_cast<st*>(qkv.data_ptr()),
|
||||
q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr())
|
||||
: nullptr,
|
||||
index_q_out.has_value()
|
||||
? reinterpret_cast<st*>(index_q_out->data_ptr())
|
||||
: nullptr,
|
||||
reinterpret_cast<st const*>(q_norm_weight.data_ptr()),
|
||||
reinterpret_cast<st const*>(k_norm_weight.data_ptr()),
|
||||
has_index
|
||||
? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr())
|
||||
: nullptr,
|
||||
has_index
|
||||
? reinterpret_cast<st const*>(index_k_norm_weight->data_ptr())
|
||||
: nullptr,
|
||||
reinterpret_cast<st const*>(cos_sin_cache.data_ptr()),
|
||||
reinterpret_cast<int64_t const*>(positions.data_ptr()),
|
||||
insert_kv
|
||||
? reinterpret_cast<int64_t const*>(slot_mapping->data_ptr())
|
||||
: nullptr,
|
||||
insert_kv ? reinterpret_cast<st*>(kv_cache->data_ptr()) : nullptr,
|
||||
(insert_kv && has_index)
|
||||
? reinterpret_cast<st*>(index_cache->data_ptr())
|
||||
: nullptr,
|
||||
static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens,
|
||||
nq, nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv,
|
||||
kv_s_token, kv_s_head, has_index, insert_kv, stream);
|
||||
});
|
||||
}
|
||||
@@ -20,7 +20,7 @@
|
||||
|
||||
#include "torch_utils.h"
|
||||
|
||||
#include "async_util.cuh"
|
||||
#include "../async_util.cuh"
|
||||
#include "../cuda_compat.h"
|
||||
#include "../type_convert.cuh"
|
||||
#include "dispatch_utils.h"
|
||||
|
||||
@@ -231,6 +231,23 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
|
||||
torch::stable::Tensor& position_ids,
|
||||
int64_t forced_token_heads_per_warp);
|
||||
|
||||
// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV /
|
||||
// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs
|
||||
// the index branch and scatters k/v/index_k into their paged caches.
|
||||
void fused_minimax_m3_qknorm_rope_kv_insert(
|
||||
torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight,
|
||||
torch::stable::Tensor const& k_norm_weight,
|
||||
torch::stable::Tensor const& cos_sin_cache,
|
||||
torch::stable::Tensor const& positions, int64_t num_heads,
|
||||
int64_t num_kv_heads, int64_t rotary_dim, double eps,
|
||||
std::optional<torch::stable::Tensor> index_q_norm_weight,
|
||||
std::optional<torch::stable::Tensor> index_k_norm_weight,
|
||||
int64_t num_index_heads, std::optional<torch::stable::Tensor> slot_mapping,
|
||||
std::optional<torch::stable::Tensor> kv_cache,
|
||||
std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
|
||||
std::optional<torch::stable::Tensor> q_out,
|
||||
std::optional<torch::stable::Tensor> index_q_out);
|
||||
|
||||
// Sampler kernels (shared CUDA/ROCm)
|
||||
void apply_repetition_penalties_(
|
||||
torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
|
||||
@@ -276,7 +293,8 @@ void selective_scan_fwd(
|
||||
// Activation kernels (shared CUDA/ROCm)
|
||||
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);
|
||||
torch::stable::Tensor& input, double limit,
|
||||
double alpha = 1.0, double beta = 0.0);
|
||||
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,
|
||||
|
||||
@@ -25,7 +25,7 @@
|
||||
#include <cuda_fp8.h>
|
||||
|
||||
#include "cuda_utils.h"
|
||||
#include "libtorch_stable/launch_bounds_utils.h"
|
||||
#include "launch_bounds_utils.h"
|
||||
|
||||
// Define before including nvfp4_utils.cuh so the header
|
||||
// can use this macro during compilation.
|
||||
|
||||
@@ -34,7 +34,7 @@
|
||||
static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
|
||||
"MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
|
||||
|
||||
#include "libtorch_stable/launch_bounds_utils.h"
|
||||
#include "launch_bounds_utils.h"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
|
||||
@@ -26,7 +26,7 @@
|
||||
|
||||
#include "cuda_utils.h"
|
||||
#include "nvfp4_utils.cuh"
|
||||
#include "libtorch_stable/launch_bounds_utils.h"
|
||||
#include "launch_bounds_utils.h"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
|
||||
@@ -26,7 +26,7 @@
|
||||
#include "../../cuda_vec_utils.cuh"
|
||||
|
||||
#include "cuda_utils.h"
|
||||
#include "libtorch_stable/launch_bounds_utils.h"
|
||||
#include "launch_bounds_utils.h"
|
||||
|
||||
// Define before including nvfp4_utils.cuh so the header
|
||||
// can use this macro during compilation.
|
||||
|
||||
@@ -7,11 +7,14 @@
|
||||
|
||||
#include <torch/csrc/stable/ops.h>
|
||||
|
||||
#include "ggml-common.h"
|
||||
#include "vecdotq.cuh"
|
||||
#include "dequantize.cuh"
|
||||
#include "mmvq.cuh"
|
||||
#include "mmq.cuh"
|
||||
// NOTE: These headers are intentionally kept in csrc/quantization/gguf/ (not
|
||||
// moved to libtorch_stable) to avoid unnecessary reformatting that would break
|
||||
// git rename detection and pollute blame history.
|
||||
#include "../../../quantization/gguf/ggml-common.h"
|
||||
#include "../../../quantization/gguf/vecdotq.cuh"
|
||||
#include "../../../quantization/gguf/dequantize.cuh"
|
||||
#include "../../../quantization/gguf/mmvq.cuh"
|
||||
#include "../../../quantization/gguf/mmq.cuh"
|
||||
#include "moe.cuh"
|
||||
#include "moe_vec.cuh"
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
#include "torch_utils.h"
|
||||
|
||||
#ifndef USE_ROCM
|
||||
#include "persistent_topk.cuh"
|
||||
#include "../persistent_topk.cuh"
|
||||
#endif
|
||||
|
||||
namespace {
|
||||
|
||||
@@ -337,6 +337,17 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"bool is_neox, Tensor position_ids, "
|
||||
"int forced_token_heads_per_warp=-1) -> ()");
|
||||
|
||||
// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert.
|
||||
ops.def(
|
||||
"fused_minimax_m3_qknorm_rope_kv_insert("
|
||||
"Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, "
|
||||
"Tensor cos_sin_cache, Tensor positions, int num_heads, "
|
||||
"int num_kv_heads, int rotary_dim, float eps, "
|
||||
"Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, "
|
||||
"int num_index_heads, "
|
||||
"Tensor? slot_mapping, Tensor!? kv_cache, Tensor!? index_cache, "
|
||||
"int block_size, Tensor!? q_out, Tensor!? index_q_out) -> ()");
|
||||
|
||||
// Apply repetition penalties to logits in-place.
|
||||
ops.def(
|
||||
"apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
|
||||
@@ -364,9 +375,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
|
||||
|
||||
// SwiGLU activation with input clamping.
|
||||
// alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to
|
||||
// the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up.
|
||||
ops.def(
|
||||
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
|
||||
"-> ()");
|
||||
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
|
||||
"float alpha=1.0, float beta=0.0) -> ()");
|
||||
|
||||
// Activation function used in GeGLU with `none` approximation.
|
||||
ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
|
||||
@@ -571,6 +584,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
// Positional encoding kernels (shared CUDA/ROCm)
|
||||
ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
|
||||
ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
|
||||
ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
|
||||
TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
|
||||
|
||||
// Sampler kernels (shared CUDA/ROCm)
|
||||
ops.impl("apply_repetition_penalties_",
|
||||
|
||||
@@ -0,0 +1,742 @@
|
||||
// CUDA C++ q2k -> k2q CSR builder.
|
||||
//
|
||||
// Five-stage pipeline. q-ascending order within each CSR row is preserved
|
||||
// by partitioning q across (CTA, warp_in_CTA) units; each unit owns a
|
||||
// contiguous q-sub-range and reserves a contiguous slot range per row via
|
||||
// a precomputed exclusive prefix scan.
|
||||
//
|
||||
// M: build_row_map -- round-robin packing of rows across batches
|
||||
// H: histogram + tile_counts
|
||||
// PR: row prefix -- single block per head, row_counts -> row_ptr
|
||||
// PT: tile prefix -- multi-block, scan tile_counts along (c, w) axis
|
||||
// S: scatter (sorted) -- per-warp slot range, q-sequential within warp
|
||||
//
|
||||
// Per-warp partitioning: each CTA has kWarps warps; warp w of CTA c owns
|
||||
// q-range [c*q_per_cta + w*q_per_warp, c*q_per_cta + (w+1)*q_per_warp).
|
||||
// tile_counts is shaped [G * kWarps, H, total_rows]; the "row" dimension
|
||||
// of the prefix scan is the flattened (c * kWarps + w) index, scanned in
|
||||
// lexicographic order so that warp-local slot ranges concatenate to the
|
||||
// global q-sorted output.
|
||||
|
||||
#include <torch/all.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
#define CHECK_CUDA(x) TORCH_CHECK((x).is_cuda(), #x " must be CUDA")
|
||||
#define CHECK_CONTIGUOUS(x) \
|
||||
TORCH_CHECK((x).is_contiguous(), #x " must be contiguous")
|
||||
#define CHECK_INT(x) \
|
||||
TORCH_CHECK((x).scalar_type() == at::kInt, #x " must be int32")
|
||||
#define CHECK_INPUT(x) \
|
||||
CHECK_CUDA(x); \
|
||||
CHECK_CONTIGUOUS(x); \
|
||||
CHECK_INT(x)
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int kWarpSize = 32;
|
||||
|
||||
__device__ __forceinline__ void advance_batch_only(int const* __restrict__ cu_q,
|
||||
int B, int q_abs, int& bi) {
|
||||
while (bi < B && cu_q[bi + 1] <= q_abs) ++bi;
|
||||
}
|
||||
|
||||
// Atomic increment of a 16-bit half within a 32-bit SMEM word; returns the
|
||||
// OLD 16-bit value (slot). Per-warp count must stay < 32768 so the low
|
||||
// half does not carry into the high half.
|
||||
// base_int32 : int32 pointer; element i holds rows 2*i (low) and 2*i+1
|
||||
// (high).
|
||||
__device__ __forceinline__ int atomic_inc_int16_packed(int* base_int32,
|
||||
int row) {
|
||||
int idx = row >> 1;
|
||||
int shift = (row & 1) << 4; // 0 or 16
|
||||
int delta = 1 << shift;
|
||||
int old = atomicAdd(&base_int32[idx], delta);
|
||||
return (old >> shift) & 0xFFFF;
|
||||
}
|
||||
|
||||
// Read 16-bit half from packed int32 storage.
|
||||
__device__ __forceinline__ int read_int16_packed(int const* base_int32,
|
||||
int row) {
|
||||
int v = base_int32[row >> 1];
|
||||
int shift = (row & 1) << 4;
|
||||
return (v >> shift) & 0xFFFF;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// M: round-robin row map.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <int kBlockK>
|
||||
__global__ void k2q_build_row_map_kernel(int const* __restrict__ cu_k,
|
||||
int* __restrict__ row_map,
|
||||
int* __restrict__ row_coords, int B,
|
||||
int max_kv_blocks) {
|
||||
int level = blockIdx.x;
|
||||
if (level >= max_kv_blocks) return;
|
||||
if (threadIdx.x != 0) return;
|
||||
int rows_before = 0;
|
||||
for (int b = 0; b < B; ++b) {
|
||||
int rb = (cu_k[b + 1] - cu_k[b] + kBlockK - 1) / kBlockK;
|
||||
rows_before += (rb < level ? rb : level);
|
||||
}
|
||||
int active_before = 0;
|
||||
for (int b = 0; b < B; ++b) {
|
||||
int rb = (cu_k[b + 1] - cu_k[b] + kBlockK - 1) / kBlockK;
|
||||
if (rb > level) {
|
||||
int row_linear = rows_before + active_before;
|
||||
row_map[(size_t)b * max_kv_blocks + level] = row_linear;
|
||||
if (row_coords != nullptr) {
|
||||
row_coords[(size_t)row_linear * 2] = b;
|
||||
row_coords[(size_t)row_linear * 2 + 1] = level;
|
||||
}
|
||||
++active_before;
|
||||
} else {
|
||||
row_map[(size_t)b * max_kv_blocks + level] = -1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// H: per-warp histogram + tile_counts.
|
||||
// kWarps warps per CTA, each owns q-sub-range = q_per_cta / kWarps.
|
||||
// SMEM hist[kWarps, total_rows] int32 (stored as packed int16 cursor:
|
||||
// 2 entries per int32 word). Each warp counts to its own row.
|
||||
// At end-of-CTA, write tile_counts[c*kWarps + w, h, r] = smem_hist[w, r]
|
||||
// and atomicAdd(row_counts[h, r], sum over w of smem_hist[w, r]).
|
||||
// ---------------------------------------------------------------------------
|
||||
template <int kTopK, int kBlockK, int kWarps>
|
||||
__global__ void k2q_hist_kernel(int const* __restrict__ q2k,
|
||||
int const* __restrict__ cu_q,
|
||||
int const* __restrict__ row_map,
|
||||
int* __restrict__ row_counts,
|
||||
int* __restrict__ tile_counts, int H, int B,
|
||||
int S_Q, int total_rows, int max_kv_blocks,
|
||||
int q_per_cta, int q_per_warp) {
|
||||
constexpr int kThreads = kWarps * kWarpSize;
|
||||
extern __shared__ int smem_hist_int[];
|
||||
int* smem_hist = smem_hist_int;
|
||||
int tid = threadIdx.x;
|
||||
int warp_id = tid >> 5;
|
||||
int lane = tid & 31;
|
||||
int c = blockIdx.x;
|
||||
int q_start_cta = c * q_per_cta;
|
||||
int q_end_cta = min(q_start_cta + q_per_cta, S_Q);
|
||||
int q_start_warp = min(q_start_cta + warp_id * q_per_warp, q_end_cta);
|
||||
int q_end_warp = min(q_start_warp + q_per_warp, q_end_cta);
|
||||
|
||||
constexpr int kInt4PerToken = kTopK / 4;
|
||||
int packed_per_warp = (total_rows + 1) >> 1;
|
||||
int* my_hist = smem_hist + warp_id * packed_per_warp;
|
||||
|
||||
for (int h = 0; h < H; ++h) {
|
||||
for (int i = lane; i < packed_per_warp; i += kWarpSize) my_hist[i] = 0;
|
||||
__syncthreads();
|
||||
|
||||
if (q_start_warp < q_end_warp) {
|
||||
int bi = 0;
|
||||
int qi = q_start_warp + lane;
|
||||
advance_batch_only(cu_q, B, qi, bi);
|
||||
|
||||
int4 const* head_topk4 =
|
||||
reinterpret_cast<int4 const*>(q2k + (size_t)h * S_Q * kTopK);
|
||||
|
||||
for (; qi < q_end_warp; qi += kWarpSize) {
|
||||
advance_batch_only(cu_q, B, qi, bi);
|
||||
int const* my_row_map = row_map + (size_t)bi * max_kv_blocks;
|
||||
|
||||
int4 buf[kInt4PerToken];
|
||||
#pragma unroll
|
||||
for (int v = 0; v < kInt4PerToken; ++v) {
|
||||
buf[v] = head_topk4[(size_t)qi * kInt4PerToken + v];
|
||||
}
|
||||
#pragma unroll
|
||||
for (int t = 0; t < kTopK; ++t) {
|
||||
int kvb_local = reinterpret_cast<int const*>(buf)[t];
|
||||
if (kvb_local >= 0 && kvb_local < max_kv_blocks) {
|
||||
int row = my_row_map[kvb_local];
|
||||
if (row >= 0 && row < total_rows) {
|
||||
atomic_inc_int16_packed(my_hist, row);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
int* head_row_counts = row_counts + (size_t)h * total_rows;
|
||||
// Each warp writes its own slice of tile_counts (full int32) by
|
||||
// unpacking int16 entries from SMEM.
|
||||
int* my_tile =
|
||||
tile_counts + ((size_t)(c * kWarps + warp_id) * H + h) * total_rows;
|
||||
for (int i = lane; i < total_rows; i += kWarpSize) {
|
||||
my_tile[i] = read_int16_packed(my_hist, i);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Sum across warps (int32 accumulator), atomicAdd to row_counts.
|
||||
for (int i = tid; i < total_rows; i += kThreads) {
|
||||
int sum = 0;
|
||||
#pragma unroll
|
||||
for (int w = 0; w < kWarps; ++w) {
|
||||
sum += read_int16_packed(smem_hist + w * packed_per_warp, i);
|
||||
}
|
||||
if (sum > 0) atomicAdd(&head_row_counts[i], sum);
|
||||
}
|
||||
if (h + 1 < H) __syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// PR: row prefix. One block per head.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <int kThreads>
|
||||
__global__ void k2q_row_prefix_kernel(int const* __restrict__ row_counts,
|
||||
int* __restrict__ row_ptr,
|
||||
int const* __restrict__ row_coords,
|
||||
int* __restrict__ scheduler_metadata,
|
||||
int* __restrict__ work_count,
|
||||
int total_rows, int target_q_per_cta,
|
||||
int work_capacity) {
|
||||
int h = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
__shared__ int scan_buf[kThreads];
|
||||
|
||||
int const* head_counts = row_counts + (size_t)h * total_rows;
|
||||
int* head_rowptr = row_ptr + (size_t)h * (total_rows + 1);
|
||||
int chunk = (total_rows + kThreads - 1) / kThreads;
|
||||
int lo = tid * chunk;
|
||||
int hi = min(lo + chunk, total_rows);
|
||||
|
||||
int local_sum = 0;
|
||||
for (int i = lo; i < hi; ++i) local_sum += head_counts[i];
|
||||
scan_buf[tid] = local_sum;
|
||||
__syncthreads();
|
||||
|
||||
for (int off = 1; off < kThreads; off <<= 1) {
|
||||
int add = (tid >= off) ? scan_buf[tid - off] : 0;
|
||||
__syncthreads();
|
||||
scan_buf[tid] += add;
|
||||
__syncthreads();
|
||||
}
|
||||
int running = scan_buf[tid] - local_sum;
|
||||
for (int i = lo; i < hi; ++i) {
|
||||
int row_count = head_counts[i];
|
||||
running += row_count;
|
||||
head_rowptr[i + 1] = running;
|
||||
if (scheduler_metadata != nullptr && work_count != nullptr &&
|
||||
row_count > 0) {
|
||||
int num_chunks = (row_count + target_q_per_cta - 1) / target_q_per_cta;
|
||||
int base = atomicAdd(work_count, num_chunks);
|
||||
int batch_idx = row_coords[(size_t)i * 2];
|
||||
int kv_block_idx = row_coords[(size_t)i * 2 + 1];
|
||||
for (int c = 0; c < num_chunks; ++c) {
|
||||
int work_idx = base + c;
|
||||
if (work_idx < work_capacity) {
|
||||
int q_begin = c * target_q_per_cta;
|
||||
int q_count = min(target_q_per_cta, row_count - q_begin);
|
||||
int* meta = scheduler_metadata + (size_t)work_idx * 6;
|
||||
meta[0] = h;
|
||||
meta[1] = i;
|
||||
meta[2] = q_begin;
|
||||
meta[3] = q_count;
|
||||
meta[4] = batch_idx;
|
||||
meta[5] = kv_block_idx;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// PT_smem: SMEM-staged tile prefix scan.
|
||||
// Each block handles kRowsPerBlock rows for one head h. Cooperative load
|
||||
// of tile_counts[*, h, base_r..base_r+M) into SMEM (better coalescing
|
||||
// than per-warp uncoalesced stride reads), then per-warp scan in SMEM,
|
||||
// then cooperative store back. Fuses row_ptr into the base.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <int kThreads, int kRowsPerBlock>
|
||||
__global__ void k2q_tile_prefix_smem_kernel(int* __restrict__ tile_counts,
|
||||
int const* __restrict__ row_ptr,
|
||||
int H, int total_rows,
|
||||
int G_total) {
|
||||
static_assert(kRowsPerBlock > 0, "kRowsPerBlock must be positive");
|
||||
extern __shared__ int smem_tprefix[];
|
||||
// smem layout: smem[r_off][g] for r_off in [0, M), g in [0, G_total).
|
||||
|
||||
int tid = threadIdx.x;
|
||||
int lane = tid & 31;
|
||||
int warp_id = tid >> 5;
|
||||
|
||||
// Grid: H * blocks_per_h. Each block stays within a single head h
|
||||
// and processes kRowsPerBlock contiguous rows starting at b_in_h *
|
||||
// kRowsPerBlock. (Earlier flat-grid mapping `h = block_job /
|
||||
// total_rows; base_r = block_job - h*total_rows` skipped rows when
|
||||
// total_rows was not a multiple of kRowsPerBlock and H > 1, because
|
||||
// the last partial block of head h-1 left blocks of head h starting
|
||||
// at a non-zero row offset.)
|
||||
int blocks_per_h = (total_rows + kRowsPerBlock - 1) / kRowsPerBlock;
|
||||
int h = blockIdx.x / blocks_per_h;
|
||||
int b_in_h = blockIdx.x - h * blocks_per_h;
|
||||
if (h >= H) return;
|
||||
int base_r = b_in_h * kRowsPerBlock;
|
||||
if (base_r >= total_rows) return;
|
||||
int actual_M = min(kRowsPerBlock, total_rows - base_r);
|
||||
|
||||
size_t stride_g = (size_t)H * total_rows;
|
||||
int* base_ptr = tile_counts + (size_t)h * total_rows + base_r;
|
||||
int total_elems = G_total * actual_M;
|
||||
|
||||
// Cooperative load. Pattern: thread tid -> (r_off=tid%M, g=tid/M),
|
||||
// then strided. 32 lanes hit M r's × (32/M) g's, giving 32/M cache
|
||||
// lines per warp (vs 32 in the naive stride-along-g pattern).
|
||||
for (int i = tid; i < total_elems; i += kThreads) {
|
||||
int r_off = i % actual_M;
|
||||
int g = i / actual_M;
|
||||
smem_tprefix[r_off * G_total + g] = base_ptr[g * stride_g + r_off];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Per-warp scan: warp w scans row (base_r + w) if w < actual_M.
|
||||
if (warp_id < actual_M) {
|
||||
int abs_r = base_r + warp_id;
|
||||
int rp = row_ptr[(size_t)h * (total_rows + 1) + abs_r];
|
||||
int* my_smem = smem_tprefix + warp_id * G_total;
|
||||
int running = rp;
|
||||
for (int g0 = 0; g0 < G_total; g0 += kWarpSize) {
|
||||
int g = g0 + lane;
|
||||
int v = (g < G_total) ? my_smem[g] : 0;
|
||||
int x = v;
|
||||
#pragma unroll
|
||||
for (int off = 1; off < kWarpSize; off <<= 1) {
|
||||
int nbr = __shfl_up_sync(0xFFFFFFFF, x, off);
|
||||
if (lane >= off) x += nbr;
|
||||
}
|
||||
int excl = running + x - v;
|
||||
if (g < G_total) my_smem[g] = excl;
|
||||
int chunk_sum = __shfl_sync(0xFFFFFFFF, x, 31);
|
||||
running += chunk_sum;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Cooperative store back.
|
||||
for (int i = tid; i < total_elems; i += kThreads) {
|
||||
int r_off = i % actual_M;
|
||||
int g = i / actual_M;
|
||||
base_ptr[g * stride_g + r_off] = smem_tprefix[r_off * G_total + g];
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// S: scatter. kWarps warps per CTA, each owns q-sub-range. Per-warp SMEM
|
||||
// cursor and per-warp tile_offset slot range. Within a warp, q's are
|
||||
// processed sequentially; lanes 0..kTopK-1 handle the topK slots in
|
||||
// lockstep. Across distinct q's in the same warp, the lockstep ordering
|
||||
// guarantees q-monotonic atomicAdd on smem_cursor[r].
|
||||
// ---------------------------------------------------------------------------
|
||||
// kQPerIter * kTopK lanes are active per warp iter; remaining lanes idle.
|
||||
// For kTopK=16, kQPerIter=2 uses all 32 lanes; for kTopK=8, kQPerIter=4.
|
||||
// CORRECTNESS NOTE: relies on lane-ordered SMEM atomicAdd return values
|
||||
// within a single warp instruction (verified on B200; tests pass).
|
||||
//
|
||||
// SMEM cursor stored as packed int16 (two cursors per int32). Per-warp
|
||||
// row count must stay < 32768 (~q_per_warp * kTopK at max sink), which
|
||||
// holds for all task.md sizes up to 1024K.
|
||||
template <int kTopK, int kBlockK, int kWarps>
|
||||
__global__ void k2q_scatter_kernel(
|
||||
int const* __restrict__ q2k, int const* __restrict__ cu_q,
|
||||
int const* __restrict__ row_map, int const* __restrict__ abs_base,
|
||||
int* __restrict__ q_idx, int* __restrict__ qsplit_idx,
|
||||
int* __restrict__ split_counts, int H, int B, int S_Q, int total_rows,
|
||||
int max_kv_blocks, int q_per_cta, int q_per_warp, int max_seqlen_q) {
|
||||
constexpr int kQPerIter = kWarpSize / kTopK > 0 ? kWarpSize / kTopK : 1;
|
||||
extern __shared__ int smem_cursor_int[];
|
||||
int* smem_cursor = smem_cursor_int;
|
||||
int tid = threadIdx.x;
|
||||
int warp_id = tid >> 5;
|
||||
int lane = tid & 31;
|
||||
int c = blockIdx.x;
|
||||
int q_start_cta = c * q_per_cta;
|
||||
int q_end_cta = min(q_start_cta + q_per_cta, S_Q);
|
||||
int q_start_warp = min(q_start_cta + warp_id * q_per_warp, q_end_cta);
|
||||
int q_end_warp = min(q_start_warp + q_per_warp, q_end_cta);
|
||||
|
||||
int q_in_iter = lane / kTopK;
|
||||
int slot_in_q = lane % kTopK;
|
||||
bool lane_active = (lane < kQPerIter * kTopK);
|
||||
|
||||
// Per-warp packed cursor: total_rows int16 entries -> ceil(total_rows/2)
|
||||
// int32.
|
||||
int packed_per_warp = (total_rows + 1) >> 1;
|
||||
int* my_cursor = smem_cursor + warp_id * packed_per_warp;
|
||||
|
||||
for (int h = 0; h < H; ++h) {
|
||||
for (int i = lane; i < packed_per_warp; i += kWarpSize) my_cursor[i] = 0;
|
||||
__syncwarp();
|
||||
|
||||
if (q_start_warp < q_end_warp) {
|
||||
int bi = 0;
|
||||
advance_batch_only(cu_q, B, q_start_warp, bi);
|
||||
|
||||
int const* head_q2k = q2k + (size_t)h * S_Q * kTopK;
|
||||
int const* my_abs_base =
|
||||
abs_base + ((size_t)(c * kWarps + warp_id) * H + h) * total_rows;
|
||||
int* head_qidx = q_idx + (size_t)h * S_Q * kTopK;
|
||||
|
||||
// (Hot-row register cache experiment showed no measurable
|
||||
// benefit; relying on L1 to keep row 0 / row total_rows-1
|
||||
// hot since they're hit every iteration in sink workloads.)
|
||||
|
||||
constexpr int kUnroll = 16;
|
||||
int qi_base = q_start_warp;
|
||||
for (; qi_base + kUnroll * kQPerIter <= q_end_warp;
|
||||
qi_base += kUnroll * kQPerIter) {
|
||||
int kvb[kUnroll];
|
||||
int qloc[kUnroll];
|
||||
int batch[kUnroll];
|
||||
int const* rmap[kUnroll];
|
||||
|
||||
#pragma unroll
|
||||
for (int u = 0; u < kUnroll; ++u) {
|
||||
int qi_u = qi_base + u * kQPerIter + q_in_iter;
|
||||
kvb[u] = -1;
|
||||
qloc[u] = 0;
|
||||
batch[u] = 0;
|
||||
if (lane_active) {
|
||||
advance_batch_only(cu_q, B, qi_u, bi);
|
||||
qloc[u] = qi_u - cu_q[bi];
|
||||
batch[u] = bi;
|
||||
kvb[u] = head_q2k[(size_t)qi_u * kTopK + slot_in_q];
|
||||
}
|
||||
rmap[u] = row_map + (size_t)bi * max_kv_blocks;
|
||||
}
|
||||
|
||||
int row[kUnroll];
|
||||
#pragma unroll
|
||||
for (int u = 0; u < kUnroll; ++u) {
|
||||
row[u] = -1;
|
||||
if (lane_active && kvb[u] >= 0 && kvb[u] < max_kv_blocks)
|
||||
row[u] = rmap[u][kvb[u]];
|
||||
}
|
||||
|
||||
// Pre-issue all kUnroll abs_base loads in parallel before
|
||||
// the atomic chain so memory pipeline runs concurrently
|
||||
// with SMEM atomic-adds.
|
||||
int abs_v[kUnroll];
|
||||
#pragma unroll
|
||||
for (int u = 0; u < kUnroll; ++u) {
|
||||
abs_v[u] =
|
||||
(row[u] >= 0 && row[u] < total_rows) ? my_abs_base[row[u]] : 0;
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int u = 0; u < kUnroll; ++u) {
|
||||
int r = row[u];
|
||||
bool valid_edge = r >= 0 && r < total_rows;
|
||||
unsigned int valid_mask = __ballot_sync(0xFFFFFFFFu, valid_edge);
|
||||
unsigned int group_mask =
|
||||
(kTopK == 32) ? 0xFFFFFFFFu
|
||||
: (((1u << kTopK) - 1u) << (q_in_iter * kTopK));
|
||||
unsigned int lower_lane_mask = lane == 0 ? 0u : ((1u << lane) - 1u);
|
||||
int split_slot = __popc(valid_mask & group_mask & lower_lane_mask);
|
||||
int valid_count = __popc(valid_mask & group_mask);
|
||||
if (split_counts != nullptr && slot_in_q == 0) {
|
||||
split_counts[((size_t)batch[u] * max_seqlen_q + qloc[u]) * H + h] =
|
||||
valid_count;
|
||||
}
|
||||
if (valid_edge) {
|
||||
int slot = atomic_inc_int16_packed(my_cursor, r);
|
||||
int out_pos = abs_v[u] + slot;
|
||||
head_qidx[out_pos] = qloc[u];
|
||||
if (qsplit_idx != nullptr) {
|
||||
qsplit_idx[(size_t)h * S_Q * kTopK + out_pos] =
|
||||
qloc[u] | ((split_slot & 0xFF) << 24);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Tail: 1-3 iters left.
|
||||
for (; qi_base < q_end_warp; qi_base += kQPerIter) {
|
||||
int my_qi = qi_base + q_in_iter;
|
||||
bool valid_q = (my_qi < q_end_warp) && lane_active;
|
||||
int kvb_local = -1;
|
||||
int q_local = 0;
|
||||
int batch_local = 0;
|
||||
if (valid_q) {
|
||||
advance_batch_only(cu_q, B, my_qi, bi);
|
||||
batch_local = bi;
|
||||
q_local = my_qi - cu_q[bi];
|
||||
kvb_local = head_q2k[(size_t)my_qi * kTopK + slot_in_q];
|
||||
}
|
||||
int const* my_row_map = row_map + (size_t)bi * max_kv_blocks;
|
||||
int row = -1;
|
||||
if (valid_q && kvb_local >= 0 && kvb_local < max_kv_blocks) {
|
||||
row = my_row_map[kvb_local];
|
||||
}
|
||||
bool valid_edge = row >= 0 && row < total_rows;
|
||||
unsigned int valid_mask = __ballot_sync(0xFFFFFFFFu, valid_edge);
|
||||
unsigned int group_mask =
|
||||
(kTopK == 32) ? 0xFFFFFFFFu
|
||||
: (((1u << kTopK) - 1u) << (q_in_iter * kTopK));
|
||||
unsigned int lower_lane_mask = lane == 0 ? 0u : ((1u << lane) - 1u);
|
||||
int split_slot = __popc(valid_mask & group_mask & lower_lane_mask);
|
||||
int valid_count = __popc(valid_mask & group_mask);
|
||||
if (split_counts != nullptr && valid_q && slot_in_q == 0) {
|
||||
split_counts[((size_t)batch_local * max_seqlen_q + q_local) * H + h] =
|
||||
valid_count;
|
||||
}
|
||||
if (valid_edge) {
|
||||
int slot = atomic_inc_int16_packed(my_cursor, row);
|
||||
int out_pos = my_abs_base[row] + slot;
|
||||
head_qidx[out_pos] = q_local;
|
||||
if (qsplit_idx != nullptr) {
|
||||
qsplit_idx[(size_t)h * S_Q * kTopK + out_pos] =
|
||||
q_local | ((split_slot & 0xFF) << 24);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if (h + 1 < H) __syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
// ===========================================================================
|
||||
// Host orchestration
|
||||
// ===========================================================================
|
||||
|
||||
template <int kTopK, int kBlockK>
|
||||
static void launch_pipeline(torch::Tensor q2k, torch::Tensor cu_q,
|
||||
torch::Tensor cu_k, torch::Tensor row_ptr,
|
||||
torch::Tensor q_idx, int total_rows,
|
||||
int max_kv_blocks,
|
||||
torch::Tensor scheduler_metadata = torch::Tensor(),
|
||||
torch::Tensor work_count = torch::Tensor(),
|
||||
torch::Tensor qsplit_idx = torch::Tensor(),
|
||||
torch::Tensor split_counts = torch::Tensor(),
|
||||
int target_q_per_cta = 1, int work_capacity = 0,
|
||||
int max_seqlen_q = 0) {
|
||||
int H = (int)q2k.size(0);
|
||||
int S_Q = (int)q2k.size(1);
|
||||
int topK = (int)q2k.size(2);
|
||||
TORCH_CHECK(topK == kTopK, "topK runtime != template kTopK");
|
||||
int B = (int)cu_q.size(0) - 1;
|
||||
auto device = q2k.device();
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
AT_CUDA_CHECK(cudaMemsetAsync(row_ptr.data_ptr<int>(), 0,
|
||||
(size_t)H * (total_rows + 1) * sizeof(int),
|
||||
stream));
|
||||
AT_CUDA_CHECK(cudaMemsetAsync(q_idx.data_ptr<int>(), 0xFF,
|
||||
(size_t)H * S_Q * kTopK * sizeof(int), stream));
|
||||
|
||||
auto opts = torch::TensorOptions().dtype(torch::kInt32).device(device);
|
||||
auto row_counts = torch::zeros({H, total_rows}, opts);
|
||||
auto row_map = torch::empty({B, max_kv_blocks}, opts);
|
||||
bool emit_schedule = scheduler_metadata.defined();
|
||||
auto row_coords =
|
||||
emit_schedule ? torch::empty({total_rows, 2}, opts) : torch::Tensor();
|
||||
int* scheduler_metadata_ptr =
|
||||
emit_schedule ? scheduler_metadata.data_ptr<int>() : nullptr;
|
||||
int* work_count_ptr = emit_schedule ? work_count.data_ptr<int>() : nullptr;
|
||||
int* qsplit_idx_ptr = emit_schedule ? qsplit_idx.data_ptr<int>() : nullptr;
|
||||
int* split_counts_ptr =
|
||||
emit_schedule ? split_counts.data_ptr<int>() : nullptr;
|
||||
int* row_coords_ptr = emit_schedule ? row_coords.data_ptr<int>() : nullptr;
|
||||
if (emit_schedule) {
|
||||
AT_CUDA_CHECK(cudaMemsetAsync(work_count_ptr, 0, sizeof(int), stream));
|
||||
AT_CUDA_CHECK(cudaMemsetAsync(scheduler_metadata_ptr, 0,
|
||||
(size_t)work_capacity * 6 * sizeof(int),
|
||||
stream));
|
||||
}
|
||||
|
||||
int dev = q2k.get_device();
|
||||
int num_sms = 0;
|
||||
AT_CUDA_CHECK(
|
||||
cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev));
|
||||
|
||||
// -- Pick kWarps per CTA based on SMEM budget for cursor/hist ---------
|
||||
// SMEM per CTA = kWarps * total_rows * sizeof(int) (for both H and S).
|
||||
// Want at least 2 CTAs/SM for memory parallelism. SM100 SMEM = 228KB.
|
||||
// Pick the largest kWarps that fits two CTAs/SM, capped at 4.
|
||||
// SMEM cursor packed as int16 (2 entries per int32 word):
|
||||
int per_warp_smem = ((total_rows + 1) >> 1) * (int)sizeof(int);
|
||||
int kWarps_pick = 4;
|
||||
while (kWarps_pick > 1 && (kWarps_pick * per_warp_smem) * 2 > 228 * 1024) {
|
||||
kWarps_pick >>= 1;
|
||||
}
|
||||
if (kWarps_pick < 1) kWarps_pick = 1;
|
||||
|
||||
// -- Pick G (CTAs) ----------------------------------------------------
|
||||
// For each (kWarps, per_warp_smem) pair, the SMEM-bound occupancy is
|
||||
// 228KB / (kWarps*per_warp_smem) CTAs/SM. We size G as
|
||||
// num_sms * occupancy so a single resident wave covers all CTAs and
|
||||
// the memory pipeline runs at peak.
|
||||
int per_cta_smem_bytes = kWarps_pick * per_warp_smem;
|
||||
int max_ctas_per_sm =
|
||||
std::max(1, (228 * 1024) / std::max(1, per_cta_smem_bytes));
|
||||
if (max_ctas_per_sm > 8) max_ctas_per_sm = 8;
|
||||
constexpr int kMinQPerCta = 256;
|
||||
// Cap target_g at num_sms * 3 — empirically this balances
|
||||
// per-CTA work-size against parallelism. Higher caps regress
|
||||
// mid-size cases due to row_counts atomicAdd contention and
|
||||
// smaller q_per_cta. SMEM-bound configurations naturally cap
|
||||
// lower if max_ctas_per_sm < 3.
|
||||
int target_g = num_sms * std::min(max_ctas_per_sm, 3);
|
||||
int max_g_for_q = (S_Q + kMinQPerCta - 1) / kMinQPerCta;
|
||||
int G = std::min({target_g, max_g_for_q, S_Q});
|
||||
if (G < 1) G = 1;
|
||||
int q_per_cta = (S_Q + G - 1) / G;
|
||||
G = (S_Q + q_per_cta - 1) / q_per_cta;
|
||||
int q_per_warp = (q_per_cta + kWarps_pick - 1) / kWarps_pick;
|
||||
int G_total = G * kWarps_pick;
|
||||
|
||||
auto tile_counts = torch::empty({G_total, H, total_rows}, opts);
|
||||
|
||||
// -- Compile-time switch on kWarps for the templated kernels ---------
|
||||
auto rmap_fn = k2q_build_row_map_kernel<kBlockK>;
|
||||
auto rprefix_fn = k2q_row_prefix_kernel<1024>;
|
||||
constexpr int kPtRowsPerBlock = 8;
|
||||
constexpr int kPtThreads = 256;
|
||||
auto tprefix_smem_fn =
|
||||
k2q_tile_prefix_smem_kernel<kPtThreads, kPtRowsPerBlock>;
|
||||
|
||||
if (max_kv_blocks > 0) {
|
||||
rmap_fn<<<max_kv_blocks, 32, 0, stream>>>(cu_k.data_ptr<int>(),
|
||||
row_map.data_ptr<int>(),
|
||||
row_coords_ptr, B, max_kv_blocks);
|
||||
}
|
||||
|
||||
auto launch_hist_scatter = [&](auto kWarps_const) {
|
||||
constexpr int W = decltype(kWarps_const)::value;
|
||||
size_t smem_bytes = (size_t)W * per_warp_smem;
|
||||
auto hist_fn = k2q_hist_kernel<kTopK, kBlockK, W>;
|
||||
auto scat_fn = k2q_scatter_kernel<kTopK, kBlockK, W>;
|
||||
AT_CUDA_CHECK(cudaFuncSetAttribute(
|
||||
hist_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem_bytes));
|
||||
AT_CUDA_CHECK(cudaFuncSetAttribute(
|
||||
scat_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem_bytes));
|
||||
|
||||
hist_fn<<<G, W * kWarpSize, smem_bytes, stream>>>(
|
||||
q2k.data_ptr<int>(), cu_q.data_ptr<int>(), row_map.data_ptr<int>(),
|
||||
row_counts.data_ptr<int>(), tile_counts.data_ptr<int>(), H, B, S_Q,
|
||||
total_rows, max_kv_blocks, q_per_cta, q_per_warp);
|
||||
|
||||
rprefix_fn<<<H, 1024, 0, stream>>>(
|
||||
row_counts.data_ptr<int>(), row_ptr.data_ptr<int>(),
|
||||
emit_schedule ? row_coords.data_ptr<int>() : nullptr,
|
||||
scheduler_metadata_ptr, work_count_ptr, total_rows, target_q_per_cta,
|
||||
work_capacity);
|
||||
|
||||
// Grid is H * blocks_per_h so each block stays within a single
|
||||
// head; flat (H*total_rows) grid would skip rows when total_rows
|
||||
// is not a multiple of kPtRowsPerBlock.
|
||||
int blocks_per_h = (total_rows + kPtRowsPerBlock - 1) / kPtRowsPerBlock;
|
||||
int pt_grid = H * blocks_per_h;
|
||||
if (pt_grid < 1) pt_grid = 1;
|
||||
size_t pt_smem = (size_t)kPtRowsPerBlock * G_total * sizeof(int);
|
||||
AT_CUDA_CHECK(cudaFuncSetAttribute(
|
||||
tprefix_smem_fn, cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
(int)pt_smem));
|
||||
tprefix_smem_fn<<<pt_grid, kPtThreads, pt_smem, stream>>>(
|
||||
tile_counts.data_ptr<int>(), row_ptr.data_ptr<int>(), H, total_rows,
|
||||
G_total);
|
||||
|
||||
scat_fn<<<G, W * kWarpSize, smem_bytes, stream>>>(
|
||||
q2k.data_ptr<int>(), cu_q.data_ptr<int>(), row_map.data_ptr<int>(),
|
||||
tile_counts.data_ptr<int>(), q_idx.data_ptr<int>(), qsplit_idx_ptr,
|
||||
split_counts_ptr, H, B, S_Q, total_rows, max_kv_blocks, q_per_cta,
|
||||
q_per_warp, max_seqlen_q);
|
||||
};
|
||||
|
||||
if (kWarps_pick == 4) {
|
||||
launch_hist_scatter(std::integral_constant<int, 4>{});
|
||||
} else if (kWarps_pick == 2) {
|
||||
launch_hist_scatter(std::integral_constant<int, 2>{});
|
||||
} else {
|
||||
launch_hist_scatter(std::integral_constant<int, 1>{});
|
||||
}
|
||||
}
|
||||
|
||||
void run_minimax_m3_build_k2q_csr_with_schedule(
|
||||
torch::Tensor q2k, torch::Tensor cu_q, torch::Tensor cu_k,
|
||||
torch::Tensor row_ptr, torch::Tensor q_idx,
|
||||
torch::Tensor scheduler_metadata, torch::Tensor work_count,
|
||||
torch::Tensor qsplit_idx, torch::Tensor split_counts, int64_t topk,
|
||||
int64_t blk_kv, int64_t total_rows, int64_t max_kv_blocks,
|
||||
int64_t target_q_per_cta, int64_t work_capacity, int64_t max_seqlen_q) {
|
||||
CHECK_INPUT(q2k);
|
||||
CHECK_INPUT(cu_q);
|
||||
CHECK_INPUT(cu_k);
|
||||
CHECK_INPUT(row_ptr);
|
||||
CHECK_INPUT(q_idx);
|
||||
CHECK_INPUT(scheduler_metadata);
|
||||
CHECK_INPUT(work_count);
|
||||
CHECK_INPUT(qsplit_idx);
|
||||
CHECK_INPUT(split_counts);
|
||||
TORCH_CHECK(blk_kv == 128, "build_k2q_csr only supports blk_kv == 128");
|
||||
int H = (int)q2k.size(0);
|
||||
int S_Q = (int)q2k.size(1);
|
||||
int tr = (int)total_rows;
|
||||
int mkv = (int)max_kv_blocks;
|
||||
int target = (int)target_q_per_cta;
|
||||
int capacity = (int)work_capacity;
|
||||
int max_sq = (int)max_seqlen_q;
|
||||
TORCH_CHECK(tr >= 0 && mkv >= 0 && target > 0 && capacity > 0 && max_sq >= 0,
|
||||
"invalid schedule sizing arguments");
|
||||
TORCH_CHECK(row_ptr.size(0) == H && row_ptr.size(1) == tr + 1,
|
||||
"row_ptr shape mismatch");
|
||||
TORCH_CHECK(q_idx.size(0) == H && q_idx.size(1) == (int64_t)S_Q * (int)topk,
|
||||
"q_idx shape mismatch");
|
||||
TORCH_CHECK(qsplit_idx.sizes() == q_idx.sizes(), "qsplit_idx shape mismatch");
|
||||
TORCH_CHECK(
|
||||
scheduler_metadata.size(0) == capacity && scheduler_metadata.size(1) == 6,
|
||||
"scheduler_metadata shape mismatch");
|
||||
TORCH_CHECK(work_count.numel() == 1,
|
||||
"work_count must have one int32 element");
|
||||
TORCH_CHECK(split_counts.dim() == 3 &&
|
||||
split_counts.size(0) == cu_q.size(0) - 1 &&
|
||||
split_counts.size(1) == max_sq && split_counts.size(2) == H,
|
||||
"split_counts shape mismatch");
|
||||
if (S_Q == 0 || tr == 0 || H == 0 || mkv == 0) {
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
AT_CUDA_CHECK(cudaMemsetAsync(row_ptr.data_ptr<int>(), 0,
|
||||
(size_t)H * (tr + 1) * sizeof(int), stream));
|
||||
AT_CUDA_CHECK(cudaMemsetAsync(q_idx.data_ptr<int>(), 0xFF,
|
||||
(size_t)H * S_Q * (int)topk * sizeof(int),
|
||||
stream));
|
||||
AT_CUDA_CHECK(
|
||||
cudaMemsetAsync(work_count.data_ptr<int>(), 0, sizeof(int), stream));
|
||||
if (split_counts.numel() > 0) {
|
||||
AT_CUDA_CHECK(cudaMemsetAsync(split_counts.data_ptr<int>(), 0,
|
||||
(size_t)split_counts.numel() * sizeof(int),
|
||||
stream));
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (topk == 16) {
|
||||
launch_pipeline<16, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
|
||||
scheduler_metadata, work_count, qsplit_idx,
|
||||
split_counts, target, capacity, max_sq);
|
||||
} else if (topk == 8) {
|
||||
launch_pipeline<8, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
|
||||
scheduler_metadata, work_count, qsplit_idx,
|
||||
split_counts, target, capacity, max_sq);
|
||||
} else if (topk == 32) {
|
||||
launch_pipeline<32, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
|
||||
scheduler_metadata, work_count, qsplit_idx,
|
||||
split_counts, target, capacity, max_sq);
|
||||
} else if (topk == 4) {
|
||||
launch_pipeline<4, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
|
||||
scheduler_metadata, work_count, qsplit_idx,
|
||||
split_counts, target, capacity, max_sq);
|
||||
} else {
|
||||
TORCH_CHECK(false, "unsupported topK ", topk,
|
||||
" (expected 4, 8, 16, or 32)");
|
||||
}
|
||||
}
|
||||
+10
-1
@@ -62,7 +62,8 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
|
||||
|
||||
void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
|
||||
|
||||
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
|
||||
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);
|
||||
@@ -146,4 +147,12 @@ std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
|
||||
torch::Tensor const& norm_weight_k, torch::Tensor workspace,
|
||||
int64_t const q_size, int64_t const kv_size, int64_t const rank,
|
||||
int64_t const nranks, double const eps);
|
||||
|
||||
void run_minimax_m3_build_k2q_csr_with_schedule(
|
||||
torch::Tensor q2k, torch::Tensor cu_q, torch::Tensor cu_k,
|
||||
torch::Tensor row_ptr, torch::Tensor q_idx,
|
||||
torch::Tensor scheduler_metadata, torch::Tensor work_count,
|
||||
torch::Tensor qsplit_idx, torch::Tensor split_counts, int64_t topk,
|
||||
int64_t blk_kv, int64_t total_rows, int64_t max_kv_blocks,
|
||||
int64_t target_q_per_cta, int64_t work_capacity, int64_t max_seqlen_q);
|
||||
#endif
|
||||
|
||||
@@ -187,6 +187,27 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"float eps) -> (Tensor, Tensor)");
|
||||
ops.impl("minimax_allreduce_rms_qk", torch::kCUDA, &minimax_allreduce_rms_qk);
|
||||
|
||||
ops.def(
|
||||
"minimax_m3_build_k2q_csr_with_schedule("
|
||||
"Tensor q2k,"
|
||||
"Tensor cu_q,"
|
||||
"Tensor cu_k,"
|
||||
"Tensor! row_ptr,"
|
||||
"Tensor! q_idx,"
|
||||
"Tensor! scheduler_metadata,"
|
||||
"Tensor! work_count,"
|
||||
"Tensor! qsplit_idx,"
|
||||
"Tensor! split_counts,"
|
||||
"int topk,"
|
||||
"int blk_kv,"
|
||||
"int total_rows,"
|
||||
"int max_kv_blocks,"
|
||||
"int target_q_per_cta,"
|
||||
"int work_capacity,"
|
||||
"int max_seqlen_q) -> ()");
|
||||
ops.impl("minimax_m3_build_k2q_csr_with_schedule", torch::kCUDA,
|
||||
&run_minimax_m3_build_k2q_csr_with_schedule);
|
||||
|
||||
// conditionally compiled so impl in source file
|
||||
#endif
|
||||
}
|
||||
|
||||
+1
-1
@@ -757,7 +757,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
|
||||
# https://docs.flashinfer.ai/installation.html
|
||||
# From versions.json: .flashinfer.version
|
||||
ARG FLASHINFER_VERSION=0.6.11.post2
|
||||
ARG FLASHINFER_VERSION=0.6.12
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
|
||||
--extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
|
||||
|
||||
@@ -256,13 +256,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
|
||||
|
||||
|
||||
# build flashinfer for torch nightly from source around 10 mins
|
||||
# release version: v0.6.11.post2
|
||||
# release version: v0.6.12
|
||||
# todo(elainewy): cache flashinfer build result for faster build
|
||||
ENV CCACHE_DIR=/root/.cache/ccache
|
||||
RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
--mount=type=cache,target=/root/.cache/uv \
|
||||
echo "git clone flashinfer..." \
|
||||
&& git clone --depth 1 --branch v0.6.11.post2 --recursive https://github.com/flashinfer-ai/flashinfer.git \
|
||||
&& git clone --depth 1 --branch v0.6.12 --recursive https://github.com/flashinfer-ai/flashinfer.git \
|
||||
&& cd flashinfer \
|
||||
&& git submodule update --init --recursive \
|
||||
&& echo "finish git clone flashinfer..." \
|
||||
|
||||
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
|
||||
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
|
||||
ARG FA_BRANCH="0e60e394"
|
||||
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
|
||||
ARG AITER_BRANCH="v0.1.13.post1"
|
||||
ARG AITER_BRANCH="v0.1.13"
|
||||
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
|
||||
ARG MORI_BRANCH="v1.1.0"
|
||||
ARG MORI_REPO="https://github.com/ROCm/mori.git"
|
||||
|
||||
@@ -68,7 +68,7 @@
|
||||
"default": "true"
|
||||
},
|
||||
"FLASHINFER_VERSION": {
|
||||
"default": "0.6.11.post2"
|
||||
"default": "0.6.12"
|
||||
},
|
||||
"GDRCOPY_CUDA_VERSION": {
|
||||
"default": "12.8"
|
||||
|
||||
@@ -246,12 +246,6 @@ Every image listed in "image_files" is added to the request in the listed order
|
||||
|
||||
The "image" shorthand accepts the same values as "image_files". The "image_url" field accepts either an OpenAI-style object with a "url" field or a URL string.
|
||||
|
||||
By default, image references are sent to the serving endpoint as provided, with local image paths converted to `file://` URLs.
|
||||
|
||||
If the benchmark client should load local and HTTP(S) images before sending requests, pass `--custom-ensure-client-side-data` to encode them as base64 data URLs on the client side.
|
||||
|
||||
Existing `data:image/...` URLs are already self-contained and are kept unchanged.
|
||||
|
||||
```bash
|
||||
# need a model with vision capability here
|
||||
vllm serve Qwen/Qwen2-VL-7B-Instruct
|
||||
@@ -259,13 +253,13 @@ vllm serve Qwen/Qwen2-VL-7B-Instruct
|
||||
|
||||
```bash
|
||||
# run benchmarking script
|
||||
vllm bench serve --save-result --save-detailed \
|
||||
vllm bench serve--save-result --save-detailed \
|
||||
--backend openai-chat \
|
||||
--model Qwen/Qwen2-VL-7B-Instruct \
|
||||
--endpoint /v1/chat/completions \
|
||||
--dataset-name custom_image \
|
||||
--dataset-path <path-to-your-image-data-jsonl> \
|
||||
--custom-ensure-client-side-data
|
||||
--allowed-local-media-path /path/to/image/folder
|
||||
```
|
||||
|
||||
Note that we need to use the `openai-chat` backend and `/v1/chat/completions` endpoint for multimodal inputs.
|
||||
|
||||
@@ -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 | 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 | 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 |
|
||||
@@ -187,6 +187,18 @@ Priority is **1 = highest** (tried first).
|
||||
>
|
||||
> **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2`, `3`, or `4`. Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), FA2 otherwise.
|
||||
|
||||
## MiniMax M3 Sparse Attention Backends
|
||||
|
||||
Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")
|
||||
layers. It is wired in directly by the model and is not part of the
|
||||
automatic priority lists above. A lightning indexer scores KV blocks, the
|
||||
top-k blocks (plus fixed init/local blocks) are selected, and attention
|
||||
attends only to those blocks; index keys live in a separate side cache.
|
||||
|
||||
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
|
||||
| `MINIMAX_M3_SPARSE` | bf16, fp16 | `bfloat16` | 128 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
|
||||
|
||||
## MLA (Multi-head Latent Attention) Backends
|
||||
|
||||
MLA uses separate backends for prefill and decode phases.
|
||||
|
||||
@@ -778,7 +778,7 @@ Then, you can use the OpenAI client as follows:
|
||||
base_url=openai_api_base,
|
||||
)
|
||||
|
||||
video_url = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
|
||||
video_url = "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerFun.mp4"
|
||||
|
||||
## Use video url in the payload
|
||||
chat_completion_from_url = client.chat.completions.create(
|
||||
|
||||
@@ -437,7 +437,6 @@ th {
|
||||
| `LongcatFlashForCausalLM` | LongCat-Flash | `meituan-longcat/LongCat-Flash-Chat`, `meituan-longcat/LongCat-Flash-Chat-FP8` | ✅︎ | ✅︎ |
|
||||
| `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ |
|
||||
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
|
||||
| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ |
|
||||
| `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
|
||||
| `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
|
||||
|
||||
@@ -100,44 +100,14 @@ For further details on renderer APIs, please refer to [this page](renderer.md).
|
||||
- `/version` - Version information
|
||||
- `/load` - Server load metrics
|
||||
|
||||
## Server in development mode
|
||||
|
||||
When using the flag VLLM_SERVER_DEV_MODE=1, you enable development endpoints.
|
||||
|
||||
**SECURITY WARNING: These endpoints should NOT be used in production!**
|
||||
|
||||
### Cache Management APIs
|
||||
|
||||
- `/reset_prefix_cache` - Reset prefix cache (can disrupt service)
|
||||
- `/reset_mm_cache` - Reset multimodal cache (can disrupt service)
|
||||
- `/reset_encoder_cache` - Reset encoder cache (can disrupt service)
|
||||
|
||||
### Weight Transfer APIs (RL Training)
|
||||
|
||||
For further details on Weight Transfer, please refer to [this page](../../training/weight_transfer/README.md).
|
||||
|
||||
- `/pause` - Pause generation (causes denial of service)
|
||||
- `/resume` - Resume generation
|
||||
- `/is_paused` - Check if generation is paused
|
||||
- `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF
|
||||
- `/update_weights` - Update model weights (can alter model behavior)
|
||||
- `/get_world_size` - Get distributed world size
|
||||
|
||||
### Collective RPC
|
||||
|
||||
- `/collective_rpc` - Execute arbitrary RPC methods on the engine (extremely dangerous)
|
||||
|
||||
### Server info
|
||||
|
||||
- `/server_info` - Get detailed server configuration
|
||||
|
||||
### Sleep Mode APIs
|
||||
## Sleep Mode APIs
|
||||
|
||||
For further details on sleep mode, please refer to [this page](../../features/sleep_mode.md).
|
||||
|
||||
- `/sleep` - Put engine to sleep (causes denial of service)
|
||||
- `/wake_up` - Wake engine from sleep
|
||||
- `/is_sleeping` - Check if engine is sleeping
|
||||
- `/collective_rpc` - Execute arbitrary RPC methods on the engine (extremely dangerous)
|
||||
|
||||
## Chat Template
|
||||
|
||||
|
||||
@@ -84,10 +84,7 @@ Both the trainer (`NCCLTrainerSendWeightsArgs`) and inference side (`NCCLWeightT
|
||||
|
||||
## Receiving Weights (Inference Side)
|
||||
|
||||
The inference side triggers weight reception using the four-phase protocol:
|
||||
`init_weight_transfer_engine`, `start_weight_update`, `update_weights`,
|
||||
`finish_weight_update`. The init phase is shown [above](#initialization). The
|
||||
remaining three steps are:
|
||||
The inference side triggers weight reception using the four-phase protocol — `init_weight_transfer_engine`, `start_weight_update`, `update_weights`, `finish_weight_update`. The init phase is shown [above](#initialization). The remaining three steps are:
|
||||
|
||||
```python
|
||||
from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
|
||||
@@ -111,24 +108,12 @@ llm.update_weights(
|
||||
llm.finish_weight_update()
|
||||
```
|
||||
|
||||
The `names`, `dtype_names`, and `shapes` lists describe each parameter. These
|
||||
must match the order in which the trainer iterates over its parameters.
|
||||
The `names`, `dtype_names`, and `shapes` lists describe each parameter. These must match the order in which the trainer iterates over its parameters.
|
||||
|
||||
`start_weight_update` must be called before `update_weights`, and
|
||||
`finish_weight_update` must be called after all weight chunks have been
|
||||
transferred. The `is_checkpoint_format` flag controls whether layerwise reload
|
||||
processing is applied (`True` for checkpoint-format weights, `False` for
|
||||
pre-processed kernel-format weights).
|
||||
|
||||
Sparse NCCL patches still use `update_kind="sparse_flat"` inside
|
||||
`update_info`, but they should be wrapped in
|
||||
`start_weight_update(is_checkpoint_format=False)` because sparse patches apply
|
||||
directly to runtime/kernel-format parameters. The current sparse MVP requires
|
||||
`TP=1` and `PP=1`.
|
||||
`start_weight_update` must be called before `update_weights`, and `finish_weight_update` must be called after all weight chunks have been transferred. The `is_checkpoint_format` flag controls whether layerwise reload processing is applied (`True` for checkpoint-format weights, `False` for pre-processed kernel-format weights).
|
||||
|
||||
## Examples
|
||||
|
||||
- [RLHF with NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_nccl.py) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
|
||||
- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `start_weight_update(is_checkpoint_format=False)` and currently require `TP=1` and `PP=1`
|
||||
- [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
|
||||
- [RLHF with NCCL weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane
|
||||
|
||||
@@ -203,7 +203,7 @@ def run_multi_image(model: str, max_completion_tokens: int) -> None:
|
||||
|
||||
# Video input inference
|
||||
def run_video(model: str, max_completion_tokens: int) -> None:
|
||||
video_url = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
|
||||
video_url = "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerFun.mp4"
|
||||
video_base64 = encode_base64_content_from_url(video_url)
|
||||
|
||||
## Use video url in the payload
|
||||
|
||||
@@ -1,526 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Demonstrates dense-vs-sparse NCCL weight syncing with a real model.
|
||||
|
||||
This example mirrors the validation story used for the sparse NCCL MVP:
|
||||
both the dense update path and the sparse patch path start from the same real
|
||||
checkpoint and apply the same deterministic trainer-side patch. The script then
|
||||
checks that greedy 1-token outputs match between the dense and sparse vLLM
|
||||
engines after the update.
|
||||
|
||||
The example performs the following steps:
|
||||
* Load a training model on one GPU via a Ray actor.
|
||||
* Launch a vLLM engine with the same real model on a second GPU.
|
||||
* Verify trainer vs vLLM baseline agreement before any update.
|
||||
* Apply a deterministic patch to ``model.embed_tokens.weight`` on the trainer.
|
||||
* Run a dense NCCL update into a fresh vLLM engine and collect post-update
|
||||
outputs.
|
||||
* Reset the trainer back to the baseline checkpoint.
|
||||
* Apply the same deterministic patch again.
|
||||
* Run a sparse NCCL update into another fresh vLLM engine and collect
|
||||
post-update outputs.
|
||||
* Compare dense vs sparse baseline outputs, dense vs sparse post-update
|
||||
outputs, estimated payload sizes, and trainer-side send times.
|
||||
|
||||
Current sparse weight transfer MVP limitations:
|
||||
* ``TP=1`` and ``PP=1`` only
|
||||
* sparse updates use runtime/kernel-format parameter names
|
||||
* sparse updates are not composable with checkpoint-format or packed updates
|
||||
|
||||
This example assumes a single-node cluster with two GPUs.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
|
||||
import ray
|
||||
import torch
|
||||
from ray.util.placement_group import placement_group
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.config import WeightTransferConfig
|
||||
from vllm.distributed.weight_transfer.base import SparseWeightPatch
|
||||
from vllm.distributed.weight_transfer.nccl_engine import (
|
||||
NCCLTrainerSendWeightsArgs,
|
||||
NCCLWeightTransferEngine,
|
||||
)
|
||||
from vllm.utils.network_utils import get_ip, get_open_port
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
PATCHED_PARAM_NAME = "model.embed_tokens.weight"
|
||||
MAX_PATCH_ROWS = 32
|
||||
PROMPTS = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
SAMPLING_PARAMS = SamplingParams(temperature=0.0, max_tokens=1)
|
||||
|
||||
|
||||
class MyLLM(LLM):
|
||||
"""Configure the vLLM worker for Ray placement group execution."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1)
|
||||
class TrainModel:
|
||||
"""Ray actor that owns the trainer-side model and deterministic patch state."""
|
||||
|
||||
def __init__(self, model_name: str):
|
||||
self.model_name = model_name
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
if self.tokenizer.pad_token_id is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
|
||||
self.model = None
|
||||
self.patched_param = None
|
||||
self.pending_sparse_patches: list[SparseWeightPatch] | None = None
|
||||
self.model_update_group = None
|
||||
self.master_address = get_ip()
|
||||
self.port = get_open_port()
|
||||
self.reset_model()
|
||||
|
||||
def reset_model(self) -> None:
|
||||
self.model = AutoModelForCausalLM.from_pretrained(
|
||||
self.model_name,
|
||||
torch_dtype=torch.bfloat16,
|
||||
).to("cuda:0")
|
||||
self.model.eval()
|
||||
|
||||
try:
|
||||
self.patched_param = self.model.get_parameter(PATCHED_PARAM_NAME)
|
||||
except AttributeError as exc:
|
||||
raise RuntimeError(
|
||||
f"Expected trainer model to expose `{PATCHED_PARAM_NAME}`"
|
||||
) from exc
|
||||
|
||||
self.pending_sparse_patches = None
|
||||
|
||||
def create_rendezvous(self) -> tuple[str, int]:
|
||||
self.port = get_open_port()
|
||||
return self.master_address, self.port
|
||||
|
||||
def init_weight_transfer_group(self, world_size: int) -> None:
|
||||
self.model_update_group = NCCLWeightTransferEngine.trainer_init(
|
||||
dict(
|
||||
master_address=self.master_address,
|
||||
master_port=self.port,
|
||||
world_size=world_size,
|
||||
)
|
||||
)
|
||||
|
||||
def get_dense_update_info(self, packed: bool = False) -> tuple[dict, int]:
|
||||
names = []
|
||||
dtype_names = []
|
||||
shapes = []
|
||||
payload_bytes = 0
|
||||
for name, param in self.model.named_parameters():
|
||||
names.append(name)
|
||||
dtype_names.append(str(param.dtype).split(".")[-1])
|
||||
shapes.append(list(param.shape))
|
||||
payload_bytes += param.numel() * param.element_size()
|
||||
|
||||
return (
|
||||
dict(
|
||||
names=names,
|
||||
dtype_names=dtype_names,
|
||||
shapes=shapes,
|
||||
packed=packed,
|
||||
),
|
||||
payload_bytes,
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate(
|
||||
self,
|
||||
prompts: Sequence[str],
|
||||
max_new_tokens: int = 1,
|
||||
) -> list[dict[str, object]]:
|
||||
generations = []
|
||||
for prompt in prompts:
|
||||
model_inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda:0")
|
||||
output = self.model.generate(
|
||||
**model_inputs,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
pad_token_id=self.tokenizer.pad_token_id,
|
||||
)
|
||||
new_token_ids = output[0, model_inputs["input_ids"].shape[1] :].tolist()
|
||||
generations.append(
|
||||
{
|
||||
"token_ids": new_token_ids,
|
||||
"text": self.tokenizer.decode(
|
||||
new_token_ids,
|
||||
skip_special_tokens=False,
|
||||
),
|
||||
}
|
||||
)
|
||||
return generations
|
||||
|
||||
def prepare_sparse_patch(
|
||||
self,
|
||||
prompts: Sequence[str],
|
||||
max_patch_rows: int = MAX_PATCH_ROWS,
|
||||
) -> tuple[dict[str, object], list[int], str, int]:
|
||||
selected_token_ids: list[int] = []
|
||||
special_ids = set(self.tokenizer.all_special_ids)
|
||||
for prompt in prompts:
|
||||
token_ids = self.tokenizer(prompt, add_special_tokens=False)["input_ids"]
|
||||
for token_id in token_ids:
|
||||
if token_id in special_ids or token_id in selected_token_ids:
|
||||
continue
|
||||
selected_token_ids.append(token_id)
|
||||
if len(selected_token_ids) == max_patch_rows:
|
||||
break
|
||||
if len(selected_token_ids) == max_patch_rows:
|
||||
break
|
||||
|
||||
if not selected_token_ids:
|
||||
raise ValueError("Could not derive any non-special token IDs to patch")
|
||||
|
||||
vocab_size = self.patched_param.shape[0]
|
||||
next_token_id = selected_token_ids[-1]
|
||||
while len(selected_token_ids) < max_patch_rows:
|
||||
next_token_id = (next_token_id + 1) % vocab_size
|
||||
if next_token_id in special_ids or next_token_id in selected_token_ids:
|
||||
continue
|
||||
selected_token_ids.append(next_token_id)
|
||||
|
||||
row_ids = torch.tensor(
|
||||
selected_token_ids,
|
||||
device=self.patched_param.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
hidden_size = self.patched_param.shape[1]
|
||||
column_offsets = torch.arange(
|
||||
hidden_size,
|
||||
device=self.patched_param.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
# Rotate the selected embedding rows instead of zeroing them so the
|
||||
# patch remains deterministic while avoiding a degenerate collapse
|
||||
# to the same special token after the update.
|
||||
replacement_rows = self.patched_param[row_ids].roll(shifts=1, dims=0)
|
||||
self.patched_param[row_ids] = replacement_rows
|
||||
|
||||
flat_indices = (
|
||||
row_ids.unsqueeze(1).mul(hidden_size).add(column_offsets).reshape(-1)
|
||||
)
|
||||
flat_values = self.patched_param[row_ids].reshape(-1).contiguous()
|
||||
self.pending_sparse_patches = [
|
||||
SparseWeightPatch(
|
||||
name=PATCHED_PARAM_NAME,
|
||||
indices=flat_indices.to(torch.int32),
|
||||
values=flat_values,
|
||||
)
|
||||
]
|
||||
patch_digest = hashlib.sha256(
|
||||
self.pending_sparse_patches[0].indices.cpu().numpy().tobytes()
|
||||
+ self.pending_sparse_patches[0]
|
||||
.values.detach()
|
||||
.float()
|
||||
.cpu()
|
||||
.numpy()
|
||||
.tobytes()
|
||||
).hexdigest()
|
||||
|
||||
sparse_payload_bytes = (
|
||||
flat_indices.numel() * torch.tensor([], dtype=torch.int32).element_size()
|
||||
+ flat_values.numel() * flat_values.element_size()
|
||||
)
|
||||
update_info = dict(
|
||||
names=[PATCHED_PARAM_NAME],
|
||||
dtype_names=[str(self.patched_param.dtype).split(".")[-1]],
|
||||
shapes=[list(self.patched_param.shape)],
|
||||
num_updates_list=[flat_indices.numel()],
|
||||
update_kind="sparse_flat",
|
||||
)
|
||||
return update_info, selected_token_ids, patch_digest, sparse_payload_bytes
|
||||
|
||||
def broadcast_weights(self, packed: bool = False) -> float:
|
||||
if self.model_update_group is None:
|
||||
raise RuntimeError("Weight transfer group is not initialized")
|
||||
|
||||
trainer_args = NCCLTrainerSendWeightsArgs(
|
||||
group=self.model_update_group,
|
||||
packed=packed,
|
||||
)
|
||||
start = time.perf_counter()
|
||||
NCCLWeightTransferEngine.trainer_send_weights(
|
||||
iterator=self.model.named_parameters(),
|
||||
trainer_args=trainer_args,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
return (time.perf_counter() - start) * 1000.0
|
||||
|
||||
def broadcast_pending_sparse_patch(self) -> float:
|
||||
if self.model_update_group is None:
|
||||
raise RuntimeError("Weight transfer group is not initialized")
|
||||
if self.pending_sparse_patches is None:
|
||||
raise RuntimeError("Sparse patch has not been prepared")
|
||||
|
||||
start = time.perf_counter()
|
||||
NCCLWeightTransferEngine.trainer_send_sparse_weights(
|
||||
iter(self.pending_sparse_patches),
|
||||
NCCLTrainerSendWeightsArgs(group=self.model_update_group),
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
self.pending_sparse_patches = None
|
||||
return (time.perf_counter() - start) * 1000.0
|
||||
|
||||
|
||||
def launch_llm(
|
||||
scheduling_inference: PlacementGroupSchedulingStrategy,
|
||||
):
|
||||
return ray.remote(
|
||||
num_cpus=0,
|
||||
num_gpus=0,
|
||||
scheduling_strategy=scheduling_inference,
|
||||
)(MyLLM).remote(
|
||||
model=MODEL_NAME,
|
||||
enforce_eager=True,
|
||||
tensor_parallel_size=1,
|
||||
distributed_executor_backend="ray",
|
||||
gpu_memory_utilization=0.7,
|
||||
weight_transfer_config=WeightTransferConfig(backend="nccl"),
|
||||
)
|
||||
|
||||
|
||||
def collect_vllm_generations(llm_handle) -> list[dict[str, object]]:
|
||||
outputs = ray.get(llm_handle.generate.remote(PROMPTS, SAMPLING_PARAMS))
|
||||
generations = []
|
||||
for output in outputs:
|
||||
generations.append(
|
||||
{
|
||||
"token_ids": output.outputs[0].token_ids,
|
||||
"text": output.outputs[0].text,
|
||||
}
|
||||
)
|
||||
return generations
|
||||
|
||||
|
||||
def token_sequences_match(
|
||||
left: Sequence[dict[str, object]],
|
||||
right: Sequence[dict[str, object]],
|
||||
) -> bool:
|
||||
return [item["token_ids"] for item in left] == [item["token_ids"] for item in right]
|
||||
|
||||
|
||||
def print_generations(label: str, prompts: Sequence[str], generations) -> None:
|
||||
print(f"\n{label}")
|
||||
print("-" * 50)
|
||||
for prompt, generation in zip(prompts, generations):
|
||||
print(f"Prompt: {prompt!r}")
|
||||
print(f"Token IDs: {generation['token_ids']}")
|
||||
print(f"Text: {generation['text']!r}")
|
||||
print("-" * 50)
|
||||
|
||||
|
||||
def run_dense_phase(
|
||||
train_model,
|
||||
scheduling_inference: PlacementGroupSchedulingStrategy,
|
||||
) -> dict[str, object]:
|
||||
ray.get(train_model.reset_model.remote())
|
||||
llm = launch_llm(scheduling_inference)
|
||||
try:
|
||||
dense_before = collect_vllm_generations(llm)
|
||||
|
||||
ray.get(llm.sleep.remote(level=0))
|
||||
master_address, master_port = ray.get(train_model.create_rendezvous.remote())
|
||||
world_size = ray.get(llm.get_world_size.remote()) + 1
|
||||
inference_init = llm.init_weight_transfer_engine.remote(
|
||||
dict(
|
||||
init_info=dict(
|
||||
master_address=master_address,
|
||||
master_port=master_port,
|
||||
rank_offset=1,
|
||||
world_size=world_size,
|
||||
)
|
||||
)
|
||||
)
|
||||
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
|
||||
ray.get([trainer_init, inference_init])
|
||||
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
|
||||
|
||||
dense_update_info, dense_payload_bytes = ray.get(
|
||||
train_model.get_dense_update_info.remote()
|
||||
)
|
||||
_, selected_token_ids, patch_digest, _ = ray.get(
|
||||
train_model.prepare_sparse_patch.remote(PROMPTS)
|
||||
)
|
||||
|
||||
inference_update = llm.update_weights.remote(
|
||||
dict(update_info=dense_update_info)
|
||||
)
|
||||
dense_send_ms, _ = ray.get(
|
||||
[
|
||||
train_model.broadcast_weights.remote(packed=False),
|
||||
inference_update,
|
||||
]
|
||||
)
|
||||
ray.get(llm.finish_weight_update.remote())
|
||||
ray.get(llm.wake_up.remote(tags=["scheduling"]))
|
||||
|
||||
dense_after = collect_vllm_generations(llm)
|
||||
|
||||
return {
|
||||
"dense_before": dense_before,
|
||||
"dense_after": dense_after,
|
||||
"selected_token_ids": selected_token_ids,
|
||||
"patch_digest": patch_digest,
|
||||
"dense_payload_bytes": dense_payload_bytes,
|
||||
"dense_send_ms": dense_send_ms,
|
||||
}
|
||||
finally:
|
||||
ray.kill(llm)
|
||||
|
||||
|
||||
def run_sparse_phase(
|
||||
train_model,
|
||||
scheduling_inference: PlacementGroupSchedulingStrategy,
|
||||
) -> dict[str, object]:
|
||||
ray.get(train_model.reset_model.remote())
|
||||
llm = launch_llm(scheduling_inference)
|
||||
try:
|
||||
sparse_before = collect_vllm_generations(llm)
|
||||
|
||||
ray.get(llm.sleep.remote(level=0))
|
||||
master_address, master_port = ray.get(train_model.create_rendezvous.remote())
|
||||
world_size = ray.get(llm.get_world_size.remote()) + 1
|
||||
inference_init = llm.init_weight_transfer_engine.remote(
|
||||
dict(
|
||||
init_info=dict(
|
||||
master_address=master_address,
|
||||
master_port=master_port,
|
||||
rank_offset=1,
|
||||
world_size=world_size,
|
||||
)
|
||||
)
|
||||
)
|
||||
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
|
||||
ray.get([trainer_init, inference_init])
|
||||
ray.get(llm.start_weight_update.remote(is_checkpoint_format=False))
|
||||
|
||||
sparse_update_info, selected_token_ids, patch_digest, sparse_payload_bytes = (
|
||||
ray.get(train_model.prepare_sparse_patch.remote(PROMPTS))
|
||||
)
|
||||
|
||||
inference_update = llm.update_weights.remote(
|
||||
dict(update_info=sparse_update_info)
|
||||
)
|
||||
sparse_send_ms, _ = ray.get(
|
||||
[
|
||||
train_model.broadcast_pending_sparse_patch.remote(),
|
||||
inference_update,
|
||||
]
|
||||
)
|
||||
ray.get(llm.finish_weight_update.remote())
|
||||
ray.get(llm.wake_up.remote(tags=["scheduling"]))
|
||||
|
||||
sparse_after = collect_vllm_generations(llm)
|
||||
|
||||
return {
|
||||
"sparse_before": sparse_before,
|
||||
"sparse_after": sparse_after,
|
||||
"selected_token_ids": selected_token_ids,
|
||||
"patch_digest": patch_digest,
|
||||
"sparse_payload_bytes": sparse_payload_bytes,
|
||||
"sparse_send_ms": sparse_send_ms,
|
||||
}
|
||||
finally:
|
||||
ray.kill(llm)
|
||||
|
||||
|
||||
ray.init()
|
||||
|
||||
try:
|
||||
train_model = TrainModel.remote(MODEL_NAME)
|
||||
|
||||
pg_inference = placement_group([{"GPU": 1, "CPU": 0}])
|
||||
ray.get(pg_inference.ready())
|
||||
scheduling_inference = PlacementGroupSchedulingStrategy(
|
||||
placement_group=pg_inference,
|
||||
placement_group_capture_child_tasks=True,
|
||||
placement_group_bundle_index=0,
|
||||
)
|
||||
|
||||
dense_results = run_dense_phase(train_model, scheduling_inference)
|
||||
sparse_results = run_sparse_phase(train_model, scheduling_inference)
|
||||
|
||||
baseline_equal = token_sequences_match(
|
||||
dense_results["dense_before"],
|
||||
sparse_results["sparse_before"],
|
||||
)
|
||||
patch_selection_equal = (
|
||||
dense_results["selected_token_ids"] == sparse_results["selected_token_ids"]
|
||||
)
|
||||
patch_digest_equal = dense_results["patch_digest"] == sparse_results["patch_digest"]
|
||||
after_equal = token_sequences_match(
|
||||
dense_results["dense_after"],
|
||||
sparse_results["sparse_after"],
|
||||
)
|
||||
any_output_changed = any(
|
||||
before["token_ids"] != after["token_ids"]
|
||||
for before, after in zip(
|
||||
dense_results["dense_before"],
|
||||
dense_results["dense_after"],
|
||||
)
|
||||
)
|
||||
dense_payload_mb = dense_results["dense_payload_bytes"] / (1024 * 1024)
|
||||
sparse_payload_mb = sparse_results["sparse_payload_bytes"] / (1024 * 1024)
|
||||
|
||||
print_generations(
|
||||
"Dense baseline outputs",
|
||||
PROMPTS,
|
||||
dense_results["dense_before"],
|
||||
)
|
||||
print_generations(
|
||||
"Sparse baseline outputs", PROMPTS, sparse_results["sparse_before"]
|
||||
)
|
||||
print_generations(
|
||||
"Dense outputs after update", PROMPTS, dense_results["dense_after"]
|
||||
)
|
||||
print_generations(
|
||||
"Sparse outputs after update",
|
||||
PROMPTS,
|
||||
sparse_results["sparse_after"],
|
||||
)
|
||||
|
||||
print(f"patched_token_ids = {dense_results['selected_token_ids']}")
|
||||
print(f"patch_selection_equal = {patch_selection_equal}")
|
||||
print(f"dense_patch_digest = {dense_results['patch_digest']}")
|
||||
print(f"sparse_patch_digest = {sparse_results['patch_digest']}")
|
||||
print(f"patch_digest_equal = {patch_digest_equal}")
|
||||
print(f"baseline_equal = {baseline_equal}")
|
||||
print(f"after_equal = {after_equal}")
|
||||
print(f"any_output_changed = {any_output_changed}")
|
||||
print(f"dense_payload_mb = {dense_payload_mb:.2f}")
|
||||
print(f"sparse_payload_mb = {sparse_payload_mb:.2f}")
|
||||
print(f"dense_send_ms = {dense_results['dense_send_ms']:.2f}")
|
||||
print(f"sparse_send_ms = {sparse_results['sparse_send_ms']:.2f}")
|
||||
|
||||
if not baseline_equal:
|
||||
raise RuntimeError(
|
||||
"Dense and sparse phases did not start from the same baseline"
|
||||
)
|
||||
if not patch_selection_equal:
|
||||
raise RuntimeError("Dense and sparse phases used different sparse patches")
|
||||
if not patch_digest_equal:
|
||||
raise RuntimeError("Dense and sparse phases produced different patch values")
|
||||
if not after_equal:
|
||||
raise RuntimeError("Dense and sparse updates produced different outputs")
|
||||
if not any_output_changed:
|
||||
raise RuntimeError("Patch did not change the observed outputs")
|
||||
finally:
|
||||
ray.shutdown()
|
||||
@@ -162,7 +162,12 @@ dout = "dout"
|
||||
Pn = "Pn"
|
||||
arange = "arange"
|
||||
thw = "thw"
|
||||
# temporal position ids (parallels hpos/wpos in vision RoPE)
|
||||
tpos = "tpos"
|
||||
subtile = "subtile"
|
||||
subtiles = "subtiles"
|
||||
reord = "reord"
|
||||
Ot = "Ot"
|
||||
HSA = "HSA"
|
||||
setp = "setp"
|
||||
CPY = "CPY"
|
||||
|
||||
@@ -29,6 +29,7 @@ xgrammar >= 0.2.0, < 1.0.0; platform_machine == "x86_64" or platform_machine ==
|
||||
typing_extensions >= 4.10
|
||||
filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
|
||||
partial-json-parser # used for parsing partial JSON outputs
|
||||
jsonschema >= 4.23.0 # required for MiniMax M3 tool schema validation
|
||||
pyzmq >= 25.0.0
|
||||
msgspec
|
||||
gguf >= 0.17.0
|
||||
|
||||
@@ -9,8 +9,8 @@ torchaudio==2.11.0
|
||||
# These must be updated alongside torch
|
||||
torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
|
||||
# FlashInfer should be updated together with the Dockerfile
|
||||
flashinfer-python==0.6.11.post2
|
||||
flashinfer-cubin==0.6.11.post2
|
||||
flashinfer-python==0.6.12
|
||||
flashinfer-cubin==0.6.12
|
||||
apache-tvm-ffi==0.1.9
|
||||
tilelang==0.1.9
|
||||
# Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to
|
||||
|
||||
@@ -360,6 +360,7 @@ jsonpointer==3.0.0
|
||||
# via jsonschema
|
||||
jsonschema==4.23.0
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# hypothesis-jsonschema
|
||||
# mistral-common
|
||||
# ray
|
||||
|
||||
@@ -440,6 +440,8 @@ jsonpointer==3.1.0
|
||||
# via jsonschema
|
||||
jsonschema==4.26.0
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/../common.txt
|
||||
# hypothesis-jsonschema
|
||||
# mcp
|
||||
# mistral-common
|
||||
|
||||
@@ -229,6 +229,7 @@ jsonlines==4.0.0
|
||||
# via lm-eval
|
||||
jsonschema==4.26.0
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# hypothesis-jsonschema
|
||||
# mistral-common
|
||||
# schemathesis
|
||||
|
||||
Generated
+87
@@ -3458,6 +3458,75 @@ version = "0.1.29"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e0c5ccf5294c6ccd63a74f1565028353830a9c2f5eb0c682c355c471726a6e3f"
|
||||
|
||||
[[package]]
|
||||
name = "pyo3"
|
||||
version = "0.28.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"once_cell",
|
||||
"portable-atomic",
|
||||
"pyo3-build-config",
|
||||
"pyo3-ffi",
|
||||
"pyo3-macros",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-build-config"
|
||||
version = "0.28.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e"
|
||||
dependencies = [
|
||||
"target-lexicon",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-ffi"
|
||||
version = "0.28.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"pyo3-build-config",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-macros"
|
||||
version = "0.28.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"pyo3-macros-backend",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pyo3-macros-backend"
|
||||
version = "0.28.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb"
|
||||
dependencies = [
|
||||
"heck",
|
||||
"proc-macro2",
|
||||
"pyo3-build-config",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pythonize"
|
||||
version = "0.28.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b79f670c9626c8b651c0581011b57b6ba6970bb69faf01a7c4c0cfc81c43f95"
|
||||
dependencies = [
|
||||
"pyo3",
|
||||
"serde",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "qoi"
|
||||
version = "0.4.1"
|
||||
@@ -4669,6 +4738,12 @@ dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "target-lexicon"
|
||||
version = "0.13.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
|
||||
|
||||
[[package]]
|
||||
name = "task-local"
|
||||
version = "0.1.1"
|
||||
@@ -5622,6 +5697,7 @@ dependencies = [
|
||||
"expect-test",
|
||||
"futures",
|
||||
"half",
|
||||
"indexmap 2.13.0",
|
||||
"itertools 0.14.0",
|
||||
"llm-multimodal",
|
||||
"minijinja",
|
||||
@@ -5900,6 +5976,17 @@ dependencies = [
|
||||
"winnow",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-tool-parser-py"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"pyo3",
|
||||
"pythonize",
|
||||
"serde_json",
|
||||
"thiserror-ext",
|
||||
"vllm-tool-parser",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "walkdir"
|
||||
version = "2.5.0"
|
||||
|
||||
+5
-1
@@ -12,6 +12,7 @@ members = [
|
||||
"src/text",
|
||||
"src/tokenizer",
|
||||
"src/tool-parser",
|
||||
"src/tool-parser/python",
|
||||
]
|
||||
resolver = "3"
|
||||
|
||||
@@ -43,6 +44,7 @@ half = { version = "2.7.1", features = ["bytemuck"] }
|
||||
hex = "0.4.3"
|
||||
hf-hub = { version = "0.5.0", features = ["tokio"] }
|
||||
http-body = "1.0.1"
|
||||
indexmap = "2.13.0"
|
||||
itertools = "0.14.0"
|
||||
libc = "0.2.177"
|
||||
llm-multimodal = { git = "https://github.com/vllm-project/llm-multimodal", rev = "5b558989844d1c7af3e43d0f604069ffd9c06320" }
|
||||
@@ -59,6 +61,8 @@ prometheus-client = "0.24.0"
|
||||
prometheus-client-derive-encode = "0.5.0"
|
||||
prost = "0.14.3"
|
||||
prost-types = "0.14.3"
|
||||
pyo3 = "0.28.3"
|
||||
pythonize = "0.28.0"
|
||||
rand = "0.9.2"
|
||||
reasoning-parser = "1.2.2"
|
||||
reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] }
|
||||
@@ -69,7 +73,7 @@ rustc-hash = "1.1.0"
|
||||
serde = { version = "1.0.228", features = ["derive"] }
|
||||
serde-json-fmt = "0.1.0"
|
||||
serde_default = "0.2.0"
|
||||
serde_json = { version = "1.0.145", features = ["arbitrary_precision", "preserve_order"] }
|
||||
serde_json = { version = "1.0.145", features = ["preserve_order"] }
|
||||
serde_repr = "0.1.20"
|
||||
serde_tuple = "1.1.3"
|
||||
serde_with = "3.18.0"
|
||||
|
||||
@@ -10,6 +10,7 @@ asynk-strim-attr.workspace = true
|
||||
easy-ext.workspace = true
|
||||
futures.workspace = true
|
||||
half.workspace = true
|
||||
indexmap.workspace = true
|
||||
itertools.workspace = true
|
||||
llm-multimodal.workspace = true
|
||||
minijinja.workspace = true
|
||||
|
||||
@@ -233,7 +233,7 @@ mod tests {
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
|
||||
expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, hy_v3, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -5,8 +5,9 @@ use std::sync::LazyLock;
|
||||
pub use vllm_reasoning_parser::{
|
||||
CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
|
||||
DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser,
|
||||
KimiReasoningParser, MiniMaxM2ReasoningParser, NemotronV3ReasoningParser, Qwen3ReasoningParser,
|
||||
ReasoningDelta, ReasoningError, ReasoningParser, Step3ReasoningParser,
|
||||
KimiReasoningParser, MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser,
|
||||
NemotronV3ReasoningParser, Qwen3ReasoningParser, ReasoningDelta, ReasoningError,
|
||||
ReasoningParser, Step3ReasoningParser,
|
||||
};
|
||||
use vllm_tokenizer::DynTokenizer;
|
||||
|
||||
@@ -23,6 +24,7 @@ pub mod names {
|
||||
pub const KIMI: &str = "kimi";
|
||||
pub const KIMI_K2: &str = "kimi_k2";
|
||||
pub const MINIMAX_M2: &str = "minimax_m2";
|
||||
pub const MINIMAX_M3: &str = "minimax_m3";
|
||||
pub const NEMOTRON_V3: &str = "nemotron_v3";
|
||||
pub const QWEN3: &str = "qwen3";
|
||||
pub const STEP3: &str = "step3";
|
||||
@@ -59,6 +61,7 @@ impl ReasoningParserFactory {
|
||||
.register_parser::<KimiReasoningParser>(names::KIMI)
|
||||
.register_parser::<KimiK2ReasoningParser>(names::KIMI_K2)
|
||||
.register_parser::<MiniMaxM2ReasoningParser>(names::MINIMAX_M2)
|
||||
.register_parser::<MiniMaxM3ReasoningParser>(names::MINIMAX_M3)
|
||||
.register_parser::<NemotronV3ReasoningParser>(names::NEMOTRON_V3)
|
||||
.register_parser::<Qwen3ReasoningParser>(names::QWEN3)
|
||||
.register_parser::<Step3ReasoningParser>(names::STEP3);
|
||||
@@ -78,6 +81,8 @@ impl ReasoningParserFactory {
|
||||
.register_pattern("kimi-k2", names::KIMI_K2)
|
||||
.register_pattern("kimi", names::KIMI)
|
||||
.register_pattern("step3", names::STEP3)
|
||||
.register_pattern("minimax-m3", names::MINIMAX_M3)
|
||||
.register_pattern("mm-m3", names::MINIMAX_M3)
|
||||
.register_pattern("minimax", names::MINIMAX_M2)
|
||||
.register_pattern("mm-m2", names::MINIMAX_M2)
|
||||
.register_pattern("cohere", names::COHERE_CMD)
|
||||
|
||||
@@ -32,8 +32,10 @@ fn factory_contains_and_lists_registered_parsers() {
|
||||
let factory = ReasoningParserFactory::new();
|
||||
assert!(factory.contains(names::QWEN3));
|
||||
assert!(factory.contains(names::DEEPSEEK_V4));
|
||||
assert!(factory.contains(names::MINIMAX_M3));
|
||||
assert!(factory.list().contains(&names::QWEN3.to_string()));
|
||||
assert!(factory.list().contains(&names::DEEPSEEK_V4.to_string()));
|
||||
assert!(factory.list().contains(&names::MINIMAX_M3.to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -49,6 +51,19 @@ fn factory_resolves_deepseek_v4_to_qwen3_alias() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn factory_resolves_minimax_m3_before_generic_minimax() {
|
||||
let factory = ReasoningParserFactory::new();
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("MiniMaxAI/Minimax-M3-preview"),
|
||||
Some(names::MINIMAX_M3)
|
||||
);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("mm-m3"),
|
||||
Some(names::MINIMAX_M3)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn factory_rejects_unknown_parser_names() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
|
||||
@@ -5,7 +5,7 @@ use std::sync::LazyLock;
|
||||
pub use vllm_tool_parser::{
|
||||
DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
|
||||
Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, HermesToolParser, HyV3ToolParser,
|
||||
Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser,
|
||||
KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, MinimaxM3ToolParser,
|
||||
MistralToolParser, Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser,
|
||||
ToolParserError, ToolParserOutput,
|
||||
};
|
||||
@@ -24,13 +24,11 @@ pub mod names {
|
||||
pub const GEMMA4: &str = "gemma4";
|
||||
pub const HERMES: &str = "hermes";
|
||||
pub const HY_V3: &str = "hy_v3";
|
||||
// Matches the Python CLI name `--tool-call-parser internlm`, which Python
|
||||
// also routes to `Internlm2ToolParser` despite the version-agnostic name.
|
||||
pub const INTERNLM: &str = "internlm";
|
||||
pub const KIMI_K2: &str = "kimi_k2";
|
||||
pub const LLAMA3_JSON: &str = "llama3_json";
|
||||
pub const LLAMA4_JSON: &str = "llama4_json";
|
||||
pub const MINIMAX_M2: &str = "minimax_m2";
|
||||
pub const MINIMAX_M3: &str = "minimax_m3";
|
||||
pub const MISTRAL: &str = "mistral";
|
||||
pub const QWEN3_CODER: &str = "qwen3_coder";
|
||||
pub const QWEN3_XML: &str = "qwen3_xml";
|
||||
@@ -65,11 +63,11 @@ impl ToolParserFactory {
|
||||
.register_parser::<Gemma4ToolParser>(names::GEMMA4)
|
||||
.register_parser::<HermesToolParser>(names::HERMES)
|
||||
.register_parser::<HyV3ToolParser>(names::HY_V3)
|
||||
.register_parser::<Internlm2ToolParser>(names::INTERNLM)
|
||||
.register_parser::<KimiK2ToolParser>(names::KIMI_K2)
|
||||
.register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
|
||||
.register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
|
||||
.register_parser::<MinimaxM2ToolParser>(names::MINIMAX_M2)
|
||||
.register_parser::<MinimaxM3ToolParser>(names::MINIMAX_M3)
|
||||
.register_parser::<MistralToolParser>(names::MISTRAL)
|
||||
.register_parser::<Qwen3XmlToolParser>(names::QWEN3_XML)
|
||||
.register_parser::<Qwen3CoderToolParser>(names::QWEN3_CODER);
|
||||
@@ -84,12 +82,6 @@ impl ToolParserFactory {
|
||||
.register_pattern("hermes", names::HERMES)
|
||||
.register_pattern("hy3", names::HY_V3)
|
||||
.register_pattern("hy_v3", names::HY_V3)
|
||||
// Narrow to `internlm2` substring so it matches `internlm2-chat-7b`
|
||||
// and `internlm2_5-7b-chat` but NOT `internlm-chat-7b` (InternLM v1,
|
||||
// routes to Llama), `internlm3-*` (also Llama-architecture per
|
||||
// vllm/model_executor/models/registry.py:146), or `Intern-S1` /
|
||||
// `Intern-S1-Pro` (separate intern-s1 parser, see PR #40115).
|
||||
.register_pattern("internlm2", names::INTERNLM)
|
||||
.register_pattern("llama-4", names::LLAMA4_JSON)
|
||||
.register_pattern("llama-3.2", names::LLAMA3_JSON)
|
||||
.register_pattern("llama-3.1", names::LLAMA3_JSON)
|
||||
@@ -106,6 +98,8 @@ impl ToolParserFactory {
|
||||
.register_pattern("gemma4", names::GEMMA4)
|
||||
.register_pattern("gemma-4", names::GEMMA4)
|
||||
.register_pattern("kimi-k2", names::KIMI_K2)
|
||||
.register_pattern("minimax-m3", names::MINIMAX_M3)
|
||||
.register_pattern("mm-m3", names::MINIMAX_M3)
|
||||
.register_pattern("minimax", names::MINIMAX_M2)
|
||||
.register_pattern("mm-m2", names::MINIMAX_M2);
|
||||
|
||||
|
||||
@@ -153,6 +153,14 @@ fn factory_new_resolves_default_patterns() {
|
||||
factory.resolve_name_for_model("tencent/Hy3-preview"),
|
||||
Some(names::HY_V3)
|
||||
);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("MiniMax/MiniMax-M3-Text"),
|
||||
Some(names::MINIMAX_M3)
|
||||
);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("org/mm-m3-base"),
|
||||
Some(names::MINIMAX_M3)
|
||||
);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("MiniMax/MiniMax-M2-01"),
|
||||
Some(names::MINIMAX_M2)
|
||||
@@ -161,33 +169,4 @@ fn factory_new_resolves_default_patterns() {
|
||||
factory.resolve_name_for_model("org/mm-m2-base"),
|
||||
Some(names::MINIMAX_M2)
|
||||
);
|
||||
|
||||
// InternLM2 positive: both dashed and underscored versioned names route.
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("internlm/internlm2-chat-7b"),
|
||||
Some(names::INTERNLM)
|
||||
);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("internlm/internlm2_5-7b-chat"),
|
||||
Some(names::INTERNLM)
|
||||
);
|
||||
|
||||
// Negative: other internlm-org models do NOT route to the InternLM2 parser,
|
||||
// since they use unrelated prompt formats.
|
||||
// - InternLM v1 (`internlm-chat-7b`) routes to Llama
|
||||
// - InternLM3 (`internlm3-8b-instruct`) routes to Llama
|
||||
// - Intern-S1 / Intern-S1-Pro have their own parser (Python PR #40115)
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("internlm/internlm-chat-7b"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("internlm/internlm3-8b-instruct"),
|
||||
None
|
||||
);
|
||||
assert_eq!(factory.resolve_name_for_model("internlm/Intern-S1"), None);
|
||||
assert_eq!(
|
||||
factory.resolve_name_for_model("internlm/Intern-S1-Pro"),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
use serde_json::Value as JsonValue;
|
||||
use thiserror_ext::AsReport as _;
|
||||
use tracing::{info, trace, warn};
|
||||
use vllm_text::Prompt;
|
||||
@@ -13,6 +13,7 @@ use self::format::{
|
||||
ChatTemplateContentFormat, ChatTemplateContentFormatOption as ContentFormatOption,
|
||||
};
|
||||
use self::template::{CompiledChatTemplate, TemplateContext};
|
||||
use self::value::{TemplateValue, to_template_value};
|
||||
use super::{ChatRenderer, RenderedPrompt};
|
||||
use crate::error::Result;
|
||||
use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest};
|
||||
@@ -24,6 +25,7 @@ mod error;
|
||||
mod format;
|
||||
mod template;
|
||||
mod tojson;
|
||||
mod value;
|
||||
|
||||
pub use template::{load_chat_template, resolve_chat_template};
|
||||
|
||||
@@ -38,7 +40,7 @@ pub struct MultimodalRenderInfo {
|
||||
/// state.
|
||||
pub struct HfChatRenderer {
|
||||
default_template: Option<CompiledChatTemplate>,
|
||||
default_template_kwargs: HashMap<String, Value>,
|
||||
default_template_kwargs: HashMap<String, JsonValue>,
|
||||
content_format: ContentFormatOption,
|
||||
special_tokens: Option<HfSpecialTokens>,
|
||||
multimodal: Option<MultimodalRenderInfo>,
|
||||
@@ -48,7 +50,7 @@ impl HfChatRenderer {
|
||||
/// Create a renderer from the given template string.
|
||||
pub fn new(
|
||||
template: Option<String>,
|
||||
default_template_kwargs: HashMap<String, Value>,
|
||||
default_template_kwargs: HashMap<String, JsonValue>,
|
||||
content_format: ContentFormatOption,
|
||||
) -> Result<Self> {
|
||||
Ok(Self {
|
||||
@@ -245,7 +247,7 @@ struct TemplateToolCall {
|
||||
#[derive(Debug, Serialize)]
|
||||
struct TemplateToolFunction {
|
||||
name: String,
|
||||
arguments: Value,
|
||||
arguments: TemplateValue,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
@@ -259,7 +261,7 @@ pub(super) struct TemplateTool {
|
||||
struct TemplateToolDefinition {
|
||||
name: String,
|
||||
description: Option<String>,
|
||||
parameters: Value,
|
||||
parameters: TemplateValue,
|
||||
strict: Option<bool>,
|
||||
}
|
||||
|
||||
@@ -345,13 +347,14 @@ fn to_template_tool_calls(
|
||||
let mut tool_calls = Vec::new();
|
||||
|
||||
for tool_call in content.tool_calls() {
|
||||
let arguments = serde_json::from_str::<Value>(&tool_call.arguments).map_err(|error| {
|
||||
let arguments = serde_json::from_str(&tool_call.arguments).map_err(|error| {
|
||||
Error::ChatTemplate(format!(
|
||||
"assistant tool call `{}` has invalid JSON arguments: {}",
|
||||
tool_call.id,
|
||||
error.as_report()
|
||||
))
|
||||
})?;
|
||||
let arguments = to_template_value(arguments);
|
||||
|
||||
tool_calls.push(TemplateToolCall {
|
||||
id: tool_call.id.clone(),
|
||||
@@ -434,7 +437,7 @@ fn to_template_tools(tools: &[ChatTool]) -> Vec<TemplateTool> {
|
||||
function: TemplateToolDefinition {
|
||||
name: tool.name.clone(),
|
||||
description: tool.description.clone(),
|
||||
parameters: tool.parameters.clone(),
|
||||
parameters: to_template_value(tool.parameters.clone()),
|
||||
strict: tool.strict,
|
||||
},
|
||||
})
|
||||
@@ -909,6 +912,29 @@ mod tests {
|
||||
assert_eq!(rendered, "get_weather|Paris|call_1|Sunny");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chat_template_tool_call_argument_items_method_is_not_shadowed_by_field() {
|
||||
let request = sample_request(vec![ChatMessage::assistant_blocks(vec![
|
||||
AssistantContentBlock::ToolCall(crate::AssistantToolCall {
|
||||
id: "call_1".to_string(),
|
||||
name: "add".to_string(),
|
||||
arguments: r#"{"items":"operands","x":2,"y":1.0}"#.to_string(),
|
||||
}),
|
||||
])]);
|
||||
|
||||
let rendered = render(
|
||||
Some(
|
||||
"{%- set arguments = messages[0].tool_calls[0].function.arguments -%}
|
||||
{%- for key, value in arguments.items() -%}{{ key }}={{ value }};{%- endfor -%}
|
||||
|{{ arguments['items'] }}",
|
||||
),
|
||||
&request,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(rendered, "items=operands;x=2;y=1.0;|operands");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn qwen35_template_renders_prefilled_reasoning_start_when_thinking_enabled() {
|
||||
let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
|
||||
|
||||
@@ -208,11 +208,27 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tojson_preserves_arbitrary_precision_number_spelling() {
|
||||
fn tojson_uses_standard_serde_json_number_spelling() {
|
||||
let payload = serde_json::from_str(r#"{"x":2,"y":1.00}"#).unwrap();
|
||||
let rendered = render("{{ payload|tojson }}", payload);
|
||||
|
||||
assert_eq!(rendered, "{\"x\": 2, \"y\": 1.00}");
|
||||
// TODO: we cannot preserve the original number precision by enabling `serde_json`'s
|
||||
// `arbitrary_precision` feature, otherwise the following test
|
||||
// `serialized_json_numbers_do_not_leak_serde_private_representation` will fail.
|
||||
// See issue: https://github.com/mitsuhiko/minijinja/issues/641
|
||||
assert_eq!(rendered, "{\"x\": 2, \"y\": 1.0}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serialized_json_numbers_do_not_leak_serde_private_representation() {
|
||||
let payload: serde_json::Value = serde_json::from_str(r#"{"x":2,"y":1.00}"#).unwrap();
|
||||
let rendered = render("{{ payload }}", payload);
|
||||
|
||||
// TODO: we cannot preserve the original number precision by enabling `serde_json`'s
|
||||
// `arbitrary_precision` feature, otherwise this will fail.
|
||||
// See issue: https://github.com/mitsuhiko/minijinja/issues/641
|
||||
assert!(!rendered.contains("$serde_json::private::Number"));
|
||||
assert_eq!(rendered, r#"{"x": 2, "y": 1.0}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use indexmap::IndexMap;
|
||||
use minijinja::value::{Enumerator, Object, ObjectExt, ObjectRepr};
|
||||
use minijinja::{Error as TemplateError, ErrorKind as TemplateErrorKind, State};
|
||||
use serde::Serialize;
|
||||
use serde_json::Value as JsonValue;
|
||||
|
||||
/// A wrapper around `minijinja::Value` that can be constructed with `to_template_value` and used
|
||||
/// as a value in the chat template.
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(transparent)]
|
||||
pub(super) struct TemplateValue(minijinja::Value);
|
||||
|
||||
pub(super) fn to_template_value(value: JsonValue) -> TemplateValue {
|
||||
TemplateValue(match value {
|
||||
JsonValue::Array(values) => values
|
||||
.into_iter()
|
||||
.map(to_template_value)
|
||||
.map(|value| value.0)
|
||||
.collect::<minijinja::Value>(),
|
||||
JsonValue::Object(values) => minijinja::Value::from_object(TemplateMap(
|
||||
values
|
||||
.into_iter()
|
||||
.map(|(key, value)| (key, to_template_value(value).0))
|
||||
.collect(),
|
||||
)),
|
||||
// For primitive values, directly convert them to `minijinja::Value` using `from_serialize`.
|
||||
value => minijinja::Value::from_serialize(value),
|
||||
})
|
||||
}
|
||||
|
||||
/// A custom map type that always returns `UnknownMethod` for method calls, so that pycompat can
|
||||
/// always handle dict methods through the unknown-method callback.
|
||||
///
|
||||
/// Use `IndexMap` to preserve the original key order when iterating.
|
||||
///
|
||||
/// MiniJinja's default map can resolve a same-named field before Python dict methods. HF templates
|
||||
/// commonly call `dict.items()`, which would fail if the map had an `items` field.
|
||||
/// See issue: https://github.com/mitsuhiko/minijinja/issues/903
|
||||
#[derive(Debug)]
|
||||
struct TemplateMap(IndexMap<String, minijinja::Value>);
|
||||
|
||||
impl Object for TemplateMap {
|
||||
fn repr(self: &Arc<Self>) -> ObjectRepr {
|
||||
ObjectRepr::Map
|
||||
}
|
||||
|
||||
fn get_value(self: &Arc<Self>, key: &minijinja::Value) -> Option<minijinja::Value> {
|
||||
self.0.get(key.as_str()?).cloned()
|
||||
}
|
||||
|
||||
fn get_value_by_str(self: &Arc<Self>, key: &str) -> Option<minijinja::Value> {
|
||||
self.0.get(key).cloned()
|
||||
}
|
||||
|
||||
fn enumerate(self: &Arc<Self>) -> Enumerator {
|
||||
self.mapped_rev_enumerator(|this| {
|
||||
Box::new(this.0.keys().map(|key| minijinja::Value::from(key.as_str())))
|
||||
})
|
||||
}
|
||||
|
||||
fn enumerator_len(self: &Arc<Self>) -> Option<usize> {
|
||||
Some(self.0.len())
|
||||
}
|
||||
|
||||
fn call_method(
|
||||
self: &Arc<Self>,
|
||||
_state: &State<'_, '_>,
|
||||
_method: &str,
|
||||
_args: &[minijinja::Value],
|
||||
) -> std::result::Result<minijinja::Value, TemplateError> {
|
||||
// Always return `UnknownMethod` for method calls,
|
||||
// so that pycompat can handle dict methods through the unknown-method callback.
|
||||
Err(TemplateError::from(TemplateErrorKind::UnknownMethod))
|
||||
}
|
||||
}
|
||||
@@ -20,7 +20,6 @@ use vllm_chat::{
|
||||
use vllm_text::{DecodedTextEvent, Finished, Prompt};
|
||||
|
||||
/// One model/parser configuration used to run the fixed roundtrip fixtures.
|
||||
#[derive(Clone)]
|
||||
struct RoundtripCase {
|
||||
/// Hugging Face model id resolved through the production backend loader.
|
||||
model_id: &'static str,
|
||||
@@ -32,45 +31,11 @@ struct RoundtripCase {
|
||||
tool_call_parser: ParserSelection,
|
||||
/// Reasoning parser selection used by the output processor.
|
||||
reasoning_parser: ParserSelection,
|
||||
/// How this model's chat template handles thinking mode.
|
||||
thinking_behavior: ThinkingBehavior,
|
||||
/// JSON formatting expected after this model's template has materialized
|
||||
/// tool-call arguments.
|
||||
json_fmt: JsonFmt,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum ThinkingBehavior {
|
||||
/// The chat template accepts explicit thinking on/off kwargs, and uses
|
||||
/// `default` when the request does not specify either kwarg.
|
||||
Toggleable { default: bool },
|
||||
/// The chat template always behaves as `value` for this fixture.
|
||||
Always { value: bool },
|
||||
}
|
||||
|
||||
impl ThinkingBehavior {
|
||||
fn default(self) -> bool {
|
||||
match self {
|
||||
Self::Toggleable { default } => default,
|
||||
Self::Always { value } => value,
|
||||
}
|
||||
}
|
||||
|
||||
fn fixtures(self) -> Vec<Option<bool>> {
|
||||
match self {
|
||||
Self::Toggleable { .. } => vec![
|
||||
Some(true), // explicitly enable thinking
|
||||
Some(false), // explicitly disable thinking
|
||||
None, // use default template behavior
|
||||
],
|
||||
Self::Always { value } => vec![
|
||||
Some(value), // explicitly request the supported thinking behavior
|
||||
None, // use default template behavior
|
||||
],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RoundtripCase {
|
||||
/// Qwen3 XML tool-call format with `qwen3` reasoning tags.
|
||||
fn qwen3() -> Self {
|
||||
@@ -79,7 +44,6 @@ impl RoundtripCase {
|
||||
assistant_stop_suffix: "<|im_end|>\n",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
|
||||
json_fmt: spaced_json_fmt(),
|
||||
}
|
||||
}
|
||||
@@ -91,7 +55,6 @@ impl RoundtripCase {
|
||||
assistant_stop_suffix: "<|im_end|>\n",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
|
||||
json_fmt: compact_json_fmt(),
|
||||
}
|
||||
}
|
||||
@@ -103,7 +66,6 @@ impl RoundtripCase {
|
||||
assistant_stop_suffix: "[e~[\n",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Always { value: true },
|
||||
json_fmt: compact_json_fmt(),
|
||||
}
|
||||
}
|
||||
@@ -115,7 +77,6 @@ impl RoundtripCase {
|
||||
assistant_stop_suffix: "<|end▁of▁sentence|>",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Toggleable { default: false },
|
||||
json_fmt: compact_json_fmt(),
|
||||
}
|
||||
}
|
||||
@@ -127,7 +88,6 @@ impl RoundtripCase {
|
||||
assistant_stop_suffix: "",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
|
||||
json_fmt: compact_json_fmt(),
|
||||
}
|
||||
}
|
||||
@@ -140,7 +100,6 @@ impl RoundtripCase {
|
||||
assistant_stop_suffix: "<|im_end|>",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
|
||||
json_fmt: spaced_json_fmt(),
|
||||
}
|
||||
}
|
||||
@@ -176,44 +135,35 @@ roundtrip_tests! {
|
||||
|
||||
/// Run the fixed reasoning+content fixture for one model/parser case.
|
||||
async fn run_roundtrip_reasoning_and_content(case: RoundtripCase) -> Result<()> {
|
||||
for thinking in case.thinking_behavior.fixtures() {
|
||||
run_roundtrip_reasoning_and_content_inner(case.clone(), thinking).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_roundtrip_reasoning_and_content_inner(
|
||||
case: RoundtripCase,
|
||||
thinking: Option<bool>,
|
||||
) -> Result<()> {
|
||||
let backends = load_roundtrip_backends(&case).await?;
|
||||
let request = roundtrip_request(
|
||||
"roundtrip-reasoning-content",
|
||||
vec![ChatMessage::text(ChatRole::User, "What is 2 + 2?")],
|
||||
Vec::new(),
|
||||
thinking,
|
||||
);
|
||||
let expected_reasoning = "Need compute 2 + 2 directly.";
|
||||
let expected_text = "The answer is 4.";
|
||||
let effective_thinking = thinking.unwrap_or(case.thinking_behavior.default());
|
||||
|
||||
let assistant = {
|
||||
let mut content = Vec::new();
|
||||
if effective_thinking {
|
||||
content.push(AssistantContentBlock::Reasoning {
|
||||
text: expected_reasoning.to_string(),
|
||||
});
|
||||
}
|
||||
content.push(AssistantContentBlock::Text {
|
||||
text: expected_text.to_string(),
|
||||
});
|
||||
AssistantMessage { content }
|
||||
};
|
||||
let result = run_roundtrip(&case, &backends, &request, assistant).await?;
|
||||
let result = run_roundtrip(
|
||||
&case,
|
||||
&backends,
|
||||
&request,
|
||||
AssistantMessage {
|
||||
content: vec![
|
||||
AssistantContentBlock::Reasoning {
|
||||
text: expected_reasoning.to_string(),
|
||||
},
|
||||
AssistantContentBlock::Text {
|
||||
text: expected_text.to_string(),
|
||||
},
|
||||
],
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
|
||||
assert_eq!(
|
||||
result.parsed_message.reasoning().as_deref().map(str::trim),
|
||||
effective_thinking.then_some(expected_reasoning)
|
||||
Some(expected_reasoning)
|
||||
);
|
||||
assert_eq!(result.parsed_message.text().trim(), expected_text);
|
||||
assert_eq!(result.parsed_message.tool_calls().count(), 0);
|
||||
@@ -233,10 +183,9 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
|
||||
"roundtrip-reasoning-tools",
|
||||
vec![ChatMessage::text(
|
||||
ChatRole::User,
|
||||
"Check Shanghai weather and add 1.00 plus 2.",
|
||||
"Check Shanghai weather and add 1.0 plus 2.",
|
||||
)],
|
||||
test_tools(),
|
||||
Some(true), // always enable thinking in this fixture
|
||||
);
|
||||
let expected_reasoning = "Need call the weather and add tools.";
|
||||
let expected_text = "I will call the tools.";
|
||||
@@ -261,9 +210,10 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
|
||||
AssistantContentBlock::ToolCall(AssistantToolCall {
|
||||
id: "functions.add:1".to_string(),
|
||||
name: "add".to_string(),
|
||||
// Intentionally use a non-lexical order of keys and a different number
|
||||
// formatting style to verify text-level fidelity of the roundtrip.
|
||||
arguments: r#"{"y":1.00,"x":2}"#.to_string(),
|
||||
// Intentionally use a non-lexical order of keys to verify text-level
|
||||
// fidelity of the roundtrip where JSON formatting remains stable. The
|
||||
// `items` key also exercises templates that call `arguments.items()`.
|
||||
arguments: r#"{"y":1.0,"x":2,"items":["left","right"]}"#.to_string(),
|
||||
}),
|
||||
],
|
||||
},
|
||||
@@ -291,7 +241,7 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
|
||||
assert_eq!(tool_calls[1].name, "add");
|
||||
assert_eq!(
|
||||
tool_calls[1].arguments,
|
||||
expected_arguments(&case, r#"{"y": 1.00, "x": 2}"#)?,
|
||||
expected_arguments(&case, r#"{"y": 1.0, "x": 2, "items": ["left", "right"]}"#)?,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
@@ -538,7 +488,6 @@ fn roundtrip_request(
|
||||
request_id: impl Into<String>,
|
||||
messages: Vec<ChatMessage>,
|
||||
tools: Vec<ChatTool>,
|
||||
thinking: Option<bool>,
|
||||
) -> ChatRequest {
|
||||
let mut request = ChatRequest {
|
||||
request_id: request_id.into(),
|
||||
@@ -552,12 +501,10 @@ fn roundtrip_request(
|
||||
..ChatRequest::for_test()
|
||||
};
|
||||
|
||||
// Explicitly enable or disable thinking so that rendering and parsing the reasoning block is
|
||||
// exercised or skipped in the roundtrip. If unspecified, use the default template behavior.
|
||||
if let Some(thinking) = thinking {
|
||||
for key in ["thinking", "enable_thinking"] {
|
||||
request.chat_options.template_kwargs.insert(key.to_string(), thinking.into());
|
||||
}
|
||||
// Enable thinking for some models so that rendering and parsing the reasoning block is
|
||||
// exercised in the roundtrip.
|
||||
for key in ["thinking", "enable_thinking"] {
|
||||
request.chat_options.template_kwargs.insert(key.to_string(), true.into());
|
||||
}
|
||||
|
||||
request
|
||||
@@ -585,9 +532,13 @@ fn test_tools() -> Vec<ChatTool> {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"y": { "type": "number" },
|
||||
"x": { "type": "number" }
|
||||
"x": { "type": "number" },
|
||||
"items": {
|
||||
"type": "array",
|
||||
"items": { "type": "string" }
|
||||
}
|
||||
},
|
||||
"required": ["y", "x"]
|
||||
"required": ["y", "x", "items"]
|
||||
}),
|
||||
strict: None,
|
||||
},
|
||||
|
||||
@@ -165,15 +165,6 @@ pub struct SharedRuntimeArgs {
|
||||
#[serde(default)]
|
||||
pub enable_log_requests: bool,
|
||||
|
||||
/// If specified, API server will add X-Request-Id header to responses.
|
||||
#[arg(
|
||||
long,
|
||||
default_missing_value = "true",
|
||||
num_args = 0..=1
|
||||
)]
|
||||
#[serde(default)]
|
||||
pub enable_request_id_headers: bool,
|
||||
|
||||
/// Disable periodic logging of engine statistics (throughput, queue depth,
|
||||
/// cache usage).
|
||||
#[arg(long)]
|
||||
@@ -247,7 +238,6 @@ impl SharedRuntimeArgs {
|
||||
default_chat_template_kwargs: self.default_chat_template_kwargs,
|
||||
chat_template_content_format: self.chat_template_content_format,
|
||||
enable_log_requests: self.enable_log_requests,
|
||||
enable_request_id_headers: self.enable_request_id_headers,
|
||||
disable_log_stats: self.disable_log_stats,
|
||||
grpc_port: self.grpc_port,
|
||||
shutdown_timeout,
|
||||
@@ -288,7 +278,6 @@ impl SharedRuntimeArgs {
|
||||
default_chat_template_kwargs: self.default_chat_template_kwargs,
|
||||
chat_template_content_format: self.chat_template_content_format,
|
||||
enable_log_requests: self.enable_log_requests,
|
||||
enable_request_id_headers: self.enable_request_id_headers,
|
||||
disable_log_stats: self.disable_log_stats,
|
||||
grpc_port: self.grpc_port,
|
||||
shutdown_timeout,
|
||||
|
||||
@@ -43,7 +43,6 @@ fn serve_args_forward_python_flags_with_separator() {
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
disable_log_stats: false,
|
||||
served_model_name: [],
|
||||
},
|
||||
@@ -117,46 +116,6 @@ fn serve_args_accept_explicit_deepseek_v32_renderer() {
|
||||
assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serve_passes_enable_request_id_headers_into_config() {
|
||||
let cli = Cli::try_parse_from([
|
||||
"vllm-rs",
|
||||
"serve",
|
||||
"Qwen/Qwen3-0.6B",
|
||||
"--enable-request-id-headers",
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
let Command::Serve(args) = cli.command else {
|
||||
panic!("expected serve args");
|
||||
};
|
||||
let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
|
||||
assert!(config.enable_request_id_headers);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn frontend_args_json_passes_enable_request_id_headers_into_config() {
|
||||
let cli = Cli::try_parse_from([
|
||||
"vllm-rs",
|
||||
"frontend",
|
||||
"--listen-fd",
|
||||
"3",
|
||||
"--input-address",
|
||||
"ipc:///tmp/input.sock",
|
||||
"--output-address",
|
||||
"ipc:///tmp/output.sock",
|
||||
"--args-json",
|
||||
r#"{"model_tag":"Qwen/Qwen3-0.6B","enable_request_id_headers":true}"#,
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
let Command::Frontend(args) = cli.command else {
|
||||
panic!("expected frontend args");
|
||||
};
|
||||
let config = args.into_config();
|
||||
assert!(config.enable_request_id_headers);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serve_args_reject_unknown_renderer_value() {
|
||||
let error = Cli::try_parse_from([
|
||||
@@ -259,7 +218,6 @@ fn frontend_args_accept_json() {
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
disable_log_stats: false,
|
||||
served_model_name: [],
|
||||
},
|
||||
@@ -658,7 +616,6 @@ fn serve_args_accept_handshake_aliases() {
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
disable_log_stats: false,
|
||||
served_model_name: [],
|
||||
},
|
||||
@@ -776,7 +733,6 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
disable_log_stats: false,
|
||||
grpc_port: None,
|
||||
shutdown_timeout: 0ns,
|
||||
@@ -839,7 +795,6 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
disable_log_stats: false,
|
||||
grpc_port: None,
|
||||
shutdown_timeout: 0ns,
|
||||
@@ -917,7 +872,6 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
disable_log_stats: false,
|
||||
grpc_port: None,
|
||||
shutdown_timeout: 0ns,
|
||||
|
||||
@@ -620,6 +620,15 @@ pub struct ServerUnsupportedArgs {
|
||||
#[arg(long)]
|
||||
pub middleware: Option<Unsupported>,
|
||||
|
||||
/// If specified, API server will add X-Request-Id header to responses.
|
||||
#[arg(
|
||||
long,
|
||||
visible_alias = "no-enable-request-id-headers",
|
||||
default_missing_value = "true",
|
||||
num_args = 0..=1
|
||||
)]
|
||||
pub enable_request_id_headers: Option<Unsupported>,
|
||||
|
||||
/// Disable FastAPI's OpenAPI schema, Swagger UI, and ReDoc endpoint.
|
||||
#[arg(
|
||||
long,
|
||||
|
||||
@@ -19,6 +19,7 @@ mod deepseek_r1;
|
||||
mod delimited;
|
||||
mod gemma4;
|
||||
mod kimi;
|
||||
mod minimax_m3;
|
||||
mod qwen3;
|
||||
|
||||
use thiserror::Error;
|
||||
@@ -29,6 +30,7 @@ pub use self::deepseek_r1::DeepSeekR1ReasoningParser;
|
||||
pub(crate) use self::delimited::DelimitedReasoningParser;
|
||||
pub use self::gemma4::Gemma4ReasoningParser;
|
||||
pub use self::kimi::KimiReasoningParser;
|
||||
pub use self::minimax_m3::MiniMaxM3ReasoningParser;
|
||||
pub use self::qwen3::Qwen3ReasoningParser;
|
||||
|
||||
/// DeepSeek V3 currently shares the standard `<think>...</think>` parser.
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
use vllm_tokenizer::DynTokenizer;
|
||||
|
||||
use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
|
||||
|
||||
/// Reasoning parser for MiniMax M3 style outputs.
|
||||
///
|
||||
/// MiniMax M3 uses `<mm:think>...</mm:think>` delimiters. Its chat template may
|
||||
/// prefill either delimiter depending on the requested thinking mode, so the
|
||||
/// shared delimited parser derives the starting state from the rendered prompt.
|
||||
pub struct MiniMaxM3ReasoningParser {
|
||||
inner: DelimitedReasoningParser,
|
||||
}
|
||||
|
||||
impl MiniMaxM3ReasoningParser {
|
||||
/// Create a MiniMax M3 parser backed by the shared delimited state machine.
|
||||
pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
|
||||
Ok(Self {
|
||||
inner: DelimitedReasoningParser::new(tokenizer, "<mm:think>", "</mm:think>", false)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ReasoningParser for MiniMaxM3ReasoningParser {
|
||||
fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
|
||||
where
|
||||
Self: Sized + 'static,
|
||||
{
|
||||
Ok(Box::new(Self::new(tokenizer)?))
|
||||
}
|
||||
|
||||
fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
|
||||
self.inner.initialize(prompt_token_ids);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
|
||||
Ok(self.inner.push(delta))
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<ReasoningDelta> {
|
||||
Ok(self.inner.finish())
|
||||
}
|
||||
}
|
||||
@@ -3,7 +3,8 @@ use std::sync::Arc;
|
||||
use vllm_tokenizer::Tokenizer;
|
||||
|
||||
use super::{
|
||||
DeepSeekR1ReasoningParser, DelimitedReasoningParser, Qwen3ReasoningParser, ReasoningParser,
|
||||
DeepSeekR1ReasoningParser, DelimitedReasoningParser, MiniMaxM3ReasoningParser,
|
||||
Qwen3ReasoningParser, ReasoningParser,
|
||||
};
|
||||
|
||||
struct FakeTokenizer;
|
||||
@@ -32,6 +33,8 @@ impl Tokenizer for FakeTokenizer {
|
||||
"<|END_THINKING|>" => Some(4),
|
||||
"◁think▷" => Some(5),
|
||||
"◁/think▷" => Some(6),
|
||||
"<mm:think>" => Some(8),
|
||||
"</mm:think>" => Some(9),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
@@ -159,3 +162,35 @@ fn deepseek_r1_stops_scanning_at_last_special_token() {
|
||||
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
|
||||
assert_eq!(delta.content.as_deref(), Some("answer"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn minimax_m3_handles_explicit_think_delimiters() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
|
||||
|
||||
let delta = parser.push("<mm:think>reason</mm:think>answer").unwrap();
|
||||
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
|
||||
assert_eq!(delta.content.as_deref(), Some("answer"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn minimax_m3_uses_prompt_prefilled_start_marker() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
|
||||
parser.initialize(&[8]).unwrap();
|
||||
|
||||
let delta = parser.push("reason</mm:think>answer").unwrap();
|
||||
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
|
||||
assert_eq!(delta.content.as_deref(), Some("answer"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn minimax_m3_uses_prompt_prefilled_end_marker() {
|
||||
let tokenizer = Arc::new(FakeTokenizer);
|
||||
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
|
||||
parser.initialize(&[9]).unwrap();
|
||||
|
||||
let delta = parser.push("answer").unwrap();
|
||||
assert_eq!(delta.reasoning, None);
|
||||
assert_eq!(delta.content.as_deref(), Some("answer"));
|
||||
}
|
||||
|
||||
@@ -68,7 +68,6 @@ async fn main() -> Result<()> {
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: ChatTemplateContentFormatOption::Auto,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
disable_log_stats: false,
|
||||
grpc_port: None,
|
||||
shutdown_timeout: Duration::ZERO,
|
||||
|
||||
@@ -61,8 +61,6 @@ pub struct Config {
|
||||
pub chat_template_content_format: ChatTemplateContentFormatOption,
|
||||
/// Log a summary line for each completed request.
|
||||
pub enable_log_requests: bool,
|
||||
/// When `true`, set `X-Request-Id` on every HTTP response.
|
||||
pub enable_request_id_headers: bool,
|
||||
/// When `true`, suppress periodic stats logging (throughput, queue depth,
|
||||
/// cache usage).
|
||||
pub disable_log_stats: bool,
|
||||
|
||||
@@ -85,9 +85,7 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
|
||||
};
|
||||
|
||||
Ok(Arc::new(
|
||||
AppState::new(served_model_names, chat)
|
||||
.with_log_requests(config.enable_log_requests)
|
||||
.with_request_id_headers(config.enable_request_id_headers),
|
||||
AppState::new(served_model_names, chat).with_log_requests(config.enable_log_requests),
|
||||
))
|
||||
}
|
||||
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
mod load;
|
||||
mod metrics;
|
||||
mod request_id;
|
||||
|
||||
pub use load::track_server_load;
|
||||
pub use metrics::track_http_metrics;
|
||||
pub use request_id::set_request_id_header;
|
||||
|
||||
@@ -1,24 +0,0 @@
|
||||
use axum::extract::Request;
|
||||
use axum::http::HeaderValue;
|
||||
use axum::http::header::HeaderName;
|
||||
use axum::middleware::Next;
|
||||
use axum::response::Response;
|
||||
use uuid::Uuid;
|
||||
|
||||
const X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");
|
||||
|
||||
/// Echo the request's `X-Request-Id` on the response, or generate a fresh
|
||||
/// `uuid4` hex if the request did not provide one.
|
||||
///
|
||||
/// Original Python:
|
||||
/// `vllm.entrypoints.openai.server_utils.XRequestIdMiddleware`.
|
||||
pub async fn set_request_id_header(req: Request, next: Next) -> Response {
|
||||
let incoming = req.headers().get(&X_REQUEST_ID).cloned();
|
||||
let mut response = next.run(req).await;
|
||||
let value = incoming.unwrap_or_else(|| {
|
||||
HeaderValue::from_str(&Uuid::new_v4().simple().to_string())
|
||||
.expect("uuid hex is valid header value")
|
||||
});
|
||||
response.headers_mut().insert(X_REQUEST_ID, value);
|
||||
response
|
||||
}
|
||||
@@ -56,18 +56,11 @@ fn build_router_with_dev_mode(state: Arc<AppState>, dev_mode_enabled: bool) -> R
|
||||
.route("/is_sleeping", get(sleep::is_sleeping))
|
||||
}
|
||||
|
||||
let enable_request_id_headers = state.enable_request_id_headers;
|
||||
let mut router = router
|
||||
router
|
||||
.with_state(state.clone())
|
||||
.layer(from_fn_with_state(state, middleware::track_server_load))
|
||||
.layer(from_fn(middleware::track_http_metrics))
|
||||
.layer(TraceLayer::new_for_http());
|
||||
|
||||
if enable_request_id_headers {
|
||||
router = router.layer(from_fn(middleware::set_request_id_header));
|
||||
}
|
||||
|
||||
router
|
||||
.layer(TraceLayer::new_for_http())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -3,33 +3,23 @@ mod types;
|
||||
mod validate;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::convert::Infallible;
|
||||
use std::result::Result;
|
||||
use std::sync::Arc;
|
||||
|
||||
use asynk_strim_attr::{TryYielder, try_stream};
|
||||
use axum::Json;
|
||||
use axum::extract::State;
|
||||
use axum::http::HeaderMap;
|
||||
use axum::response::sse::{Event, Sse};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use futures::{Stream, StreamExt as _, pin_mut};
|
||||
use thiserror_ext::AsReport as _;
|
||||
use tracing::{error, info, trace};
|
||||
use tracing::info;
|
||||
use tracing_futures::Instrument as _;
|
||||
use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs};
|
||||
use vllm_llm::{
|
||||
CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _,
|
||||
};
|
||||
use vllm_llm::{CollectedGenerateOutput, GenerateOutputStreamExt as _};
|
||||
|
||||
use self::convert::prepare_generate_request;
|
||||
use self::types::{
|
||||
GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice,
|
||||
GenerateResponseStreamChoice, GenerateStreamResponse,
|
||||
};
|
||||
use crate::error::{ApiError, bail_server_error, server_error};
|
||||
use self::types::{GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice};
|
||||
use crate::error::{ApiError, server_error};
|
||||
use crate::routes::openai::utils::logprobs::clamp_logprob;
|
||||
use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb, Usage};
|
||||
use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb};
|
||||
use crate::routes::openai::utils::validated_json::ValidatedJson;
|
||||
use crate::state::AppState;
|
||||
use crate::utils::resolve_request_context;
|
||||
@@ -56,7 +46,6 @@ pub async fn generate(
|
||||
let log_request = state.enable_log_requests;
|
||||
let include_logprobs = prepared.include_logprobs;
|
||||
let include_prompt_logprobs = prepared.include_prompt_logprobs;
|
||||
let stream = prepared.stream;
|
||||
|
||||
let raw_stream = match state
|
||||
.chat
|
||||
@@ -75,20 +64,6 @@ pub async fn generate(
|
||||
}
|
||||
};
|
||||
|
||||
if stream {
|
||||
let chunk_stream = generate_chunk_stream(
|
||||
raw_stream,
|
||||
prepared.request_id,
|
||||
log_request,
|
||||
prepared.include_usage,
|
||||
prepared.include_continuous_usage,
|
||||
include_logprobs,
|
||||
);
|
||||
let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span);
|
||||
|
||||
return Sse::new(sse_stream).into_response();
|
||||
}
|
||||
|
||||
let collected = match raw_stream.collect_output().instrument(request_span.clone()).await {
|
||||
Ok(collected) => collected,
|
||||
Err(error) => {
|
||||
@@ -123,102 +98,6 @@ pub async fn generate(
|
||||
Json(response).into_response()
|
||||
}
|
||||
|
||||
#[try_stream]
|
||||
async fn generate_chunk_stream(
|
||||
stream: impl Stream<Item = vllm_llm::Result<GenerateOutput>>,
|
||||
request_id: String,
|
||||
log_request: bool,
|
||||
include_usage: bool,
|
||||
include_continuous_usage: bool,
|
||||
include_logprobs: bool,
|
||||
mut y: TryYielder<GenerateStreamResponse, ApiError>,
|
||||
) -> Result<(), ApiError> {
|
||||
pin_mut!(stream);
|
||||
let mut prompt_tokens: Option<u32> = None;
|
||||
let mut output_tokens = 0_u32;
|
||||
|
||||
while let Some(next) = stream.next().await {
|
||||
match next {
|
||||
Ok(output) => {
|
||||
if prompt_tokens.is_none() {
|
||||
prompt_tokens =
|
||||
output.prompt_info.as_ref().map(|info| info.prompt_token_ids.len() as u32);
|
||||
}
|
||||
let usage_prompt_tokens = prompt_tokens.unwrap_or_default();
|
||||
|
||||
let token_ids = output.token_ids;
|
||||
output_tokens = output_tokens.saturating_add(token_ids.len() as u32);
|
||||
let finish_reason = output.finish_reason;
|
||||
|
||||
if matches!(finish_reason.as_ref(), Some(FinishReason::Error)) {
|
||||
bail_server_error!("Internal server error");
|
||||
}
|
||||
|
||||
if let Some(finish_reason) = finish_reason.as_ref()
|
||||
&& log_request
|
||||
{
|
||||
info!(
|
||||
stream = true,
|
||||
prompt_tokens = usage_prompt_tokens,
|
||||
output_tokens,
|
||||
finish_reason = finish_reason.as_str(),
|
||||
"generate finished"
|
||||
);
|
||||
}
|
||||
|
||||
if token_ids.is_empty() && finish_reason.is_none() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let logprobs = if include_logprobs && !token_ids.is_empty() {
|
||||
let logprobs = output.logprobs.as_ref().ok_or_else(|| {
|
||||
server_error!(
|
||||
"raw generate stream requested logprobs but generation returned none"
|
||||
)
|
||||
})?;
|
||||
Some(raw_logprobs_to_openai_chat(logprobs)?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
y.yield_ok(GenerateStreamResponse {
|
||||
request_id: request_id.clone(),
|
||||
choices: vec![GenerateResponseStreamChoice {
|
||||
index: 0,
|
||||
logprobs,
|
||||
finish_reason: finish_reason.map(|reason| reason.as_str().to_string()),
|
||||
token_ids,
|
||||
}],
|
||||
usage: include_continuous_usage
|
||||
.then(|| Usage::from_counts(usage_prompt_tokens, output_tokens)),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
Err(error) => {
|
||||
error!(
|
||||
error = %error.as_report(),
|
||||
"raw generate stream failed"
|
||||
);
|
||||
bail_server_error!("{}", error.to_report_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if include_usage {
|
||||
y.yield_ok(GenerateStreamResponse {
|
||||
request_id,
|
||||
choices: Vec::new(),
|
||||
usage: Some(Usage::from_counts(
|
||||
prompt_tokens.unwrap_or_default(),
|
||||
output_tokens,
|
||||
)),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn collect_generate(
|
||||
collected: CollectedGenerateOutput,
|
||||
request_id: String,
|
||||
@@ -334,94 +213,3 @@ fn position_to_logprob_map(position: &PositionLogprobs) -> HashMap<u32, Generate
|
||||
fn format_token_id(token_id: u32) -> String {
|
||||
format!("token_id:{token_id}")
|
||||
}
|
||||
|
||||
/// Convert one raw-generate chunk stream into SSE events.
|
||||
#[try_stream]
|
||||
async fn generate_sse_stream(
|
||||
stream: impl Stream<Item = Result<GenerateStreamResponse, ApiError>>,
|
||||
mut y: TryYielder<Event, Infallible>,
|
||||
) -> Result<(), Infallible> {
|
||||
pin_mut!(stream);
|
||||
|
||||
while let Some(next) = stream.next().await {
|
||||
match next {
|
||||
Ok(chunk) => y.yield_ok(to_sse_event(&chunk)).await,
|
||||
Err(error) => {
|
||||
y.yield_ok(to_error_sse_event(&error)).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
y.yield_ok(done_sse_event()).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn to_sse_event(chunk: &GenerateStreamResponse) -> Event {
|
||||
let payload = serde_json::to_string(chunk).expect("generate chunk must serialize to JSON");
|
||||
trace!(payload, "generate emitting chunk");
|
||||
Event::default().data(payload)
|
||||
}
|
||||
|
||||
fn to_error_sse_event(error: &ApiError) -> Event {
|
||||
let payload = serde_json::to_string(&error.to_error_response())
|
||||
.expect("ErrorResponse must serialize to JSON");
|
||||
trace!(payload, "generate emitting error");
|
||||
Event::default().data(payload)
|
||||
}
|
||||
|
||||
fn done_sse_event() -> Event {
|
||||
trace!("generate emitting done");
|
||||
Event::default().data("[DONE]")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use futures::{TryStreamExt as _, stream};
|
||||
use vllm_llm::GeneratePromptInfo;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn generate_chunk_stream_captures_late_prompt_info() {
|
||||
let stream = stream::iter(vec![
|
||||
Ok(GenerateOutput {
|
||||
request_id: String::new(),
|
||||
prompt_info: None,
|
||||
token_ids: Vec::new(),
|
||||
logprobs: None,
|
||||
finish_reason: None,
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
Ok(GenerateOutput {
|
||||
request_id: String::new(),
|
||||
prompt_info: Some(GeneratePromptInfo {
|
||||
prompt_token_ids: Arc::from([11_u32, 22_u32]),
|
||||
prompt_logprobs: None,
|
||||
}),
|
||||
token_ids: vec![33],
|
||||
logprobs: None,
|
||||
finish_reason: Some(FinishReason::stop_eos()),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
let chunks: Vec<_> =
|
||||
generate_chunk_stream(stream, "raw-stream".to_string(), false, true, true, false)
|
||||
.try_collect()
|
||||
.await
|
||||
.expect("collect chunks");
|
||||
|
||||
assert_eq!(chunks.len(), 2);
|
||||
assert_eq!(
|
||||
chunks[0].usage.as_ref().expect("chunk usage").prompt_tokens,
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
chunks[1].usage.as_ref().expect("final usage").prompt_tokens,
|
||||
2
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,9 +10,6 @@ use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params};
|
||||
pub struct PreparedRequest {
|
||||
pub request_id: String,
|
||||
pub text_request: TextRequest,
|
||||
pub stream: bool,
|
||||
pub include_usage: bool,
|
||||
pub include_continuous_usage: bool,
|
||||
pub include_logprobs: bool,
|
||||
pub include_prompt_logprobs: bool,
|
||||
}
|
||||
@@ -26,18 +23,6 @@ pub fn prepare_generate_request(
|
||||
) -> Result<PreparedRequest, ApiError> {
|
||||
validate::validate_request_compat(&request, served_model_names)?;
|
||||
|
||||
let stream = request.stream;
|
||||
let include_usage = request
|
||||
.stream_options
|
||||
.as_ref()
|
||||
.and_then(|options| options.include_usage)
|
||||
.unwrap_or(false);
|
||||
let include_continuous_usage = include_usage
|
||||
&& request
|
||||
.stream_options
|
||||
.as_ref()
|
||||
.and_then(|options| options.continuous_usage_stats)
|
||||
.unwrap_or(false);
|
||||
let include_logprobs = request.sampling_params.logprobs.is_some();
|
||||
let include_prompt_logprobs = request.sampling_params.prompt_logprobs.is_some();
|
||||
let mut sampling_params = request.sampling_params;
|
||||
@@ -62,9 +47,6 @@ pub fn prepare_generate_request(
|
||||
Ok(PreparedRequest {
|
||||
request_id: ctx.request_id,
|
||||
text_request,
|
||||
stream,
|
||||
include_usage,
|
||||
include_continuous_usage,
|
||||
include_logprobs,
|
||||
include_prompt_logprobs,
|
||||
})
|
||||
@@ -127,28 +109,4 @@ mod tests {
|
||||
Some(json!({"connector": "x"}))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_generate_request_gates_continuous_usage_on_include_usage() {
|
||||
let request: GenerateRequest = serde_json::from_value(json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"token_ids": [11, 22],
|
||||
"stream": true,
|
||||
"stream_options": {
|
||||
"continuous_usage_stats": true
|
||||
},
|
||||
"sampling_params": {}
|
||||
}))
|
||||
.expect("parse request");
|
||||
|
||||
let prepared = prepare_generate_request(
|
||||
request,
|
||||
&["Qwen/Qwen1.5-0.5B-Chat".to_string()],
|
||||
ResolvedRequestContext::default(),
|
||||
)
|
||||
.expect("prepare");
|
||||
|
||||
assert!(!prepared.include_usage);
|
||||
assert!(!prepared.include_continuous_usage);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,7 +5,7 @@ use serde_json::{Map, Value};
|
||||
use validator::Validate;
|
||||
use vllm_text::SamplingParams;
|
||||
|
||||
use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable, StreamOptions, Usage};
|
||||
use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable};
|
||||
|
||||
/// vLLM-compatible request type for the token-in/token-out generate API.
|
||||
#[serde_with::skip_serializing_none]
|
||||
@@ -17,7 +17,6 @@ pub struct GenerateRequest {
|
||||
pub sampling_params: SamplingParams,
|
||||
#[serde(default)]
|
||||
pub stream: bool,
|
||||
pub stream_options: Option<StreamOptions>,
|
||||
pub cache_salt: Option<String>,
|
||||
#[serde(default)]
|
||||
pub priority: i32,
|
||||
@@ -38,25 +37,6 @@ pub(super) struct GenerateResponseChoice {
|
||||
pub token_ids: Vec<u32>,
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `GenerateResponseStreamChoice` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub(super) struct GenerateResponseStreamChoice {
|
||||
pub index: u32,
|
||||
pub logprobs: Option<ChatLogProbs>,
|
||||
pub finish_reason: Option<String>,
|
||||
pub token_ids: Vec<u32>,
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `GenerateStreamResponse` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub(super) struct GenerateStreamResponse {
|
||||
pub request_id: String,
|
||||
pub choices: Vec<GenerateResponseStreamChoice>,
|
||||
pub usage: Option<Usage>,
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `GenerateResponse` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
|
||||
@@ -13,11 +13,8 @@ pub(super) fn validate_request_compat(
|
||||
return Err(ApiError::model_not_found(model.clone()));
|
||||
}
|
||||
|
||||
if request.stream_options.is_some() && !request.stream {
|
||||
bail_invalid_request!(
|
||||
param = "stream_options",
|
||||
"stream_options are only supported when stream=true."
|
||||
);
|
||||
if request.stream {
|
||||
bail_invalid_request!(param = "stream", "stream=true is not supported.");
|
||||
}
|
||||
|
||||
if request.token_ids.is_empty() {
|
||||
@@ -68,24 +65,11 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_request_compat_accepts_streaming() {
|
||||
fn validate_request_compat_rejects_streaming() {
|
||||
let request = GenerateRequest {
|
||||
stream: true,
|
||||
..base_request()
|
||||
};
|
||||
assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_request_compat_rejects_stream_options_without_streaming() {
|
||||
let request: GenerateRequest = serde_json::from_value(json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"token_ids": [11, 22],
|
||||
"stream": false,
|
||||
"stream_options": {"include_usage": true},
|
||||
"sampling_params": {}
|
||||
}))
|
||||
.expect("parse request");
|
||||
assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
|
||||
}
|
||||
|
||||
|
||||
@@ -746,20 +746,6 @@ async fn test_app() -> axum::Router {
|
||||
)))
|
||||
}
|
||||
|
||||
async fn test_app_with_request_id_headers() -> (axum::Router, MockEngineTask) {
|
||||
let (chat, engine_task) = test_models_with_engine_outputs_and_backend(
|
||||
b"engine-openai-request-id",
|
||||
default_stream_output_specs(),
|
||||
Arc::new(FakeChatBackend::new()),
|
||||
)
|
||||
.await;
|
||||
let app = build_router(Arc::new(
|
||||
AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat)
|
||||
.with_request_id_headers(true),
|
||||
));
|
||||
(app, engine_task)
|
||||
}
|
||||
|
||||
async fn test_health_app_with_engine_script<F>(
|
||||
script: F,
|
||||
) -> (axum::Router, Arc<AppState>, MockEngineTask)
|
||||
@@ -961,18 +947,6 @@ async fn health_status(app: &axum::Router) -> (StatusCode, Bytes) {
|
||||
(status, body)
|
||||
}
|
||||
|
||||
async fn health_response(app: &axum::Router, request_id: Option<&str>) -> axum::response::Response {
|
||||
let mut builder = Request::builder().method("GET").uri("/health");
|
||||
if let Some(request_id) = request_id {
|
||||
builder = builder.header("X-Request-Id", request_id);
|
||||
}
|
||||
|
||||
app.clone()
|
||||
.call(builder.body(Body::empty()).expect("build request"))
|
||||
.await
|
||||
.expect("call app")
|
||||
}
|
||||
|
||||
fn metric_value(rendered: &str, metric: &str, labels: Option<&str>) -> Option<f64> {
|
||||
rendered.lines().find_map(|line| {
|
||||
let rest = line.strip_prefix(metric)?;
|
||||
@@ -1020,43 +994,6 @@ async fn list_models_returns_configured_model() {
|
||||
assert_eq!(json["data"][0]["id"], "Qwen/Qwen1.5-0.5B-Chat");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn request_id_header_is_absent_by_default() {
|
||||
let app = test_app().await;
|
||||
let response = health_response(&app, None).await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert!(!response.headers().contains_key("x-request-id"));
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn request_id_header_generates_uuid_hex_when_enabled() {
|
||||
let (app, _engine_task) = test_app_with_request_id_headers().await;
|
||||
let response = health_response(&app, None).await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let request_id = response
|
||||
.headers()
|
||||
.get("x-request-id")
|
||||
.expect("x-request-id header")
|
||||
.to_str()
|
||||
.expect("header is ascii");
|
||||
assert_eq!(request_id.len(), 32);
|
||||
assert!(request_id.chars().all(|ch| ch.is_ascii_hexdigit()));
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn request_id_header_echoes_incoming_header_when_enabled() {
|
||||
let (app, _engine_task) = test_app_with_request_id_headers().await;
|
||||
let response = health_response(&app, Some("req-123")).await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(response.headers().get("x-request-id").unwrap(), "req-123");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn version_returns_engine_vllm_version() {
|
||||
@@ -2488,72 +2425,8 @@ async fn non_stream_raw_generate_returns_token_output_envelope() {
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn stream_raw_generate_returns_sse_chunks_and_usage() {
|
||||
let ipc = IpcNamespace::new().expect("create ipc namespace");
|
||||
let handshake_address = ipc.handshake_endpoint();
|
||||
let engine_id = b"engine-raw-generate-stream".to_vec();
|
||||
|
||||
let engine_task = MockEngineTask::new(spawn_mock_engine_task(
|
||||
handshake_address.clone(),
|
||||
engine_id.clone(),
|
||||
|dealer, push| {
|
||||
boxed_test_future(async move {
|
||||
let add = recv_engine_message(dealer).await;
|
||||
let request: EngineCoreRequest =
|
||||
rmp_serde::from_slice(&add[1]).expect("decode request");
|
||||
assert_eq!(request.prompt_token_ids.as_deref(), Some(&[11, 22][..]));
|
||||
assert_eq!(request.external_req_id.as_deref(), Some("raw-stream"));
|
||||
|
||||
send_outputs(
|
||||
push,
|
||||
EngineCoreOutputs {
|
||||
engine_index: 0,
|
||||
outputs: vec![
|
||||
request_output_with_logprobs(
|
||||
&request.request_id,
|
||||
vec![33],
|
||||
None,
|
||||
None,
|
||||
Some(sample_logprobs_for_token(33, 34)),
|
||||
None,
|
||||
),
|
||||
request_output_with_logprobs(
|
||||
&request.request_id,
|
||||
vec![44],
|
||||
Some(EngineCoreFinishReason::Stop),
|
||||
None,
|
||||
Some(sample_logprobs_for_token(44, 45)),
|
||||
None,
|
||||
),
|
||||
],
|
||||
scheduler_stats: None,
|
||||
timestamp: 0.0,
|
||||
utility_output: None,
|
||||
finished_requests: None,
|
||||
wave_complete: None,
|
||||
start_wave: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
})
|
||||
},
|
||||
));
|
||||
|
||||
let client = EngineCoreClient::connect(
|
||||
EngineCoreClientConfig::new_single(handshake_address)
|
||||
.with_model_name("test-model")
|
||||
.with_local_input_output_addresses(
|
||||
Some(ipc.input_endpoint()),
|
||||
Some(ipc.output_endpoint()),
|
||||
),
|
||||
)
|
||||
.await
|
||||
.expect("connect client");
|
||||
let chat = ChatLlm::from_shared_backend(Llm::new(client), Arc::new(FakeChatBackend::new()));
|
||||
let mut app = build_router(Arc::new(AppState::new(
|
||||
vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
|
||||
chat,
|
||||
)));
|
||||
async fn raw_generate_rejects_streaming() {
|
||||
let mut app = test_app().await;
|
||||
|
||||
let response = app
|
||||
.call(
|
||||
@@ -2564,17 +2437,9 @@ async fn stream_raw_generate_returns_sse_chunks_and_usage() {
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"request_id": "raw-stream",
|
||||
"token_ids": [11, 22],
|
||||
"stream": true,
|
||||
"stream_options": {
|
||||
"include_usage": true,
|
||||
"continuous_usage_stats": true
|
||||
},
|
||||
"sampling_params": {
|
||||
"max_tokens": 2,
|
||||
"logprobs": 1
|
||||
}
|
||||
"sampling_params": {}
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
@@ -2583,196 +2448,10 @@ async fn stream_raw_generate_returns_sse_chunks_and_usage() {
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.headers().get("content-type").and_then(|value| value.to_str().ok()),
|
||||
Some("text/event-stream")
|
||||
);
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
|
||||
engine_task.await.expect("mock engine task");
|
||||
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
|
||||
let payloads = sse_data_payloads(&text);
|
||||
assert_eq!(payloads.len(), 4, "{text}");
|
||||
|
||||
let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
|
||||
assert_eq!(first["request_id"], "raw-stream");
|
||||
assert_eq!(first["choices"][0]["index"], 0);
|
||||
assert_eq!(first["choices"][0]["token_ids"], json!([33]));
|
||||
assert_eq!(
|
||||
first["choices"][0]["logprobs"]["content"][0]["token"],
|
||||
"token_id:33"
|
||||
);
|
||||
assert_eq!(first["usage"]["prompt_tokens"], 2);
|
||||
assert_eq!(first["usage"]["completion_tokens"], 1);
|
||||
|
||||
let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
|
||||
assert_eq!(second["choices"][0]["token_ids"], json!([44]));
|
||||
assert_eq!(second["choices"][0]["finish_reason"], "stop");
|
||||
assert_eq!(second["usage"]["completion_tokens"], 2);
|
||||
|
||||
let usage: serde_json::Value = serde_json::from_str(payloads[2]).expect("usage chunk json");
|
||||
assert_eq!(usage["choices"], json!([]));
|
||||
assert_eq!(usage["usage"]["prompt_tokens"], 2);
|
||||
assert_eq!(usage["usage"]["completion_tokens"], 2);
|
||||
assert_eq!(usage["usage"]["total_tokens"], 4);
|
||||
assert_eq!(payloads[3], "[DONE]");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn stream_raw_generate_emits_final_usage_without_continuous_usage() {
|
||||
let (mut app, engine_task) = test_app_with_stream_output_specs(vec![
|
||||
(vec![33], None),
|
||||
(vec![44], Some(EngineCoreFinishReason::Stop)),
|
||||
])
|
||||
.await;
|
||||
|
||||
let response = app
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/inference/v1/generate")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"request_id": "raw-stream-final-usage",
|
||||
"token_ids": [11, 22],
|
||||
"stream": true,
|
||||
"stream_options": {
|
||||
"include_usage": true
|
||||
},
|
||||
"sampling_params": {
|
||||
"max_tokens": 2
|
||||
}
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
|
||||
engine_task.await.expect("mock engine task");
|
||||
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
|
||||
let payloads = sse_data_payloads(&text);
|
||||
assert_eq!(payloads.len(), 4, "{text}");
|
||||
|
||||
let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
|
||||
assert_eq!(first["choices"][0]["token_ids"], json!([33]));
|
||||
assert!(first.get("usage").is_none());
|
||||
|
||||
let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
|
||||
assert_eq!(second["choices"][0]["token_ids"], json!([44]));
|
||||
assert_eq!(second["choices"][0]["finish_reason"], "stop");
|
||||
assert!(second.get("usage").is_none());
|
||||
|
||||
let usage: serde_json::Value = serde_json::from_str(payloads[2]).expect("usage chunk json");
|
||||
assert_eq!(usage["choices"], json!([]));
|
||||
assert_eq!(usage["usage"]["prompt_tokens"], 2);
|
||||
assert_eq!(usage["usage"]["completion_tokens"], 2);
|
||||
assert_eq!(usage["usage"]["total_tokens"], 4);
|
||||
assert_eq!(payloads[3], "[DONE]");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn stream_raw_generate_emits_empty_finish_chunk() {
|
||||
let (mut app, engine_task) = test_app_with_stream_output_specs(vec![
|
||||
(vec![33], None),
|
||||
(vec![], Some(EngineCoreFinishReason::Stop)),
|
||||
])
|
||||
.await;
|
||||
|
||||
let response = app
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/inference/v1/generate")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"request_id": "raw-stream-empty-finish",
|
||||
"token_ids": [11, 22],
|
||||
"stream": true,
|
||||
"sampling_params": {
|
||||
"max_tokens": 2
|
||||
}
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
|
||||
engine_task.await.expect("mock engine task");
|
||||
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
|
||||
let payloads = sse_data_payloads(&text);
|
||||
assert_eq!(payloads.len(), 3, "{text}");
|
||||
|
||||
let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
|
||||
assert_eq!(first["choices"][0]["token_ids"], json!([33]));
|
||||
assert!(first["choices"][0].get("finish_reason").is_none());
|
||||
|
||||
let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
|
||||
assert_eq!(second["choices"][0]["token_ids"], json!([]));
|
||||
assert_eq!(second["choices"][0]["finish_reason"], "stop");
|
||||
assert_eq!(payloads[2], "[DONE]");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn stream_raw_generate_error_finish_returns_sse_error() {
|
||||
let (mut app, engine_task) =
|
||||
test_app_with_stream_output_specs(vec![(vec![], Some(EngineCoreFinishReason::Error))])
|
||||
.await;
|
||||
|
||||
let response = app
|
||||
.call(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/inference/v1/generate")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"request_id": "raw-stream-error",
|
||||
"token_ids": [11, 22],
|
||||
"stream": true,
|
||||
"stream_options": {
|
||||
"include_usage": true
|
||||
},
|
||||
"sampling_params": {
|
||||
"max_tokens": 2
|
||||
}
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.expect("build request"),
|
||||
)
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
|
||||
engine_task.await.expect("mock engine task");
|
||||
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
|
||||
|
||||
assert!(text.contains("\"type\":\"server_error\""), "{text}");
|
||||
assert!(text.contains("Internal server error"), "{text}");
|
||||
assert!(!text.contains("\"finish_reason\":\"error\""), "{text}");
|
||||
assert!(!text.contains("\"usage\":"), "{text}");
|
||||
assert!(text.trim_end().ends_with("data: [DONE]"), "{text}");
|
||||
let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
|
||||
assert_eq!(json["error"]["param"], "stream");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
|
||||
@@ -17,8 +17,6 @@ pub struct AppState {
|
||||
pub chat: ChatLlm,
|
||||
/// Whether to log a summary line for each completed request.
|
||||
pub enable_log_requests: bool,
|
||||
/// Whether to set X-Request-Id on every HTTP response.
|
||||
pub enable_request_id_headers: bool,
|
||||
/// Number of in-flight inference requests currently owned by this frontend.
|
||||
server_load: AtomicU64,
|
||||
}
|
||||
@@ -41,7 +39,6 @@ impl AppState {
|
||||
served_model_names,
|
||||
chat,
|
||||
enable_log_requests: false,
|
||||
enable_request_id_headers: false,
|
||||
server_load: AtomicU64::new(0),
|
||||
}
|
||||
}
|
||||
@@ -52,12 +49,6 @@ impl AppState {
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable X-Request-Id response headers.
|
||||
pub fn with_request_id_headers(mut self, enabled: bool) -> Self {
|
||||
self.enable_request_id_headers = enabled;
|
||||
self
|
||||
}
|
||||
|
||||
/// The primary model name echoed back in API responses (the first served
|
||||
/// name).
|
||||
pub fn primary_model_name(&self) -> &str {
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
[package]
|
||||
name = "vllm-tool-parser-py"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
[lib]
|
||||
name = "_rust_tool_parser"
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
|
||||
[features]
|
||||
default = []
|
||||
extension-module = ["pyo3/extension-module"]
|
||||
|
||||
[dependencies]
|
||||
pyo3.workspace = true
|
||||
pythonize = { workspace = true, features = ["serde_json"] }
|
||||
serde_json.workspace = true
|
||||
thiserror-ext.workspace = true
|
||||
vllm-tool-parser.workspace = true
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,367 @@
|
||||
//! Thin PyO3 bindings for `vllm_tool_parser`.
|
||||
//!
|
||||
//! This crate exposes the Rust tool parser trait and data shapes to Python
|
||||
//! while keeping parser state, grammar, and schema-aware argument conversion in
|
||||
//! Rust. Python callers should use this module as a typed bridge and keep any
|
||||
//! vLLM protocol adaptation outside the binding.
|
||||
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::{PyAny, PyModule};
|
||||
use pythonize::{depythonize, pythonize};
|
||||
use serde_json::Value;
|
||||
use thiserror_ext::AsReport as _;
|
||||
use vllm_tool_parser::{Tool, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
|
||||
macro_rules! tool_parser_factory {
|
||||
($($parser:ident),+ $(,)?) => {
|
||||
fn create_tool_parser(
|
||||
name: &str,
|
||||
tools: &[Tool],
|
||||
) -> PyResult<Box<dyn ToolParser>> {
|
||||
match name {
|
||||
$(
|
||||
stringify!($parser) => {
|
||||
<vllm_tool_parser::$parser as ToolParser>::create(tools)
|
||||
}
|
||||
)+
|
||||
_ => {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"unsupported tool parser `{name}`"
|
||||
)));
|
||||
}
|
||||
}
|
||||
.map_err(|error| PyValueError::new_err(error.to_report_string()))
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
// Export a tool parser to Python by registering it here.
|
||||
tool_parser_factory! {
|
||||
DeepSeekV4ToolParser,
|
||||
MinimaxM3ToolParser,
|
||||
}
|
||||
|
||||
#[pyclass(name = "Tool", module = "vllm._rust_tool_parser", skip_from_py_object)]
|
||||
#[derive(Clone)]
|
||||
struct PyTool(Tool);
|
||||
|
||||
#[pymethods]
|
||||
impl PyTool {
|
||||
#[new]
|
||||
#[pyo3(signature = (name, description, parameters, strict=None))]
|
||||
fn new(
|
||||
name: String,
|
||||
description: Option<String>,
|
||||
parameters: &Bound<'_, PyAny>,
|
||||
strict: Option<bool>,
|
||||
) -> PyResult<Self> {
|
||||
let parameters = depythonize::<Value>(parameters).map_err(|error| {
|
||||
PyValueError::new_err(format!(
|
||||
"failed to convert tool parameters from Python to JSON: {error}"
|
||||
))
|
||||
})?;
|
||||
Ok(Self(Tool {
|
||||
name,
|
||||
description,
|
||||
parameters,
|
||||
strict,
|
||||
}))
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn name(&self) -> &str {
|
||||
&self.0.name
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn description(&self) -> Option<&str> {
|
||||
self.0.description.as_deref()
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn parameters(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
pythonize(py, &self.0.parameters).map(Bound::unbind).map_err(|error| {
|
||||
PyValueError::new_err(format!(
|
||||
"failed to convert tool parameters from JSON to Python: {error}"
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn strict(&self) -> Option<bool> {
|
||||
self.0.strict
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(
|
||||
name = "ToolCallDelta",
|
||||
module = "vllm._rust_tool_parser",
|
||||
skip_from_py_object
|
||||
)]
|
||||
#[derive(Clone)]
|
||||
struct PyToolCallDelta(ToolCallDelta);
|
||||
|
||||
#[pymethods]
|
||||
impl PyToolCallDelta {
|
||||
#[new]
|
||||
#[pyo3(signature = (tool_index, name, arguments))]
|
||||
fn new(tool_index: usize, name: Option<String>, arguments: String) -> Self {
|
||||
Self(ToolCallDelta {
|
||||
tool_index,
|
||||
name,
|
||||
arguments,
|
||||
})
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn tool_index(&self) -> usize {
|
||||
self.0.tool_index
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn name(&self) -> Option<&str> {
|
||||
self.0.name.as_deref()
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn arguments(&self) -> &str {
|
||||
&self.0.arguments
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(
|
||||
name = "ToolParserOutput",
|
||||
module = "vllm._rust_tool_parser",
|
||||
skip_from_py_object
|
||||
)]
|
||||
#[derive(Clone)]
|
||||
struct PyToolParserOutput(ToolParserOutput);
|
||||
|
||||
#[pymethods]
|
||||
impl PyToolParserOutput {
|
||||
#[new]
|
||||
#[pyo3(signature = (normal_text="", calls=None))]
|
||||
fn new(py: Python<'_>, normal_text: &str, calls: Option<Vec<Py<PyToolCallDelta>>>) -> Self {
|
||||
let calls =
|
||||
calls.unwrap_or_default().iter().map(|call| call.borrow(py).0.clone()).collect();
|
||||
Self(ToolParserOutput {
|
||||
normal_text: normal_text.to_owned(),
|
||||
calls,
|
||||
})
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn normal_text(&self) -> &str {
|
||||
&self.0.normal_text
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn calls(&self) -> Vec<PyToolCallDelta> {
|
||||
self.0.calls.iter().cloned().map(PyToolCallDelta).collect()
|
||||
}
|
||||
|
||||
fn append(&mut self, other: PyRef<'_, PyToolParserOutput>) {
|
||||
self.0.append(other.0.clone());
|
||||
}
|
||||
|
||||
fn coalesce_calls(&self) -> Self {
|
||||
Self(self.0.clone().coalesce_calls())
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(name = "ToolParser", module = "vllm._rust_tool_parser", unsendable)]
|
||||
struct PyToolParser(Box<dyn ToolParser>);
|
||||
|
||||
impl PyToolParser {
|
||||
fn parse_into_output(&mut self, chunk: &str, output: &mut PyToolParserOutput) -> PyResult<()> {
|
||||
self.0
|
||||
.parse_into(chunk, &mut output.0)
|
||||
.map_err(|error| PyValueError::new_err(error.to_report_string()))
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyToolParser {
|
||||
#[new]
|
||||
fn new(py: Python<'_>, parser_name: &str, tools: Vec<Py<PyTool>>) -> PyResult<Self> {
|
||||
let tools = tools.iter().map(|tool| tool.borrow(py).0.clone()).collect::<Vec<_>>();
|
||||
create_tool_parser(parser_name, &tools).map(Self)
|
||||
}
|
||||
|
||||
fn parse_into(
|
||||
&mut self,
|
||||
chunk: &str,
|
||||
mut output: PyRefMut<'_, PyToolParserOutput>,
|
||||
) -> PyResult<()> {
|
||||
self.parse_into_output(chunk, &mut output)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> PyResult<PyToolParserOutput> {
|
||||
self.0
|
||||
.finish()
|
||||
.map(PyToolParserOutput)
|
||||
.map_err(|error| PyValueError::new_err(error.to_report_string()))
|
||||
}
|
||||
|
||||
fn reset(&mut self) -> String {
|
||||
self.0.reset()
|
||||
}
|
||||
|
||||
fn preserve_special_tokens(&self) -> bool {
|
||||
self.0.preserve_special_tokens()
|
||||
}
|
||||
}
|
||||
|
||||
#[pymodule]
|
||||
fn _rust_tool_parser(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<PyTool>()?;
|
||||
m.add_class::<PyToolCallDelta>()?;
|
||||
m.add_class::<PyToolParserOutput>()?;
|
||||
m.add_class::<PyToolParser>()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn with_python<R>(f: impl for<'py> FnOnce(Python<'py>) -> R) -> R {
|
||||
Python::initialize();
|
||||
Python::attach(f)
|
||||
}
|
||||
|
||||
fn tool_schema() -> Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"user_id": {"type": "integer"},
|
||||
"shipping": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"city": {"type": "string"},
|
||||
"zip": {"type": "integer"}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn build_call() -> String {
|
||||
r#"<|DSML|tool_calls>
|
||||
<|DSML|invoke name="create_order">
|
||||
<|DSML|parameter name="user_id" string="false">42</|DSML|parameter>
|
||||
<|DSML|parameter name="shipping" string="false">{"city":"Singapore","zip":18956}</|DSML|parameter>
|
||||
</|DSML|invoke>
|
||||
</|DSML|tool_calls>"#
|
||||
.to_owned()
|
||||
}
|
||||
|
||||
fn make_py_tool(py: Python<'_>) -> PyResult<Py<PyTool>> {
|
||||
let parameters = pythonize(py, &tool_schema()).map_err(|error| {
|
||||
PyValueError::new_err(format!(
|
||||
"failed to convert test schema from JSON to Python: {error}"
|
||||
))
|
||||
})?;
|
||||
Py::new(
|
||||
py,
|
||||
PyTool::new(
|
||||
"create_order".to_owned(),
|
||||
Some("Create an order".to_owned()),
|
||||
¶meters,
|
||||
None,
|
||||
)?,
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_round_trips_typed_fields() {
|
||||
with_python(|py| {
|
||||
let tool = make_py_tool(py)?;
|
||||
let borrowed = tool.borrow(py);
|
||||
assert_eq!(borrowed.name(), "create_order");
|
||||
assert_eq!(borrowed.description(), Some("Create an order"));
|
||||
assert_eq!(borrowed.strict(), None);
|
||||
|
||||
let parameters = borrowed.parameters(py)?;
|
||||
let parameters = depythonize::<Value>(parameters.bind(py))?;
|
||||
assert_eq!(parameters, tool_schema());
|
||||
PyResult::Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn output_append_and_coalesce_calls() {
|
||||
with_python(|py| {
|
||||
let first = Py::new(
|
||||
py,
|
||||
PyToolCallDelta::new(0, Some("create_order".to_owned()), "{\"a\"".to_owned()),
|
||||
)?;
|
||||
let second = Py::new(py, PyToolCallDelta::new(0, None, ":1}".to_owned()))?;
|
||||
let mut output = PyToolParserOutput::new(py, "text", Some(vec![first]));
|
||||
let other = Py::new(py, PyToolParserOutput::new(py, "", Some(vec![second])))?;
|
||||
output.append(other.borrow(py));
|
||||
|
||||
let coalesced = output.coalesce_calls();
|
||||
assert_eq!(coalesced.normal_text(), "text");
|
||||
let calls = coalesced.calls();
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].tool_index(), 0);
|
||||
assert_eq!(calls[0].name(), Some("create_order"));
|
||||
assert_eq!(calls[0].arguments(), "{\"a\":1}");
|
||||
PyResult::Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parser_parse_finish_and_preserve_special_tokens() {
|
||||
with_python(|py| {
|
||||
let tool = make_py_tool(py)?;
|
||||
let mut parser = PyToolParser::new(py, "DeepSeekV4ToolParser", vec![tool])?;
|
||||
assert!(parser.preserve_special_tokens());
|
||||
|
||||
let mut output = PyToolParserOutput::new(py, "", None);
|
||||
parser.parse_into_output(&build_call(), &mut output)?;
|
||||
let finish = Py::new(py, parser.finish()?)?;
|
||||
output.append(finish.borrow(py));
|
||||
let output = output.coalesce_calls();
|
||||
|
||||
assert_eq!(output.normal_text(), "");
|
||||
let calls = output.calls();
|
||||
assert_eq!(calls.len(), 1);
|
||||
assert_eq!(calls[0].name(), Some("create_order"));
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(calls[0].arguments()).unwrap(),
|
||||
json!({
|
||||
"user_id": 42,
|
||||
"shipping": {
|
||||
"city": "Singapore",
|
||||
"zip": 18956
|
||||
}
|
||||
})
|
||||
);
|
||||
|
||||
assert_eq!(parser.reset(), "");
|
||||
PyResult::Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parser_errors_for_unknown_name() {
|
||||
with_python(|py| {
|
||||
let tool = make_py_tool(py)?;
|
||||
let error = match PyToolParser::new(py, "missing", vec![tool]) {
|
||||
Ok(_) => panic!("missing parser name unexpectedly succeeded"),
|
||||
Err(error) => error,
|
||||
};
|
||||
let message = format!("{error}");
|
||||
assert!(message.contains("unsupported tool parser `missing`"));
|
||||
PyResult::Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
@@ -8,7 +8,7 @@ const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
||||
marker_whitespace: JsonToolCallWhitespace::Optional,
|
||||
delimiter: None,
|
||||
name_key: "name",
|
||||
arguments_key: &["arguments"],
|
||||
arguments_key: "arguments",
|
||||
};
|
||||
|
||||
/// Tool parser for Hermes XML-wrapped JSON tool calls.
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user