forked from Karylab-cklius/vllm
Compare commits
125
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0fc695fc6d | ||
|
|
91df0fad4d | ||
|
|
78743ab5bf | ||
|
|
b2d7294b0f | ||
|
|
741ba421d8 | ||
|
|
ac94893da3 | ||
|
|
967c5c3bc3 | ||
|
|
54c660c3a6 | ||
|
|
8fb0274415 | ||
|
|
eebce65756 | ||
|
|
303916e93d | ||
|
|
5633405964 | ||
|
|
6124a98a9b | ||
|
|
2ed0a9627b | ||
|
|
4dcd10eb0d | ||
|
|
228bcc436b | ||
|
|
3d3ba460a2 | ||
|
|
66ecfd0568 | ||
|
|
f0f6805d8a | ||
|
|
15652a6b70 | ||
|
|
51ef688831 | ||
|
|
6ac69203e8 | ||
|
|
1505b3d8a1 | ||
|
|
32f34d3935 | ||
|
|
9c7f7741d4 | ||
|
|
6181e80fe0 | ||
|
|
3bb46975bd | ||
|
|
810966453a | ||
|
|
2a983c79ac | ||
|
|
bc5745a00f | ||
|
|
3b3d5287fa | ||
|
|
062b05ff3a | ||
|
|
fa27d4e9cf | ||
|
|
67d3792d99 | ||
|
|
00d1fb7747 | ||
|
|
c9b4b184b4 | ||
|
|
f87df1df9e | ||
|
|
eafbb06331 | ||
|
|
ec0a31d4aa | ||
|
|
c8beda4cc3 | ||
|
|
2f27c9a150 | ||
|
|
4765f0f189 | ||
|
|
a50e675b0d | ||
|
|
f6a708ab2b | ||
|
|
4200f62147 | ||
|
|
c73b0d0db9 | ||
|
|
e28e369f78 | ||
|
|
703fb17b13 | ||
|
|
b593396c7a | ||
|
|
91e17d4315 | ||
|
|
aa6fb8a329 | ||
|
|
6a894574bf | ||
|
|
7f003a1285 | ||
|
|
ef0df7dbd6 | ||
|
|
a80af24356 | ||
|
|
c66b19800b | ||
|
|
6a11d72df7 | ||
|
|
02d2da0748 | ||
|
|
bbb6c274c8 | ||
|
|
62215e72c6 | ||
|
|
7fe7800fa4 | ||
|
|
8a83e6f2d7 | ||
|
|
efc347f1b2 | ||
|
|
d98b8f371c | ||
|
|
e64237ae82 | ||
|
|
d61d8566ec | ||
|
|
d2f70da116 | ||
|
|
6542d48964 | ||
|
|
ca73293fa6 | ||
|
|
ef3af56a97 | ||
|
|
b4a6f26c90 | ||
|
|
165b7864d0 | ||
|
|
c505cd93ef | ||
|
|
96229fa99e | ||
|
|
da1daf40bf | ||
|
|
4efd6ffde0 | ||
|
|
56aff0dd15 | ||
|
|
063ce98fb7 | ||
|
|
62d6f06e3d | ||
|
|
b7c5baf63d | ||
|
|
a55fccfc7c | ||
|
|
41a4829f22 | ||
|
|
38fd2405f3 | ||
|
|
a947f7a420 | ||
|
|
439203d32c | ||
|
|
8d9536a775 | ||
|
|
3da29aa4a5 | ||
|
|
06f94633e7 | ||
|
|
99ef652907 | ||
|
|
4cc78c9d5d | ||
|
|
3dbb4e0ace | ||
|
|
b21443e23c | ||
|
|
06ee2d8433 | ||
|
|
b5235fca2e | ||
|
|
3e77036768 | ||
|
|
6f68ca3e91 | ||
|
|
0c96dd64fb | ||
|
|
68f5e565c9 | ||
|
|
9354fb1ba5 | ||
|
|
f35b557239 | ||
|
|
e68988a248 | ||
|
|
4b87b3e845 | ||
|
|
90619351e3 | ||
|
|
d0975a4b50 | ||
|
|
1bdc60ed53 | ||
|
|
a6183563b6 | ||
|
|
22c2e87555 | ||
|
|
d01d0b4646 | ||
|
|
b4b4aaa70e | ||
|
|
5e2af28838 | ||
|
|
4f423bd5bc | ||
|
|
f0cd590d62 | ||
|
|
e6018c644a | ||
|
|
f25952e59b | ||
|
|
b58e082d95 | ||
|
|
0c1e6f63f5 | ||
|
|
ceb0111a90 | ||
|
|
0414d75410 | ||
|
|
128adabfe0 | ||
|
|
bdbf08fc02 | ||
|
|
6bad553f4e | ||
|
|
91945b6e4a | ||
|
|
2b237c7a41 | ||
|
|
dad95e34d8 | ||
|
|
a248b45d05 |
@@ -28,18 +28,19 @@ steps:
|
||||
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"
|
||||
|
||||
- label: CPU-Compatibility Tests
|
||||
depends_on: []
|
||||
device: intel_cpu
|
||||
no_plugin: true
|
||||
source_file_dependencies:
|
||||
- cmake/cpu_extension.cmake
|
||||
- setup.py
|
||||
- vllm/platforms/cpu.py
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-compatibility-test.sh"
|
||||
# Note: SDE can't be downloaded from CI host because of AWS WAF
|
||||
# - label: CPU-Compatibility Tests
|
||||
# depends_on: []
|
||||
# device: intel_cpu
|
||||
# no_plugin: true
|
||||
# source_file_dependencies:
|
||||
# - cmake/cpu_extension.cmake
|
||||
# - setup.py
|
||||
# - vllm/platforms/cpu.py
|
||||
# commands:
|
||||
# - |
|
||||
# bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
|
||||
# bash .buildkite/scripts/hardware_ci/run-cpu-compatibility-test.sh"
|
||||
|
||||
- label: CPU-Language Generation and Pooling Model Tests
|
||||
depends_on: []
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
group: Basic Correctness
|
||||
depends_on:
|
||||
- image-build-xpu
|
||||
steps:
|
||||
- label: XPU Sleep Mode
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/basic_correctness/test_cumem.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
|
||||
pytest -v -s basic_correctness/test_mem.py::test_end_to_end'
|
||||
@@ -53,3 +53,26 @@ steps:
|
||||
pytest -v -s v1/test_oracle.py &&
|
||||
pytest -v -s v1/test_request.py &&
|
||||
pytest -v -s v1/test_outputs.py'
|
||||
|
||||
- label: XPU CPU Offload
|
||||
timeout_in_minutes: 60
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- vllm/v1/kv_offload/
|
||||
- vllm/v1/kv_connector/
|
||||
- tests/v1/kv_offload/
|
||||
- tests/v1/kv_connector/unit/test_offloading_connector.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
|
||||
cd tests &&
|
||||
pytest -v -s v1/kv_offload &&
|
||||
pytest -v -s v1/kv_connector/unit/test_offloading_connector.py'
|
||||
|
||||
@@ -40,7 +40,9 @@ steps:
|
||||
python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192 &&
|
||||
python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 &&
|
||||
python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel &&
|
||||
python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192
|
||||
python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192 &&
|
||||
VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192 &&
|
||||
python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC --enforce-eager -tp 2 --max-model-len 8192
|
||||
'
|
||||
- label: "XPU V1 test"
|
||||
depends_on:
|
||||
|
||||
@@ -28,8 +28,10 @@
|
||||
###############################################################################
|
||||
set -o pipefail
|
||||
|
||||
# Export Python path
|
||||
export PYTHONPATH=".."
|
||||
# Export Python path for commands that run directly on the host. Containerized
|
||||
# tests set this to /vllm-workspace below so spawned Python processes do not
|
||||
# depend on their current working directory.
|
||||
export PYTHONPATH="${PYTHONPATH:-..}"
|
||||
|
||||
###############################################################################
|
||||
# Helper Functions
|
||||
@@ -377,6 +379,14 @@ HF_CACHE="$(realpath ~)/huggingface"
|
||||
mkdir -p "${HF_CACHE}"
|
||||
HF_MOUNT="/root/.cache/huggingface"
|
||||
|
||||
# Hugging Face Hub defaults to 10s request/download timeouts, while the ROCm
|
||||
# CI image currently raises downloads to 60s. AMD model-test jobs routinely
|
||||
# start from a cold or partially-populated shared cache, and the 60s read cap
|
||||
# has still timed out before pytest reached the vLLM behavior under test.
|
||||
# Keep the CI default explicit and overridable from the Buildkite environment.
|
||||
: "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
|
||||
: "${HF_HUB_ETAG_TIMEOUT:=60}"
|
||||
|
||||
# ---- Command source selection ----
|
||||
# Prefer VLLM_TEST_COMMANDS (preserves all inner quoting intact).
|
||||
# Fall back to $* for backward compatibility, but warn that inner
|
||||
@@ -416,7 +426,14 @@ fi
|
||||
|
||||
echo "Final commands: $commands"
|
||||
|
||||
MYPYTHONPATH=".."
|
||||
MYPYTHONPATH="/vllm-workspace"
|
||||
|
||||
container_job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-0}}"
|
||||
container_job_id="${container_job_id//[^A-Za-z0-9_.-]/_}"
|
||||
container_job_id_short="${container_job_id:0:8}"
|
||||
CONTAINER_TMPDIR="/tmp/vllm-${container_job_id_short}"
|
||||
CONTAINER_CACHE_ROOT="/tmp/vllm-buildkite-${container_job_id}/cache"
|
||||
CONTAINER_PREFLIGHT="mkdir -p \"\$TMPDIR\" \"\$TORCHINDUCTOR_CACHE_DIR\" \"\$TRITON_CACHE_DIR\" \"\$VLLM_CACHE_ROOT\" \"\$XDG_CACHE_HOME\" && python -c \"import encodings, importlib.metadata as im, importlib.util as iu; [im.version(d) for d in ('transformers', 'torch', 'ray', 'sympy', 'markupsafe', 'vllm')]; missing=[m for m in ('torch.utils.model_zoo', 'transformers.models.nomic_bert', 'ray.dag', 'sympy.physics', 'markupsafe._speedups') if iu.find_spec(m) is None]; assert not missing, missing\""
|
||||
|
||||
# Verify GPU access
|
||||
render_gid=$(getent group render | cut -d: -f3)
|
||||
@@ -493,6 +510,8 @@ else
|
||||
--group-add "$render_gid" \
|
||||
--rm \
|
||||
-e HF_TOKEN \
|
||||
-e "HF_HUB_DOWNLOAD_TIMEOUT=${HF_HUB_DOWNLOAD_TIMEOUT}" \
|
||||
-e "HF_HUB_ETAG_TIMEOUT=${HF_HUB_ETAG_TIMEOUT}" \
|
||||
-e AWS_ACCESS_KEY_ID \
|
||||
-e AWS_SECRET_ACCESS_KEY \
|
||||
-e BUILDKITE_PARALLEL_JOB \
|
||||
@@ -500,10 +519,15 @@ else
|
||||
-v "${HF_CACHE}:${HF_MOUNT}" \
|
||||
-e "HF_HOME=${HF_MOUNT}" \
|
||||
-e "PYTHONPATH=${MYPYTHONPATH}" \
|
||||
-e "TMPDIR=${CONTAINER_TMPDIR}/tmp" \
|
||||
-e "TORCHINDUCTOR_CACHE_DIR=${CONTAINER_CACHE_ROOT}/torchinductor" \
|
||||
-e "TRITON_CACHE_DIR=${CONTAINER_CACHE_ROOT}/triton" \
|
||||
-e "VLLM_CACHE_ROOT=${CONTAINER_CACHE_ROOT}/vllm" \
|
||||
-e "XDG_CACHE_HOME=${CONTAINER_CACHE_ROOT}/xdg" \
|
||||
-e "PYTORCH_ROCM_ARCH=" \
|
||||
--name "${container_name}" \
|
||||
"${image_name}" \
|
||||
/bin/bash -c "${commands}"
|
||||
/bin/bash -c "${CONTAINER_PREFLIGHT} && ${commands}"
|
||||
|
||||
exit_code=$?
|
||||
handle_pytest_exit "$exit_code"
|
||||
|
||||
+73
-26
@@ -315,24 +315,6 @@ steps:
|
||||
- pytest -v -s distributed/test_pp_cudagraph.py
|
||||
- pytest -v -s distributed/test_pipeline_parallel.py
|
||||
|
||||
#---------------------------------------------------------- mi250 · engine -----------------------------------------------------------#
|
||||
|
||||
- label: Engine # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
|
||||
agent_pool: mi250_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/engine
|
||||
- tests/test_sequence
|
||||
- tests/test_config
|
||||
- tests/test_logger
|
||||
- tests/test_vllm_port
|
||||
commands:
|
||||
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
|
||||
|
||||
#----------------------------------------------------------- mi250 · evals -----------------------------------------------------------#
|
||||
|
||||
- label: Multi-Modal Accuracy Eval (Small Models) # TBD
|
||||
@@ -908,10 +890,10 @@ steps:
|
||||
- vllm/
|
||||
- tests/basic_correctness/test_basic_correctness
|
||||
- tests/basic_correctness/test_cpu_offload
|
||||
- tests/basic_correctness/test_cumem.py
|
||||
- tests/basic_correctness/test_mem.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s basic_correctness/test_cumem.py
|
||||
- pytest -v -s basic_correctness/test_mem.py
|
||||
- pytest -v -s basic_correctness/test_basic_correctness.py
|
||||
- pytest -v -s basic_correctness/test_cpu_offload.py
|
||||
|
||||
@@ -1299,12 +1281,11 @@ steps:
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/llm
|
||||
- tests/entrypoints/offline_mode
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
|
||||
- pytest -v -s entrypoints/llm/test_generate.py
|
||||
- pytest -v -s entrypoints/offline_mode
|
||||
- pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode
|
||||
- pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
|
||||
- pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
|
||||
|
||||
- label: Entrypoints Integration (Pooling) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -1346,12 +1327,13 @@ 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/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
|
||||
|
||||
- label: OpenAI API correctness # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
@@ -1410,6 +1392,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
@@ -2184,10 +2167,72 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
|
||||
|
||||
- label: Speculators Correctness # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/config/speculative.py
|
||||
- vllm/engine/arg_utils.py
|
||||
- vllm/transformers_utils/config.py
|
||||
- vllm/transformers_utils/configs/speculators/
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/v1/worker/gpu/spec_decode/
|
||||
- vllm/v1/worker/gpu_model_runner.py
|
||||
- vllm/v1/sample/
|
||||
- vllm/v1/attention/backends/
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/model_executor/model_loader/
|
||||
- vllm/model_executor/layers/
|
||||
- vllm/model_executor/models/llama_eagle3.py
|
||||
- vllm/model_executor/models/qwen3.py
|
||||
- vllm/model_executor/models/qwen3_dflash.py
|
||||
- vllm/model_executor/models/registry.py
|
||||
- vllm/_aiter_ops.py
|
||||
- tests/evals/gsm8k/
|
||||
- tests/v1/spec_decode/test_speculators_correctness.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
|
||||
- pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
|
||||
|
||||
- label: Extract Hidden States Integration # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/config/speculative.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/
|
||||
- vllm/model_executor/layers/attention/
|
||||
- vllm/model_executor/layers/mamba/
|
||||
- vllm/model_executor/model_loader/
|
||||
- vllm/model_executor/models/extract_hidden_states.py
|
||||
- vllm/model_executor/models/llama.py
|
||||
- vllm/model_executor/models/qwen3_5.py
|
||||
- vllm/model_executor/models/qwen3_next.py
|
||||
- vllm/model_executor/models/registry.py
|
||||
- vllm/transformers_utils/configs/extract_hidden_states.py
|
||||
- vllm/transformers_utils/configs/qwen3_5.py
|
||||
- vllm/v1/attention/backends/
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/v1/kv_cache_interface.py
|
||||
- vllm/v1/spec_decode/extract_hidden_states.py
|
||||
- vllm/v1/worker/gpu_model_runner.py
|
||||
- vllm/_aiter_ops.py
|
||||
- tests/v1/kv_connector/extract_hidden_states_integration/
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s v1/kv_connector/extract_hidden_states_integration
|
||||
|
||||
- label: V1 attention (H100-MI300) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/config/attention.py
|
||||
@@ -2546,6 +2591,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2555,7 +2601,7 @@ steps:
|
||||
- tests/test_logger
|
||||
- tests/test_vllm_port
|
||||
commands:
|
||||
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
|
||||
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
|
||||
|
||||
#----------------------------------------------------------- mi325 · evals -----------------------------------------------------------#
|
||||
|
||||
@@ -2637,6 +2683,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -2,8 +2,8 @@ group: Attention
|
||||
depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: V1 attention (H100)
|
||||
key: v1-attention-h100
|
||||
- label: V1 attention (H100-MI300)
|
||||
key: v1-attention-h100-mi300
|
||||
timeout_in_minutes: 30
|
||||
device: h100
|
||||
source_file_dependencies:
|
||||
@@ -13,6 +13,20 @@ steps:
|
||||
- tests/v1/attention
|
||||
commands:
|
||||
- pytest -v -s v1/attention
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 70
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/config/attention.py
|
||||
- vllm/model_executor/layers/attention
|
||||
- vllm/v1/attention
|
||||
- tests/v1/attention
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/envs.py
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
- label: V1 attention (B200)
|
||||
key: v1-attention-b200
|
||||
|
||||
@@ -10,9 +10,9 @@ steps:
|
||||
- vllm/
|
||||
- tests/basic_correctness/test_basic_correctness
|
||||
- tests/basic_correctness/test_cpu_offload
|
||||
- tests/basic_correctness/test_cumem.py
|
||||
- tests/basic_correctness/test_mem.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s basic_correctness/test_cumem.py
|
||||
- pytest -v -s basic_correctness/test_mem.py
|
||||
- pytest -v -s basic_correctness/test_basic_correctness.py
|
||||
- pytest -v -s basic_correctness/test_cpu_offload.py
|
||||
|
||||
@@ -26,6 +26,12 @@ steps:
|
||||
- tests/test_jit_monitor.py
|
||||
commands:
|
||||
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Engine (1 GPU)
|
||||
key: engine-1-gpu
|
||||
@@ -38,7 +44,7 @@ steps:
|
||||
- pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 40
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -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/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
|
||||
|
||||
- label: Entrypoints Integration (LLM)
|
||||
key: entrypoints-integration-llm
|
||||
@@ -20,16 +20,14 @@ steps:
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/llm
|
||||
- tests/entrypoints/offline_mode
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
|
||||
- pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode
|
||||
- pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
|
||||
- pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
|
||||
- pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
soft_fail: true
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -46,8 +44,7 @@ steps:
|
||||
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
soft_fail: true
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 80
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
@@ -69,8 +66,7 @@ steps:
|
||||
- pytest -v -s tool_use
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
soft_fail: true
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
@@ -90,7 +86,6 @@ steps:
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
soft_fail: true
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
@@ -109,8 +104,7 @@ steps:
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
soft_fail: true
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -156,3 +150,21 @@ steps:
|
||||
- vllm/entrypoints/openai/
|
||||
commands: # LMEval
|
||||
- pytest -s entrypoints/openai/correctness/
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/model_executor/models/whisper.py
|
||||
- vllm/model_executor/layers/
|
||||
- vllm/v1/attention/backends/
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
- vllm/model_executor/model_loader/
|
||||
commands:
|
||||
- bash ../tools/install_torchcodec_rocm.sh || exit 1
|
||||
- pytest -s entrypoints/openai/correctness/
|
||||
|
||||
@@ -86,7 +86,7 @@ steps:
|
||||
parallelism: 2
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/
|
||||
- vllm/model_executor/layers/quantization
|
||||
|
||||
@@ -12,6 +12,21 @@ steps:
|
||||
autorun_on_main: true
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 55
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/model_executor/layers/quantization
|
||||
- vllm/model_executor/models/
|
||||
- vllm/model_executor/model_loader/
|
||||
- vllm/v1/attention/backends/
|
||||
- vllm/v1/attention/selector.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
# - label: LM Eval Large Models (4 GPUs)(A100)
|
||||
# device: a100
|
||||
|
||||
@@ -51,6 +51,7 @@ steps:
|
||||
torch_nightly: {}
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 90
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
|
||||
@@ -15,7 +15,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -33,7 +33,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -50,7 +50,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -153,3 +153,12 @@ steps:
|
||||
- tests/models/multimodal/pooling
|
||||
commands:
|
||||
- pytest -v -s models/multimodal/pooling -m 'not core_model'
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal/pooling
|
||||
|
||||
@@ -37,6 +37,21 @@ steps:
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 65
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/v1/worker/gpu/spec_decode/
|
||||
- vllm/model_executor/model_loader/
|
||||
- vllm/v1/sample/
|
||||
- vllm/model_executor/layers/
|
||||
- vllm/transformers_utils/configs/speculators/
|
||||
- tests/v1/e2e/spec_decode/
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
- label: Spec Decode Speculators + MTP Nightly B200
|
||||
key: spec-decode-speculators-mtp-nightly-b200
|
||||
@@ -61,6 +76,20 @@ steps:
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 65
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/v1/worker/gpu/spec_decode/
|
||||
- vllm/model_executor/model_loader/
|
||||
- vllm/v1/sample/
|
||||
- vllm/model_executor/layers/
|
||||
- tests/v1/e2e/spec_decode/
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
- label: Spec Decode Draft Model
|
||||
key: spec-decode-draft-model
|
||||
@@ -72,6 +101,20 @@ steps:
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 50
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/v1/worker/gpu/spec_decode/
|
||||
- vllm/model_executor/model_loader/
|
||||
- vllm/v1/sample/
|
||||
- vllm/model_executor/layers/
|
||||
- tests/v1/e2e/spec_decode/
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
- label: Spec Decode Draft Model Nightly B200
|
||||
key: spec-decode-draft-model-nightly-b200
|
||||
|
||||
@@ -33,10 +33,3 @@ share/python-wheels/
|
||||
*.egg
|
||||
MANIFEST
|
||||
rust/target/
|
||||
# Not needed in Docker builds
|
||||
docs/
|
||||
.github/
|
||||
.pre-commit-config.yaml
|
||||
.clang-format
|
||||
.gitattributes
|
||||
format.sh
|
||||
|
||||
+2
-1
@@ -34,10 +34,11 @@
|
||||
/vllm/entrypoints/speech_to_text/realtime @njhill
|
||||
/vllm/entrypoints/speech_to_text @NickLucche
|
||||
/vllm/entrypoints/pooling @noooop
|
||||
/vllm/entrypoints/sagemaker @DarkLight1337
|
||||
/vllm/entrypoints/serve/sagemaker @DarkLight1337
|
||||
/vllm/entrypoints/serve @njhill
|
||||
/vllm/entrypoints/*.py @njhill
|
||||
/vllm/entrypoints/chat_utils.py @DarkLight1337
|
||||
/vllm/entrypoints/offline_utils.py @DarkLight1337
|
||||
/vllm/entrypoints/llm.py @DarkLight1337
|
||||
|
||||
# Rust Frontend
|
||||
|
||||
+3
-12
@@ -21,6 +21,9 @@ pull_request_rules:
|
||||
- check-failure=pre-commit
|
||||
- -closed
|
||||
- -draft
|
||||
- or:
|
||||
- label=ready
|
||||
- label=verified
|
||||
actions:
|
||||
comment:
|
||||
message: |
|
||||
@@ -36,18 +39,6 @@ pull_request_rules:
|
||||
|
||||
For future commits, `pre-commit` will run automatically on changed files before each commit.
|
||||
|
||||
> [!TIP]
|
||||
> <details>
|
||||
> <summary>Is <code>mypy</code> failing?</summary>
|
||||
> <br/>
|
||||
> <code>mypy</code> is run differently in CI. If the failure is related to this check, please use the following command to run it locally:
|
||||
>
|
||||
> ```bash
|
||||
> # For mypy (substitute "3.10" with the failing version if needed)
|
||||
> pre-commit run --hook-stage manual mypy-3.10
|
||||
> ```
|
||||
> </details>
|
||||
|
||||
- name: comment-dco-failure
|
||||
description: Comment on PR when DCO check fails
|
||||
conditions:
|
||||
|
||||
@@ -15,7 +15,7 @@ jobs:
|
||||
actions: write
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/stale@997185467fa4f803885201cee163a9f38240193d # v10.1.1
|
||||
- uses: actions/stale@eb5cf3af3ac0a1aa4c9c45633dd1ae542a27a899 # v10.3.0
|
||||
with:
|
||||
# Increasing this value ensures that changes to this workflow
|
||||
# propagate to all issues and PRs in days rather than months
|
||||
|
||||
+7
-13
@@ -148,33 +148,27 @@ repos:
|
||||
language: python
|
||||
entry: python tools/pre_commit/generate_nightly_torch_test.py
|
||||
files: ^requirements/test/cuda\.(in|txt)$
|
||||
- id: mypy-local
|
||||
name: Run mypy locally for lowest supported Python version
|
||||
entry: python tools/pre_commit/mypy.py 0 "3.10"
|
||||
stages: [pre-commit] # Don't run in CI
|
||||
- id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
|
||||
name: Run mypy for Python 3.10
|
||||
entry: python tools/pre_commit/mypy.py "3.10"
|
||||
<<: &mypy_common
|
||||
language: python
|
||||
types_or: [python, pyi]
|
||||
require_serial: true
|
||||
additional_dependencies: ["mypy[faster-cache]==1.19.1", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
|
||||
- id: mypy-3.10 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
|
||||
name: Run mypy for Python 3.10
|
||||
entry: python tools/pre_commit/mypy.py 1 "3.10"
|
||||
<<: *mypy_common
|
||||
stages: [manual] # Only run in CI
|
||||
additional_dependencies: ["mypy==1.20.2", regex, types-cachetools, types-setuptools, types-PyYAML, types-requests, types-torch, pydantic]
|
||||
- id: mypy-3.11 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
|
||||
name: Run mypy for Python 3.11
|
||||
entry: python tools/pre_commit/mypy.py 1 "3.11"
|
||||
entry: python tools/pre_commit/mypy.py "3.11"
|
||||
<<: *mypy_common
|
||||
stages: [manual] # Only run in CI
|
||||
- id: mypy-3.12 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
|
||||
name: Run mypy for Python 3.12
|
||||
entry: python tools/pre_commit/mypy.py 1 "3.12"
|
||||
entry: python tools/pre_commit/mypy.py "3.12"
|
||||
<<: *mypy_common
|
||||
stages: [manual] # Only run in CI
|
||||
- id: mypy-3.13 # TODO: Use https://github.com/pre-commit/mirrors-mypy when mypy setup is less awkward
|
||||
name: Run mypy for Python 3.13
|
||||
entry: python tools/pre_commit/mypy.py 1 "3.13"
|
||||
entry: python tools/pre_commit/mypy.py "3.13"
|
||||
<<: *mypy_common
|
||||
stages: [manual] # Only run in CI
|
||||
- id: shellcheck
|
||||
|
||||
@@ -98,11 +98,13 @@ pre-commit run --all-files
|
||||
pre-commit run ruff-check --all-files
|
||||
|
||||
# Run mypy as it is in CI:
|
||||
pre-commit run mypy-3.10 --all-files --hook-stage manual
|
||||
pre-commit run mypy-3.12 --all-files --hook-stage manual
|
||||
```
|
||||
|
||||
The line length limit for Python code is 88 characters. If you are not sure, use pre-commit to check.
|
||||
|
||||
Use [Google-style docstrings](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) (`Args:`/`Returns:`/`Raises:` sections), not reStructuredText/Sphinx fields (`:param:`, `:return:`, `:rtype:`).
|
||||
|
||||
### Commit messages
|
||||
|
||||
Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:
|
||||
|
||||
+22
-19
@@ -307,10 +307,7 @@ endif()
|
||||
#
|
||||
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/cuda_view.cu"
|
||||
"csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
|
||||
"csrc/quantization/activation_kernels.cu"
|
||||
"csrc/cuda_utils_kernels.cu"
|
||||
"csrc/torch_bindings.cpp")
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
@@ -346,9 +343,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
FetchContent_MakeAvailable(cutlass)
|
||||
|
||||
list(APPEND VLLM_EXT_SRC
|
||||
"csrc/cutlass_extensions/common.cpp")
|
||||
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${VLLM_EXT_SRC}"
|
||||
CUDA_ARCHS "${CUDA_ARCHS}")
|
||||
@@ -596,8 +590,11 @@ endif()
|
||||
|
||||
if (VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
|
||||
# TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
|
||||
list(APPEND VLLM_EXT_SRC
|
||||
"csrc/custom_quickreduce.cu"
|
||||
"csrc/cuda_view.cu"
|
||||
"csrc/libtorch_stable/cuda_utils_kernels.cu"
|
||||
)
|
||||
# if ROCM endif
|
||||
endif()
|
||||
@@ -632,6 +629,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
"csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
|
||||
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
|
||||
"csrc/libtorch_stable/permute_cols.cu"
|
||||
"csrc/libtorch_stable/quantization/gptq/q_gemm.cu"
|
||||
"csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"
|
||||
"csrc/libtorch_stable/pos_encoding_kernels.cu"
|
||||
@@ -639,6 +637,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
"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"
|
||||
"csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
|
||||
"csrc/libtorch_stable/attention/merge_attn_states.cu"
|
||||
"csrc/libtorch_stable/sampler.cu"
|
||||
"csrc/libtorch_stable/topk.cu"
|
||||
@@ -653,14 +652,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC
|
||||
"csrc/cuda_utils_kernels.cu"
|
||||
"csrc/cutlass_extensions/common.cpp"
|
||||
"csrc/libtorch_stable/cuda_view.cu"
|
||||
"csrc/libtorch_stable/cuda_utils_kernels.cu"
|
||||
"csrc/libtorch_stable/cutlass_extensions/common.cpp"
|
||||
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
|
||||
"csrc/libtorch_stable/permute_cols.cu"
|
||||
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/awq/gemm_kernels.cu"
|
||||
"csrc/libtorch_stable/minimax_reduce_rms_kernel.cu")
|
||||
|
||||
@@ -1075,20 +1072,25 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
USE_SABI 3
|
||||
WITH_SOABI)
|
||||
|
||||
# Set TORCH_TARGET_VERSION for stable ABI compatibility.
|
||||
# This ensures we only use C-shim APIs available in PyTorch 2.10.
|
||||
# _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
|
||||
# which is currently set to 2.10.
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
TORCH_TARGET_VERSION=0x020A000000000000ULL)
|
||||
|
||||
# Needed to use cuda/hip APIs from C-shim
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
# Set TORCH_TARGET_VERSION for stable ABI compatibility.
|
||||
# This ensures we only use C-shim APIs available in PyTorch 2.11.
|
||||
# _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
|
||||
# which is currently set to 2.11.
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
TORCH_TARGET_VERSION=0x020B000000000000ULL)
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
|
||||
# Needed by CUTLASS kernels
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
|
||||
elseif(VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# Set TORCH_TARGET_VERSION for stable ABI compatibility.
|
||||
# This ensures we only use C-shim APIs available in PyTorch 2.10.
|
||||
# _C_stable_libtorch is abi compatible with PyTorch >= TORCH_TARGET_VERSION
|
||||
# which is currently set to 2.10.
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
TORCH_TARGET_VERSION=0x020A000000000000ULL)
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE USE_ROCM)
|
||||
endif()
|
||||
|
||||
@@ -1306,7 +1308,8 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set(VLLM_ROCM_HAS_GFX1100 ON)
|
||||
list(APPEND VLLM_ROCM_EXT_SRC
|
||||
"csrc/rocm/q_gemm_rdna3.cu"
|
||||
"csrc/rocm/q_gemm_rdna3_wmma.cu")
|
||||
"csrc/rocm/q_gemm_rdna3_wmma.cu"
|
||||
"csrc/rocm/moe_q_gemm_rdna3.cu")
|
||||
endif()
|
||||
|
||||
define_extension_target(
|
||||
|
||||
@@ -271,7 +271,6 @@ def benchmark_config(
|
||||
moe_config=moe_config,
|
||||
quant_config=quant_config,
|
||||
),
|
||||
inplace=not disable_inplace(),
|
||||
)
|
||||
|
||||
with override_config(config):
|
||||
@@ -279,7 +278,6 @@ def benchmark_config(
|
||||
x, input_gating, topk, renormalize=not use_deep_gemm
|
||||
)
|
||||
|
||||
inplace = not disable_inplace()
|
||||
if use_deep_gemm:
|
||||
return deep_gemm_experts.apply(
|
||||
x,
|
||||
@@ -298,7 +296,6 @@ def benchmark_config(
|
||||
w2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
inplace=inplace,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
#include <torch/all.h>
|
||||
#include <torch/cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
@@ -1,9 +1,6 @@
|
||||
#ifndef CONCAT_MLA_Q_CUH_
|
||||
#define CONCAT_MLA_Q_CUH_
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
#include "cuda_vec_utils.cuh"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include <torch/csrc/stable/ops.h>
|
||||
#include <torch/csrc/stable/accelerator.h>
|
||||
#include <torch/headeronly/core/ScalarType.h>
|
||||
#include <torch/csrc/stable/device.h>
|
||||
#include <torch/csrc/stable/c/shim.h>
|
||||
#include <torch/headeronly/version.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <array>
|
||||
#include <optional>
|
||||
|
||||
// This function assumes that `cpu_tensor` is a CPU tensor,
|
||||
// and that UVA (Unified Virtual Addressing) is enabled.
|
||||
torch::stable::Tensor get_cuda_view_from_cpu_tensor(
|
||||
torch::stable::Tensor& cpu_tensor) {
|
||||
STD_TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU");
|
||||
|
||||
const auto dtype = cpu_tensor.scalar_type();
|
||||
const auto layout = cpu_tensor.layout();
|
||||
const torch::stable::Device cuda_dev(torch::headeronly::DeviceType::CUDA);
|
||||
|
||||
// handle empty tensor
|
||||
if (cpu_tensor.numel() == 0) {
|
||||
return torch::stable::empty(cpu_tensor.sizes(), dtype, layout, cuda_dev);
|
||||
}
|
||||
|
||||
std::array<StableIValue, 2> is_pinned_stack{
|
||||
torch::stable::detail::from(cpu_tensor),
|
||||
torch::stable::detail::from(std::nullopt)};
|
||||
TORCH_ERROR_CODE_CHECK(torch_call_dispatcher(
|
||||
"aten::is_pinned", "", is_pinned_stack.data(), TORCH_ABI_VERSION));
|
||||
if (torch::stable::detail::to<bool>(is_pinned_stack[0])) {
|
||||
// If CPU tensor is pinned, directly get the device pointer.
|
||||
void* host_ptr = const_cast<void*>(cpu_tensor.mutable_data_ptr());
|
||||
void* device_ptr = nullptr;
|
||||
cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
|
||||
STD_TORCH_CHECK(err == cudaSuccess, "cudaHostGetDevicePointer failed: ",
|
||||
cudaGetErrorString(err));
|
||||
|
||||
return torch::stable::from_blob(
|
||||
device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(), cuda_dev, dtype,
|
||||
[base = cpu_tensor](void*) {}); // keep cpu tensor alive
|
||||
}
|
||||
|
||||
// If CPU tensor is not pinned, allocate a new pinned memory buffer.
|
||||
torch::stable::Tensor contiguous_cpu = torch::stable::contiguous(cpu_tensor);
|
||||
size_t nbytes = contiguous_cpu.numel() * contiguous_cpu.element_size();
|
||||
|
||||
void* host_ptr = nullptr;
|
||||
cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped);
|
||||
if (err != cudaSuccess) {
|
||||
STD_TORCH_CHECK(false, "cudaHostAlloc failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
|
||||
err = cudaMemcpy(host_ptr, contiguous_cpu.const_data_ptr(), nbytes,
|
||||
cudaMemcpyDefault);
|
||||
if (err != cudaSuccess) {
|
||||
cudaFreeHost(host_ptr);
|
||||
STD_TORCH_CHECK(false, "cudaMemcpy failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
|
||||
void* device_ptr = nullptr;
|
||||
err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
|
||||
if (err != cudaSuccess) {
|
||||
cudaFreeHost(host_ptr);
|
||||
STD_TORCH_CHECK(
|
||||
false, "cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
|
||||
auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); };
|
||||
|
||||
return torch::stable::from_blob(device_ptr, contiguous_cpu.sizes(),
|
||||
contiguous_cpu.strides(), cuda_dev,
|
||||
contiguous_cpu.scalar_type(), deleter);
|
||||
}
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "common.hpp"
|
||||
|
||||
int32_t get_sm_version_num() {
|
||||
int32_t major_capability, minor_capability;
|
||||
@@ -102,6 +102,35 @@ constexpr float NUM_TOKEN_CUTOFF = 1024;
|
||||
constexpr int kNumLanes = 32;
|
||||
constexpr int kElemsPerLane = kHeadDim / kNumLanes; // 16
|
||||
|
||||
// Pack this lane's 16 fp32 elements into per-tensor E4M3 FP8 (one uint4 = 16
|
||||
// B), scaling by `scale` (a reciprocal scale) and saturating to ±448. Used by
|
||||
// the FlashInfer full-cache path for both the Q and KV stores.
|
||||
__device__ __forceinline__ uint4 packFp8E4M3x16(float const* values,
|
||||
float const scale) {
|
||||
#ifndef USE_ROCM
|
||||
uint4 out;
|
||||
auto* out2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&out);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane / 2; i++) {
|
||||
float2 scaled =
|
||||
make_float2(values[2 * i] * scale, values[2 * i + 1] * scale);
|
||||
scaled.x = fminf(fmaxf(scaled.x, -kFp8Max), kFp8Max);
|
||||
scaled.y = fminf(fmaxf(scaled.y, -kFp8Max), kFp8Max);
|
||||
out2[i] = __nv_cvt_float2_to_fp8x2(scaled, __NV_SATFINITE, __NV_E4M3);
|
||||
}
|
||||
return out;
|
||||
#else
|
||||
uint8_t out_bytes[kElemsPerLane];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
float scaled = values[i] * scale;
|
||||
scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
|
||||
out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
|
||||
}
|
||||
return *reinterpret_cast<uint4 const*>(out_bytes);
|
||||
#endif
|
||||
}
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Small inline helpers
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
@@ -649,6 +678,257 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
|
||||
#undef DISPATCH
|
||||
}
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// FlashInfer full-cache kernel
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
//
|
||||
// Sibling to the FlashMLA kernel above, used by the FlashInfer V4 sparse-MLA
|
||||
// backend. Differences from the legacy path:
|
||||
// * No Q head padding — output Q layout matches the input num_heads_q.
|
||||
// * KV is written as a *contiguous* 512-wide row per token (token-strided),
|
||||
// not the legacy UE8M0 paged layout with a separate scale tail.
|
||||
// * Q/KV are stored either as bf16 or as per-tensor E4M3 FP8 (one global
|
||||
// scale), selected by the STORE_Q_FP8 / STORE_KV_FP8 template flags.
|
||||
//
|
||||
// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) / warps).
|
||||
// Each warp handles one (token, slot): slot < num_heads_q → Q, slot ==
|
||||
// num_heads_q → KV.
|
||||
template <typename scalar_t_in, bool STORE_Q_FP8, bool STORE_KV_FP8>
|
||||
__global__ void fusedDeepseekV4FullCacheKernel(
|
||||
scalar_t_in* __restrict__ q_inout, // [N, H, 512], in place (bf16)
|
||||
uint8_t* __restrict__ q_fp8_out, // [N, H, 512] fp8, optional
|
||||
int64_t const q_fp8_stride0, // elements (fp8 == bytes)
|
||||
int64_t const q_fp8_stride1, // elements (fp8 == bytes)
|
||||
scalar_t_in const* __restrict__ kv_in, // [N, 512] bf16
|
||||
uint8_t* __restrict__ k_cache, // contiguous bf16 or fp8 cache
|
||||
int64_t const* __restrict__ slot_mapping, // [num_tokens_insert] i64
|
||||
int64_t const* __restrict__ position_ids, // [N] i64
|
||||
float const* __restrict__ cos_sin_cache, // [max_pos, 64] fp32
|
||||
float const* __restrict__ fp8_scale_ptr, // scalar, KV fp8 only
|
||||
float const* __restrict__ q_fp8_scale_inv, // scalar, Q fp8 only
|
||||
float const eps,
|
||||
int const num_tokens_full, // = q.size(0) = kv.size(0)
|
||||
int const num_tokens_insert, // = slot_mapping.size(0)
|
||||
int const num_heads_q, // H (no padding)
|
||||
int const cache_block_size, // tokens per cache block
|
||||
int64_t const kv_block_stride, // bytes per cache block
|
||||
int64_t const kv_token_stride) { // bytes per cache token
|
||||
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
|
||||
if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
|
||||
return;
|
||||
} else {
|
||||
#endif
|
||||
using Converter = vllm::_typeConvert<scalar_t_in>;
|
||||
int const warpsPerBlock = blockDim.x / 32;
|
||||
int const warpId = threadIdx.x / 32;
|
||||
int const laneId = threadIdx.x % 32;
|
||||
int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
|
||||
|
||||
int const slotsPerToken = num_heads_q + 1;
|
||||
int const tokenIdx = globalWarpIdx / slotsPerToken;
|
||||
int const slotIdx = globalWarpIdx % slotsPerToken;
|
||||
if (tokenIdx >= num_tokens_full) return;
|
||||
bool const isKV = (slotIdx == num_heads_q);
|
||||
// KV branch: skip DP-padded tokens (no slot reserved for them).
|
||||
if (isKV && tokenIdx >= num_tokens_insert) return;
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaGridDependencySynchronize();
|
||||
#endif
|
||||
|
||||
int const dim_base = laneId * kElemsPerLane; // in [0, 512) step 16
|
||||
scalar_t_in const* src_ptr;
|
||||
if (isKV) {
|
||||
src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
|
||||
} else {
|
||||
src_ptr = q_inout +
|
||||
(static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) *
|
||||
kHeadDim +
|
||||
dim_base;
|
||||
}
|
||||
uint4 const v0 = *reinterpret_cast<uint4 const*>(src_ptr);
|
||||
uint4 const v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
|
||||
|
||||
// ── Decode bf16 → 16 fp32 registers ───────────────────────────────────
|
||||
float elements[kElemsPerLane];
|
||||
{
|
||||
auto const* p0 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
|
||||
auto const* p1 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
float2 f2 = Converter::convert(p0[i]);
|
||||
elements[2 * i] = f2.x;
|
||||
elements[2 * i + 1] = f2.y;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
float2 f2 = Converter::convert(p1[i]);
|
||||
elements[8 + 2 * i] = f2.x;
|
||||
elements[8 + 2 * i + 1] = f2.y;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Q branch: RMSNorm (no weight) ─────────────────────────────────────
|
||||
if (!isKV) {
|
||||
float sumOfSquares = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
sumOfSquares += elements[i] * elements[i];
|
||||
}
|
||||
sumOfSquares = warpSum<float>(sumOfSquares);
|
||||
float const rms_rcp =
|
||||
rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
elements[i] = elements[i] * rms_rcp;
|
||||
}
|
||||
}
|
||||
|
||||
// ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ───────────────────────────
|
||||
bool const is_rope_lane = dim_base >= kNopeDim;
|
||||
if (is_rope_lane) {
|
||||
int64_t const pos = position_ids[tokenIdx];
|
||||
constexpr int kHalfRope = kRopeDim / 2;
|
||||
float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
|
||||
float const* sin_ptr = cos_ptr + kHalfRope;
|
||||
int const rope_local_base = dim_base - kNopeDim;
|
||||
int const half_base = rope_local_base >> 1;
|
||||
float4 const c0 = *reinterpret_cast<float4 const*>(cos_ptr + half_base);
|
||||
float4 const c1 = *reinterpret_cast<float4 const*>(cos_ptr + half_base + 4);
|
||||
float4 const s0 = *reinterpret_cast<float4 const*>(sin_ptr + half_base);
|
||||
float4 const s1 = *reinterpret_cast<float4 const*>(sin_ptr + half_base + 4);
|
||||
float const cos_arr[8] = {c0.x, c0.y, c0.z, c0.w, c1.x, c1.y, c1.z, c1.w};
|
||||
float const sin_arr[8] = {s0.x, s0.y, s0.z, s0.w, s1.x, s1.y, s1.z, s1.w};
|
||||
#pragma unroll
|
||||
for (int p = 0; p < kElemsPerLane / 2; p++) {
|
||||
float const x_even = elements[2 * p];
|
||||
float const x_odd = elements[2 * p + 1];
|
||||
elements[2 * p] = x_even * cos_arr[p] - x_odd * sin_arr[p];
|
||||
elements[2 * p + 1] = x_even * sin_arr[p] + x_odd * cos_arr[p];
|
||||
}
|
||||
}
|
||||
|
||||
// ── Store ─────────────────────────────────────────────────────────────
|
||||
if (!isKV) {
|
||||
if constexpr (STORE_Q_FP8) {
|
||||
float const scale_inv = VLLM_LDG(q_fp8_scale_inv);
|
||||
uint4 const out = packFp8E4M3x16(elements, scale_inv);
|
||||
uint8_t* dst = q_fp8_out +
|
||||
static_cast<int64_t>(tokenIdx) * q_fp8_stride0 +
|
||||
static_cast<int64_t>(slotIdx) * q_fp8_stride1 + dim_base;
|
||||
*reinterpret_cast<uint4*>(dst) = out;
|
||||
} else {
|
||||
uint4 out0, out1;
|
||||
auto* po0 = reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
|
||||
auto* po1 = reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po0[i] = Converter::convert(
|
||||
make_float2(elements[2 * i], elements[2 * i + 1]));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po1[i] = Converter::convert(
|
||||
make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
|
||||
}
|
||||
scalar_t_in* dst =
|
||||
q_inout +
|
||||
(static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
|
||||
dim_base;
|
||||
*reinterpret_cast<uint4*>(dst) = out0;
|
||||
*reinterpret_cast<uint4*>(dst + 8) = out1;
|
||||
}
|
||||
} else {
|
||||
int64_t const slot_id = slot_mapping[tokenIdx];
|
||||
if (slot_id >= 0) {
|
||||
int64_t const block_idx = slot_id / cache_block_size;
|
||||
int64_t const pos_in_block = slot_id % cache_block_size;
|
||||
uint8_t* cache_row =
|
||||
k_cache + block_idx * kv_block_stride + pos_in_block * kv_token_stride;
|
||||
if constexpr (STORE_KV_FP8) {
|
||||
float const inv_scale = 1.0f / VLLM_LDG(fp8_scale_ptr);
|
||||
uint4 const out = packFp8E4M3x16(elements, inv_scale);
|
||||
*reinterpret_cast<uint4*>(cache_row + dim_base) = out;
|
||||
} else {
|
||||
uint4 out0, out1;
|
||||
auto* po0 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
|
||||
auto* po1 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po0[i] = Converter::convert(
|
||||
make_float2(elements[2 * i], elements[2 * i + 1]));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po1[i] = Converter::convert(
|
||||
make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
|
||||
}
|
||||
scalar_t_in* dst = reinterpret_cast<scalar_t_in*>(cache_row) + dim_base;
|
||||
*reinterpret_cast<uint4*>(dst) = out0;
|
||||
*reinterpret_cast<uint4*>(dst + 8) = out1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// Configure + launch helper shared by the bf16 and fp8 full-cache launchers.
|
||||
template <typename scalar_t_in, bool STORE_Q_FP8, bool STORE_KV_FP8>
|
||||
static void launchFullCacheKernel(
|
||||
scalar_t_in* q_inout, uint8_t* q_fp8_out, int64_t q_fp8_stride0,
|
||||
int64_t q_fp8_stride1, scalar_t_in const* kv_in, uint8_t* k_cache,
|
||||
int64_t const* slot_mapping, int64_t const* position_ids,
|
||||
float const* cos_sin_cache, float const* fp8_scale,
|
||||
float const* q_fp8_scale_inv, float const eps, int const num_tokens_full,
|
||||
int const num_tokens_insert, int const num_heads_q,
|
||||
int const cache_block_size, int64_t const kv_block_stride,
|
||||
int64_t const kv_token_stride, char const* op_name, cudaStream_t stream) {
|
||||
constexpr int kBlockSize = 256;
|
||||
constexpr int kWarpsPerBlock = kBlockSize / 32;
|
||||
int64_t const total_warps =
|
||||
static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
|
||||
int const grid =
|
||||
static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
|
||||
auto* kernel =
|
||||
fusedDeepseekV4FullCacheKernel<scalar_t_in, STORE_Q_FP8, STORE_KV_FP8>;
|
||||
#ifndef USE_ROCM
|
||||
static int const sm_version = getSMVersion();
|
||||
STD_TORCH_CHECK(sm_version >= 80, op_name,
|
||||
" requires sm_80+ (Ampere or newer); got sm_", sm_version);
|
||||
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;
|
||||
cudaLaunchKernelEx(&config, kernel, q_inout, q_fp8_out, q_fp8_stride0,
|
||||
q_fp8_stride1, kv_in, k_cache, slot_mapping, position_ids,
|
||||
cos_sin_cache, fp8_scale, q_fp8_scale_inv, eps,
|
||||
num_tokens_full, num_tokens_insert, num_heads_q,
|
||||
cache_block_size, kv_block_stride, kv_token_stride);
|
||||
#else
|
||||
kernel<<<grid, kBlockSize, 0, stream>>>(
|
||||
q_inout, q_fp8_out, q_fp8_stride0, q_fp8_stride1, kv_in, k_cache,
|
||||
slot_mapping, position_ids, cos_sin_cache, fp8_scale, q_fp8_scale_inv,
|
||||
eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
|
||||
kv_block_stride, kv_token_stride);
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace deepseek_v4_fused_ops
|
||||
} // namespace vllm
|
||||
|
||||
@@ -735,3 +1015,167 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
|
||||
});
|
||||
return q_out;
|
||||
}
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// FlashInfer full-cache torch ops
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
|
||||
torch::stable::Tensor& q, // [N, H, 512] bf16, in place
|
||||
torch::stable::Tensor const& kv, // [N, 512] bf16, read-only
|
||||
torch::stable::Tensor& k_cache, // [num_blocks, bs, 512] bf16
|
||||
torch::stable::Tensor const& slot_mapping, // [num_tokens_insert] int64
|
||||
torch::stable::Tensor const& position_ids, // [N] int64
|
||||
torch::stable::Tensor const& cos_sin_cache, // [max_pos, 64] float32
|
||||
double eps, int64_t cache_block_size) {
|
||||
using torch::headeronly::ScalarType;
|
||||
STD_TORCH_CHECK(q.device().is_cuda() && q.is_contiguous(),
|
||||
"q must be contiguous CUDA");
|
||||
STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
|
||||
"kv must be contiguous CUDA");
|
||||
STD_TORCH_CHECK(k_cache.device().is_cuda(), "k_cache must be CUDA");
|
||||
STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
|
||||
slot_mapping.scalar_type() == ScalarType::Long,
|
||||
"slot_mapping must be int64 CUDA");
|
||||
STD_TORCH_CHECK(position_ids.device().is_cuda() &&
|
||||
position_ids.scalar_type() == ScalarType::Long,
|
||||
"position_ids must be int64 CUDA");
|
||||
STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() &&
|
||||
cos_sin_cache.scalar_type() == ScalarType::Float &&
|
||||
cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
|
||||
"cos_sin_cache shape [max_pos, 64] float32");
|
||||
STD_TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
|
||||
STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
|
||||
STD_TORCH_CHECK(q.scalar_type() == ScalarType::BFloat16 &&
|
||||
kv.scalar_type() == ScalarType::BFloat16,
|
||||
"q and kv must be bfloat16");
|
||||
STD_TORCH_CHECK(k_cache.dim() == 3 && k_cache.size(1) == cache_block_size &&
|
||||
k_cache.size(2) == 512 && k_cache.stride(2) == 1,
|
||||
"k_cache shape [num_blocks, cache_block_size, 512] contiguous");
|
||||
STD_TORCH_CHECK(k_cache.scalar_type() == ScalarType::BFloat16,
|
||||
"k_cache must be bfloat16");
|
||||
|
||||
int const num_tokens_full = static_cast<int>(q.size(0));
|
||||
int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
|
||||
STD_TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
|
||||
static_cast<int>(position_ids.size(0)) == num_tokens_full,
|
||||
"q/kv/position_ids row counts must match");
|
||||
STD_TORCH_CHECK(num_tokens_insert <= num_tokens_full,
|
||||
"slot_mapping must not exceed q row count");
|
||||
int const num_heads_q = static_cast<int>(q.size(1));
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
q.get_device_index());
|
||||
const cudaStream_t stream = get_current_cuda_stream(q.get_device_index());
|
||||
|
||||
// bf16 cache: 2 bytes/element -> byte strides for the uint8-addressed kernel.
|
||||
int64_t const kv_block_stride = k_cache.stride(0) * 2;
|
||||
int64_t const kv_token_stride = k_cache.stride(1) * 2;
|
||||
|
||||
VLLM_STABLE_DISPATCH_HALF_TYPES(
|
||||
q.scalar_type(),
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert", [&] {
|
||||
vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, false,
|
||||
false>(
|
||||
reinterpret_cast<scalar_t*>(q.mutable_data_ptr()), nullptr, 0, 0,
|
||||
reinterpret_cast<scalar_t const*>(kv.const_data_ptr()),
|
||||
reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
|
||||
slot_mapping.const_data_ptr<int64_t>(),
|
||||
position_ids.const_data_ptr<int64_t>(),
|
||||
cos_sin_cache.const_data_ptr<float>(), nullptr, nullptr,
|
||||
static_cast<float>(eps), num_tokens_full, num_tokens_insert,
|
||||
num_heads_q, static_cast<int>(cache_block_size), kv_block_stride,
|
||||
kv_token_stride,
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert",
|
||||
stream);
|
||||
});
|
||||
}
|
||||
|
||||
void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
|
||||
torch::stable::Tensor const& q, // [N, H, 512] bf16, read-only
|
||||
torch::stable::Tensor const& kv, // [N, 512] bf16, read-only
|
||||
torch::stable::Tensor& q_fp8, // [N, H, 512] fp8 e4m3
|
||||
torch::stable::Tensor& k_cache, // [num_blocks, bs, 512] fp8
|
||||
torch::stable::Tensor const& slot_mapping, // [num_tokens_insert] int64
|
||||
torch::stable::Tensor const& position_ids, // [N] int64
|
||||
torch::stable::Tensor const& cos_sin_cache, // [max_pos, 64] float32
|
||||
torch::stable::Tensor const& fp8_scale, // scalar float32 (KV scale)
|
||||
torch::stable::Tensor const& q_fp8_scale_inv, // scalar float32 (1 / Q scale)
|
||||
double eps, int64_t cache_block_size) {
|
||||
using torch::headeronly::ScalarType;
|
||||
STD_TORCH_CHECK(q.device().is_cuda() && q.is_contiguous(),
|
||||
"q must be contiguous CUDA");
|
||||
STD_TORCH_CHECK(kv.device().is_cuda() && kv.is_contiguous(),
|
||||
"kv must be contiguous CUDA");
|
||||
STD_TORCH_CHECK(q_fp8.device().is_cuda() && q_fp8.is_contiguous() &&
|
||||
q_fp8.scalar_type() == ScalarType::Float8_e4m3fn &&
|
||||
q_fp8.dim() == 3 && q_fp8.size(0) == q.size(0) &&
|
||||
q_fp8.size(1) == q.size(1) && q_fp8.size(2) == q.size(2),
|
||||
"q_fp8 must be a contiguous float8_e4m3fn tensor matching q");
|
||||
STD_TORCH_CHECK(k_cache.device().is_cuda(), "k_cache must be CUDA");
|
||||
STD_TORCH_CHECK(slot_mapping.device().is_cuda() &&
|
||||
slot_mapping.scalar_type() == ScalarType::Long,
|
||||
"slot_mapping must be int64 CUDA");
|
||||
STD_TORCH_CHECK(position_ids.device().is_cuda() &&
|
||||
position_ids.scalar_type() == ScalarType::Long,
|
||||
"position_ids must be int64 CUDA");
|
||||
STD_TORCH_CHECK(cos_sin_cache.device().is_cuda() &&
|
||||
cos_sin_cache.scalar_type() == ScalarType::Float &&
|
||||
cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
|
||||
"cos_sin_cache shape [max_pos, 64] float32");
|
||||
STD_TORCH_CHECK(fp8_scale.device().is_cuda() &&
|
||||
fp8_scale.scalar_type() == ScalarType::Float &&
|
||||
fp8_scale.size(0) == 1,
|
||||
"fp8_scale must be a scalar float32 CUDA tensor");
|
||||
STD_TORCH_CHECK(q_fp8_scale_inv.device().is_cuda() &&
|
||||
q_fp8_scale_inv.scalar_type() == ScalarType::Float &&
|
||||
q_fp8_scale_inv.size(0) == 1,
|
||||
"q_fp8_scale_inv must be a scalar float32 CUDA tensor");
|
||||
STD_TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
|
||||
STD_TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
|
||||
STD_TORCH_CHECK(q.scalar_type() == kv.scalar_type(),
|
||||
"q and kv dtype must match");
|
||||
STD_TORCH_CHECK(k_cache.dim() == 3 && k_cache.size(1) == cache_block_size &&
|
||||
k_cache.size(2) == 512 && k_cache.stride(2) == 1,
|
||||
"k_cache shape [num_blocks, cache_block_size, 512] contiguous");
|
||||
STD_TORCH_CHECK(k_cache.scalar_type() == ScalarType::Float8_e4m3fn,
|
||||
"k_cache must be float8_e4m3fn");
|
||||
|
||||
int const num_tokens_full = static_cast<int>(q.size(0));
|
||||
int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
|
||||
STD_TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
|
||||
static_cast<int>(position_ids.size(0)) == num_tokens_full,
|
||||
"q/kv/position_ids row counts must match");
|
||||
STD_TORCH_CHECK(num_tokens_insert <= num_tokens_full,
|
||||
"slot_mapping must not exceed q row count");
|
||||
int const num_heads_q = static_cast<int>(q.size(1));
|
||||
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
q.get_device_index());
|
||||
const cudaStream_t stream = get_current_cuda_stream(q.get_device_index());
|
||||
|
||||
VLLM_STABLE_DISPATCH_HALF_TYPES(
|
||||
q.scalar_type(),
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert", [&] {
|
||||
vllm::deepseek_v4_fused_ops::launchFullCacheKernel<scalar_t, true,
|
||||
true>(
|
||||
// q is read-only in the fp8 path (the kernel writes q_fp8); the
|
||||
// launcher signature is non-const, so cast away const on the ptr.
|
||||
reinterpret_cast<scalar_t*>(
|
||||
const_cast<void*>(q.const_data_ptr())),
|
||||
reinterpret_cast<uint8_t*>(q_fp8.mutable_data_ptr()),
|
||||
q_fp8.stride(0), q_fp8.stride(1),
|
||||
reinterpret_cast<scalar_t const*>(kv.const_data_ptr()),
|
||||
reinterpret_cast<uint8_t*>(k_cache.mutable_data_ptr()),
|
||||
slot_mapping.const_data_ptr<int64_t>(),
|
||||
position_ids.const_data_ptr<int64_t>(),
|
||||
cos_sin_cache.const_data_ptr<float>(),
|
||||
fp8_scale.const_data_ptr<float>(),
|
||||
q_fp8_scale_inv.const_data_ptr<float>(), static_cast<float>(eps),
|
||||
num_tokens_full, num_tokens_insert, num_heads_q,
|
||||
static_cast<int>(cache_block_size),
|
||||
// fp8 cache: 1 byte/element -> stride already in bytes.
|
||||
k_cache.stride(0), k_cache.stride(1),
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
|
||||
stream);
|
||||
});
|
||||
}
|
||||
|
||||
@@ -24,10 +24,10 @@ void per_token_group_quant_int8(const torch::stable::Tensor& input,
|
||||
int64_t group_size, double eps, double int8_min,
|
||||
double int8_max);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
|
||||
torch::stable::Tensor const& perm);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
|
||||
bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
|
||||
bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
|
||||
@@ -162,6 +162,11 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
|
||||
// AllSpark ops: declarations are in the source files
|
||||
// (allspark_repack.cu and allspark_qgemm_w8a16.cu)
|
||||
|
||||
// TODO: Move this out once ROCm upgrade their torch to 2.11.
|
||||
// CPU tensor -> CUDA UVA view (shared CUDA)
|
||||
torch::stable::Tensor get_cuda_view_from_cpu_tensor(
|
||||
torch::stable::Tensor& cpu_tensor);
|
||||
|
||||
#endif
|
||||
|
||||
// Attention kernels (shared CUDA/ROCm)
|
||||
@@ -215,6 +220,13 @@ void rms_norm_per_block_quant(torch::stable::Tensor& out,
|
||||
std::optional<torch::stable::Tensor> residual,
|
||||
int64_t group_size, bool is_scale_transposed);
|
||||
|
||||
void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
|
||||
torch::stable::Tensor const& input,
|
||||
torch::stable::Tensor& scales,
|
||||
int64_t group_size,
|
||||
std::optional<torch::stable::Tensor> scale_ub,
|
||||
bool is_scale_transposed);
|
||||
|
||||
// Positional encoding kernels (shared CUDA/ROCm)
|
||||
void rotary_embedding(torch::stable::Tensor& positions,
|
||||
torch::stable::Tensor& query,
|
||||
@@ -238,6 +250,23 @@ torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
|
||||
torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
|
||||
double eps, int64_t cache_block_size);
|
||||
|
||||
void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
|
||||
torch::stable::Tensor& q, torch::stable::Tensor const& kv,
|
||||
torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
|
||||
torch::stable::Tensor const& position_ids,
|
||||
torch::stable::Tensor const& cos_sin_cache, double eps,
|
||||
int64_t cache_block_size);
|
||||
|
||||
void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
|
||||
torch::stable::Tensor const& q, torch::stable::Tensor const& kv,
|
||||
torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache,
|
||||
torch::stable::Tensor const& slot_mapping,
|
||||
torch::stable::Tensor const& position_ids,
|
||||
torch::stable::Tensor const& cos_sin_cache,
|
||||
torch::stable::Tensor const& fp8_scale,
|
||||
torch::stable::Tensor const& q_fp8_scale_inv, double eps,
|
||||
int64_t cache_block_size);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
torch::stable::Tensor minimax_allreduce_rms(
|
||||
torch::stable::Tensor const& input,
|
||||
|
||||
@@ -18,7 +18,7 @@
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
#include "cutlass_extensions/torch_utils.hpp"
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
#include "get_group_starts.cuh"
|
||||
#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
|
||||
|
||||
@@ -21,7 +21,7 @@
|
||||
#include "cutlass/util/packed_stride.hpp"
|
||||
#include "cutlass/util/mixed_dtype_utils.hpp"
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
@@ -12,7 +12,7 @@
|
||||
|
||||
#include <cutlass/arch/arch.h>
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
|
||||
@@ -20,7 +20,7 @@
|
||||
|
||||
#include <cutlass/arch/arch.h>
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
|
||||
@@ -18,7 +18,7 @@
|
||||
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
#include "nvfp4_utils.cuh"
|
||||
|
||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||
|
||||
@@ -18,7 +18,7 @@
|
||||
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
|
||||
void cutlass_scaled_fp4_mm_sm100a(torch::stable::Tensor& D,
|
||||
|
||||
@@ -18,7 +18,7 @@
|
||||
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
|
||||
@@ -18,7 +18,7 @@
|
||||
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
|
||||
+44
-39
@@ -1,11 +1,10 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include "../../torch_utils.h"
|
||||
|
||||
#include "../../dispatch_utils.h"
|
||||
#include "libtorch_stable/quantization/fused_kernels/quant_conversions.cuh"
|
||||
#include "quant_conversions.cuh"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
@@ -105,64 +104,70 @@ __global__ void silu_and_mul_per_block_quant_kernel(
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
void silu_and_mul_per_block_quant(torch::Tensor& out,
|
||||
torch::Tensor const& input,
|
||||
torch::Tensor& scales, int64_t group_size,
|
||||
std::optional<torch::Tensor> scale_ub,
|
||||
void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
|
||||
torch::stable::Tensor const& input,
|
||||
torch::stable::Tensor& scales,
|
||||
int64_t group_size,
|
||||
std::optional<torch::stable::Tensor> scale_ub,
|
||||
bool is_scale_transposed) {
|
||||
static c10::ScalarType kFp8Type = is_fp8_ocp()
|
||||
? c10::ScalarType::Float8_e4m3fn
|
||||
: c10::ScalarType::Float8_e4m3fnuz;
|
||||
static torch::headeronly::ScalarType kFp8Type =
|
||||
is_fp8_ocp() ? torch::headeronly::ScalarType::Float8_e4m3fn
|
||||
: torch::headeronly::ScalarType::Float8_e4m3fnuz;
|
||||
|
||||
TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
|
||||
TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
|
||||
TORCH_CHECK(
|
||||
input.dtype() == torch::kFloat16 || input.dtype() == torch::kBFloat16,
|
||||
STD_TORCH_CHECK(out.scalar_type() == kFp8Type ||
|
||||
out.scalar_type() == torch::headeronly::ScalarType::Char);
|
||||
STD_TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
|
||||
STD_TORCH_CHECK(
|
||||
input.scalar_type() == torch::headeronly::ScalarType::Half ||
|
||||
input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
|
||||
"Input must be FP16 or BF16");
|
||||
TORCH_CHECK(scales.dtype() == torch::kFloat32, "Scales must be FP32");
|
||||
TORCH_CHECK(group_size == 128 || group_size == 64,
|
||||
"Unsupported group size: ", group_size);
|
||||
STD_TORCH_CHECK(scales.scalar_type() == torch::headeronly::ScalarType::Float);
|
||||
STD_TORCH_CHECK(group_size == 128 || group_size == 64,
|
||||
"Unsupported group size: ", group_size);
|
||||
|
||||
if (scale_ub.has_value()) {
|
||||
TORCH_CHECK(out.dtype() == kFp8Type);
|
||||
STD_TORCH_CHECK(out.scalar_type() == kFp8Type);
|
||||
}
|
||||
|
||||
int32_t hidden_size = out.size(-1);
|
||||
auto num_tokens = input.size(0);
|
||||
int32_t num_groups = hidden_size / group_size;
|
||||
|
||||
TORCH_CHECK(input.size(-1) == hidden_size * 2,
|
||||
"input last dim must be 2x output hidden_size");
|
||||
TORCH_CHECK(hidden_size % group_size == 0,
|
||||
"hidden_size must be divisible by group_size");
|
||||
STD_TORCH_CHECK(input.size(-1) == hidden_size * 2,
|
||||
"input last dim must be 2x output hidden_size");
|
||||
STD_TORCH_CHECK(hidden_size % group_size == 0,
|
||||
"hidden_size must be divisible by group_size");
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
input.get_device_index());
|
||||
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
|
||||
|
||||
dim3 grid(num_tokens, num_groups);
|
||||
dim3 block(group_size);
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
|
||||
input.scalar_type(), "silu_and_mul_per_block_quant", [&] {
|
||||
using scalar_in_t = scalar_t;
|
||||
|
||||
VLLM_DISPATCH_QUANT_TYPES(
|
||||
VLLM_STABLE_DISPATCH_QUANT_TYPES(
|
||||
out.scalar_type(), "silu_and_mul_per_block_quant", [&] {
|
||||
using scalar_out_t = scalar_t;
|
||||
|
||||
VLLM_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
|
||||
VLLM_DISPATCH_BOOL(is_scale_transposed, transpose_scale, [&] {
|
||||
vllm::silu_and_mul_per_block_quant_kernel<
|
||||
scalar_in_t, scalar_out_t, transpose_scale, gs>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
out.data_ptr<scalar_out_t>(),
|
||||
scales.data_ptr<float>(),
|
||||
input.data_ptr<scalar_in_t>(),
|
||||
scale_ub.has_value() ? scale_ub->data_ptr<float>()
|
||||
: nullptr,
|
||||
hidden_size);
|
||||
});
|
||||
VLLM_STABLE_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
|
||||
VLLM_STABLE_DISPATCH_BOOL(
|
||||
is_scale_transposed, transpose_scale, [&] {
|
||||
vllm::silu_and_mul_per_block_quant_kernel<
|
||||
scalar_in_t, scalar_out_t, transpose_scale, gs>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
out.mutable_data_ptr<scalar_out_t>(),
|
||||
scales.mutable_data_ptr<float>(),
|
||||
input.const_data_ptr<scalar_in_t>(),
|
||||
scale_ub.has_value()
|
||||
? scale_ub->const_data_ptr<float>()
|
||||
: nullptr,
|
||||
hidden_size);
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -20,7 +20,7 @@
|
||||
#include "cutlass/util/packed_stride.hpp"
|
||||
|
||||
#include "core/math.hpp"
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
// clang-format on
|
||||
|
||||
namespace vllm::c3x {
|
||||
|
||||
@@ -15,7 +15,7 @@
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "core/math.hpp"
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
// clang-format on
|
||||
|
||||
/*
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include <torch/headeronly/core/ScalarType.h>
|
||||
#include "cuda_utils.h"
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
template <typename Fp8Func, typename Int8Func, typename BlockwiseFunc>
|
||||
void dispatch_scaled_mm(torch::stable::Tensor& c,
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
|
||||
#include <torch/csrc/stable/ops.h>
|
||||
#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
#include "get_group_starts.cuh"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
@@ -23,7 +23,7 @@
|
||||
#include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"
|
||||
|
||||
#include "core/math.hpp"
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
// clang-format on
|
||||
|
||||
using namespace cute;
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
|
||||
void cutlass_scaled_mm_sm75(torch::stable::Tensor& c,
|
||||
torch::stable::Tensor const& a,
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
#include "ops.h"
|
||||
#include "cuda_utils.h"
|
||||
#include "core/registration.h"
|
||||
|
||||
#include <torch/csrc/stable/library.h>
|
||||
@@ -26,9 +27,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! "
|
||||
"output_s, int group_size, float eps, float int8_min, float int8_max) -> "
|
||||
"()");
|
||||
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
|
||||
|
||||
#ifndef USE_ROCM
|
||||
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
|
||||
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
|
||||
#endif
|
||||
|
||||
#ifndef USE_ROCM
|
||||
@@ -321,6 +325,16 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"Tensor? scale_ub, Tensor!? residual, int group_size, "
|
||||
"bool is_scale_transposed) -> ()");
|
||||
|
||||
// Fused SiLU+Mul + per-block quantization
|
||||
ops.def(
|
||||
"silu_and_mul_per_block_quant("
|
||||
"Tensor! out, "
|
||||
"Tensor input, "
|
||||
"Tensor! scales, "
|
||||
"int group_size, "
|
||||
"Tensor? scale_ub=None, "
|
||||
"bool is_scale_transposed=False) -> ()");
|
||||
|
||||
// Rotary embedding
|
||||
// Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
|
||||
ops.def(
|
||||
@@ -343,6 +357,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
|
||||
"int q_head_padded, float eps, int cache_block_size) -> Tensor");
|
||||
|
||||
// FlashInfer V4 full-cache variants: write Q in place (bf16) or to a separate
|
||||
// FP8 tensor, and KV into a contiguous 512-wide token-strided cache.
|
||||
ops.def(
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert("
|
||||
"Tensor! q, Tensor kv, Tensor! k_cache, Tensor slot_mapping, "
|
||||
"Tensor position_ids, Tensor cos_sin_cache, float eps, "
|
||||
"int cache_block_size) -> ()");
|
||||
ops.def(
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert("
|
||||
"Tensor q, Tensor kv, Tensor! q_fp8, Tensor! k_cache, "
|
||||
"Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
|
||||
"Tensor fp8_scale, Tensor q_fp8_scale_inv, float eps, "
|
||||
"int cache_block_size) -> ()");
|
||||
|
||||
#ifndef USE_ROCM
|
||||
ops.def(
|
||||
"minimax_allreduce_rms("
|
||||
@@ -533,9 +561,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
ops.impl("per_token_group_quant_int8",
|
||||
TORCH_BOX(&per_token_group_quant_int8));
|
||||
|
||||
#ifndef USE_ROCM
|
||||
ops.impl("permute_cols", TORCH_BOX(&permute_cols));
|
||||
#endif
|
||||
|
||||
#ifndef USE_ROCM
|
||||
// CUTLASS scaled_mm ops
|
||||
@@ -585,12 +611,20 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
ops.impl("rms_norm_dynamic_per_token_quant",
|
||||
TORCH_BOX(&rms_norm_dynamic_per_token_quant));
|
||||
ops.impl("rms_norm_per_block_quant", TORCH_BOX(&rms_norm_per_block_quant));
|
||||
ops.impl("silu_and_mul_per_block_quant",
|
||||
TORCH_BOX(&silu_and_mul_per_block_quant));
|
||||
|
||||
// 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_deepseek_v4_qnorm_rope_kv_rope_quant_insert",
|
||||
TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert));
|
||||
ops.impl(
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert",
|
||||
TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert));
|
||||
ops.impl(
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert",
|
||||
TORCH_BOX(&fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert));
|
||||
#ifndef USE_ROCM
|
||||
ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms));
|
||||
ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
|
||||
@@ -641,6 +675,28 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2));
|
||||
}
|
||||
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
#ifndef USE_ROCM
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) {
|
||||
ops.impl("get_cuda_view_from_cpu_tensor",
|
||||
TORCH_BOX(&get_cuda_view_from_cpu_tensor));
|
||||
}
|
||||
|
||||
STABLE_TORCH_LIBRARY_FRAGMENT(_C_cuda_utils, cuda_utils) {
|
||||
cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
|
||||
cuda_utils.def(
|
||||
"get_max_shared_memory_per_block_device_attribute(int device_id) -> int");
|
||||
}
|
||||
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
|
||||
cuda_utils) {
|
||||
cuda_utils.impl("get_device_attribute", TORCH_BOX(&get_device_attribute));
|
||||
cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
|
||||
TORCH_BOX(&get_max_shared_memory_per_block_device_attribute));
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
// These capability-check functions take only primitive args (no tensors), so
|
||||
// there is no device to dispatch on. CompositeExplicitAutograd makes them
|
||||
// available for all backends. This is the stable ABI equivalent of calling
|
||||
|
||||
+3
-8
@@ -40,12 +40,6 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
|
||||
void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
|
||||
torch::Tensor& weight, double epsilon);
|
||||
|
||||
void silu_and_mul_per_block_quant(torch::Tensor& out,
|
||||
torch::Tensor const& input,
|
||||
torch::Tensor& scales, int64_t group_size,
|
||||
std::optional<torch::Tensor> scale_ub,
|
||||
bool is_scale_transposed);
|
||||
|
||||
// rotary_embedding also exist in csrc/libtorch_stable/ops.h (torch::stable
|
||||
// ABI for CUDA). It remains here because the CPU build still uses these
|
||||
// torch::Tensor declarations.
|
||||
@@ -84,8 +78,6 @@ void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
|
||||
torch::Tensor const& seq_lens,
|
||||
torch::Tensor const& page_table, double scale);
|
||||
|
||||
torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
|
||||
|
||||
void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
|
||||
torch::Tensor const& scale,
|
||||
std::optional<torch::Tensor> const& azp);
|
||||
@@ -110,4 +102,7 @@ void qr_open_handles(fptr_t _fa, const std::vector<torch::Tensor>& handles);
|
||||
void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
|
||||
int64_t quant_level, bool cast_bf2half = false);
|
||||
int64_t qr_max_size();
|
||||
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
|
||||
#endif
|
||||
|
||||
@@ -0,0 +1,639 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Fused MoE W4A16 GPTQ kernel for RDNA3 (gfx1100).
|
||||
//
|
||||
// Combines expert routing (sorted_token_ids / expert_ids) with the RDNA3
|
||||
// W4A16 dequant+dot from q_gemm_rdna3.cu into a single kernel launch.
|
||||
// Each block processes BLOCK_SIZE_M tokens assigned to one expert, covering
|
||||
// a tile of N output columns and K input positions.
|
||||
//
|
||||
// Weight format: same as the dense kernel — [E, K/8, N] uint32 shuffled,
|
||||
// [E, groups, N] scales, [E, groups, N/8] packed zeros.
|
||||
//
|
||||
// Design: THREADS_X=256 (8 waves on wave32), BLOCK_KN_SIZE=256, each thread
|
||||
// handles 4 N columns. Output via 64-bit packed CAS atomic-add directly to
|
||||
// the pre-zeroed output tensor (no FP32 scratch buffer).
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
#include <torch/all.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <hip/hip_bf16.h>
|
||||
#include <hip/hip_fp16.h>
|
||||
|
||||
#include "qdq_4_rdna3.cuh"
|
||||
|
||||
#if defined(__HIPCC__) && defined(__gfx1100__)
|
||||
#define __HIP__RDNA3__
|
||||
#endif
|
||||
|
||||
namespace vllm {
|
||||
namespace moe_gptq_rdna3 {
|
||||
|
||||
#define BLOCK_KN_SIZE 256
|
||||
#define THREADS_X 256
|
||||
|
||||
#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
|
||||
|
||||
using gptq_rdna3::bf162_t;
|
||||
using gptq_rdna3::bf16_t;
|
||||
|
||||
// --- Helpers (same as q_gemm_rdna3.cu) ---
|
||||
|
||||
template <typename T>
|
||||
__forceinline__ __device__ T tzero();
|
||||
|
||||
template <>
|
||||
__forceinline__ __device__ half tzero<half>() {
|
||||
return __float2half_rn(0.0f);
|
||||
}
|
||||
|
||||
template <>
|
||||
__forceinline__ __device__ bf16_t tzero<bf16_t>() {
|
||||
return __float2bfloat16(0.0f);
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) {
|
||||
float result = 0.0f;
|
||||
const half2* a2_ptr = (const half2*)a_ptr;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ float dot22_8_f(float (&dq)[8],
|
||||
const bf16_t* a_ptr) {
|
||||
float result = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
uint32_t aw;
|
||||
__builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
|
||||
float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
|
||||
float a_y = __uint_as_float(aw & 0xFFFF0000u);
|
||||
result = __fmaf_rn(dq[2 * i + 0], a_x, result);
|
||||
result = __fmaf_rn(dq[2 * i + 1], a_y, result);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01,
|
||||
half2 v23) {
|
||||
unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
|
||||
unsigned long long old = *addr_u;
|
||||
while (true) {
|
||||
union {
|
||||
unsigned long long u;
|
||||
half2 h2[2];
|
||||
} cur, sum;
|
||||
cur.u = old;
|
||||
sum.h2[0] = __hadd2(cur.h2[0], v01);
|
||||
sum.h2[1] = __hadd2(cur.h2[1], v23);
|
||||
unsigned long long prev = atomicCAS(addr_u, old, sum.u);
|
||||
if (prev == old) break;
|
||||
old = prev;
|
||||
}
|
||||
}
|
||||
|
||||
__forceinline__ __device__ void atomic_add_pk4_bf16(bf16_t* addr, bf162_t v01,
|
||||
bf162_t v23) {
|
||||
unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
|
||||
unsigned long long old = *addr_u;
|
||||
while (true) {
|
||||
union {
|
||||
unsigned long long u;
|
||||
bf162_t b2[2];
|
||||
} cur, sum;
|
||||
cur.u = old;
|
||||
sum.b2[0] = __hadd2(cur.b2[0], v01);
|
||||
sum.b2[1] = __hadd2(cur.b2[1], v23);
|
||||
unsigned long long prev = atomicCAS(addr_u, old, sum.u);
|
||||
if (prev == old) break;
|
||||
old = prev;
|
||||
}
|
||||
}
|
||||
|
||||
__forceinline__ __device__ void load4_zeros(const uint32_t* qzeros_row, int n,
|
||||
int (&zeros)[4]) {
|
||||
int qcol = n / 8;
|
||||
int shift = (n & 0x07) * 4;
|
||||
uint32_t d = qzeros_row[qcol] >> shift;
|
||||
zeros[0] = (int)(d & 0xF);
|
||||
zeros[1] = (int)((d >> 4) & 0xF);
|
||||
zeros[2] = (int)((d >> 8) & 0xF);
|
||||
zeros[3] = (int)((d >> 12) & 0xF);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__forceinline__ __device__ void load4_scales(const T* scales_row, int n,
|
||||
T (&scales)[4]) {
|
||||
scales[0] = scales_row[n + 0];
|
||||
scales[1] = scales_row[n + 1];
|
||||
scales[2] = scales_row[n + 2];
|
||||
scales[3] = scales_row[n + 3];
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Fused MoE kernel.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
template <typename T, int BLOCK_SIZE_M>
|
||||
__global__ void moe_gemm_q4_kernel_rdna3(
|
||||
const T* __restrict__ a, // [size_m, size_k] or [M*topk, K]
|
||||
T* __restrict__ c, // [M*topk, size_n] pre-zeroed
|
||||
const uint32_t* __restrict__ b_q_weight, // [E, K/8, N] packed
|
||||
const T* __restrict__ b_scales, // [E, groups, N]
|
||||
const uint32_t* __restrict__ b_qzeros, // [E, groups, N/8] packed
|
||||
const float* __restrict__ topk_weights, // [M*topk] or nullptr
|
||||
const int32_t* __restrict__ sorted_token_ids,
|
||||
const int32_t* __restrict__ expert_ids,
|
||||
const int32_t* __restrict__ num_tokens_post_padded,
|
||||
const int size_m, // total tokens (original M, or M*topk for w2)
|
||||
const int size_n, // output features per expert
|
||||
const int size_k, // input features
|
||||
const int groups, // K / group_size
|
||||
const int top_k, // routing top-k (1 for w2 pass)
|
||||
// Per-expert strides (in elements, not bytes)
|
||||
const int expert_weight_stride, // (K/8) * N
|
||||
const int expert_scales_stride, // groups * N
|
||||
const int expert_zeros_stride, // groups * (N/8)
|
||||
const bool mul_topk_weight,
|
||||
const int output_topk) { // >0: reduce output by token_id/output_topk
|
||||
const int t = threadIdx.x;
|
||||
const int token_block = blockIdx.x;
|
||||
const int offset_n = blockIdx.y * BLOCK_KN_SIZE * 4;
|
||||
const int offset_k = blockIdx.z * BLOCK_KN_SIZE;
|
||||
const int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
|
||||
const int n = offset_n + t * 4;
|
||||
|
||||
// Early exit for padding blocks or invalid experts (expert_map = -1)
|
||||
if (token_block * BLOCK_SIZE_M >= num_tokens_post_padded[0]) return;
|
||||
|
||||
const int expert_id = expert_ids[token_block];
|
||||
if (expert_id == -1) return;
|
||||
|
||||
// Expert-specific pointers
|
||||
const uint32_t* expert_weights =
|
||||
b_q_weight + (int64_t)expert_id * expert_weight_stride;
|
||||
const T* expert_scales = b_scales + (int64_t)expert_id * expert_scales_stride;
|
||||
const uint32_t* expert_qzeros =
|
||||
b_qzeros + (int64_t)expert_id * expert_zeros_stride;
|
||||
|
||||
// LDS for activations
|
||||
constexpr int LDS_PAD = 8;
|
||||
__shared__ T block_a[BLOCK_SIZE_M][BLOCK_KN_SIZE + LDS_PAD];
|
||||
|
||||
static_assert(BLOCK_KN_SIZE == THREADS_X,
|
||||
"BLOCK_KN_SIZE must equal THREADS_X");
|
||||
|
||||
// For bf16 M=1, we can skip LDS and read A from global (same as dense).
|
||||
// fp16 always needs LDS due to the dot22_8_f indexing pattern.
|
||||
constexpr bool USE_LDS_A = (BLOCK_SIZE_M > 1) || std::is_same<T, half>::value;
|
||||
|
||||
const int offset_m_base = token_block * BLOCK_SIZE_M;
|
||||
|
||||
if constexpr (USE_LDS_A) {
|
||||
if (offset_k + t < end_k) {
|
||||
#pragma unroll
|
||||
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
|
||||
int32_t token_id = sorted_token_ids[offset_m_base + m];
|
||||
int token_row = token_id / top_k;
|
||||
T av;
|
||||
if (token_row < size_m) {
|
||||
av = a[(int64_t)token_row * size_k + offset_k + t];
|
||||
} else {
|
||||
av = tzero<T>();
|
||||
}
|
||||
block_a[m][t] = av;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (n >= size_n) return;
|
||||
|
||||
// Group bookkeeping
|
||||
const int groupsize = size_k / groups;
|
||||
int group = offset_k / groupsize;
|
||||
int nextgroup = (group + 1) * groupsize;
|
||||
|
||||
// Weight pointer for this expert
|
||||
int qk = offset_k / 8;
|
||||
const uint32_t* b_ptr = expert_weights + qk * size_n + n;
|
||||
|
||||
// Per-column dequant constants (4 columns per thread)
|
||||
half2 z1z16_h[4][2], y1y16_h[4][2];
|
||||
float z_b_f[4], y_b_f[4];
|
||||
|
||||
// GPTQv1: zero_offset = 1
|
||||
constexpr int zero_offset = 1;
|
||||
|
||||
auto refresh_group = [&](int g) {
|
||||
const uint32_t* qz_row = expert_qzeros + g * (size_n / 8);
|
||||
const T* sc_row = expert_scales + g * size_n;
|
||||
int zeros[4];
|
||||
T scales[4];
|
||||
load4_zeros(qz_row, n, zeros);
|
||||
load4_scales<T>(sc_row, n, scales);
|
||||
if constexpr (std::is_same<T, half>::value) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
gptq_rdna3::prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset),
|
||||
scales[i], z1z16_h[i], y1y16_h[i]);
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
gptq_rdna3::prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset),
|
||||
scales[i], z_b_f[i], y_b_f[i]);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
refresh_group(group);
|
||||
|
||||
float block_c[BLOCK_SIZE_M][4];
|
||||
#pragma unroll
|
||||
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j) block_c[m][j] = 0.0f;
|
||||
}
|
||||
|
||||
// --- Main K-loop ---
|
||||
int k = offset_k;
|
||||
while (k < end_k) {
|
||||
if (k == nextgroup) {
|
||||
group++;
|
||||
nextgroup += groupsize;
|
||||
refresh_group(group);
|
||||
}
|
||||
|
||||
// Prefetch 4 weight words (128 bytes)
|
||||
int4 b_w[4];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
b_w[j] = *(const int4*)(b_ptr + j * size_n);
|
||||
}
|
||||
b_ptr += 4 * size_n;
|
||||
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
const int a_off = (k - offset_k) + 8 * j;
|
||||
|
||||
if constexpr (std::is_same<T, half>::value) {
|
||||
// fp16 path: dequant via bit-trick, dot via v_dot2_f32_f16
|
||||
half2 dq[4][4];
|
||||
gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0],
|
||||
y1y16_h[0]);
|
||||
gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1],
|
||||
y1y16_h[1]);
|
||||
gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2],
|
||||
y1y16_h[2]);
|
||||
gptq_rdna3::dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3],
|
||||
y1y16_h[3]);
|
||||
|
||||
#pragma unroll
|
||||
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
|
||||
const half* a_ptr = reinterpret_cast<const half*>(&block_a[m][a_off]);
|
||||
block_c[m][0] += dot22_8_f(dq[0], a_ptr);
|
||||
block_c[m][1] += dot22_8_f(dq[1], a_ptr);
|
||||
block_c[m][2] += dot22_8_f(dq[2], a_ptr);
|
||||
block_c[m][3] += dot22_8_f(dq[3], a_ptr);
|
||||
}
|
||||
} else if constexpr (BLOCK_SIZE_M == 1) {
|
||||
// bf16 M=1: v_dot2_f32_bf16 with InstCombine-defeating opacity
|
||||
typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
|
||||
constexpr uint32_t BF16_MAGIC = 0x43004300u;
|
||||
constexpr uint32_t BF16_ONES = 0x3F803F80u;
|
||||
union pack4 {
|
||||
float f[4];
|
||||
uint32_t u[4];
|
||||
};
|
||||
|
||||
uint32_t w[4];
|
||||
__builtin_memcpy(w, &b_w[j], sizeof(int4));
|
||||
|
||||
// Load activations — read from global (no LDS for bf16 M=1)
|
||||
pack4 a_pack;
|
||||
{
|
||||
int32_t token_id = sorted_token_ids[offset_m_base];
|
||||
int token_row = token_id / top_k;
|
||||
if (token_row < size_m) {
|
||||
const uint32_t* a_words = reinterpret_cast<const uint32_t*>(
|
||||
a + (int64_t)token_row * size_k + offset_k + a_off);
|
||||
a_pack.u[0] = a_words[0];
|
||||
a_pack.u[1] = a_words[1];
|
||||
a_pack.u[2] = a_words[2];
|
||||
a_pack.u[3] = a_words[3];
|
||||
} else {
|
||||
a_pack.u[0] = 0;
|
||||
a_pack.u[1] = 0;
|
||||
a_pack.u[2] = 0;
|
||||
a_pack.u[3] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
// sum_a for bias correction
|
||||
float sum_a = 0.0f;
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; ++b) {
|
||||
sum_a = __builtin_amdgcn_fdot2_f32_bf16(
|
||||
*((bf16x2_t*)(&a_pack.f[b])), *((const bf16x2_t*)&BF16_ONES),
|
||||
sum_a, /*clamp=*/false);
|
||||
}
|
||||
|
||||
#pragma unroll 1
|
||||
for (int col = 0; col < 4; ++col) {
|
||||
pack4 q_pack;
|
||||
const uint32_t qa = w[col];
|
||||
q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
|
||||
q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
|
||||
q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
|
||||
q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
|
||||
|
||||
float partial = 0.0f;
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; ++b) {
|
||||
partial = __builtin_amdgcn_fdot2_f32_bf16(
|
||||
*((bf16x2_t*)(&a_pack.f[b])), *((bf16x2_t*)(&q_pack.f[b])),
|
||||
partial, /*clamp=*/false);
|
||||
}
|
||||
|
||||
block_c[0][col] =
|
||||
__fmaf_rn(y_b_f[col], partial,
|
||||
__fmaf_rn(z_b_f[col], sum_a, block_c[0][col]));
|
||||
}
|
||||
} else {
|
||||
// bf16 M>1: v_dot2_f32_bf16 with LDS-staged activations
|
||||
typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
|
||||
constexpr uint32_t BF16_MAGIC = 0x43004300u;
|
||||
constexpr uint32_t BF16_ONES = 0x3F803F80u;
|
||||
union pack4 {
|
||||
float f[4];
|
||||
uint32_t u[4];
|
||||
};
|
||||
|
||||
uint32_t w[4];
|
||||
__builtin_memcpy(w, &b_w[j], sizeof(int4));
|
||||
|
||||
pack4 a_pack[BLOCK_SIZE_M];
|
||||
#pragma unroll
|
||||
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
|
||||
const uint32_t* a_words =
|
||||
reinterpret_cast<const uint32_t*>(&block_a[m][a_off]);
|
||||
a_pack[m].u[0] = a_words[0];
|
||||
a_pack[m].u[1] = a_words[1];
|
||||
a_pack[m].u[2] = a_words[2];
|
||||
a_pack[m].u[3] = a_words[3];
|
||||
}
|
||||
|
||||
float sum_a[BLOCK_SIZE_M];
|
||||
#pragma unroll
|
||||
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
|
||||
float s = 0.0f;
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; ++b) {
|
||||
s = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&a_pack[m].f[b])),
|
||||
*((const bf16x2_t*)&BF16_ONES),
|
||||
s, /*clamp=*/false);
|
||||
}
|
||||
sum_a[m] = s;
|
||||
}
|
||||
|
||||
#pragma unroll 1
|
||||
for (int col = 0; col < 4; ++col) {
|
||||
pack4 q_pack;
|
||||
const uint32_t qa = w[col];
|
||||
q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
|
||||
q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
|
||||
q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
|
||||
q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
|
||||
|
||||
#pragma unroll
|
||||
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
|
||||
float partial = 0.0f;
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; ++b) {
|
||||
partial = __builtin_amdgcn_fdot2_f32_bf16(
|
||||
*((bf16x2_t*)(&a_pack[m].f[b])), *((bf16x2_t*)(&q_pack.f[b])),
|
||||
partial, /*clamp=*/false);
|
||||
}
|
||||
block_c[m][col] =
|
||||
__fmaf_rn(y_b_f[col], partial,
|
||||
__fmaf_rn(z_b_f[col], sum_a[m], block_c[m][col]));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
k += 32;
|
||||
}
|
||||
|
||||
// --- Epilogue: apply topk_weight and atomic-add to output ---
|
||||
#pragma unroll
|
||||
for (int m = 0; m < BLOCK_SIZE_M; ++m) {
|
||||
int32_t token_id = sorted_token_ids[offset_m_base + m];
|
||||
if (token_id / top_k >= size_m) continue;
|
||||
|
||||
// Apply router weight
|
||||
if (mul_topk_weight && topk_weights != nullptr) {
|
||||
float tw = topk_weights[token_id];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j) block_c[m][j] *= tw;
|
||||
}
|
||||
|
||||
// output_topk > 0: reduce by mapping token_id back to original token
|
||||
// (multiple experts write to the same row via atomics)
|
||||
int64_t out_row = (output_topk > 0) ? (int64_t)(token_id / output_topk)
|
||||
: (int64_t)token_id;
|
||||
T* out = c + out_row * size_n + n;
|
||||
if constexpr (std::is_same<T, half>::value) {
|
||||
half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]),
|
||||
__float2half_rn(block_c[m][1]));
|
||||
half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]),
|
||||
__float2half_rn(block_c[m][3]));
|
||||
atomic_add_pk4_f16(out, r01, r23);
|
||||
} else {
|
||||
bf162_t r01;
|
||||
r01.x = __float2bfloat16(block_c[m][0]);
|
||||
r01.y = __float2bfloat16(block_c[m][1]);
|
||||
bf162_t r23;
|
||||
r23.x = __float2bfloat16(block_c[m][2]);
|
||||
r23.y = __float2bfloat16(block_c[m][3]);
|
||||
atomic_add_pk4_bf16(out, r01, r23);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#else // non-RDNA3: empty stub for symbol parity
|
||||
|
||||
template <typename T, int BLOCK_SIZE_M>
|
||||
__global__ void moe_gemm_q4_kernel_rdna3(
|
||||
const T*, T*, const uint32_t*, const T*, const uint32_t*, const float*,
|
||||
const int32_t*, const int32_t*, const int32_t*, const int, const int,
|
||||
const int, const int, const int, const int, const int, const int,
|
||||
const bool, const int) {}
|
||||
|
||||
#endif // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Launcher
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
template <typename T, int BLOCK_SIZE_M>
|
||||
void launch_moe_gemm_q4(
|
||||
const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales,
|
||||
const uint32_t* b_qzeros, const float* topk_weights,
|
||||
const int32_t* sorted_token_ids, const int32_t* expert_ids,
|
||||
const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m,
|
||||
int size_n, int size_k, int groups, int top_k, int expert_weight_stride,
|
||||
int expert_scales_stride, int expert_zeros_stride, bool mul_topk_weight,
|
||||
int output_topk, cudaStream_t stream) {
|
||||
dim3 block(THREADS_X);
|
||||
dim3 grid(num_token_blocks,
|
||||
(size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4),
|
||||
(size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE);
|
||||
|
||||
moe_gemm_q4_kernel_rdna3<T, BLOCK_SIZE_M><<<grid, block, 0, stream>>>(
|
||||
a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
|
||||
expert_ids, num_tokens_post_padded, size_m, size_n, size_k, groups, top_k,
|
||||
expert_weight_stride, expert_scales_stride, expert_zeros_stride,
|
||||
mul_topk_weight, output_topk);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void dispatch_moe_gemm_q4(
|
||||
const T* a, T* c, const uint32_t* b_q_weight, const T* b_scales,
|
||||
const uint32_t* b_qzeros, const float* topk_weights,
|
||||
const int32_t* sorted_token_ids, const int32_t* expert_ids,
|
||||
const int32_t* num_tokens_post_padded, int num_token_blocks, int size_m,
|
||||
int size_n, int size_k, int groups, int top_k, int block_size_m,
|
||||
int expert_weight_stride, int expert_scales_stride, int expert_zeros_stride,
|
||||
bool mul_topk_weight, int output_topk, cudaStream_t stream) {
|
||||
// Dispatch to template instantiation based on block_size_m
|
||||
switch (block_size_m) {
|
||||
case 1:
|
||||
launch_moe_gemm_q4<T, 1>(
|
||||
a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
|
||||
expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
|
||||
size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
|
||||
expert_zeros_stride, mul_topk_weight, output_topk, stream);
|
||||
break;
|
||||
case 2:
|
||||
launch_moe_gemm_q4<T, 2>(
|
||||
a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
|
||||
expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
|
||||
size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
|
||||
expert_zeros_stride, mul_topk_weight, output_topk, stream);
|
||||
break;
|
||||
case 4:
|
||||
launch_moe_gemm_q4<T, 4>(
|
||||
a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
|
||||
expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
|
||||
size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
|
||||
expert_zeros_stride, mul_topk_weight, output_topk, stream);
|
||||
break;
|
||||
case 8:
|
||||
launch_moe_gemm_q4<T, 8>(
|
||||
a, c, b_q_weight, b_scales, b_qzeros, topk_weights, sorted_token_ids,
|
||||
expert_ids, num_tokens_post_padded, num_token_blocks, size_m, size_n,
|
||||
size_k, groups, top_k, expert_weight_stride, expert_scales_stride,
|
||||
expert_zeros_stride, mul_topk_weight, output_topk, stream);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false,
|
||||
"moe_gptq_gemm_rdna3: block_size_m must be 1, 2, 4, or 8, "
|
||||
"got ",
|
||||
block_size_m);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace moe_gptq_rdna3
|
||||
} // namespace vllm
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Public entry point
|
||||
// ---------------------------------------------------------------------------
|
||||
//
|
||||
// Inputs:
|
||||
// a [M, K] or [M*top_k, K] half or bfloat16
|
||||
// c [M*top_k, N] same dtype (pre-zeroed!)
|
||||
// b_q_weight [E, K/8, N] uint32 (shuffled)
|
||||
// b_scales [E, groups, N] same dtype as a
|
||||
// b_qzeros [E, groups, N/8] uint32 (packed 4-bit)
|
||||
// topk_weights [M*top_k] or empty float32
|
||||
// sorted_token_ids [num_blocks * block_m] int32
|
||||
// expert_ids [num_blocks] int32
|
||||
// num_tokens_post_padded [1] int32
|
||||
// top_k int
|
||||
// block_size_m int (1, 2, 4, or 8)
|
||||
// mul_topk_weight bool
|
||||
|
||||
void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c,
|
||||
torch::Tensor b_q_weight, torch::Tensor b_scales,
|
||||
torch::Tensor b_qzeros, torch::Tensor topk_weights,
|
||||
torch::Tensor sorted_token_ids,
|
||||
torch::Tensor expert_ids,
|
||||
torch::Tensor num_tokens_post_padded, int64_t top_k,
|
||||
int64_t block_size_m, bool mul_topk_weight,
|
||||
int64_t output_topk) {
|
||||
TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor");
|
||||
TORCH_CHECK(c.is_cuda(), "c must be a CUDA/HIP tensor");
|
||||
TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor");
|
||||
TORCH_CHECK(a.dim() == 2, "a must be 2D");
|
||||
TORCH_CHECK(c.dim() == 2, "c must be 2D");
|
||||
TORCH_CHECK(b_q_weight.dim() == 3, "b_q_weight must be 3D [E, K/8, N]");
|
||||
TORCH_CHECK(b_scales.dim() == 3, "b_scales must be 3D [E, groups, N]");
|
||||
TORCH_CHECK(b_qzeros.dim() == 3, "b_qzeros must be 3D [E, groups, N/8]");
|
||||
TORCH_CHECK(
|
||||
a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16,
|
||||
"a must be half or bfloat16");
|
||||
TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(),
|
||||
"b_scales dtype must match a");
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
int size_m = (int)a.size(0);
|
||||
int size_k = (int)a.size(1);
|
||||
int size_n = (int)b_q_weight.size(2);
|
||||
int groups = (int)b_scales.size(1);
|
||||
|
||||
// Per-expert strides
|
||||
int expert_weight_stride = (int)(b_q_weight.size(1) * b_q_weight.size(2));
|
||||
int expert_scales_stride = (int)(b_scales.size(1) * b_scales.size(2));
|
||||
int expert_zeros_stride = (int)(b_qzeros.size(1) * b_qzeros.size(2));
|
||||
|
||||
int num_token_blocks = (int)(sorted_token_ids.size(0) / block_size_m);
|
||||
|
||||
const float* topk_w_ptr =
|
||||
(topk_weights.numel() > 0) ? topk_weights.data_ptr<float>() : nullptr;
|
||||
|
||||
// Manual dtype dispatch using HIP native types (c10::Half/BFloat16 don't
|
||||
// implicitly convert to half/__hip_bfloat16 in device code).
|
||||
using vllm::gptq_rdna3::bf16_t;
|
||||
|
||||
auto dispatch = [&](auto* a_ptr, auto* c_ptr, const auto* s_ptr) {
|
||||
using T = std::remove_const_t<std::remove_pointer_t<decltype(a_ptr)>>;
|
||||
vllm::moe_gptq_rdna3::dispatch_moe_gemm_q4<T>(
|
||||
a_ptr, c_ptr, (const uint32_t*)b_q_weight.data_ptr<int32_t>(), s_ptr,
|
||||
(const uint32_t*)b_qzeros.data_ptr<int32_t>(), topk_w_ptr,
|
||||
sorted_token_ids.data_ptr<int32_t>(), expert_ids.data_ptr<int32_t>(),
|
||||
num_tokens_post_padded.data_ptr<int32_t>(), num_token_blocks, size_m,
|
||||
size_n, size_k, groups, (int)top_k, (int)block_size_m,
|
||||
expert_weight_stride, expert_scales_stride, expert_zeros_stride,
|
||||
mul_topk_weight, (int)output_topk, stream);
|
||||
};
|
||||
|
||||
if (a.scalar_type() == torch::kHalf) {
|
||||
dispatch((const half*)a.data_ptr(), (half*)c.data_ptr(),
|
||||
(const half*)b_scales.data_ptr());
|
||||
} else {
|
||||
dispatch((const bf16_t*)a.data_ptr(), (bf16_t*)c.data_ptr(),
|
||||
(const bf16_t*)b_scales.data_ptr());
|
||||
}
|
||||
}
|
||||
@@ -27,6 +27,15 @@ torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
|
||||
torch::Tensor b_scales,
|
||||
torch::Tensor b_g_idx, bool use_v2_format);
|
||||
|
||||
void moe_gptq_gemm_rdna3(torch::Tensor a, torch::Tensor c,
|
||||
torch::Tensor b_q_weight, torch::Tensor b_scales,
|
||||
torch::Tensor b_qzeros, torch::Tensor topk_weights,
|
||||
torch::Tensor sorted_token_ids,
|
||||
torch::Tensor expert_ids,
|
||||
torch::Tensor num_tokens_post_padded, int64_t top_k,
|
||||
int64_t block_size_m, bool mul_topk_weight,
|
||||
int64_t output_topk);
|
||||
|
||||
void paged_attention(
|
||||
torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
|
||||
torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
|
||||
|
||||
@@ -50,6 +50,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
|
||||
"gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
|
||||
"Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor");
|
||||
rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma);
|
||||
|
||||
rocm_ops.def(
|
||||
"moe_gptq_gemm_rdna3(Tensor a, Tensor! c, Tensor b_q_weight, "
|
||||
"Tensor b_scales, Tensor b_qzeros, Tensor topk_weights, "
|
||||
"Tensor sorted_token_ids, Tensor expert_ids, "
|
||||
"Tensor num_tokens_post_padded, "
|
||||
"int top_k, int block_size_m, bool mul_topk_weight, "
|
||||
"int output_topk) -> ()");
|
||||
rocm_ops.impl("moe_gptq_gemm_rdna3", torch::kCUDA, &moe_gptq_gemm_rdna3);
|
||||
#endif
|
||||
|
||||
// Custom attention op
|
||||
|
||||
+9
-15
@@ -32,30 +32,24 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
|
||||
ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
|
||||
|
||||
#ifdef USE_ROCM
|
||||
// TODO: Remove this once we upgrade to torch 2.11.
|
||||
// ROCm still uses torch 2.10,
|
||||
// So we still need to use unstable torch ABI for now.
|
||||
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
|
||||
ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
|
||||
&get_cuda_view_from_cpu_tensor);
|
||||
#endif
|
||||
|
||||
// Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
|
||||
ops.def(
|
||||
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
|
||||
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
|
||||
|
||||
// Fused SiLU+Mul + per-block quantization
|
||||
ops.def(
|
||||
"silu_and_mul_per_block_quant("
|
||||
"Tensor! out, "
|
||||
"Tensor input, "
|
||||
"Tensor! scales, "
|
||||
"int group_size, "
|
||||
"Tensor? scale_ub=None, "
|
||||
"bool is_scale_transposed=False) -> ()");
|
||||
ops.impl("silu_and_mul_per_block_quant", torch::kCUDA,
|
||||
&silu_and_mul_per_block_quant);
|
||||
|
||||
// Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
|
||||
// GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
|
||||
// kernel launch. Registered in _C_stable_libtorch.
|
||||
// kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
|
||||
// full-cache bf16/fp8 variants).
|
||||
|
||||
// Quantization ops
|
||||
#ifndef USE_ROCM
|
||||
@@ -175,11 +169,10 @@ TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
|
||||
|
||||
custom_ar.def("qr_max_size", &qr_max_size);
|
||||
}
|
||||
#endif
|
||||
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
|
||||
// Cuda utils
|
||||
|
||||
// Gets the specified device attribute.
|
||||
cuda_utils.def("get_device_attribute(int attribute, int device_id) -> int");
|
||||
cuda_utils.impl("get_device_attribute", &get_device_attribute);
|
||||
@@ -190,5 +183,6 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
|
||||
cuda_utils.impl("get_max_shared_memory_per_block_device_attribute",
|
||||
&get_max_shared_memory_per_block_device_attribute);
|
||||
}
|
||||
#endif
|
||||
|
||||
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
|
||||
|
||||
@@ -218,6 +218,10 @@ COPY requirements/common.txt requirements/common.txt
|
||||
COPY requirements/cuda.txt requirements/cuda.txt
|
||||
COPY use_existing_torch.py use_existing_torch.py
|
||||
COPY pyproject.toml pyproject.toml
|
||||
# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that
|
||||
# share paths with different content. uv can extract them in either order,
|
||||
# leaving base files that break CUDA 13 CuTe DSL JIT.
|
||||
# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed.
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
|
||||
sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' requirements/cuda.txt; \
|
||||
@@ -234,6 +238,13 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
else \
|
||||
uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
|
||||
fi \
|
||||
&& if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \
|
||||
CUTLASS_DSL_VERSION=$(uv pip show --python /opt/venv/bin/python3 nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \
|
||||
if [ -n "$CUTLASS_DSL_VERSION" ]; then \
|
||||
uv pip install --python /opt/venv/bin/python3 --force-reinstall --no-deps \
|
||||
"nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \
|
||||
fi; \
|
||||
fi
|
||||
|
||||
# Track PyTorch lib versions used during build and match in downstream instances.
|
||||
@@ -745,6 +756,10 @@ ENV VLLM_ENABLE_CUDA_COMPATIBILITY=0
|
||||
ARG PYTORCH_CUDA_INDEX_BASE_URL
|
||||
COPY requirements/common.txt /tmp/common.txt
|
||||
COPY requirements/cuda.txt /tmp/requirements-cuda.txt
|
||||
# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that
|
||||
# share paths with different content. uv can extract them in either order,
|
||||
# leaving base files that break CUDA 13 CuTe DSL JIT.
|
||||
# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed.
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
|
||||
sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' /tmp/requirements-cuda.txt; \
|
||||
@@ -752,6 +767,13 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
fi && \
|
||||
uv pip install --system -r /tmp/requirements-cuda.txt \
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \
|
||||
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \
|
||||
CUTLASS_DSL_VERSION=$(uv pip show --system nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \
|
||||
if [ -n "$CUTLASS_DSL_VERSION" ]; then \
|
||||
uv pip install --system --force-reinstall --no-deps \
|
||||
"nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \
|
||||
fi; \
|
||||
fi && \
|
||||
rm /tmp/requirements-cuda.txt /tmp/common.txt
|
||||
|
||||
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
|
||||
@@ -842,6 +864,19 @@ RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm
|
||||
uv pip install --system ep_kernels/dist/*.whl --verbose \
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
|
||||
|
||||
# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that
|
||||
# share paths with different content. Force -libs-cu13 last after runtime
|
||||
# dependency installs so uv cannot leave base files behind.
|
||||
# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed.
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \
|
||||
CUTLASS_DSL_VERSION=$(uv pip show --system nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \
|
||||
if [ -n "$CUTLASS_DSL_VERSION" ]; then \
|
||||
uv pip install --system --force-reinstall --no-deps \
|
||||
"nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \
|
||||
fi; \
|
||||
fi
|
||||
|
||||
# Download FlashInfer precompiled cubins AFTER all pip installs are done.
|
||||
# This must run after the vLLM wheel and EP kernels installs above, because
|
||||
# those can reinstall/touch flashinfer packages. Downloading cubins earlier
|
||||
|
||||
@@ -168,6 +168,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
######################### TRITON-CPU BUILD IMAGE #########################
|
||||
FROM base AS vllm-triton-cpu-build
|
||||
|
||||
# Support for cross-compilation with x86 ISA including AVX2 and AVX512: docker build --build-arg VLLM_CPU_X86="true" ...
|
||||
# Re-declared here because this stage is `FROM base` (not `vllm-build`), so it
|
||||
# does not inherit the ARG/ENV defined there. Without it, the guard below would
|
||||
# see an empty value and build triton-cpu on non-x86 targets (e.g. arm64).
|
||||
ARG VLLM_CPU_X86=0
|
||||
|
||||
WORKDIR /vllm-workspace
|
||||
|
||||
RUN mkdir dist
|
||||
@@ -269,6 +275,11 @@ ENV HF_HUB_DOWNLOAD_TIMEOUT 60
|
||||
######################### RELEASE IMAGE #########################
|
||||
FROM base AS vllm-openai
|
||||
|
||||
# Re-declared here because this stage is `FROM base` (not `vllm-build`), so the
|
||||
# RUN below that gates the triton-cpu wheel install on $VLLM_CPU_X86 would
|
||||
# otherwise see an empty value and try to install it on non-x86 targets.
|
||||
ARG VLLM_CPU_X86=0
|
||||
|
||||
WORKDIR /vllm-workspace
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
|
||||
@@ -98,7 +98,6 @@ RUN if [ "$USE_SCCACHE" = "1" ]; then \
|
||||
ARG USE_SCCACHE
|
||||
ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
|
||||
ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
|
||||
ENV SCCACHE_ENDPOINT=${USE_SCCACHE:+${SCCACHE_ENDPOINT}}
|
||||
ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
|
||||
ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
|
||||
|
||||
@@ -231,6 +230,7 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
|
||||
@@ -514,6 +514,7 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requir
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
|
||||
|
||||
@@ -296,7 +296,7 @@ llm = LLM(model="Qwen/Qwen3-8B")
|
||||
The `fastokens` Python package (>= 0.2.0) must be installed; if it isn't,
|
||||
vLLM raises a clear `ImportError` at tokenizer load. The override applies to
|
||||
any `--tokenizer-mode` that ends up loading an HF fast tokenizer (`hf`,
|
||||
`deepseek_v32`, `deepseek_v4`, `qwen_vl`, …). Modes that don't use the HF
|
||||
`deepseek_v32`, `deepseek_v4`, `qwen_vl`, …). Models that don't use the HF
|
||||
fast tokenizer (`mistral`, `grok2`, `kimi_audio`) ignore the flag.
|
||||
|
||||
Tokenizer-bound workloads — long shared prefixes, bursty short prompts,
|
||||
|
||||
@@ -101,7 +101,7 @@ vLLM's `pre-commit` hooks will now run automatically every time you commit.
|
||||
Some `pre-commit` hooks only run in CI. If you need to, you can run them locally with:
|
||||
|
||||
```bash
|
||||
pre-commit run --hook-stage manual mypy-3.10
|
||||
pre-commit run --hook-stage manual mypy-3.11
|
||||
```
|
||||
|
||||
### Documentation
|
||||
|
||||
@@ -228,3 +228,17 @@ MLA decode backends are selected using the standard
|
||||
| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
|
||||
|
||||
### DeepSeek V4 Decode Backends
|
||||
|
||||
DeepSeek V4 sparse MLA uses its own decode backends, selected via
|
||||
`--attention-backend=<BACKEND>` (e.g., `FLASHMLA_SPARSE_DSV4`,
|
||||
`FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index
|
||||
pipeline (compressor + SWA + indexer, 256-token blocks, head 512);
|
||||
default on NVIDIA is `FLASHMLA_SPARSE_DSV4`.
|
||||
|
||||
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
|
||||
| `FLASHINFER_MLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
|
||||
| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla`, `fp8` | 256 | 512 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
|
||||
| `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
|
||||
@@ -82,6 +82,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
|
||||
|
||||
| Architecture | Models | CG for Image | CG for Video |
|
||||
| ------------ | ------ | ------------ | ------------ |
|
||||
| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ |
|
||||
| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ |
|
||||
| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
|
||||
| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -128,7 +128,7 @@ The lease mechanism is controlled through `kv_connector_extra_config` in `--kv-t
|
||||
vllm serve <MODEL> \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "NixlConnector",
|
||||
"kv_role": "kv_both",
|
||||
"kv_role": "kv_producer",
|
||||
"kv_connector_extra_config": {"kv_lease_duration": 60}
|
||||
}'
|
||||
```
|
||||
|
||||
@@ -88,7 +88,7 @@ If compilation fails for a multimodal model:
|
||||
|
||||
1. **Disable and test**: First verify the model works without compilation:
|
||||
```bash
|
||||
VLLM_TORCH_COMPILE_LEVEL=0 vllm serve <model> --compilation-config='{"compile_mm_encoder":"false"}'
|
||||
vllm serve <model> --compilation-config='{"mode":0,"compile_mm_encoder":"false"}'
|
||||
```
|
||||
|
||||
2. **Check logs**: Enable debug logging to see compilation details:
|
||||
|
||||
@@ -17,10 +17,7 @@ Batch invariance is crucial for several use cases:
|
||||
|
||||
## Hardware Requirements
|
||||
|
||||
Batch invariance currently requires NVIDIA GPUs with compute capability 9.0 or higher:
|
||||
|
||||
- **H-series**: H100, H200
|
||||
- **B-series**: B100, B200
|
||||
Batch invariance requires NVIDIA GPUs with compute capability 8.0 or higher.
|
||||
|
||||
## Enabling Batch Invariance
|
||||
|
||||
@@ -107,7 +104,7 @@ Batch invariance has been tested and verified on the following models:
|
||||
- **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`, `Qwen/Qwen3-4B-AWQ`, `Qwen/Qwen3-8B-AWQ`
|
||||
- **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct`, `Qwen/Qwen3-30B-A3B-Thinking-2507-FP8`
|
||||
- **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct`
|
||||
- **Llama 3**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct`
|
||||
- **Llama 3**: Llama3.1 and 3.2 series, `meta-llama/Llama-3.2-3B-Instruct` for example
|
||||
- **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b`
|
||||
- **Mistral**: `mistralai/Mistral-7B-v0.3`
|
||||
|
||||
|
||||
@@ -50,7 +50,7 @@ To select a different backend, set `kv_connector_extra_config.backends` in `--kv
|
||||
vllm serve <MODEL> \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector":"NixlConnector",
|
||||
"kv_role":"kv_both",
|
||||
"kv_role":"kv_producer",
|
||||
"kv_connector_extra_config":{"backends":["LIBFABRIC"]}
|
||||
}'
|
||||
```
|
||||
@@ -60,7 +60,7 @@ You can also pass JSON keys individually using dotted arguments, and you can app
|
||||
```bash
|
||||
vllm serve <MODEL> \
|
||||
--kv-transfer-config.kv_connector NixlConnector \
|
||||
--kv-transfer-config.kv_role kv_both \
|
||||
--kv-transfer-config.kv_role kv_producer \
|
||||
--kv-transfer-config.kv_connector_extra_config.backends+ LIBFABRIC
|
||||
```
|
||||
|
||||
@@ -81,7 +81,7 @@ VLLM_NIXL_SIDE_CHANNEL_PORT=5600 \
|
||||
vllm serve Qwen/Qwen3-0.6B \
|
||||
--port 8100 \
|
||||
--enforce-eager \
|
||||
--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail"}'
|
||||
--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail"}'
|
||||
```
|
||||
|
||||
### Consumer (Decoder) Configuration
|
||||
@@ -96,7 +96,7 @@ VLLM_NIXL_SIDE_CHANNEL_PORT=5601 \
|
||||
vllm serve Qwen/Qwen3-0.6B \
|
||||
--port 8200 \
|
||||
--enforce-eager \
|
||||
--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_load_failure_policy":"fail"}'
|
||||
--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail"}'
|
||||
```
|
||||
|
||||
### Proxy Server
|
||||
@@ -212,10 +212,21 @@ sequenceDiagram
|
||||
Enable bidirectional KV transfer by setting `bidirectional_kv_xfer` in `kv_connector_extra_config` on **both** P and D instances:
|
||||
|
||||
```bash
|
||||
# Prefill instance
|
||||
vllm serve <MODEL> \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "NixlConnector",
|
||||
"kv_role": "kv_both",
|
||||
"kv_role": "kv_producer",
|
||||
"kv_connector_extra_config": {
|
||||
"bidirectional_kv_xfer": true
|
||||
}
|
||||
}'
|
||||
|
||||
# Decode instance
|
||||
vllm serve <MODEL> \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "NixlConnector",
|
||||
"kv_role": "kv_consumer",
|
||||
"kv_connector_extra_config": {
|
||||
"bidirectional_kv_xfer": true
|
||||
}
|
||||
@@ -359,11 +370,10 @@ For multi-host DP deployment, only need to provide the host/port of the head ins
|
||||
|
||||
- **kv_producer**: For prefiller instances that generate KV caches
|
||||
- **kv_consumer**: For decoder instances that consume KV caches from prefiller
|
||||
- **kv_both**: Enables symmetric functionality where the connector can act as both producer and consumer. This provides flexibility for experimental setups and scenarios where the role distinction is not predetermined.
|
||||
- **kv_both** (deprecated): Previously used as a catch-all when the role was not predetermined. This value is now deprecated for NixlConnector and will be removed in a future release.
|
||||
|
||||
!!! tip
|
||||
NixlConnector currently does not distinguish `kv_role`; the actual prefiller/decoder roles are determined by the upper-level proxy (e.g., `toy_proxy_server.py` using `--prefiller-hosts` and `--decoder-hosts`).
|
||||
Therefore, `kv_role` in `--kv-transfer-config` is effectively a placeholder and does not affect NixlConnector's behavior.
|
||||
!!! warning
|
||||
`kv_role="kv_both"` is deprecated for NixlConnector. Please set `kv_role="kv_producer"` for prefill instances and `kv_role="kv_consumer"` for decode instances. See [#33702](https://github.com/vllm-project/vllm/issues/33702) for details.
|
||||
|
||||
### KV Load Failure Policy
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
Quantization trades off model precision for smaller memory footprint, allowing large models to be run on a wider range of devices.
|
||||
|
||||
!!! tip
|
||||
To get started with quantization, see [LLM Compressor](llm_compressor.md), a library for optimizing models for deployment with vLLM that supports FP8, INT8, INT4, and other quantization formats.
|
||||
To get started with quantization, see [LLM Compressor](llm_compressor/README.md), a library for optimizing models for deployment with vLLM that supports FP8, INT8, INT4, and other quantization formats.
|
||||
|
||||
The following are the supported quantization formats for vLLM:
|
||||
|
||||
@@ -12,9 +12,11 @@ The following are the supported quantization formats for vLLM:
|
||||
- [GGUF](gguf.md)
|
||||
- [GPTQModel](gptqmodel.md)
|
||||
- [Intel Neural Compressor](inc.md)
|
||||
- [INT4 W4A16](int4.md)
|
||||
- [INT8 W8A8](int8.md)
|
||||
- [FP8 W8A8](fp8.md)
|
||||
- [LLM Compressor](llm_compressor/README.md)
|
||||
- [FP8 W8A8](llm_compressor/fp8.md)
|
||||
- [INT4 W4A16](llm_compressor/int4.md)
|
||||
- [INT8 W4A8](llm_compressor/int8_w4a8.md)
|
||||
- [INT8 W8A8](llm_compressor/int8_w8a8.md)
|
||||
- [NVIDIA Model Optimizer](modelopt.md)
|
||||
- [Online Quantization](online.md)
|
||||
- [AMD Quark](quark.md)
|
||||
@@ -46,16 +48,17 @@ th:not(:first-child) {
|
||||
}
|
||||
</style>
|
||||
|
||||
| Implementation | Volta | Turing | Ampere | Ada | Hopper | AMD GPU | Intel GPU | x86 CPU |
|
||||
| ------------------------- | ----- | ------ | ------ | --- | ------ | ------- | --------- | ------- |
|
||||
| AWQ | ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ✅︎ |
|
||||
| GPTQ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ✅︎ |
|
||||
| Marlin (GPTQ/AWQ/FP8/FP4) | ❌ | ✅︎* | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ |
|
||||
| INT8 (W8A8) | ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ✅︎ |
|
||||
| FP8 (W8A8) | ❌ | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ |
|
||||
| bitsandbytes | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ |
|
||||
| DeepSpeedFP | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ |
|
||||
| GGUF | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ |
|
||||
| Implementation | Volta | Turing | Ampere | Ada | Hopper | AMD GPU | Intel GPU | x86 CPU | Arm CPU |
|
||||
| ------------------------- | ----- | ------ | ------ | --- | ------ | ------- | --------- | ------- | ------- |
|
||||
| AWQ | ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ✅︎ | ❌ |
|
||||
| GPTQ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ✅︎ | ✅︎ | ❌ |
|
||||
| Marlin (GPTQ/AWQ/FP8/FP4) | ❌ | ✅︎* | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
|
||||
| llm-compressor INT8 (W8A8)| ❌ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ✅︎ | ✅︎ |
|
||||
| llm-compressor INT8 (W4A8)| ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅︎ |
|
||||
| llm-compressor FP8 (W8A8) | ❌ | ❌ | ❌ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ |
|
||||
| bitsandbytes | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
|
||||
| DeepSpeedFP | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ | ❌ |
|
||||
| GGUF | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ✅︎ | ❌ | ❌ | ❌ |
|
||||
|
||||
- Volta refers to SM 7.0, Turing to SM 7.5, Ampere to SM 8.0/8.6, Ada to SM 8.9, and Hopper to SM 9.0.
|
||||
- ✅︎ indicates that the quantization method is supported on the specified hardware.
|
||||
|
||||
+25
-25
@@ -21,9 +21,17 @@ The FP8 types typically supported in hardware have two distinct representations,
|
||||
To produce performant FP8 quantized models with vLLM, you'll need to install the [llm-compressor](https://github.com/vllm-project/llm-compressor/) library:
|
||||
|
||||
```bash
|
||||
pip install llmcompressor
|
||||
(venv-llm-compressor) pip install llmcompressor
|
||||
```
|
||||
|
||||
Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
(venv-vllm) pip install vllm "lm-eval[api]>=0.4.12"
|
||||
```
|
||||
|
||||
Please use separate environments for vLLM and llm-compressor as they might not work together.
|
||||
|
||||
## Quantization Process
|
||||
|
||||
The quantization process involves three main steps:
|
||||
@@ -57,36 +65,28 @@ For FP8 quantization, we can recover accuracy with simple RTN quantization. We r
|
||||
|
||||
Since simple RTN does not require data for weight quantization and the activations are quantized dynamically, we do not need any calibration data for this quantization flow.
|
||||
|
||||
??? code
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import QuantizationModifier
|
||||
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import QuantizationModifier
|
||||
# Configure the simple PTQ quantization
|
||||
recipe = QuantizationModifier(
|
||||
targets="Linear",
|
||||
scheme="FP8_DYNAMIC",
|
||||
ignore=["lm_head"],
|
||||
)
|
||||
|
||||
# Configure the simple PTQ quantization
|
||||
recipe = QuantizationModifier(
|
||||
targets="Linear",
|
||||
scheme="FP8_DYNAMIC",
|
||||
ignore=["lm_head"],
|
||||
)
|
||||
# Apply the quantization algorithm.
|
||||
oneshot(model=model, recipe=recipe)
|
||||
|
||||
# Apply the quantization algorithm.
|
||||
oneshot(model=model, recipe=recipe)
|
||||
|
||||
# Save the model: Meta-Llama-3-8B-Instruct-FP8-Dynamic
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-Dynamic"
|
||||
model.save_pretrained(SAVE_DIR)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
# Save the model: Meta-Llama-3-8B-Instruct-FP8-Dynamic
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-Dynamic"
|
||||
model.save_pretrained(SAVE_DIR)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
|
||||
### 3. Evaluating Accuracy
|
||||
|
||||
Install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
pip install vllm "lm-eval[api]>=0.4.12"
|
||||
```
|
||||
|
||||
Load and run the model in `vllm`:
|
||||
|
||||
```python
|
||||
+64
-68
@@ -12,15 +12,17 @@ Please visit the HF collection of [quantized INT4 checkpoints of popular LLMs re
|
||||
To use INT4 quantization with vLLM, you'll need to install the [llm-compressor](https://github.com/vllm-project/llm-compressor/) library:
|
||||
|
||||
```bash
|
||||
pip install llmcompressor
|
||||
(venv-llm-compressor) pip install llmcompressor
|
||||
```
|
||||
|
||||
Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
pip install vllm "lm-eval[api]>=0.4.12"
|
||||
(venv-vllm) pip install vllm "lm-eval[api]>=0.4.12"
|
||||
```
|
||||
|
||||
Please use separate environments for vLLM and llm-compressor as they might not work together.
|
||||
|
||||
## Quantization Process
|
||||
|
||||
The quantization process involves four main steps:
|
||||
@@ -52,55 +54,51 @@ When quantizing weights to INT4, you need sample data to estimate the weight upd
|
||||
It's best to use calibration data that closely matches your deployment data.
|
||||
For a general-purpose instruction-tuned model, you can use a dataset like `ultrachat`:
|
||||
|
||||
??? code
|
||||
```python
|
||||
from datasets import load_dataset
|
||||
|
||||
```python
|
||||
from datasets import load_dataset
|
||||
NUM_CALIBRATION_SAMPLES = 512
|
||||
MAX_SEQUENCE_LENGTH = 2048
|
||||
|
||||
NUM_CALIBRATION_SAMPLES = 512
|
||||
MAX_SEQUENCE_LENGTH = 2048
|
||||
# Load and preprocess the dataset
|
||||
ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft")
|
||||
ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
|
||||
|
||||
# Load and preprocess the dataset
|
||||
ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft")
|
||||
ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
|
||||
def preprocess(example):
|
||||
return {"text": tokenizer.apply_chat_template(example["messages"], tokenize=False)}
|
||||
ds = ds.map(preprocess)
|
||||
|
||||
def preprocess(example):
|
||||
return {"text": tokenizer.apply_chat_template(example["messages"], tokenize=False)}
|
||||
ds = ds.map(preprocess)
|
||||
|
||||
def tokenize(sample):
|
||||
return tokenizer(sample["text"], padding=False, max_length=MAX_SEQUENCE_LENGTH, truncation=True, add_special_tokens=False)
|
||||
ds = ds.map(tokenize, remove_columns=ds.column_names)
|
||||
```
|
||||
def tokenize(sample):
|
||||
return tokenizer(sample["text"], padding=False, max_length=MAX_SEQUENCE_LENGTH, truncation=True, add_special_tokens=False)
|
||||
ds = ds.map(tokenize, remove_columns=ds.column_names)
|
||||
```
|
||||
|
||||
### 3. Applying Quantization
|
||||
|
||||
Now, apply the quantization algorithms:
|
||||
|
||||
??? code
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import GPTQModifier
|
||||
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
|
||||
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import GPTQModifier
|
||||
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
|
||||
# Configure the quantization algorithms
|
||||
recipe = GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"])
|
||||
|
||||
# Configure the quantization algorithms
|
||||
recipe = GPTQModifier(targets="Linear", scheme="W4A16", ignore=["lm_head"])
|
||||
# Apply quantization
|
||||
oneshot(
|
||||
model=model,
|
||||
dataset=ds,
|
||||
recipe=recipe,
|
||||
max_seq_length=MAX_SEQUENCE_LENGTH,
|
||||
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
|
||||
)
|
||||
|
||||
# Apply quantization
|
||||
oneshot(
|
||||
model=model,
|
||||
dataset=ds,
|
||||
recipe=recipe,
|
||||
max_seq_length=MAX_SEQUENCE_LENGTH,
|
||||
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
|
||||
)
|
||||
|
||||
# Save the compressed model: Meta-Llama-3-8B-Instruct-W4A16-G128
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-W4A16-G128"
|
||||
model.save_pretrained(SAVE_DIR, save_compressed=True)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
# Save the compressed model: Meta-Llama-3-8B-Instruct-W4A16-G128
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-W4A16-G128"
|
||||
model.save_pretrained(SAVE_DIR, save_compressed=True)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
|
||||
This process creates a W4A16 model with weights quantized to 4-bit integers.
|
||||
|
||||
@@ -141,36 +139,34 @@ lm_eval --model vllm \
|
||||
|
||||
The following is an example of an expanded quantization recipe you can tune to your own use case:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from compressed_tensors.quantization import (
|
||||
QuantizationArgs,
|
||||
QuantizationScheme,
|
||||
QuantizationStrategy,
|
||||
QuantizationType,
|
||||
)
|
||||
recipe = GPTQModifier(
|
||||
targets="Linear",
|
||||
config_groups={
|
||||
"config_group": QuantizationScheme(
|
||||
targets=["Linear"],
|
||||
weights=QuantizationArgs(
|
||||
num_bits=4,
|
||||
type=QuantizationType.INT,
|
||||
strategy=QuantizationStrategy.GROUP,
|
||||
group_size=128,
|
||||
symmetric=True,
|
||||
dynamic=False,
|
||||
actorder="weight",
|
||||
),
|
||||
```python
|
||||
from compressed_tensors.quantization import (
|
||||
QuantizationArgs,
|
||||
QuantizationScheme,
|
||||
QuantizationStrategy,
|
||||
QuantizationType,
|
||||
)
|
||||
recipe = GPTQModifier(
|
||||
targets="Linear",
|
||||
config_groups={
|
||||
"config_group": QuantizationScheme(
|
||||
targets=["Linear"],
|
||||
weights=QuantizationArgs(
|
||||
num_bits=4,
|
||||
type=QuantizationType.INT,
|
||||
strategy=QuantizationStrategy.GROUP,
|
||||
group_size=128,
|
||||
symmetric=True,
|
||||
dynamic=False,
|
||||
actorder="weight",
|
||||
),
|
||||
},
|
||||
ignore=["lm_head"],
|
||||
update_size=NUM_CALIBRATION_SAMPLES,
|
||||
dampening_frac=0.01,
|
||||
)
|
||||
```
|
||||
),
|
||||
},
|
||||
ignore=["lm_head"],
|
||||
update_size=NUM_CALIBRATION_SAMPLES,
|
||||
dampening_frac=0.01,
|
||||
)
|
||||
```
|
||||
|
||||
## Troubleshooting and Support
|
||||
|
||||
@@ -0,0 +1,217 @@
|
||||
# INT8 W4A8
|
||||
|
||||
vLLM supports quantizing weights to INT4 and activations to INT8 for memory savings and inference acceleration.
|
||||
This quantization method is particularly useful for reducing model size while maintaining good performance.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
To use INT8 W4A8 quantization with vLLM, you'll need to install the [llm-compressor](https://github.com/vllm-project/llm-compressor/) library.
|
||||
|
||||
```bash
|
||||
(venv-llm-compressor) pip install llmcompressor
|
||||
```
|
||||
|
||||
Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
(venv-vllm) pip install vllm "lm-eval[api]>=0.4.12"
|
||||
```
|
||||
|
||||
Please use separate environments for vLLM and llm-compressor as they might not work together.
|
||||
|
||||
## Quantization Process
|
||||
|
||||
The quantization process involves four main steps:
|
||||
|
||||
1. Loading the model
|
||||
2. Preparing calibration data
|
||||
3. Applying quantization
|
||||
4. Evaluating accuracy in vLLM
|
||||
|
||||
### 1. Loading the Model
|
||||
|
||||
Load your model and tokenizer using the standard `transformers` AutoModel classes:
|
||||
|
||||
```python
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
|
||||
MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
MODEL_ID,
|
||||
dtype="auto",
|
||||
)
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
||||
```
|
||||
|
||||
### 2. Preparing Calibration Data
|
||||
|
||||
When quantizing activations to INT8 and weights to INT4, you need sample data to estimate the activation scales.
|
||||
It's best to use calibration data that closely matches your deployment data.
|
||||
For a general-purpose instruction-tuned model, you can use a dataset like `ultrachat`:
|
||||
|
||||
```python
|
||||
from datasets import load_dataset
|
||||
|
||||
NUM_CALIBRATION_SAMPLES = 512
|
||||
MAX_SEQUENCE_LENGTH = 2048
|
||||
|
||||
# Load and preprocess the dataset
|
||||
ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft")
|
||||
ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
|
||||
|
||||
def preprocess(example):
|
||||
return {"text": tokenizer.apply_chat_template(example["messages"], tokenize=False)}
|
||||
ds = ds.map(preprocess)
|
||||
|
||||
def tokenize(sample):
|
||||
return tokenizer(sample["text"], padding=False, max_length=MAX_SEQUENCE_LENGTH, truncation=True, add_special_tokens=False)
|
||||
ds = ds.map(tokenize, remove_columns=ds.column_names)
|
||||
```
|
||||
|
||||
### 3. Applying Quantization
|
||||
|
||||
Now, apply the quantization algorithms.
|
||||
|
||||
The following recipes create W4A8 models (int4 weights, int8 activations). On Arm® CPUs, this is accelerated through [KleidiAI](https://github.com/ARM-software/kleidiai).
|
||||
|
||||
Use groupwise for best accuracy, and channelwise for best inference performance.
|
||||
|
||||
=== "Groupwise"
|
||||
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import GPTQModifier
|
||||
|
||||
# Configure the quantization algorithms
|
||||
recipe = [
|
||||
GPTQModifier(
|
||||
targets="Linear",
|
||||
scheme="W4A8",
|
||||
ignore=["lm_head"],
|
||||
dampening_frac=0.01
|
||||
),
|
||||
]
|
||||
|
||||
# Apply quantization
|
||||
oneshot(
|
||||
model=model,
|
||||
dataset=ds,
|
||||
recipe=recipe,
|
||||
max_seq_length=MAX_SEQUENCE_LENGTH,
|
||||
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
|
||||
)
|
||||
|
||||
# Save the compressed model: Meta-Llama-3-8B-Instruct-W4A8-G128-Dynamic-Per-Token
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-W4A8-G128-Dynamic-Per-Token"
|
||||
model.save_pretrained(SAVE_DIR, save_compressed=True)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
|
||||
=== "Channelwise"
|
||||
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import GPTQModifier
|
||||
from compressed_tensors.quantization import QuantizationStrategy, QuantizationType
|
||||
|
||||
scheme = {
|
||||
"targets": ["Linear"],
|
||||
"weights": {
|
||||
"num_bits": 4,
|
||||
"type": QuantizationType.INT,
|
||||
"strategy": QuantizationStrategy.CHANNEL,
|
||||
"symmetric": True,
|
||||
"dynamic": False,
|
||||
"group_size": None,
|
||||
},
|
||||
"input_activations": {
|
||||
"num_bits": 8,
|
||||
"type": QuantizationType.INT,
|
||||
"strategy": QuantizationStrategy.TOKEN,
|
||||
"dynamic": True,
|
||||
"symmetric": False,
|
||||
"observer": None,
|
||||
},
|
||||
"output_activations": None,
|
||||
}
|
||||
|
||||
recipe = [
|
||||
GPTQModifier(
|
||||
targets="Linear",
|
||||
config_groups={"group_0": scheme},
|
||||
ignore=["lm_head"],
|
||||
dampening_frac=0.01,
|
||||
),
|
||||
]
|
||||
|
||||
oneshot(
|
||||
model=model,
|
||||
dataset=ds,
|
||||
recipe=recipe,
|
||||
max_seq_length=MAX_SEQUENCE_LENGTH,
|
||||
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
|
||||
)
|
||||
|
||||
# Save the compressed model: Meta-Llama-3-8B-Instruct-W4A8-Channelwise-Dynamic-Per-Token
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-W4A8-Channelwise-Dynamic-Per-Token"
|
||||
model.save_pretrained(SAVE_DIR, save_compressed=True)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
|
||||
### 4. Evaluating Accuracy
|
||||
|
||||
=== "Groupwise"
|
||||
|
||||
After quantization, you can load and run the model in vLLM:
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
|
||||
llm = LLM("./Meta-Llama-3-8B-Instruct-W4A8-G128-Dynamic-Per-Token")
|
||||
```
|
||||
|
||||
To evaluate accuracy, you can use `lm_eval`:
|
||||
|
||||
```bash
|
||||
lm_eval --model vllm \
|
||||
--model_args pretrained="./Meta-Llama-3-8B-Instruct-W4A8-G128-Dynamic-Per-Token",add_bos_token=true \
|
||||
--tasks gsm8k \
|
||||
--num_fewshot 5 \
|
||||
--limit 250 \
|
||||
--batch_size 'auto'
|
||||
```
|
||||
|
||||
=== "Channelwise"
|
||||
|
||||
After quantization, you can load and run the model in vLLM:
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
|
||||
llm = LLM("./Meta-Llama-3-8B-Instruct-W4A8-Channelwise-Dynamic-Per-Token")
|
||||
```
|
||||
|
||||
To evaluate accuracy, you can use `lm_eval`:
|
||||
|
||||
```bash
|
||||
lm_eval --model vllm \
|
||||
--model_args pretrained="./Meta-Llama-3-8B-Instruct-W4A8-Channelwise-Dynamic-Per-Token",add_bos_token=true \
|
||||
--tasks gsm8k \
|
||||
--num_fewshot 5 \
|
||||
--limit 250 \
|
||||
--batch_size 'auto'
|
||||
```
|
||||
|
||||
!!! note
|
||||
Quantized models can be sensitive to the presence of the `bos` token. Make sure to include the `add_bos_token=True` argument when running evaluations.
|
||||
|
||||
## Best Practices
|
||||
|
||||
- Start with 512 samples for calibration data (increase if accuracy drops)
|
||||
- Use a sequence length of 2048 as a starting point
|
||||
- Employ the chat template or instruction template that the model was trained with
|
||||
- If you've fine-tuned a model, consider using a sample of your training data for calibration
|
||||
|
||||
## Troubleshooting and Support
|
||||
|
||||
If you encounter any issues or have feature requests, please open an issue on the [vllm-project/llm-compressor](https://github.com/vllm-project/llm-compressor/issues) GitHub repository.
|
||||
+40
-42
@@ -17,15 +17,17 @@ Please visit the HF collection of [quantized INT8 checkpoints of popular LLMs re
|
||||
To use INT8 quantization with vLLM, you'll need to install the [llm-compressor](https://github.com/vllm-project/llm-compressor/) library:
|
||||
|
||||
```bash
|
||||
pip install llmcompressor
|
||||
(venv-llm-compressor) pip install llmcompressor
|
||||
```
|
||||
|
||||
Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
pip install vllm "lm-eval[api]>=0.4.12"
|
||||
(venv-vllm) pip install vllm "lm-eval[api]>=0.4.12"
|
||||
```
|
||||
|
||||
Please use separate environments for vLLM and llm-compressor as they might not work together.
|
||||
|
||||
## Quantization Process
|
||||
|
||||
The quantization process involves four main steps:
|
||||
@@ -57,26 +59,24 @@ When quantizing activations to INT8, you need sample data to estimate the activa
|
||||
It's best to use calibration data that closely matches your deployment data.
|
||||
For a general-purpose instruction-tuned model, you can use a dataset like `ultrachat`:
|
||||
|
||||
??? code
|
||||
```python
|
||||
from datasets import load_dataset
|
||||
|
||||
```python
|
||||
from datasets import load_dataset
|
||||
NUM_CALIBRATION_SAMPLES = 512
|
||||
MAX_SEQUENCE_LENGTH = 2048
|
||||
|
||||
NUM_CALIBRATION_SAMPLES = 512
|
||||
MAX_SEQUENCE_LENGTH = 2048
|
||||
# Load and preprocess the dataset
|
||||
ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft")
|
||||
ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
|
||||
|
||||
# Load and preprocess the dataset
|
||||
ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft")
|
||||
ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
|
||||
def preprocess(example):
|
||||
return {"text": tokenizer.apply_chat_template(example["messages"], tokenize=False)}
|
||||
ds = ds.map(preprocess)
|
||||
|
||||
def preprocess(example):
|
||||
return {"text": tokenizer.apply_chat_template(example["messages"], tokenize=False)}
|
||||
ds = ds.map(preprocess)
|
||||
|
||||
def tokenize(sample):
|
||||
return tokenizer(sample["text"], padding=False, max_length=MAX_SEQUENCE_LENGTH, truncation=True, add_special_tokens=False)
|
||||
ds = ds.map(tokenize, remove_columns=ds.column_names)
|
||||
```
|
||||
def tokenize(sample):
|
||||
return tokenizer(sample["text"], padding=False, max_length=MAX_SEQUENCE_LENGTH, truncation=True, add_special_tokens=False)
|
||||
ds = ds.map(tokenize, remove_columns=ds.column_names)
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
@@ -84,33 +84,31 @@ For a general-purpose instruction-tuned model, you can use a dataset like `ultra
|
||||
|
||||
Now, apply the quantization algorithms:
|
||||
|
||||
??? code
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import GPTQModifier
|
||||
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
|
||||
|
||||
```python
|
||||
from llmcompressor import oneshot
|
||||
from llmcompressor.modifiers.quantization import GPTQModifier
|
||||
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
|
||||
# Configure the quantization algorithms
|
||||
recipe = [
|
||||
SmoothQuantModifier(smoothing_strength=0.8),
|
||||
GPTQModifier(targets="Linear", scheme="W8A8", ignore=["lm_head"]),
|
||||
]
|
||||
|
||||
# Configure the quantization algorithms
|
||||
recipe = [
|
||||
SmoothQuantModifier(smoothing_strength=0.8),
|
||||
GPTQModifier(targets="Linear", scheme="W8A8", ignore=["lm_head"]),
|
||||
]
|
||||
# Apply quantization
|
||||
oneshot(
|
||||
model=model,
|
||||
dataset=ds,
|
||||
recipe=recipe,
|
||||
max_seq_length=MAX_SEQUENCE_LENGTH,
|
||||
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
|
||||
)
|
||||
|
||||
# Apply quantization
|
||||
oneshot(
|
||||
model=model,
|
||||
dataset=ds,
|
||||
recipe=recipe,
|
||||
max_seq_length=MAX_SEQUENCE_LENGTH,
|
||||
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
|
||||
)
|
||||
|
||||
# Save the compressed model: Meta-Llama-3-8B-Instruct-W8A8-Dynamic-Per-Token
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-W8A8-Dynamic-Per-Token"
|
||||
model.save_pretrained(SAVE_DIR, save_compressed=True)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
# Save the compressed model: Meta-Llama-3-8B-Instruct-W8A8-Dynamic-Per-Token
|
||||
SAVE_DIR = MODEL_ID.split("/")[1] + "-W8A8-Dynamic-Per-Token"
|
||||
model.save_pretrained(SAVE_DIR, save_compressed=True)
|
||||
tokenizer.save_pretrained(SAVE_DIR)
|
||||
```
|
||||
|
||||
This process creates a W8A8 model with weights and activations quantized to 8-bit integers.
|
||||
|
||||
@@ -169,7 +169,7 @@ speculative decoding, breaking down the guarantees into three key areas:
|
||||
> distribution. [View Test Code](https://github.com/vllm-project/vllm/blob/47b65a550866c7ffbd076ecb74106714838ce7da/tests/samplers/test_rejection_sampler.py#L252)
|
||||
> - **Greedy Sampling Equality**: Confirms that greedy sampling with speculative decoding matches greedy sampling
|
||||
> without it. This verifies that vLLM's speculative decoding framework, when integrated with the vLLM forward pass and the vLLM rejection sampler,
|
||||
> provides a lossless guarantee. Almost all of the tests in [tests/spec_decode/e2e](/tests/v1/spec_decode).
|
||||
> provides a lossless guarantee. Almost all of the tests in [tests/spec_decode/e2e](../../../tests/v1/spec_decode).
|
||||
> verify this property using [this assertion implementation](https://github.com/vllm-project/vllm/blob/b67ae00cdbbe1a58ffc8ff170f0c8d79044a684a/tests/spec_decode/e2e/conftest.py#L291)
|
||||
|
||||
3. **vLLM Logprob Stability**
|
||||
|
||||
@@ -24,10 +24,11 @@ vllm serve google/gemma-4-E2B-it \
|
||||
--speculative-config '{"method":"mtp","model":"gg-hf-am/gemma-4-E2B-it-assistant","num_speculative_tokens":1}'
|
||||
```
|
||||
|
||||
The E2B, E4B, 26B-A4B, and 31B Gemma 4 IT assistant checkpoints are supported
|
||||
when their configuration uses `model_type: gemma4_assistant`. vLLM maps those
|
||||
checkpoints to `Gemma4MTPModel` internally and wires the assistant layers to
|
||||
share KV cache with the target model.
|
||||
The E2B, E4B, 12B, 26B-A4B, and 31B Gemma 4 IT assistant checkpoints are supported.
|
||||
Tower-based variants use `model_type: gemma4_assistant` and the encoder-free
|
||||
Gemma 4 Unified variant (12B) uses `model_type: gemma4_unified_assistant`.
|
||||
vLLM maps both to `Gemma4MTPModel` internally and wires the assistant layers
|
||||
to share KV cache with the target model.
|
||||
|
||||
If an older vLLM release logs `SpeculativeConfig(method='draft_model', ...)`
|
||||
for a Gemma 4 assistant checkpoint, that release is treating the assistant as a
|
||||
|
||||
@@ -378,7 +378,7 @@ th {
|
||||
| `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ |
|
||||
| `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |
|
||||
| `CohereForCausalLM`, `Cohere2ForCausalLM` | Command-R, Command-A | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025`, `CohereLabs/command-a-reasoning-08-2025`, etc. | ✅︎ | ✅︎ |
|
||||
| `Cohere2MoeForCausalLM` | Command-A+ | `CohereLabs/command-a-plus-05-2026`, etc. | ✅︎ | ✅︎ |
|
||||
| `Cohere2MoeForCausalLM` | North-Mini-Code | `CohereLabs/North-Mini-Code`, etc. | ✅︎ | ✅︎ |
|
||||
| `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ |
|
||||
| `DbrxForCausalLM` | DBRX | `databricks/dbrx-base`, `databricks/dbrx-instruct`, etc. | | ✅︎ |
|
||||
| `DeciLMForCausalLM` | DeciLM | `nvidia/Llama-3_3-Nemotron-Super-49B-v1`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -550,7 +550,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `Blip2ForConditionalGeneration` | BLIP-2 | T + I<sup>E</sup> | `Salesforce/blip2-opt-2.7b`, `Salesforce/blip2-opt-6.7b`, etc. | ✅︎ | ✅︎ |
|
||||
| `ChameleonForConditionalGeneration` | Chameleon | T + I | `facebook/chameleon-7b`, etc. | | ✅︎ |
|
||||
| `CheersForConditionalGeneration` | Cheers | T + I | `ai9stars/Cheers` | | ✅︎ |
|
||||
| `Cohere2VisionForConditionalGeneration` | Command A Vision | T + I<sup>+</sup> | `CohereLabs/command-a-vision-07-2025`, etc. | | ✅︎ |
|
||||
| `Cohere2VisionForConditionalGeneration` | Command A Vision, Command-A+ | T + I<sup>+</sup> | `CohereLabs/command-a-vision-07-2025`, `CohereLabs/command-a-plus-05-2026`, etc. | | ✅︎ |
|
||||
| `Cosmos3ForConditionalGeneration` | Cosmos3 (understanding tower) | T + I<sup>E+</sup> + V<sup>E+</sup> | `nvidia/Cosmos3-Nano` | | ✅︎ |
|
||||
| `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |
|
||||
| `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -562,12 +562,14 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `Gemma3ForConditionalGeneration` | Gemma 3 | T + I<sup>E+</sup> | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ |
|
||||
| `Gemma3nForConditionalGeneration` | Gemma 3n | T + I + A | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
|
||||
| `Gemma4ForConditionalGeneration` | Gemma 4 | T + I<sup>+</sup> + V + A<sup>*</sup> | `google/gemma-4-E2B-it`, etc. | | ✅︎ |
|
||||
| `Gemma4UnifiedForConditionalGeneration` | Gemma 4 Unified | T + I<sup>+</sup> + V + A | `google/gemma-4-12B-it`, etc. | | ✅︎ |
|
||||
| `GLM4VForCausalLM`<sup>^</sup> | GLM-4V | T + I | `zai-org/glm-4v-9b`, `zai-org/cogagent-9b-20241220`, etc. | ✅︎ | ✅︎ |
|
||||
| `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ |
|
||||
| `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ |
|
||||
| `GlmOcrForConditionalGeneration` | GLM-OCR | T + I<sup>E+</sup> | `zai-org/GLM-OCR`, etc. | ✅︎ | ✅︎ |
|
||||
| `Granite4VisionForConditionalGeneration` | Granite 4 Vision | T + I<sup>E+</sup> | `ibm-granite/granite-4.1-3b-vision`, etc. | ✅︎ | ✅︎ |
|
||||
| `GraniteSpeechForConditionalGeneration` | Granite Speech | T + A | `ibm-granite/granite-speech-3.3-8b` | ✅︎ | ✅︎ |
|
||||
| `GraniteSpeechPlusForConditionalGeneration` | Granite Speech Plus | T + A | `ibm-granite/granite-speech-4.1-2b-plus` | ✅︎ | ✅︎ |
|
||||
| `HCXVisionForCausalLM` | HyperCLOVAX-SEED-Vision-Instruct-3B | T + I<sup>+</sup> + V<sup>+</sup> | `naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B` | | |
|
||||
| `HCXVisionV2ForCausalLM` | HyperCLOVAX-SEED-Think-32B | T + I<sup>+</sup> + V<sup>+</sup> | `naver-hyperclovax/HyperCLOVAX-SEED-Think-32B` | | |
|
||||
| `H2OVLChatModel` | H2OVL | T + I<sup>E+</sup> | `h2oai/h2ovl-mississippi-800m`, `h2oai/h2ovl-mississippi-2b`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -664,10 +666,16 @@ Some models are supported only via the [Transformers modeling backend](#transfor
|
||||
For `Gemma4ForConditionalGeneration`:
|
||||
- audio input is only supported by the `gemma-4-E2B` and `gemma-4-E4B` variants.
|
||||
- The model does not ingest videos directly. However, vLLM’s Gemma 4 implementation supports video inputs by handling video processing internally. Users can send videos directly in the message structure to vLLM, where they are converted into text and image frames before being passed to the model.
|
||||
- Gemma 4 assistant checkpoints for speculative decoding use vLLM's Gemma
|
||||
- Gemma 4 assistant checkpoints for speculative decoding use vLLM’s Gemma
|
||||
4 MTP path, not generic draft-model speculative decoding. See the
|
||||
[Gemma 4 assistant model MTP example](../features/speculative_decoding/mtp.md#gemma-4-assistant-models).
|
||||
|
||||
!!! note
|
||||
For `Gemma4UnifiedForConditionalGeneration`:
|
||||
- This is the encoder-free Gemma 4 variant (e.g. `gemma-4-12B-it`). Unlike the tower-based `Gemma4ForConditionalGeneration`, it has **no SigLIP vision encoder** and **no audio encoder**. Raw pixel patches are projected directly into LM space via a Dense+LayerNorm pipeline with factorized positional embeddings, and raw audio waveform frames are projected directly through a multimodal embedder.
|
||||
- All modalities (image, video, audio) are supported.
|
||||
- Gemma 4 Unified assistant checkpoints (`model_type: gemma4_unified_assistant`) use the same MTP path as the tower-based variant. See the [Gemma 4 assistant model MTP example](../features/speculative_decoding/mtp.md#gemma-4-assistant-models).
|
||||
|
||||
!!! note
|
||||
For `InternVLChatModel`, only InternVL2.5 with Qwen2.5 text backbone (`OpenGVLab/InternVL2.5-1B` etc.), InternVL3 and InternVL3.5 have video inputs support currently.
|
||||
|
||||
@@ -702,6 +710,7 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
|
||||
| `Gemma3nForConditionalGeneration` | Gemma3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
|
||||
| `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ |
|
||||
| `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-4.0-1b-speech`, `ibm-granite/granite-speech-3.3-2b`, etc. | ✅︎ | ✅︎ |
|
||||
| `GraniteSpeechPlusForConditionalGeneration` | Granite Speech Plus | `ibm-granite/granite-speech-4.1-2b-plus` | ✅︎ | ✅︎ |
|
||||
| `Qwen3ASRForConditionalGeneration` | Qwen3-ASR | `Qwen/Qwen3-ASR-1.7B`, etc. | ✅︎ | ✅︎ |
|
||||
| `Qwen3OmniMoeThinkerForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct`, etc. | | ✅︎ |
|
||||
| `VoxtralForConditionalGeneration` | Voxtral (Mistral format) | `mistralai/Voxtral-Mini-3B-2507`, `mistralai/Voxtral-Small-24B-2507`, etc. | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -24,8 +24,14 @@ echo "Checking pre-commit/pre-run-check status..."
|
||||
MAX_WAIT=300
|
||||
INTERVAL=60
|
||||
ELAPSED=0
|
||||
# Use a GitHub token if provided to raise the API rate limit (60 -> 5000
|
||||
# requests/hour). Set GITHUB_TOKEN in the Read the Docs environment variables.
|
||||
CURL_AUTH=()
|
||||
if [ -n "$GITHUB_TOKEN" ]; then
|
||||
CURL_AUTH=(-H "Authorization: Bearer $GITHUB_TOKEN")
|
||||
fi
|
||||
while :; do
|
||||
RAW=$(curl -sS -w "\n%{http_code}" "https://api.github.com/repos/vllm-project/vllm/commits/${READTHEDOCS_GIT_COMMIT_HASH}/check-runs?check_name=pre-run-check&filter=latest")
|
||||
RAW=$(curl -sS "${CURL_AUTH[@]}" -w "\n%{http_code}" "https://api.github.com/repos/vllm-project/vllm/commits/${READTHEDOCS_GIT_COMMIT_HASH}/check-runs?check_name=pre-run-check&filter=latest")
|
||||
HTTP_CODE=$(printf %s "$RAW" | tail -n1)
|
||||
BODY=$(printf %s "$RAW" | sed '$d')
|
||||
if [ "$HTTP_CODE" != "200" ]; then
|
||||
|
||||
@@ -2554,6 +2554,7 @@ MODELS_NEED_VIDEO_METADATA = [
|
||||
|
||||
|
||||
MODELS_SUPPORT_VIT_CUDA_GRAPH = [
|
||||
"internvl_chat",
|
||||
"qwen2_5_vl",
|
||||
"qwen3_vl",
|
||||
"qwen3_vl_moe",
|
||||
|
||||
+12
-9
@@ -83,22 +83,22 @@ plugins:
|
||||
- "re:vllm\\._.*" # Internal modules
|
||||
- "vllm.third_party"
|
||||
- "vllm.vllm_flash_attn"
|
||||
- "re:vllm\\.grpc\\..*_pb2.*" # Auto-generated protobuf files
|
||||
- "vllm.transformers_utils.configs"
|
||||
- "vllm.transformers_utils.processors"
|
||||
- !ENV [API_AUTONAV_EXCLUDE, "re:^$"] # Match nothing by default
|
||||
- mkdocstrings:
|
||||
handlers:
|
||||
python:
|
||||
options:
|
||||
show_symbol_type_heading: true
|
||||
show_symbol_type_toc: true
|
||||
filters:
|
||||
- "!.*_pb2_grpc" # Exclude auto-generated gRPC stubs
|
||||
summary:
|
||||
modules: true
|
||||
show_signature_annotations: true
|
||||
separate_signature: true
|
||||
filters: []
|
||||
show_overloads: true
|
||||
signature_crossrefs: true
|
||||
# Recommendations from api-autonav
|
||||
docstring_section_style: list
|
||||
parameter_headings: true
|
||||
show_symbol_type_heading: true
|
||||
show_symbol_type_toc: true
|
||||
summary: true
|
||||
inventories:
|
||||
- https://docs.python.org/3/objects.inv
|
||||
- https://typing-extensions.readthedocs.io/en/latest/objects.inv
|
||||
@@ -110,6 +110,9 @@ plugins:
|
||||
redirect_maps:
|
||||
features/spec_decode/README.md: features/speculative_decoding/README.md
|
||||
features/spec_decode/speculators.md: features/speculative_decoding/speculators.md
|
||||
features/quantization/fp8.md: features/quantization/llm_compressor/fp8.md
|
||||
features/quantization/int4.md: features/quantization/llm_compressor/int4.md
|
||||
features/quantization/int8.md: features/quantization/llm_compressor/int8_w8a8.md
|
||||
serving/openai_compatible_server.md: serving/online_serving/README.md
|
||||
|
||||
markdown_extensions:
|
||||
|
||||
@@ -11,7 +11,7 @@ transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5
|
||||
tokenizers >= 0.21.1 # Required for fast incremental detokenization.
|
||||
safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611
|
||||
protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
|
||||
fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
|
||||
fastapi[standard] >= 0.115.0, < 0.137 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint; <0.137 due to prometheus-fastapi-instrumentator#370.
|
||||
aiohttp >= 3.13.3
|
||||
openai >= 2.0.0 # For Responses API with reasoning content
|
||||
pydantic >= 2.12.0
|
||||
@@ -32,13 +32,13 @@ partial-json-parser # used for parsing partial JSON outputs
|
||||
pyzmq >= 25.0.0
|
||||
msgspec
|
||||
gguf >= 0.17.0
|
||||
mistral_common[image] >= 1.11.2
|
||||
mistral_common[image] >= 1.11.3
|
||||
opencv-python-headless >= 4.13.0 # required for video IO
|
||||
pyyaml
|
||||
six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
|
||||
setuptools>=77.0.3,<81.0.0; python_version > '3.11' # Setuptools is used by triton, we need to ensure a modern version is installed for 3.12+ so that it does not try to import distutils, which was removed in 3.12
|
||||
einops # Required for Qwen2-VL.
|
||||
compressed-tensors == 0.15.0.1 # required for compressed-tensors
|
||||
compressed-tensors == 0.17.0 # required for compressed-tensors
|
||||
depyf==0.20.0 # required for profiling and debugging with compilation config
|
||||
cloudpickle # allows pickling lambda functions in model_executor/models/registry.py
|
||||
watchfiles # required for http server to monitor the updates of TLS files
|
||||
|
||||
@@ -13,12 +13,10 @@ 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
|
||||
# breaking changes in 1.19.0
|
||||
nvidia-cudnn-frontend>=1.13.0,<1.19.0
|
||||
nvidia-cudnn-frontend>=1.19.1
|
||||
|
||||
# Required for faster safetensors model loading
|
||||
fastsafetensors >= 0.2.2
|
||||
fastsafetensors >= 0.3.2
|
||||
|
||||
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
|
||||
nvidia-cutlass-dsl[cu13]==4.5.2
|
||||
@@ -28,4 +26,4 @@ quack-kernels>=0.3.3
|
||||
tokenspeed-mla==0.1.2
|
||||
|
||||
# Humming kernels for quantization gemm
|
||||
humming-kernels[cu13]==0.1.2
|
||||
humming-kernels[cu13]==0.1.4
|
||||
|
||||
@@ -19,7 +19,11 @@ setuptools-rust>=1.9.0
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
conch-triton-kernels==1.2.1
|
||||
timm>=1.0.17
|
||||
# amd-quark: required for Quark quantization on ROCm
|
||||
# amd-quark: required for Quark quantization on ROCm
|
||||
# To be consistent with test_quark.py
|
||||
amd-quark>=0.8.99
|
||||
tilelang==0.1.10
|
||||
# Required apache-tvm-ffi matching tilelang version
|
||||
apache-tvm-ffi==0.1.10
|
||||
# Required for faster safetensors model loading
|
||||
fastsafetensors >= 0.3.2
|
||||
|
||||
@@ -31,7 +31,7 @@ torchaudio==2.11.0
|
||||
torchvision==0.26.0
|
||||
transformers_stream_generator # required for qwen-vl test
|
||||
matplotlib # required for qwen-vl test
|
||||
mistral_common[image,audio] >= 1.11.2 # required for voxtral test
|
||||
mistral_common[image,audio] >= 1.11.3 # required for voxtral test
|
||||
num2words # required for smolvlm test
|
||||
open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py
|
||||
opencv-python-headless >= 4.13.0 # required for video test
|
||||
@@ -57,7 +57,7 @@ arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix de
|
||||
numba == 0.65.0 # Required for N-gram speculative decoding
|
||||
numpy
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors>=0.2.2; platform_machine == "x86_64" # 0.2.2 contains important fixes for multi-GPU mem usage
|
||||
fastsafetensors>=0.3.2
|
||||
instanttensor>=0.1.5; platform_machine == "x86_64"
|
||||
pydantic>=2.12 # 2.11 leads to error on python 3.13
|
||||
decord==0.6.0; platform_machine == "x86_64"
|
||||
|
||||
@@ -191,7 +191,7 @@ fastparquet==2024.11.0
|
||||
# via genai-perf
|
||||
fastrlock==0.8.2
|
||||
# via cupy-cuda12x
|
||||
fastsafetensors==0.2.2
|
||||
fastsafetensors==0.3.2
|
||||
# via
|
||||
# -c requirements/cuda.txt
|
||||
# -r requirements/test/cuda.in
|
||||
@@ -409,7 +409,7 @@ mbstrdecoder==1.1.3
|
||||
# typepy
|
||||
mdurl==0.1.2
|
||||
# via markdown-it-py
|
||||
mistral-common==1.11.2
|
||||
mistral-common==1.11.3
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/cuda.in
|
||||
|
||||
@@ -23,7 +23,7 @@ jiwer # required for audio tests
|
||||
timm # required for internvl test
|
||||
transformers_stream_generator # required for qwen-vl test
|
||||
matplotlib # required for qwen-vl test
|
||||
mistral_common[image,audio] >= 1.11.2 # required for voxtral test
|
||||
mistral_common[image,audio] >= 1.11.3 # required for voxtral test
|
||||
num2words # required for smolvlm test
|
||||
opencv-python-headless >= 4.13.0 # required for video test
|
||||
datamodel_code_generator # required for minicpm3 test
|
||||
@@ -43,6 +43,6 @@ tritonclient>=2.51.0
|
||||
numba == 0.65.0 # Required for N-gram speculative decoding
|
||||
numpy
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors>=0.2.2
|
||||
fastsafetensors>=0.3.2
|
||||
instanttensor>=0.1.5
|
||||
pydantic>=2.12 # 2.11 leads to error on python 3.13
|
||||
|
||||
@@ -30,7 +30,7 @@ tblib # for pickling test exceptions
|
||||
timm>=1.0.17 # required for internvl and gemma3n-mm test
|
||||
transformers_stream_generator # required for qwen-vl test
|
||||
matplotlib # required for qwen-vl test
|
||||
mistral_common[image,audio]>=1.11.2 # required for voxtral test
|
||||
mistral_common[image,audio]>=1.11.3 # required for voxtral test
|
||||
num2words # required for smolvlm test
|
||||
open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py
|
||||
opencv-python-headless>=4.13.0 # required for video test
|
||||
@@ -56,7 +56,7 @@ arctic-inference==0.1.1 # Required for suffix decoding test
|
||||
numba==0.65.0 # Required for N-gram speculative decoding
|
||||
numpy
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2 # PyPI only ships CUDA wheels
|
||||
fastsafetensors>=0.3.2
|
||||
instanttensor>=0.1.5
|
||||
pydantic>=2.12 # 2.11 leads to error on python 3.13
|
||||
decord==0.6.0
|
||||
|
||||
@@ -44,6 +44,7 @@ anyio==4.13.0
|
||||
# watchfiles
|
||||
apache-tvm-ffi==0.1.10
|
||||
# via
|
||||
# -c requirements/rocm.txt
|
||||
# tilelang
|
||||
# xgrammar
|
||||
arctic-inference==0.1.1
|
||||
@@ -143,7 +144,7 @@ colorful==0.5.8
|
||||
# via ray
|
||||
colorlog==6.10.1
|
||||
# via optuna
|
||||
compressed-tensors==0.15.0.1
|
||||
compressed-tensors==0.17.0
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/../common.txt
|
||||
@@ -240,8 +241,10 @@ fastar==0.10.0
|
||||
# via fastapi-cloud-cli
|
||||
fastparquet==2026.3.0
|
||||
# via genai-perf
|
||||
fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@65d80088fca7a8f567fba30415fbcc80f7d2259c
|
||||
# via -r requirements/test/rocm.in
|
||||
fastsafetensors==0.3.2
|
||||
# via
|
||||
# -c requirements/rocm.txt
|
||||
# -r requirements/test/rocm.in
|
||||
filelock==3.25.2
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
@@ -510,7 +513,7 @@ mcp==1.27.0
|
||||
# via -r requirements/test/../common.txt
|
||||
mdurl==0.1.2
|
||||
# via markdown-it-py
|
||||
mistral-common==1.11.2
|
||||
mistral-common==1.11.3
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/../common.txt
|
||||
|
||||
@@ -266,7 +266,7 @@ mbstrdecoder==1.1.4
|
||||
# typepy
|
||||
mdurl==0.1.2
|
||||
# via markdown-it-py
|
||||
mistral-common==1.11.2
|
||||
mistral-common==1.11.3
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/xpu.in
|
||||
|
||||
@@ -12,4 +12,4 @@ ray[data]
|
||||
setuptools==78.1.0
|
||||
setuptools-rust>=1.9.0
|
||||
nixl==0.3.0
|
||||
tpu-inference==0.20.0
|
||||
tpu-inference==0.21.0
|
||||
|
||||
@@ -12,7 +12,7 @@ jinja2>=3.1.6
|
||||
datasets # for benchmark scripts
|
||||
numba == 0.65.0 # Required for N-gram speculative decoding
|
||||
--extra-index-url=https://download.pytorch.org/whl/xpu
|
||||
torch==2.11.0+xpu
|
||||
torch==2.11.0
|
||||
torchaudio
|
||||
torchvision
|
||||
|
||||
|
||||
@@ -38,13 +38,17 @@ impl HfChatBackend {
|
||||
) -> Result<Self> {
|
||||
let model_config = load_model_config(files.config_path.as_deref())?;
|
||||
let model_type = model_config.model_type().unwrap_or_default();
|
||||
let multimodal_model_info = MultimodalModelInfo::from_paths(
|
||||
model_id.clone(),
|
||||
(!model_type.is_empty()).then_some(model_type.to_string()),
|
||||
files.config_path.as_deref(),
|
||||
files.preprocessor_config_path.as_deref(),
|
||||
tokenizer.clone(),
|
||||
)?;
|
||||
let multimodal_model_info = if options.language_model_only {
|
||||
None
|
||||
} else {
|
||||
MultimodalModelInfo::from_paths(
|
||||
model_id.clone(),
|
||||
(!model_type.is_empty()).then_some(model_type.to_string()),
|
||||
files.config_path.as_deref(),
|
||||
files.preprocessor_config_path.as_deref(),
|
||||
tokenizer.clone(),
|
||||
)?
|
||||
};
|
||||
let multimodal_render_info = resolve_multimodal_render_info(multimodal_model_info.as_ref());
|
||||
|
||||
let renderer = options.renderer.resolve(model_type);
|
||||
@@ -225,6 +229,7 @@ mod tests {
|
||||
"test-model".to_string(),
|
||||
LoadModelBackendsOptions {
|
||||
renderer,
|
||||
language_model_only: false,
|
||||
chat_template_content_format: Default::default(),
|
||||
chat_template: None,
|
||||
default_chat_template_kwargs: HashMap::new(),
|
||||
@@ -267,6 +272,54 @@ mod tests {
|
||||
assert_eq!(prompt, "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn language_model_only_skips_multimodal_preprocessor_config() {
|
||||
let mut files = resolved_files(
|
||||
r#"{"model_type":"deepseek_v0_vl"}"#,
|
||||
r#"{"chat_template":"{{ messages[0].content }}"}"#,
|
||||
);
|
||||
let preprocessor_config_path = files
|
||||
.config_path
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.parent()
|
||||
.unwrap()
|
||||
.join("preprocessor_config.json");
|
||||
write_json(&preprocessor_config_path, r#"{"size":[672,672]}"#);
|
||||
files.preprocessor_config_path = Some(preprocessor_config_path);
|
||||
|
||||
let backend = HfChatBackend::from_resolved_model_files(
|
||||
files.clone(),
|
||||
"test-model".to_string(),
|
||||
LoadModelBackendsOptions {
|
||||
language_model_only: true,
|
||||
chat_template_content_format: Default::default(),
|
||||
chat_template: None,
|
||||
default_chat_template_kwargs: HashMap::new(),
|
||||
..Default::default()
|
||||
},
|
||||
test_tokenizer(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(backend.multimodal_model_info().is_none());
|
||||
|
||||
let error = HfChatBackend::from_resolved_model_files(
|
||||
files,
|
||||
"test-model".to_string(),
|
||||
LoadModelBackendsOptions {
|
||||
chat_template_content_format: Default::default(),
|
||||
chat_template: None,
|
||||
default_chat_template_kwargs: HashMap::new(),
|
||||
..Default::default()
|
||||
},
|
||||
test_tokenizer(),
|
||||
)
|
||||
.err()
|
||||
.expect("invalid preprocessor config should fail without language_model_only");
|
||||
assert!(error.to_string().contains("failed to parse preprocessor_config.json"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_deepseek_renderer_overrides_generic_model_type() {
|
||||
let prompt = render_prompt(
|
||||
|
||||
@@ -60,6 +60,9 @@ pub type DynChatTextBackend = Arc<dyn ChatTextBackend>;
|
||||
pub struct LoadModelBackendsOptions {
|
||||
/// Which chat renderer implementation to use.
|
||||
pub renderer: RendererSelection,
|
||||
/// Disable frontend-side multimodal preprocessing and render the model as
|
||||
/// language-only.
|
||||
pub language_model_only: bool,
|
||||
/// How to serialize `message.content` when rendering the chat template.
|
||||
pub chat_template_content_format: ChatTemplateContentFormatOption,
|
||||
/// Optional server-default chat template override, provided either as an
|
||||
|
||||
@@ -234,7 +234,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, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, phi4_mini_json, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -6,8 +6,8 @@ pub use vllm_tool_parser::{
|
||||
DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
|
||||
Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, HermesToolParser, HyV3ToolParser,
|
||||
Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser,
|
||||
MistralToolParser, Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser,
|
||||
ToolParserError, ToolParserOutput,
|
||||
MistralToolParser, Phi4MiniJsonToolParser, Qwen3CoderToolParser, Qwen3XmlToolParser,
|
||||
ToolCallDelta, ToolParser, ToolParserError, ToolParserOutput,
|
||||
};
|
||||
|
||||
use crate::parser::ParserFactory;
|
||||
@@ -32,6 +32,7 @@ pub mod names {
|
||||
pub const LLAMA4_JSON: &str = "llama4_json";
|
||||
pub const MINIMAX_M2: &str = "minimax_m2";
|
||||
pub const MISTRAL: &str = "mistral";
|
||||
pub const PHI4_MINI_JSON: &str = "phi4_mini_json";
|
||||
pub const QWEN3_CODER: &str = "qwen3_coder";
|
||||
pub const QWEN3_XML: &str = "qwen3_xml";
|
||||
}
|
||||
@@ -71,6 +72,7 @@ impl ToolParserFactory {
|
||||
.register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
|
||||
.register_parser::<MinimaxM2ToolParser>(names::MINIMAX_M2)
|
||||
.register_parser::<MistralToolParser>(names::MISTRAL)
|
||||
.register_parser::<Phi4MiniJsonToolParser>(names::PHI4_MINI_JSON)
|
||||
.register_parser::<Qwen3XmlToolParser>(names::QWEN3_XML)
|
||||
.register_parser::<Qwen3CoderToolParser>(names::QWEN3_CODER);
|
||||
|
||||
|
||||
@@ -191,3 +191,14 @@ fn factory_new_resolves_default_patterns() {
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn factory_new_registers_phi4_mini_json_by_name() {
|
||||
// phi-4-mini is registered by explicit name only (matching Python's
|
||||
// `--tool-call-parser phi4_mini_json`); it is intentionally not mapped to
|
||||
// any model-name pattern.
|
||||
let factory = ToolParserFactory::new();
|
||||
|
||||
assert!(factory.contains(names::PHI4_MINI_JSON));
|
||||
factory.create(names::PHI4_MINI_JSON, &[]).unwrap();
|
||||
}
|
||||
|
||||
@@ -116,6 +116,10 @@ pub struct SharedRuntimeArgs {
|
||||
#[arg(long = "tokenizer-mode", default_value_t)]
|
||||
#[serde(default, rename = "tokenizer_mode")]
|
||||
pub renderer: RendererSelection,
|
||||
/// Disable multimodal inputs and treat the model as language-only.
|
||||
#[arg(long)]
|
||||
#[serde(default)]
|
||||
pub language_model_only: bool,
|
||||
/// Override the maximum model context length. When set, the frontend uses
|
||||
/// this value instead of the model's `max_position_embeddings` from
|
||||
/// `config.json`.
|
||||
@@ -243,6 +247,7 @@ impl SharedRuntimeArgs {
|
||||
tool_call_parser: self.tool_call_parser,
|
||||
reasoning_parser: self.reasoning_parser,
|
||||
renderer: self.renderer,
|
||||
language_model_only: self.language_model_only,
|
||||
chat_template: self.chat_template,
|
||||
default_chat_template_kwargs: self.default_chat_template_kwargs,
|
||||
chat_template_content_format: self.chat_template_content_format,
|
||||
@@ -284,6 +289,7 @@ impl SharedRuntimeArgs {
|
||||
tool_call_parser: self.tool_call_parser,
|
||||
reasoning_parser: self.reasoning_parser,
|
||||
renderer: self.renderer,
|
||||
language_model_only: self.language_model_only,
|
||||
chat_template: self.chat_template,
|
||||
default_chat_template_kwargs: self.default_chat_template_kwargs,
|
||||
chat_template_content_format: self.chat_template_content_format,
|
||||
@@ -419,6 +425,7 @@ impl ServeArgs {
|
||||
self.managed_engine.clone().into_config(
|
||||
self.runtime.model.clone(),
|
||||
self.runtime.max_model_len,
|
||||
self.runtime.language_model_only,
|
||||
handshake_port,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -34,6 +34,7 @@ fn serve_args_forward_python_flags_with_separator() {
|
||||
tool_call_parser: Auto,
|
||||
reasoning_parser: Auto,
|
||||
renderer: Auto,
|
||||
language_model_only: false,
|
||||
max_model_len: Some(
|
||||
512,
|
||||
),
|
||||
@@ -263,6 +264,7 @@ fn frontend_args_accept_json() {
|
||||
tool_call_parser: Auto,
|
||||
reasoning_parser: Auto,
|
||||
renderer: Auto,
|
||||
language_model_only: false,
|
||||
max_model_len: None,
|
||||
grpc_port: None,
|
||||
shutdown_timeout: 0,
|
||||
@@ -321,7 +323,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() {
|
||||
"--output-address",
|
||||
"ipc:///tmp/output.sock",
|
||||
"--args-json",
|
||||
r#"{"model_tag":"Qwen/Qwen3-0.6B","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","max_model_len":8192,"shutdown_timeout":3}"#,
|
||||
r#"{"model_tag":"Qwen/Qwen3-0.6B","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_model_len":8192,"shutdown_timeout":3}"#,
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
@@ -338,6 +340,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() {
|
||||
ParserSelection::Explicit("qwen3_thinking".to_string())
|
||||
);
|
||||
assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32);
|
||||
assert!(args.runtime.language_model_only);
|
||||
assert_eq!(args.runtime.max_model_len, Some(8192));
|
||||
assert_eq!(args.runtime.shutdown_timeout, 3);
|
||||
}
|
||||
@@ -662,6 +665,7 @@ fn serve_args_accept_handshake_aliases() {
|
||||
tool_call_parser: Auto,
|
||||
reasoning_parser: Auto,
|
||||
renderer: Auto,
|
||||
language_model_only: false,
|
||||
max_model_len: None,
|
||||
grpc_port: None,
|
||||
shutdown_timeout: 0,
|
||||
@@ -783,6 +787,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
|
||||
tool_call_parser: Auto,
|
||||
reasoning_parser: Auto,
|
||||
renderer: Auto,
|
||||
language_model_only: false,
|
||||
chat_template: None,
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: Auto,
|
||||
@@ -846,6 +851,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
|
||||
tool_call_parser: Auto,
|
||||
reasoning_parser: Auto,
|
||||
renderer: Auto,
|
||||
language_model_only: false,
|
||||
chat_template: None,
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: Auto,
|
||||
@@ -924,6 +930,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
|
||||
tool_call_parser: Auto,
|
||||
reasoning_parser: Auto,
|
||||
renderer: Auto,
|
||||
language_model_only: false,
|
||||
chat_template: None,
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: Auto,
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::slice;
|
||||
use std::sync::Arc;
|
||||
|
||||
use arc_swap::ArcSwapOption;
|
||||
@@ -253,33 +252,47 @@ pub(crate) async fn run_abort_loop(
|
||||
inner: Arc<ClientInner>,
|
||||
mut abort_rx: mpsc::UnboundedReceiver<AbortRequest>,
|
||||
) {
|
||||
// TODO: receive and abort requests in batch
|
||||
while let Some(AbortRequest { request_id, cause }) = abort_rx.recv().await {
|
||||
let Some(engine_id) = inner.take_auto_abort_target(&request_id) else {
|
||||
debug!(request_id, "skip auto-abort for inactive request");
|
||||
continue;
|
||||
};
|
||||
// Coalesce bursts of auto-aborts into a single Abort message per engine.
|
||||
// A dropped-stream storm (e.g. many clients disconnecting at once under
|
||||
// high concurrency) would otherwise issue one engine round-trip per
|
||||
// request. `recv_many` returns as soon as at least one item is ready, so a
|
||||
// lone abort is still forwarded promptly.
|
||||
const MAX_DRAIN: usize = 1024;
|
||||
let mut batch: Vec<AbortRequest> = Vec::new();
|
||||
|
||||
match cause {
|
||||
AbortCause::DroppedStream => {
|
||||
info!(request_id, "auto-aborting request due to dropped stream")
|
||||
}
|
||||
AbortCause::StopStringMatched => {
|
||||
debug!(
|
||||
request_id,
|
||||
"auto-aborting request due to stop string matched"
|
||||
)
|
||||
while abort_rx.recv_many(&mut batch, MAX_DRAIN).await > 0 {
|
||||
let mut by_engine: BTreeMap<EngineId, Vec<String>> = BTreeMap::new();
|
||||
|
||||
for AbortRequest { request_id, cause } in batch.drain(..) {
|
||||
let Some(engine_id) = inner.take_auto_abort_target(&request_id) else {
|
||||
debug!(request_id, "skip auto-abort for inactive request");
|
||||
continue;
|
||||
};
|
||||
|
||||
match cause {
|
||||
AbortCause::DroppedStream => {
|
||||
info!(request_id, "auto-aborting request due to dropped stream")
|
||||
}
|
||||
AbortCause::StopStringMatched => {
|
||||
debug!(
|
||||
request_id,
|
||||
"auto-aborting request due to stop string matched"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
by_engine.entry(engine_id).or_default().push(request_id);
|
||||
}
|
||||
|
||||
if let Err(error) = inner.do_abort_requests(&engine_id, slice::from_ref(&request_id)).await
|
||||
{
|
||||
warn!(
|
||||
request_id,
|
||||
?engine_id,
|
||||
error = %error.as_report(),
|
||||
"failed to auto-abort dropped request stream"
|
||||
);
|
||||
for (engine_id, request_ids) in by_engine {
|
||||
if let Err(error) = inner.do_abort_requests(&engine_id, &request_ids).await {
|
||||
warn!(
|
||||
?engine_id,
|
||||
?request_ids,
|
||||
error = %error.as_report(),
|
||||
"failed to auto-abort request streams"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1225,6 +1225,86 @@ async fn dropping_a_live_stream_triggers_abort() {
|
||||
client.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dropping_multiple_live_streams_aborts_all_in_a_burst() {
|
||||
init_tracing();
|
||||
let ipc = IpcNamespace::new().unwrap();
|
||||
let handshake_address = ipc.handshake_endpoint();
|
||||
let engine_id = b"engine-burst".to_vec();
|
||||
let request_ids = ["req-1", "req-2", "req-3"];
|
||||
|
||||
let (shutdown_tx, engine_task) = spawn_mock_engine_task(
|
||||
handshake_address.clone(),
|
||||
engine_id.clone(),
|
||||
|dealer, push| {
|
||||
Box::pin(async move {
|
||||
for _ in 0..3 {
|
||||
let add = recv_engine_message(dealer).await;
|
||||
assert_eq!(add[0].as_ref(), &[0x00]);
|
||||
}
|
||||
send_outputs(
|
||||
push,
|
||||
EngineCoreOutputs {
|
||||
outputs: vec![
|
||||
request_output("req-1", vec![99], None),
|
||||
request_output("req-2", vec![99], None),
|
||||
request_output("req-3", vec![99], None),
|
||||
],
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let abort =
|
||||
timeout(Duration::from_secs(1), recv_engine_message(dealer)).await.unwrap();
|
||||
assert_eq!(abort[0].as_ref(), &[0x01]);
|
||||
let ids: Vec<String> = rmp_serde::from_slice(&abort[1]).unwrap();
|
||||
assert_eq!(
|
||||
ids,
|
||||
vec![
|
||||
"req-1".to_string(),
|
||||
"req-2".to_string(),
|
||||
"req-3".to_string()
|
||||
]
|
||||
);
|
||||
assert!(
|
||||
timeout(Duration::from_millis(100), recv_engine_message(dealer)).await.is_err()
|
||||
);
|
||||
})
|
||||
},
|
||||
);
|
||||
|
||||
let client = connect_client_with_ipc(
|
||||
handshake_test_config(
|
||||
handshake_address,
|
||||
1,
|
||||
"test-model",
|
||||
Duration::from_secs(2),
|
||||
0,
|
||||
None,
|
||||
),
|
||||
&ipc,
|
||||
)
|
||||
.await;
|
||||
|
||||
// Open every request first so all three adds reach the engine before it
|
||||
// emits outputs, then drain the first token from each stream.
|
||||
let mut streams = Vec::new();
|
||||
for id in request_ids {
|
||||
streams.push(client.call(sample_request_with_id(id)).await.unwrap());
|
||||
}
|
||||
for stream in streams.iter_mut() {
|
||||
let first = timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap();
|
||||
assert_eq!(first.new_token_ids, vec![99]);
|
||||
}
|
||||
// Drop the whole burst back-to-back so the abort worker can batch them.
|
||||
drop(streams);
|
||||
|
||||
let _ = shutdown_tx.send(());
|
||||
engine_task.await.unwrap();
|
||||
client.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn dispatcher_failure_propagates_to_streams_and_future_calls() {
|
||||
init_tracing();
|
||||
|
||||
@@ -71,6 +71,7 @@ impl ManagedEngineArgs {
|
||||
self,
|
||||
model: String,
|
||||
max_model_len: Option<u32>,
|
||||
language_model_only: bool,
|
||||
handshake_port: u16,
|
||||
) -> ManagedEngineConfig {
|
||||
let mut python_args = self.python_args;
|
||||
@@ -79,6 +80,9 @@ impl ManagedEngineArgs {
|
||||
python_args.push("--max-model-len".to_string());
|
||||
python_args.push(max_model_len.to_string());
|
||||
}
|
||||
if language_model_only {
|
||||
python_args.push("--language-model-only".to_string());
|
||||
}
|
||||
if let Some(data_parallel_size_local) = self.data_parallel_size_local {
|
||||
python_args.push("--data-parallel-size-local".to_string());
|
||||
python_args.push(data_parallel_size_local.to_string());
|
||||
|
||||
@@ -64,6 +64,7 @@ async fn main() -> Result<()> {
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
renderer: RendererSelection::Auto,
|
||||
language_model_only: false,
|
||||
chat_template: None,
|
||||
default_chat_template_kwargs: None,
|
||||
chat_template_content_format: ChatTemplateContentFormatOption::Auto,
|
||||
|
||||
@@ -53,6 +53,9 @@ pub struct Config {
|
||||
pub reasoning_parser: ParserSelection,
|
||||
/// Chat renderer selection.
|
||||
pub renderer: RendererSelection,
|
||||
/// Disable frontend-side multimodal preprocessing and render the model as
|
||||
/// language-only.
|
||||
pub language_model_only: bool,
|
||||
/// Server-default chat template override, as a file path or inline
|
||||
/// template.
|
||||
pub chat_template: Option<String>,
|
||||
|
||||
@@ -42,6 +42,7 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
|
||||
&config.model,
|
||||
LoadModelBackendsOptions {
|
||||
renderer: config.renderer,
|
||||
language_model_only: config.language_model_only,
|
||||
chat_template: config.chat_template.clone(),
|
||||
chat_template_content_format: config.chat_template_content_format,
|
||||
default_chat_template_kwargs: config
|
||||
|
||||
@@ -85,6 +85,7 @@ pub async fn chat_completions(
|
||||
log_request,
|
||||
prepared.include_usage,
|
||||
prepared.requested_logprobs,
|
||||
prepared.include_reasoning,
|
||||
prepared.echo,
|
||||
prepared.return_token_ids,
|
||||
prepared.return_tokens_as_token_ids,
|
||||
@@ -100,6 +101,7 @@ pub async fn chat_completions(
|
||||
created,
|
||||
prepared.requested_logprobs,
|
||||
prepared.include_prompt_logprobs,
|
||||
prepared.include_reasoning,
|
||||
prepared.echo,
|
||||
prepared.return_token_ids,
|
||||
prepared.return_tokens_as_token_ids,
|
||||
@@ -134,6 +136,7 @@ async fn collect_chat_completion(
|
||||
created: u64,
|
||||
requested_logprobs: bool,
|
||||
include_prompt_logprobs: bool,
|
||||
include_reasoning: bool,
|
||||
echo: Option<String>,
|
||||
return_token_ids: bool,
|
||||
return_tokens_as_token_ids: bool,
|
||||
@@ -157,6 +160,11 @@ async fn collect_chat_completion(
|
||||
} = collected;
|
||||
let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json);
|
||||
let saw_tool_calls = message.tool_calls().next().is_some();
|
||||
let reasoning = message.reasoning();
|
||||
// Output logprobs and token IDs cover the complete generated token stream.
|
||||
// When reasoning is hidden, omit them rather than leaking hidden reasoning
|
||||
// tokens through per-token metadata.
|
||||
let include_output_metadata = include_reasoning || reasoning.is_none();
|
||||
let finish_reason = chat_finish_reason_to_openai(&finish_reason, saw_tool_calls)?.to_string();
|
||||
let tool_calls = message
|
||||
.tool_calls()
|
||||
@@ -169,7 +177,7 @@ async fn collect_chat_completion(
|
||||
},
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let logprobs = if requested_logprobs {
|
||||
let logprobs = if requested_logprobs && include_output_metadata {
|
||||
Some(decoded_logprobs_to_openai_chat(
|
||||
logprobs.as_ref().ok_or_else(|| {
|
||||
server_error!("chat response requested logprobs but generation returned none")
|
||||
@@ -207,12 +215,12 @@ async fn collect_chat_completion(
|
||||
None => Some(message.text()).filter(|t| !t.is_empty()),
|
||||
},
|
||||
tool_calls: Some(tool_calls).filter(|calls| !calls.is_empty()),
|
||||
reasoning: message.reasoning(),
|
||||
reasoning: if include_reasoning { reasoning } else { None },
|
||||
},
|
||||
logprobs,
|
||||
finish_reason: Some(finish_reason),
|
||||
stop_reason,
|
||||
token_ids: return_token_ids.then_some(token_ids),
|
||||
token_ids: (return_token_ids && include_output_metadata).then_some(token_ids),
|
||||
}],
|
||||
usage: Some(usage),
|
||||
system_fingerprint: None,
|
||||
@@ -232,12 +240,18 @@ async fn chat_completion_chunk_stream(
|
||||
log_request: bool,
|
||||
include_usage: bool,
|
||||
requested_logprobs: bool,
|
||||
include_reasoning: bool,
|
||||
echo: Option<String>,
|
||||
return_token_ids: bool,
|
||||
return_tokens_as_token_ids: bool,
|
||||
mut y: TryYielder<ChatCompletionStreamResponse, ApiError>,
|
||||
) -> Result<(), ApiError> {
|
||||
let mut saw_tool_calls = false;
|
||||
// `LogprobsDelta` is emitted after all chat events for one decoded update.
|
||||
// If that update contains hidden reasoning, including delimiter-only block
|
||||
// starts or ends, omit its token metadata as well as its visible delta.
|
||||
let mut inside_hidden_reasoning = false;
|
||||
let mut suppress_current_update_metadata = false;
|
||||
|
||||
// If the client requested logprobs or token_ids, we need to buffer chunks until
|
||||
// we receive the separate `LogprobsDelta` event, so that we can emit one
|
||||
@@ -268,29 +282,44 @@ async fn chat_completion_chunk_stream(
|
||||
}
|
||||
}
|
||||
Ok(ChatEvent::BlockDelta { kind, delta, .. }) => {
|
||||
if let Some(pending_chunk) = pending_chunk.as_mut() {
|
||||
pending_chunk.push_block_delta(kind, delta);
|
||||
let include_delta =
|
||||
include_reasoning || !matches!(kind, AssistantBlockKind::Reasoning);
|
||||
if include_delta {
|
||||
if let Some(pending_chunk) = pending_chunk.as_mut() {
|
||||
pending_chunk.push_block_delta(kind, delta);
|
||||
} else {
|
||||
y.yield_ok(block_delta_chunk(
|
||||
&request_id,
|
||||
&response_model,
|
||||
created,
|
||||
kind,
|
||||
delta,
|
||||
))
|
||||
.await;
|
||||
}
|
||||
} else {
|
||||
y.yield_ok(block_delta_chunk(
|
||||
&request_id,
|
||||
&response_model,
|
||||
created,
|
||||
kind,
|
||||
delta,
|
||||
))
|
||||
.await;
|
||||
suppress_current_update_metadata = true;
|
||||
}
|
||||
}
|
||||
Ok(ChatEvent::LogprobsDelta {
|
||||
logprobs,
|
||||
token_ids,
|
||||
}) => {
|
||||
let openai_logprobs = logprobs
|
||||
.as_ref()
|
||||
.map(|lp| decoded_logprobs_to_openai_chat(lp, return_tokens_as_token_ids))
|
||||
.transpose()?;
|
||||
let openai_token_ids =
|
||||
return_token_ids.then_some(token_ids).filter(|t| !t.is_empty());
|
||||
let include_metadata =
|
||||
!suppress_current_update_metadata && !inside_hidden_reasoning;
|
||||
suppress_current_update_metadata = false;
|
||||
let openai_logprobs = if include_metadata {
|
||||
logprobs
|
||||
.as_ref()
|
||||
.map(|lp| decoded_logprobs_to_openai_chat(lp, return_tokens_as_token_ids))
|
||||
.transpose()?
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let openai_token_ids = include_metadata
|
||||
.then_some(token_ids)
|
||||
.and_then(|token_ids| return_token_ids.then_some(token_ids))
|
||||
.filter(|t| !t.is_empty());
|
||||
if let Some(pending_chunk) = pending_chunk.as_mut() {
|
||||
pending_chunk.logprobs = openai_logprobs;
|
||||
pending_chunk.token_ids = openai_token_ids;
|
||||
@@ -311,9 +340,17 @@ async fn chat_completion_chunk_stream(
|
||||
}
|
||||
Ok(ChatEvent::BlockStart { kind, .. }) => {
|
||||
debug!(?kind, "starting new block");
|
||||
if !include_reasoning && matches!(kind, AssistantBlockKind::Reasoning) {
|
||||
inside_hidden_reasoning = true;
|
||||
suppress_current_update_metadata = true;
|
||||
}
|
||||
}
|
||||
Ok(ChatEvent::BlockEnd { .. }) => {
|
||||
debug!("ending current block");
|
||||
if inside_hidden_reasoning {
|
||||
inside_hidden_reasoning = false;
|
||||
suppress_current_update_metadata = true;
|
||||
}
|
||||
}
|
||||
Ok(ChatEvent::ToolCallStart { index, id, name }) => {
|
||||
let tool_index = index as u32;
|
||||
@@ -763,7 +800,9 @@ fn stop_reason_to_json(stop_reason: &StopReason) -> Value {
|
||||
mod tests {
|
||||
use futures::{StreamExt as _, stream};
|
||||
use serde_json::json;
|
||||
use vllm_chat::{AssistantBlockKind, AssistantToolCall, ChatEvent, FinishReason};
|
||||
use vllm_chat::{
|
||||
AssistantBlockKind, AssistantContentBlock, AssistantToolCall, ChatEvent, FinishReason,
|
||||
};
|
||||
use vllm_engine_core_client::protocol::StopReason;
|
||||
use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
|
||||
|
||||
@@ -895,6 +934,7 @@ mod tests {
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
true,
|
||||
None,
|
||||
false,
|
||||
false,
|
||||
@@ -958,6 +998,7 @@ mod tests {
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
true,
|
||||
None,
|
||||
false,
|
||||
false,
|
||||
@@ -976,6 +1017,292 @@ mod tests {
|
||||
assert!(chunks[1].choices[0].logprobs.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn chunk_stream_omits_reasoning_delta_when_disabled() {
|
||||
let stream = stream::iter(vec![
|
||||
Ok(ChatEvent::Start {
|
||||
prompt_token_ids: vec![].into(),
|
||||
prompt_logprobs: None,
|
||||
}),
|
||||
Ok(ChatEvent::BlockDelta {
|
||||
index: 0,
|
||||
kind: AssistantBlockKind::Reasoning,
|
||||
delta: "think".to_string(),
|
||||
}),
|
||||
Ok(ChatEvent::BlockDelta {
|
||||
index: 1,
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "answer".to_string(),
|
||||
}),
|
||||
Ok(ChatEvent::Done {
|
||||
message: Default::default(),
|
||||
prompt_token_count: 1,
|
||||
output_token_count: 2,
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
let chunks = chat_completion_chunk_stream(
|
||||
stream,
|
||||
"chatcmpl-1".to_string(),
|
||||
"model".to_string(),
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
false,
|
||||
false,
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.expect("stream chunks");
|
||||
|
||||
assert_eq!(chunks.len(), 3);
|
||||
assert_eq!(
|
||||
chunks[1].choices[0].delta.content.as_deref(),
|
||||
Some("answer")
|
||||
);
|
||||
assert!(
|
||||
chunks
|
||||
.iter()
|
||||
.all(|chunk| chunk.choices.iter().all(|choice| choice.delta.reasoning.is_none()))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn chunk_stream_omits_logprobs_for_suppressed_reasoning() {
|
||||
let stream = stream::iter(vec![
|
||||
Ok(ChatEvent::Start {
|
||||
prompt_token_ids: vec![].into(),
|
||||
prompt_logprobs: None,
|
||||
}),
|
||||
Ok(ChatEvent::BlockDelta {
|
||||
index: 0,
|
||||
kind: AssistantBlockKind::Reasoning,
|
||||
delta: "think".to_string(),
|
||||
}),
|
||||
Ok(ChatEvent::LogprobsDelta {
|
||||
logprobs: Some(DecodedLogprobs {
|
||||
positions: vec![DecodedPositionLogprobs {
|
||||
entries: vec![DecodedTokenLogprob {
|
||||
token_id: 11,
|
||||
token: "think".to_string(),
|
||||
logprob: -0.1,
|
||||
rank: 1,
|
||||
}],
|
||||
}],
|
||||
}),
|
||||
token_ids: vec![11],
|
||||
}),
|
||||
Ok(ChatEvent::BlockDelta {
|
||||
index: 1,
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "answer".to_string(),
|
||||
}),
|
||||
Ok(ChatEvent::LogprobsDelta {
|
||||
logprobs: Some(DecodedLogprobs {
|
||||
positions: vec![DecodedPositionLogprobs {
|
||||
entries: vec![DecodedTokenLogprob {
|
||||
token_id: 22,
|
||||
token: "answer".to_string(),
|
||||
logprob: -0.2,
|
||||
rank: 1,
|
||||
}],
|
||||
}],
|
||||
}),
|
||||
token_ids: vec![22],
|
||||
}),
|
||||
Ok(ChatEvent::Done {
|
||||
message: Default::default(),
|
||||
prompt_token_count: 1,
|
||||
output_token_count: 2,
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
let chunks = chat_completion_chunk_stream(
|
||||
stream,
|
||||
"chatcmpl-1".to_string(),
|
||||
"model".to_string(),
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
true,
|
||||
false,
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.expect("stream chunks");
|
||||
|
||||
assert_eq!(chunks.len(), 3);
|
||||
let choice = &chunks[1].choices[0];
|
||||
assert_eq!(choice.delta.content.as_deref(), Some("answer"));
|
||||
assert_eq!(choice.token_ids.as_deref(), Some(&[22][..]));
|
||||
let logprobs = choice.logprobs.as_ref().expect("answer logprobs");
|
||||
let content = logprobs.content.as_ref().expect("logprobs content");
|
||||
assert_eq!(content[0].token, "answer");
|
||||
assert!(chunks.iter().all(|chunk| {
|
||||
chunk.choices.iter().all(|choice| {
|
||||
choice.delta.reasoning.is_none()
|
||||
&& choice.token_ids.as_deref() != Some(&[11][..])
|
||||
&& choice
|
||||
.logprobs
|
||||
.as_ref()
|
||||
.and_then(|logprobs| logprobs.content.as_ref())
|
||||
.is_none_or(|content| content.iter().all(|entry| entry.token != "think"))
|
||||
})
|
||||
}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn chunk_stream_omits_logprobs_for_hidden_reasoning_delimiters() {
|
||||
let stream = stream::iter(vec![
|
||||
Ok(ChatEvent::Start {
|
||||
prompt_token_ids: vec![].into(),
|
||||
prompt_logprobs: None,
|
||||
}),
|
||||
Ok(ChatEvent::BlockStart {
|
||||
index: 0,
|
||||
kind: AssistantBlockKind::Reasoning,
|
||||
}),
|
||||
Ok(ChatEvent::LogprobsDelta {
|
||||
logprobs: Some(DecodedLogprobs {
|
||||
positions: vec![DecodedPositionLogprobs {
|
||||
entries: vec![DecodedTokenLogprob {
|
||||
token_id: 11,
|
||||
token: "<think>".to_string(),
|
||||
logprob: -0.1,
|
||||
rank: 1,
|
||||
}],
|
||||
}],
|
||||
}),
|
||||
token_ids: vec![11],
|
||||
}),
|
||||
Ok(ChatEvent::BlockDelta {
|
||||
index: 0,
|
||||
kind: AssistantBlockKind::Reasoning,
|
||||
delta: "think".to_string(),
|
||||
}),
|
||||
Ok(ChatEvent::LogprobsDelta {
|
||||
logprobs: Some(DecodedLogprobs {
|
||||
positions: vec![DecodedPositionLogprobs {
|
||||
entries: vec![DecodedTokenLogprob {
|
||||
token_id: 12,
|
||||
token: "think".to_string(),
|
||||
logprob: -0.2,
|
||||
rank: 1,
|
||||
}],
|
||||
}],
|
||||
}),
|
||||
token_ids: vec![12],
|
||||
}),
|
||||
Ok(ChatEvent::BlockEnd {
|
||||
index: 0,
|
||||
block: AssistantContentBlock::Reasoning {
|
||||
text: "think".to_string(),
|
||||
},
|
||||
}),
|
||||
Ok(ChatEvent::LogprobsDelta {
|
||||
logprobs: Some(DecodedLogprobs {
|
||||
positions: vec![DecodedPositionLogprobs {
|
||||
entries: vec![DecodedTokenLogprob {
|
||||
token_id: 13,
|
||||
token: "</think>".to_string(),
|
||||
logprob: -0.3,
|
||||
rank: 1,
|
||||
}],
|
||||
}],
|
||||
}),
|
||||
token_ids: vec![13],
|
||||
}),
|
||||
Ok(ChatEvent::BlockStart {
|
||||
index: 1,
|
||||
kind: AssistantBlockKind::Text,
|
||||
}),
|
||||
Ok(ChatEvent::BlockDelta {
|
||||
index: 1,
|
||||
kind: AssistantBlockKind::Text,
|
||||
delta: "answer".to_string(),
|
||||
}),
|
||||
Ok(ChatEvent::LogprobsDelta {
|
||||
logprobs: Some(DecodedLogprobs {
|
||||
positions: vec![DecodedPositionLogprobs {
|
||||
entries: vec![DecodedTokenLogprob {
|
||||
token_id: 22,
|
||||
token: "answer".to_string(),
|
||||
logprob: -0.4,
|
||||
rank: 1,
|
||||
}],
|
||||
}],
|
||||
}),
|
||||
token_ids: vec![22],
|
||||
}),
|
||||
Ok(ChatEvent::Done {
|
||||
message: Default::default(),
|
||||
prompt_token_count: 1,
|
||||
output_token_count: 4,
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
kv_transfer_params: None,
|
||||
}),
|
||||
]);
|
||||
|
||||
let chunks = chat_completion_chunk_stream(
|
||||
stream,
|
||||
"chatcmpl-1".to_string(),
|
||||
"model".to_string(),
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
true,
|
||||
false,
|
||||
)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.expect("stream chunks");
|
||||
|
||||
assert_eq!(chunks.len(), 3);
|
||||
let choice = &chunks[1].choices[0];
|
||||
assert_eq!(choice.delta.content.as_deref(), Some("answer"));
|
||||
assert_eq!(choice.token_ids.as_deref(), Some(&[22][..]));
|
||||
let logprobs = choice.logprobs.as_ref().expect("answer logprobs");
|
||||
let content = logprobs.content.as_ref().expect("logprobs content");
|
||||
assert_eq!(content[0].token, "answer");
|
||||
assert!(chunks.iter().all(|chunk| {
|
||||
chunk.choices.iter().all(|choice| {
|
||||
choice.delta.reasoning.is_none()
|
||||
&& !choice
|
||||
.token_ids
|
||||
.as_ref()
|
||||
.is_some_and(|ids| matches!(ids.as_slice(), [11] | [12] | [13]))
|
||||
&& choice
|
||||
.logprobs
|
||||
.as_ref()
|
||||
.and_then(|logprobs| logprobs.content.as_ref())
|
||||
.is_none_or(|content| {
|
||||
content.iter().all(|entry| {
|
||||
!matches!(entry.token.as_str(), "<think>" | "think" | "</think>")
|
||||
})
|
||||
})
|
||||
})
|
||||
}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn chunk_stream_preserves_tool_call_index_and_omits_id_from_arguments_delta() {
|
||||
let stream = stream::iter(vec![
|
||||
@@ -1017,6 +1344,7 @@ mod tests {
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
None,
|
||||
false,
|
||||
false,
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user