forked from Karylab-cklius/vllm
Compare commits
76
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5c846ccbda | ||
|
|
f2069b005b | ||
|
|
ccd49f6821 | ||
|
|
1c7bc18318 | ||
|
|
9a938df64e | ||
|
|
3da4a1b124 | ||
|
|
6871738777 | ||
|
|
aa4990a9a2 | ||
|
|
a4610da0c6 | ||
|
|
09cdcf34aa | ||
|
|
d2c671c29b | ||
|
|
b5a2adec4b | ||
|
|
78739e3bda | ||
|
|
89accad2cc | ||
|
|
3c8e49596c | ||
|
|
cec2ec1176 | ||
|
|
435f82d61a | ||
|
|
1c4b51b990 | ||
|
|
2e2c47928b | ||
|
|
80abe0de7d | ||
|
|
a9f7b2d41c | ||
|
|
d14e551a53 | ||
|
|
68567ef2df | ||
|
|
6bc6f2d86d | ||
|
|
1eb2cc961e | ||
|
|
31124749d1 | ||
|
|
9037498c22 | ||
|
|
db32b53e30 | ||
|
|
b529bfd6c5 | ||
|
|
f3df7a7231 | ||
|
|
485bbe1c6f | ||
|
|
a19ff2218a | ||
|
|
4f0d0049a0 | ||
|
|
13b83d77ad | ||
|
|
50241602fd | ||
|
|
12fe2a9aac | ||
|
|
89bd2c14d3 | ||
|
|
9c450b1027 | ||
|
|
635c38338a | ||
|
|
c441ad1c07 | ||
|
|
745bba5ea8 | ||
|
|
2cac89f9da | ||
|
|
3e6e33526d | ||
|
|
b91b7726e0 | ||
|
|
d3ad8e8bcd | ||
|
|
b80ce9dd2f | ||
|
|
b5495cc5f9 | ||
|
|
183a430c13 | ||
|
|
a346d589f5 | ||
|
|
7df3d7dada | ||
|
|
8dd1b702f2 | ||
|
|
f57ac274b2 | ||
|
|
6e919960af | ||
|
|
c88d3d4775 | ||
|
|
ab7fcbdd5d | ||
|
|
3b4a76b63f | ||
|
|
cc22621b51 | ||
|
|
77148992cf | ||
|
|
891cc4b9c5 | ||
|
|
1bdf9810aa | ||
|
|
ebfbcfe46a | ||
|
|
e9de72fe6c | ||
|
|
d272418f45 | ||
|
|
7ff7f5c8eb | ||
|
|
dced290769 | ||
|
|
93bad11912 | ||
|
|
0fbf42af84 | ||
|
|
e6cd8913dd | ||
|
|
859e4d436b | ||
|
|
4a083cc858 | ||
|
|
ca7e1f2c43 | ||
|
|
dec860fb19 | ||
|
|
0a49fb2b13 | ||
|
|
4a8abf37c7 | ||
|
|
01192139bf | ||
|
|
b9a7cd464c |
@@ -21,6 +21,10 @@ steps:
|
||||
timeout_in_minutes: 30
|
||||
optional: true
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
@@ -38,6 +42,10 @@ steps:
|
||||
timeout_in_minutes: 30
|
||||
optional: true
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
@@ -55,6 +63,10 @@ steps:
|
||||
timeout_in_minutes: 30
|
||||
optional: true
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
|
||||
@@ -5,6 +5,10 @@ steps:
|
||||
- label: XPU Sleep Mode
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
|
||||
@@ -5,6 +5,10 @@ steps:
|
||||
- label: Engine (1 GPU)
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
|
||||
@@ -6,6 +6,10 @@ steps:
|
||||
key: eplb-algorithm
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
|
||||
@@ -5,6 +5,10 @@ steps:
|
||||
- label: vLLM IR Tests
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
|
||||
@@ -5,6 +5,10 @@ steps:
|
||||
- label: LoRA Runtime + Utils
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -34,6 +38,10 @@ steps:
|
||||
- label: LoRA Fused/MoE Kernels
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -54,6 +62,10 @@ steps:
|
||||
- label: LoRA Punica Kernels
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -74,6 +86,10 @@ steps:
|
||||
- label: LoRA Punica FP8/XPU Ops
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -94,6 +110,10 @@ steps:
|
||||
- label: LoRA Models
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -117,6 +137,10 @@ steps:
|
||||
- label: LoRA Multimodal
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
|
||||
@@ -5,6 +5,10 @@ steps:
|
||||
- label: V1 Core + KV + Metrics
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -31,6 +35,10 @@ steps:
|
||||
- label: V1 Sample + Logits
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -71,6 +79,10 @@ steps:
|
||||
- label: XPU CPU Offload
|
||||
timeout_in_minutes: 60
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -95,6 +107,10 @@ steps:
|
||||
key: regression
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -126,6 +142,10 @@ steps:
|
||||
timeout_in_minutes: 30
|
||||
num_devices: 2
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -157,6 +177,10 @@ steps:
|
||||
key: async-engine-inputs-utils-worker
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
|
||||
@@ -5,6 +5,10 @@ steps:
|
||||
- label: Model Runner V2 Core Tests (Intel)
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -30,6 +34,10 @@ steps:
|
||||
- label: Model Runner V2 Examples (Intel)
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
|
||||
@@ -6,6 +6,10 @@ steps:
|
||||
key: multi-modal-models-standard-1-qwen2
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -27,6 +31,10 @@ steps:
|
||||
key: multi-modal-models-standard-2-qwen3-gemma
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -47,6 +55,10 @@ steps:
|
||||
key: multi-modal-models-standard-3-llava-qwen2-vl
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -68,6 +80,10 @@ steps:
|
||||
key: multi-modal-models-standard-4-other-whisper
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -88,6 +104,10 @@ steps:
|
||||
key: multi-modal-processor
|
||||
timeout_in_minutes: 45
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
|
||||
@@ -19,6 +19,10 @@ steps:
|
||||
- image-build-xpu
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
@@ -49,6 +53,10 @@ steps:
|
||||
- image-build-xpu
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
@@ -74,6 +82,10 @@ steps:
|
||||
- image-build-xpu
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
@@ -93,6 +105,10 @@ steps:
|
||||
- image-build-xpu
|
||||
timeout_in_minutes: 30
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
|
||||
@@ -34,7 +34,7 @@ case "${test_suite}" in
|
||||
pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py
|
||||
pytest -v -s v1/structured_output
|
||||
pytest -v -s v1/test_serial_utils.py
|
||||
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py
|
||||
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py
|
||||
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py
|
||||
;;
|
||||
server)
|
||||
|
||||
@@ -4,6 +4,11 @@
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
if python3 -c "import torch; raise SystemExit(0 if torch.version.hip is not None else 1)"; then
|
||||
uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
exit 0
|
||||
fi
|
||||
|
||||
REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
|
||||
|
||||
uv pip install --system -r "${REQUIREMENTS_FILE}"
|
||||
|
||||
+12
-23
@@ -647,7 +647,7 @@ steps:
|
||||
- pytest -v -s v1/cudagraph/test_cudagraph_mode.py
|
||||
|
||||
- label: e2e Core (1 GPU) # TBD
|
||||
timeout_in_minutes: 180
|
||||
timeout_in_minutes: 35
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
|
||||
agent_pool: mi250_1
|
||||
optional: true
|
||||
@@ -1594,9 +1594,10 @@ steps:
|
||||
#---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------#
|
||||
|
||||
- label: Kernels Attention Test %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
timeout_in_minutes: 55
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
parallelism: 2
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -1627,10 +1628,11 @@ steps:
|
||||
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py
|
||||
|
||||
- label: Kernels MoE Test %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
parallelism: 4
|
||||
optional: true
|
||||
parallelism: 5
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/cutlass_w8a8/moe/
|
||||
@@ -2075,19 +2077,6 @@ steps:
|
||||
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
|
||||
- pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
|
||||
|
||||
- label: e2e Core (1 GPU) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/
|
||||
- tests/v1/e2e/
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
|
||||
|
||||
- label: e2e Scheduling (1 GPU) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
@@ -2133,9 +2122,10 @@ steps:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
|
||||
|
||||
- label: Spec Decode Eagle # TBD
|
||||
timeout_in_minutes: 180
|
||||
timeout_in_minutes: 45
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
@@ -3053,7 +3043,7 @@ steps:
|
||||
#---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------#
|
||||
|
||||
- label: Kernels (B200-MI355) # TBD
|
||||
timeout_in_minutes: 180
|
||||
timeout_in_minutes: 15
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
working_dir: "/vllm-workspace/"
|
||||
@@ -3077,11 +3067,10 @@ steps:
|
||||
- pytest -v -s tests/kernels/attention/test_attention_selector.py
|
||||
|
||||
- label: Kernels Attention Test %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
timeout_in_minutes: 60
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
parallelism: 2
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- csrc/attention/
|
||||
@@ -3095,10 +3084,10 @@ steps:
|
||||
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
|
||||
- label: Kernels MoE Test %N # TBD
|
||||
timeout_in_minutes: 180
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
parallelism: 4
|
||||
parallelism: 5
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/cutlass_w8a8/moe/
|
||||
|
||||
@@ -12,7 +12,6 @@ steps:
|
||||
- tests/basic_correctness/test_cpu_offload
|
||||
- tests/basic_correctness/test_mem.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- 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
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s basic_correctness/test_mem.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s basic_correctness/test_basic_correctness.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s basic_correctness/test_cpu_offload.py
|
||||
|
||||
@@ -14,8 +14,7 @@ steps:
|
||||
- vllm/v1/cudagraph_dispatcher.py
|
||||
- tests/compile/correctness_e2e/test_sequence_parallel.py
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
|
||||
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
|
||||
|
||||
- label: Sequence Parallel Correctness Tests (2xH100)
|
||||
key: sequence-parallel-correctness-tests-2xh100
|
||||
@@ -25,8 +24,7 @@ steps:
|
||||
optional: true
|
||||
num_devices: 2
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
|
||||
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
|
||||
|
||||
- label: AsyncTP Correctness Tests (2xH100)
|
||||
key: asynctp-correctness-tests-2xh100
|
||||
@@ -36,8 +34,7 @@ steps:
|
||||
optional: true
|
||||
num_devices: 2
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
|
||||
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
|
||||
|
||||
- label: AsyncTP Correctness Tests (B200)
|
||||
key: asynctp-correctness-tests-b200
|
||||
@@ -47,8 +44,7 @@ steps:
|
||||
optional: true
|
||||
num_devices: 2
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
|
||||
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
|
||||
|
||||
- label: Distributed Compile Unit Tests (2xH100)
|
||||
key: distributed-compile-unit-tests-2xh100
|
||||
@@ -61,8 +57,7 @@ steps:
|
||||
- vllm/model_executor/layers
|
||||
- tests/compile/passes/distributed/
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- pytest -s -v tests/compile/passes/distributed
|
||||
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -s -v tests/compile/passes/distributed
|
||||
|
||||
- label: Fusion and Compile Unit Tests (2xB200)
|
||||
key: fusion-and-compile-unit-tests-2xb200
|
||||
|
||||
@@ -32,11 +32,10 @@ steps:
|
||||
- tests/entrypoints/openai/test_multi_api_servers.py
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
|
||||
- DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
|
||||
|
||||
- label: Distributed Compile + RPC Tests (2 GPUs)
|
||||
key: distributed-compile-rpc-tests-2-gpus
|
||||
@@ -56,10 +55,9 @@ steps:
|
||||
- tests/entrypoints/llm/test_collective_rpc.py
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- pytest -v -s entrypoints/llm/test_collective_rpc.py
|
||||
- pytest -v -s ./compile/fullgraph/test_basic_correctness.py
|
||||
- pytest -v -s ./compile/test_wrapper.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s entrypoints/llm/test_collective_rpc.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s ./compile/fullgraph/test_basic_correctness.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s ./compile/test_wrapper.py
|
||||
|
||||
- label: Distributed Torchrun + Shutdown Tests (2 GPUs)
|
||||
key: distributed-torchrun-shutdown-tests-2-gpus
|
||||
@@ -78,11 +76,10 @@ steps:
|
||||
- tests/v1/worker/test_worker_memory_snapshot.py
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
|
||||
- VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
|
||||
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
|
||||
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s v1/worker/test_worker_memory_snapshot.py
|
||||
|
||||
- label: Distributed Torchrun + Examples (4 GPUs)
|
||||
key: distributed-torchrun-examples-4-gpus
|
||||
@@ -97,24 +94,23 @@ steps:
|
||||
- tests/examples/features/data_parallel/data_parallel_offline.py
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
# test with torchrun tp=2 and external_dp=2
|
||||
- torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
|
||||
# test with torchrun tp=2 and pp=2
|
||||
- PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
|
||||
# test with torchrun tp=4 and dp=1
|
||||
- TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
# test with torchrun tp=2, pp=2 and dp=1
|
||||
- PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
# test with torchrun tp=1 and dp=4 with ep
|
||||
- DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
# test with torchrun tp=2 and dp=2 with ep
|
||||
- TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
# test with internal dp
|
||||
- python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager
|
||||
# rlhf examples
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py
|
||||
|
||||
- label: Distributed DP Tests (4 GPUs)
|
||||
key: distributed-dp-tests-4-gpus
|
||||
@@ -128,14 +124,13 @@ steps:
|
||||
- tests/distributed/test_utils
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
|
||||
- TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
|
||||
- TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
|
||||
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
|
||||
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
|
||||
- pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
|
||||
- pytest -v -s distributed/test_utils.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_utils.py
|
||||
|
||||
- label: Distributed Compile + Comm (4 GPUs)
|
||||
key: distributed-compile-comm-4-gpus
|
||||
@@ -151,13 +146,12 @@ steps:
|
||||
- tests/distributed/test_multiproc_executor.py
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- pytest -v -s compile/fullgraph/test_basic_correctness.py
|
||||
- pytest -v -s distributed/test_pynccl.py
|
||||
- pytest -v -s distributed/test_events.py
|
||||
- pytest -v -s distributed/test_symm_mem_allreduce.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s compile/fullgraph/test_basic_correctness.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_pynccl.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_events.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_symm_mem_allreduce.py
|
||||
# test multi-node TP with multiproc executor (simulated on single node)
|
||||
- pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
|
||||
|
||||
- label: Distributed Tests (8 GPUs)(H100)
|
||||
key: distributed-tests-8-gpus-h100
|
||||
@@ -176,9 +170,8 @@ steps:
|
||||
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
# test with torchrun tp=2 and dp=4 with ep
|
||||
- torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
|
||||
|
||||
- label: Distributed Tests (4 GPUs)(A100)
|
||||
key: distributed-tests-4-gpus-a100
|
||||
@@ -271,9 +264,7 @@ steps:
|
||||
- tests/distributed/test_pipeline_parallel.py
|
||||
- tests/basic_correctness/test_basic_correctness.py
|
||||
commands:
|
||||
- export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- pytest -v -s distributed/test_ray_v2_executor.py
|
||||
- pytest -v -s distributed/test_ray_v2_executor_e2e.py
|
||||
- pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
|
||||
- TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -k "ray"
|
||||
- VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_ray_v2_executor.py
|
||||
- VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_ray_v2_executor_e2e.py
|
||||
- VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
|
||||
- VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -k "ray"
|
||||
|
||||
@@ -74,6 +74,16 @@ steps:
|
||||
- tests/v1/e2e/general/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi250_1
|
||||
timeout_in_minutes: 35
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/v1/
|
||||
- tests/v1/e2e/general/
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
- label: V1 e2e (2 GPUs)
|
||||
key: v1-e2e-2-gpus
|
||||
|
||||
@@ -22,10 +22,9 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/llm
|
||||
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 --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
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -41,9 +40,8 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/serve
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -59,8 +57,7 @@ steps:
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -77,9 +74,8 @@ steps:
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/chat_completion
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -128,8 +124,7 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/speech_to_text
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/speech_to_text
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/speech_to_text
|
||||
|
||||
- label: Entrypoints Integration (Multimodal)
|
||||
device: h200_35gb
|
||||
@@ -140,8 +135,7 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/multimodal
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/multimodal
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/multimodal
|
||||
|
||||
- label: Entrypoints Integration (Pooling)
|
||||
key: entrypoints-integration-pooling
|
||||
@@ -151,8 +145,7 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/pooling
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/pooling
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/pooling
|
||||
|
||||
- label: OpenAI API Correctness
|
||||
key: openai-api-correctness
|
||||
|
||||
@@ -74,6 +74,20 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
parallelism: 2
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 55
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- csrc/attention/
|
||||
- vllm/v1/attention
|
||||
- vllm/model_executor/layers/attention
|
||||
- tests/kernels/attention
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/envs.py
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
- label: Kernels Attention DiffKV Test (H100)
|
||||
key: kernels-attention-diffkv-test-h100
|
||||
@@ -104,6 +118,7 @@ steps:
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/
|
||||
- vllm/model_executor/layers/quantization
|
||||
- vllm/config/
|
||||
- tests/kernels/quantization
|
||||
- tests/kernels/quantization/test_rocm_skinny_gemms.py
|
||||
- vllm/_aiter_ops.py
|
||||
@@ -127,6 +142,22 @@ steps:
|
||||
- pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
- pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
parallelism: 5
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 50
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/cutlass_w8a8/moe/
|
||||
- csrc/moe/
|
||||
- tests/kernels/moe
|
||||
- vllm/model_executor/layers/fused_moe/
|
||||
- vllm/distributed/device_communicators/
|
||||
- vllm/envs.py
|
||||
- vllm/config
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Kernels Mamba Test
|
||||
key: kernels-mamba-test
|
||||
|
||||
@@ -50,8 +50,7 @@ steps:
|
||||
- csrc/
|
||||
- vllm/model_executor/layers/quantization
|
||||
commands:
|
||||
- export VLLM_USE_DEEP_GEMM=0 # We found Triton is faster than DeepGEMM for H100
|
||||
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4
|
||||
- VLLM_USE_DEEP_GEMM=0 pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4 # Triton is faster than DeepGEMM for H100
|
||||
|
||||
- label: LM Eval Small Models (B200)
|
||||
key: lm-eval-small-models-b200
|
||||
@@ -101,6 +100,14 @@ steps:
|
||||
num_devices: 8
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_8
|
||||
timeout_in_minutes: 180
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_ROCM_ARCH=gfx942 pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt # Limit Quark compilation to save time
|
||||
|
||||
- label: MoE Refactor Integration Test (H100 - TEMPORARY)
|
||||
key: moe-refactor-integration-test-h100-temporary
|
||||
|
||||
@@ -36,14 +36,14 @@ steps:
|
||||
commands:
|
||||
# FIXIT: find out which code initialize cuda before running the test
|
||||
# before the fix, we need to use spawn to test it
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
#
|
||||
# Alot of these tests are on the edge of OOMing
|
||||
- export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
||||
#
|
||||
# There is some Tensor Parallelism related processing logic in LoRA that
|
||||
# requires multi-GPU testing for validation.
|
||||
- pytest -v -s -x lora/test_chatglm3_tp.py
|
||||
- pytest -v -s -x lora/test_llama_tp.py
|
||||
- pytest -v -s -x lora/test_qwen3_with_multi_loras.py
|
||||
- pytest -v -s -x lora/test_olmoe_tp.py
|
||||
- pytest -v -s -x lora/test_gptoss_tp.py
|
||||
- pytest -v -s -x lora/test_qwen35_densemodel_lora.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_chatglm3_tp.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_llama_tp.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_qwen3_with_multi_loras.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_olmoe_tp.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_gptoss_tp.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_qwen35_densemodel_lora.py
|
||||
@@ -18,9 +18,8 @@ steps:
|
||||
- vllm/v1/
|
||||
- tests/v1/spec_decode
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# TODO: create another `optional` test group for slow tests
|
||||
- pytest -v -s -m 'not slow_test' v1/spec_decode
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not slow_test' v1/spec_decode
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
@@ -50,12 +49,11 @@ steps:
|
||||
- tests/v1/test_request.py
|
||||
- tests/v1/test_outputs.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s v1/sample
|
||||
- pytest -v -s v1/logits_processors
|
||||
- pytest -v -s v1/test_oracle.py
|
||||
- pytest -v -s v1/test_request.py
|
||||
- pytest -v -s v1/test_outputs.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/sample
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/logits_processors
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/test_oracle.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/test_request.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/test_outputs.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -93,18 +91,23 @@ steps:
|
||||
- tests/entrypoints/openai/correctness/test_lmeval.py
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# split the test to avoid interference
|
||||
- pytest -v -s -m 'not cpu_test' v1/core
|
||||
- pytest -v -s v1/executor
|
||||
- pytest -v -s v1/kv_offload
|
||||
- pytest -v -s v1/simple_kv_offload
|
||||
- pytest -v -s v1/worker
|
||||
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'not cpu_test' v1/metrics
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not cpu_test' v1/core
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/executor
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/kv_offload
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/simple_kv_offload
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/worker
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not cpu_test' v1/metrics
|
||||
# Integration test for streaming correctness (requires special branch).
|
||||
- pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
|
||||
- pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 60
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: V1 Others (CPU)
|
||||
key: v1-others-cpu
|
||||
@@ -147,8 +150,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
|
||||
- tests/v1/kv_connector/extract_hidden_states_integration
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s v1/kv_connector/extract_hidden_states_integration
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/kv_connector/extract_hidden_states_integration
|
||||
|
||||
- label: Extract Hidden States Integration (2 GPUs)
|
||||
key: extract-hidden-states-integration-2-gpus
|
||||
@@ -161,8 +163,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
|
||||
- tests/v1/kv_connector/extract_hidden_states_integration
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration
|
||||
|
||||
- label: Regression
|
||||
key: regression
|
||||
@@ -354,10 +355,9 @@ steps:
|
||||
- vllm/model_executor/layers
|
||||
- tests/v1/determinism/
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pip install pytest-timeout pytest-forked
|
||||
- pytest -v -s v1/determinism/test_batch_invariance.py
|
||||
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_batch_invariance.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
|
||||
|
||||
- label: Batch Invariance (H100)
|
||||
key: batch-invariance-h100
|
||||
@@ -368,12 +368,11 @@ steps:
|
||||
- vllm/model_executor/layers
|
||||
- tests/v1/determinism/
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pip install pytest-timeout pytest-forked
|
||||
- pytest -v -s v1/determinism/test_batch_invariance.py
|
||||
- pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
|
||||
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
|
||||
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_batch_invariance.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
|
||||
|
||||
- label: Batch Invariance (B200)
|
||||
key: batch-invariance-b200
|
||||
@@ -384,14 +383,13 @@ steps:
|
||||
- vllm/model_executor/layers
|
||||
- tests/v1/determinism/
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pip install pytest-timeout pytest-forked
|
||||
- pytest -v -s v1/determinism/test_batch_invariance.py
|
||||
- pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
|
||||
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
|
||||
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
|
||||
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
|
||||
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_batch_invariance.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
|
||||
|
||||
- label: Acceptance Length Test (Large Models) # optional
|
||||
device: h200_35gb
|
||||
@@ -406,5 +404,4 @@ steps:
|
||||
- vllm/model_executor/models/mlp_speculator.py
|
||||
- tests/v1/spec_decode/test_acceptance_length.py
|
||||
commands:
|
||||
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
|
||||
- pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
|
||||
|
||||
@@ -13,13 +13,12 @@ steps:
|
||||
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
commands:
|
||||
- apt-get update && apt-get install -y curl libsodium23
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# Dump tracebacks of all threads if a test hangs, so a wedged GPU/CUDA
|
||||
# init surfaces a stack instead of silently stalling.
|
||||
- export PYTHONFAULTHANDLER=1
|
||||
# Per-test watchdog: a single hung test (e.g. stuck during engine/CUDA
|
||||
# init) fails fast with a traceback instead of running until the global
|
||||
# build timeout. The `thread` method also handles hangs inside C/CUDA
|
||||
# calls that the signal method cannot interrupt.
|
||||
- pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread
|
||||
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
|
||||
#
|
||||
# Env vars are inlined because CONTINUE_ON_FAILURE wraps each command
|
||||
# in a subshell, so a standalone `export` would be lost.
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTHONFAULTHANDLER=1 pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTHONFAULTHANDLER=1 pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
|
||||
|
||||
@@ -16,15 +16,14 @@ steps:
|
||||
- tests/entrypoints/llm/test_struct_output_generate.py
|
||||
commands:
|
||||
- set -x
|
||||
- export VLLM_USE_V2_MODEL_RUNNER=1
|
||||
- pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
|
||||
# This requires eager until we sort out CG correctness issues.
|
||||
# TODO: remove ENFORCE_EAGER here after https://github.com/vllm-project/vllm/pull/32936 is merged.
|
||||
- ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
|
||||
- pytest -v -s v1/e2e/general/test_context_length.py
|
||||
- pytest -v -s v1/e2e/general/test_min_tokens.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/e2e/general/test_context_length.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/e2e/general/test_min_tokens.py
|
||||
# Temporary hack filter to exclude ngram spec decoding based tests.
|
||||
- pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
|
||||
|
||||
- label: Model Runner V2 Examples
|
||||
device: h200_35gb
|
||||
@@ -42,26 +41,25 @@ steps:
|
||||
- examples/features/tensorize_vllm_model.py
|
||||
commands:
|
||||
- set -x
|
||||
- export VLLM_USE_V2_MODEL_RUNNER=1
|
||||
- pip install tensorizer # for tensorizer test
|
||||
- python3 basic/offline_inference/chat.py # for basic
|
||||
- python3 basic/offline_inference/generate.py --model facebook/opt-125m
|
||||
#- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO
|
||||
#- python3 basic/offline_inference/embed.py # TODO
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/chat.py # for basic
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/generate.py --model facebook/opt-125m
|
||||
#- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO
|
||||
#- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/embed.py # TODO
|
||||
# for multi-modal models
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# for pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# for features demo
|
||||
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
|
||||
- python3 deployment/llm_engine_example.py
|
||||
- python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
|
||||
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 features/automatic_prefix_caching/prefix_caching_offline.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 deployment/llm_engine_example.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && VLLM_USE_V2_MODEL_RUNNER=1 python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
|
||||
# https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
|
||||
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
|
||||
|
||||
- label: Model Runner V2 Distributed (2 GPUs)
|
||||
key: model-runner-v2-distributed-2-gpus
|
||||
@@ -76,13 +74,11 @@ steps:
|
||||
- tests/v1/distributed/test_eagle_dp.py
|
||||
commands:
|
||||
- set -x
|
||||
- export VLLM_USE_V2_MODEL_RUNNER=1
|
||||
# The "and not True" here is a hacky way to exclude the prompt_embeds cases which aren't yet supported.
|
||||
- TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True"
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
|
||||
|
||||
- label: Model Runner V2 Pipeline Parallelism (4 GPUs)
|
||||
key: model-runner-v2-pipeline-parallelism-4-gpus
|
||||
@@ -97,10 +93,9 @@ steps:
|
||||
- tests/v1/distributed/test_pp_dp_v2.py
|
||||
commands:
|
||||
- set -x
|
||||
- export VLLM_USE_V2_MODEL_RUNNER=1
|
||||
- pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
|
||||
- pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
|
||||
- pytest -v -s v1/distributed/test_pp_dp_v2.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/distributed/test_pp_dp_v2.py
|
||||
|
||||
- label: Model Runner V2 Spec Decode
|
||||
device: h200_35gb
|
||||
@@ -115,8 +110,7 @@ steps:
|
||||
- tests/v1/e2e/spec_decode/test_spec_decode.py
|
||||
commands:
|
||||
- set -x
|
||||
- export VLLM_USE_V2_MODEL_RUNNER=1
|
||||
- pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
|
||||
- pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
|
||||
- pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
|
||||
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
|
||||
|
||||
@@ -68,7 +68,6 @@ steps:
|
||||
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work
|
||||
mirror:
|
||||
amd:
|
||||
soft_fail: true
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -107,6 +107,12 @@ steps:
|
||||
- tests/compile/passes
|
||||
commands:
|
||||
- pytest -s -v compile/passes --ignore compile/passes/distributed
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
timeout_in_minutes: 180
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: PyTorch Fullgraph Smoke Test
|
||||
key: pytorch-fullgraph-smoke-test
|
||||
|
||||
@@ -23,17 +23,15 @@ steps:
|
||||
# - tests/entrypoints/openai/test_uds.py
|
||||
- tests/v1/sample/test_logprobs_e2e.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
|
||||
- pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py
|
||||
# - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py
|
||||
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"
|
||||
|
||||
# - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
|
||||
- pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
|
||||
# - pytest -v -s entrypoints/openai/test_return_token_ids.py
|
||||
# - pytest -v -s entrypoints/openai/test_uds.py
|
||||
- pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"
|
||||
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
|
||||
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/test_return_token_ids.py
|
||||
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/test_uds.py
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"
|
||||
|
||||
- label: Rust Frontend Serve/Admin Coverage
|
||||
timeout_in_minutes: 60
|
||||
@@ -51,13 +49,11 @@ steps:
|
||||
- tests/entrypoints/serve/instrumentator/test_metrics.py
|
||||
# - tests/entrypoints/serve/dev/test_sleep.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
- pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
|
||||
- pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
|
||||
- pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
|
||||
# - pytest -v -s entrypoints/serve/dev/test_sleep.py
|
||||
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
|
||||
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/dev/test_sleep.py
|
||||
|
||||
- label: Rust Frontend Core Correctness
|
||||
timeout_in_minutes: 30
|
||||
@@ -69,9 +65,7 @@ steps:
|
||||
- tests/utils.py
|
||||
- tests/entrypoints/openai/correctness/test_lmeval.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
|
||||
- label: Rust Frontend Tool Use
|
||||
timeout_in_minutes: 60
|
||||
@@ -83,9 +77,7 @@ steps:
|
||||
- tests/utils.py
|
||||
- tests/tool_use/
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
|
||||
|
||||
- label: Rust Frontend Distributed
|
||||
timeout_in_minutes: 30
|
||||
@@ -103,9 +95,6 @@ steps:
|
||||
- tests/v1/distributed/test_hybrid_lb_dp.py
|
||||
- tests/v1/distributed/test_internal_lb_dp.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info"
|
||||
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info"
|
||||
|
||||
@@ -12,6 +12,20 @@ steps:
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
timeout_in_minutes: 45
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/v1/worker/gpu/spec_decode/
|
||||
- vllm/model_executor/model_loader/
|
||||
- vllm/v1/sample/
|
||||
- vllm/model_executor/layers/
|
||||
- tests/v1/e2e/spec_decode/
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
- label: Spec Decode Eagle Nightly B200
|
||||
key: spec-decode-eagle-nightly-b200
|
||||
@@ -139,8 +153,7 @@ steps:
|
||||
- vllm/model_executor/models/qwen3_dflash.py
|
||||
- tests/v1/spec_decode/test_speculators_correctness.py
|
||||
commands:
|
||||
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
|
||||
- pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
|
||||
|
||||
- label: Spec Decode MTP hybrid (B200)
|
||||
timeout_in_minutes: 30
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
---
|
||||
name: ci-fails-buildkite
|
||||
description: Fetch and diagnose vLLM Buildkite CI failure logs. Use when investigating failing CI jobs on a PR or build, when the user pastes a buildkite.com URL, or asks to fetch/diagnose CI logs.
|
||||
---
|
||||
|
||||
# Diagnosing vLLM Buildkite CI Failures
|
||||
|
||||
Buildkite logs are public; no login needed.
|
||||
|
||||
`.buildkite/scripts/ci-fetch-log.sh` saves each log as `ci-<build>-<job-name>.log`, stripped of timestamps and ANSI codes. Existing files are kept; set `CI_FETCH_LOG_FORCE=1` to refetch.
|
||||
|
||||
## Fetching logs
|
||||
|
||||
```bash
|
||||
# All failed jobs in a PR's latest build (current branch's PR if omitted):
|
||||
.buildkite/scripts/ci-fetch-log.sh --pr <PR>
|
||||
|
||||
# All failed jobs in a build (--soft also includes soft-failed jobs;
|
||||
# --all fetches every finished job):
|
||||
.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>"
|
||||
|
||||
# One job — `gh pr checks` URLs (#<job_uuid>) and web UI URLs (?sid=) both
|
||||
# work; pass "-" as a second argument to stream to stdout:
|
||||
.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>#<job_uuid>"
|
||||
```
|
||||
|
||||
To clean an already-downloaded log with `.buildkite/scripts/ci-clean-log.sh`:
|
||||
|
||||
```bash
|
||||
./ci-clean-log.sh ci.log
|
||||
```
|
||||
|
||||
## Reference
|
||||
|
||||
See [docs/contributing/ci/failures.md](../../../docs/contributing/ci/failures.md) for the full guide: filing CI failure issues, investigating/bisecting, reproducing flaky tests, and daily triage.
|
||||
+2
-3
@@ -2,15 +2,14 @@
|
||||
# for more info about CODEOWNERS file
|
||||
|
||||
# This lists cover the "core" components of vLLM that require careful review
|
||||
/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng @vadiklyutiy
|
||||
/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng
|
||||
/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi
|
||||
/vllm/lora @jeejeelee
|
||||
/vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
|
||||
/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye
|
||||
/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye
|
||||
/vllm/model_executor/layers/mamba @tdoublep @tomeras91
|
||||
/vllm/model_executor/layers/mamba/gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
|
||||
/vllm/model_executor/layers/rotary_embedding.py @vadiklyutiy
|
||||
/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
|
||||
/vllm/model_executor/model_loader @22quinn
|
||||
/vllm/model_executor/layers/batch_invariant.py @yewentao256
|
||||
/vllm/ir @ProExpertProg
|
||||
|
||||
+3
-1
@@ -199,7 +199,9 @@ cython_debug/
|
||||
.vscode/
|
||||
|
||||
# Claude
|
||||
.claude/
|
||||
.claude/*
|
||||
!.claude/skills/
|
||||
!.claude/skills/**
|
||||
|
||||
# Codex
|
||||
.codex/
|
||||
|
||||
@@ -114,17 +114,6 @@ Follow these rules for all code changes in this repository:
|
||||
- Keep comments and docstrings minimal and concise.
|
||||
- Assume the reader is familiar with vLLM.
|
||||
|
||||
### Diagnosing CI failures
|
||||
|
||||
Buildkite logs are public; no login needed. Details: [docs/contributing/ci/failures.md](docs/contributing/ci/failures.md).
|
||||
|
||||
```bash
|
||||
# All failed-job logs for a PR's latest build (current branch's PR if omitted):
|
||||
.buildkite/scripts/ci-fetch-log.sh --pr <PR>
|
||||
# Any Buildkite build or job URL also works:
|
||||
.buildkite/scripts/ci-fetch-log.sh "<buildkite_url>"
|
||||
```
|
||||
|
||||
### Commit messages
|
||||
|
||||
Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:
|
||||
@@ -146,6 +135,12 @@ Do not modify code in these areas without first reading and following the
|
||||
linked guide. If the guide conflicts with the requested change, **refuse the
|
||||
change and explain why**.
|
||||
|
||||
Security reviewers should start with [`SECURITY.md`](SECURITY.md),
|
||||
[`docs/usage/security.md`](docs/usage/security.md), and
|
||||
[`docs/contributing/vulnerability_management.md`](docs/contributing/vulnerability_management.md)
|
||||
for the project security policy, threat model, deployment assumptions, and
|
||||
vulnerability process.
|
||||
|
||||
- **Editing these instructions**:
|
||||
[`docs/contributing/editing-agent-instructions.md`](docs/contributing/editing-agent-instructions.md)
|
||||
— Rules for modifying AGENTS.md or any domain-specific guide it references.
|
||||
|
||||
+56
-72
@@ -319,82 +319,35 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
endif()
|
||||
|
||||
#
|
||||
# _C extension
|
||||
# Legacy _C extension (ROCm only — CUDA ops migrated to _C_stable_libtorch)
|
||||
#
|
||||
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/quantization/activation_kernels.cu"
|
||||
"csrc/torch_bindings.cpp")
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
|
||||
|
||||
# Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
|
||||
set(CUTLASS_REVISION "v4.4.2")
|
||||
|
||||
# Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
|
||||
if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
|
||||
set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
|
||||
endif()
|
||||
|
||||
if(VLLM_CUTLASS_SRC_DIR)
|
||||
if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
|
||||
get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
|
||||
endif()
|
||||
message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
|
||||
FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
|
||||
else()
|
||||
FetchContent_Declare(
|
||||
cutlass
|
||||
GIT_REPOSITORY https://github.com/nvidia/cutlass.git
|
||||
# Please keep this in sync with CUTLASS_REVISION line above.
|
||||
GIT_TAG ${CUTLASS_REVISION}
|
||||
GIT_PROGRESS TRUE
|
||||
|
||||
# Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
|
||||
# Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
|
||||
# So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
|
||||
GIT_SHALLOW TRUE
|
||||
)
|
||||
endif()
|
||||
FetchContent_MakeAvailable(cutlass)
|
||||
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${VLLM_EXT_SRC}"
|
||||
CUDA_ARCHS "${CUDA_ARCHS}")
|
||||
|
||||
# if CUDA endif
|
||||
endif()
|
||||
|
||||
if (VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
|
||||
# TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
|
||||
list(APPEND VLLM_EXT_SRC
|
||||
if(VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set(VLLM_EXT_SRC
|
||||
"csrc/torch_bindings.cpp"
|
||||
"csrc/custom_quickreduce.cu"
|
||||
"csrc/cuda_view.cu"
|
||||
"csrc/libtorch_stable/cuda_utils_kernels.cu"
|
||||
)
|
||||
# if ROCM endif
|
||||
endif()
|
||||
"csrc/libtorch_stable/cuda_utils_kernels.cu")
|
||||
|
||||
message(STATUS "Enabling C extension.")
|
||||
define_extension_target(
|
||||
_C
|
||||
DESTINATION vllm
|
||||
LANGUAGE ${VLLM_GPU_LANG}
|
||||
SOURCES ${VLLM_EXT_SRC}
|
||||
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
|
||||
ARCHITECTURES ${VLLM_GPU_ARCHES}
|
||||
INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
|
||||
INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
|
||||
USE_SABI 3
|
||||
WITH_SOABI)
|
||||
message(STATUS "Enabling C extension.")
|
||||
define_extension_target(
|
||||
_C
|
||||
DESTINATION vllm
|
||||
LANGUAGE ${VLLM_GPU_LANG}
|
||||
SOURCES ${VLLM_EXT_SRC}
|
||||
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
|
||||
ARCHITECTURES ${VLLM_GPU_ARCHES}
|
||||
INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
|
||||
INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
|
||||
USE_SABI 3
|
||||
WITH_SOABI)
|
||||
|
||||
# If CUTLASS is compiled on NVCC >= 12.5, it by default uses
|
||||
# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
|
||||
# driver API. This causes problems when linking with earlier versions of CUDA.
|
||||
# Setting this variable sidesteps the issue by calling the driver directly.
|
||||
target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
|
||||
# If CUTLASS is compiled on NVCC >= 12.5, it by default uses
|
||||
# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
|
||||
# driver API. This causes problems when linking with earlier versions of CUDA.
|
||||
# Setting this variable sidesteps the issue by calling the driver directly.
|
||||
target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
|
||||
endif() # _C HIP endif
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
#
|
||||
@@ -403,6 +356,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set(VLLM_STABLE_EXT_SRC
|
||||
"csrc/libtorch_stable/torch_bindings.cpp"
|
||||
"csrc/libtorch_stable/activation_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/activation_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
|
||||
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
|
||||
@@ -429,6 +383,38 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
"csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
|
||||
|
||||
# Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
|
||||
set(CUTLASS_REVISION "v4.4.2")
|
||||
|
||||
# Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
|
||||
if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
|
||||
set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
|
||||
endif()
|
||||
|
||||
if(VLLM_CUTLASS_SRC_DIR)
|
||||
if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
|
||||
get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
|
||||
endif()
|
||||
message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
|
||||
FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
|
||||
else()
|
||||
FetchContent_Declare(
|
||||
cutlass
|
||||
GIT_REPOSITORY https://github.com/nvidia/cutlass.git
|
||||
# Please keep this in sync with CUTLASS_REVISION line above.
|
||||
GIT_TAG ${CUTLASS_REVISION}
|
||||
GIT_PROGRESS TRUE
|
||||
|
||||
# Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
|
||||
# Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
|
||||
# So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
|
||||
GIT_SHALLOW TRUE
|
||||
)
|
||||
endif()
|
||||
FetchContent_MakeAvailable(cutlass)
|
||||
|
||||
list(APPEND VLLM_STABLE_EXT_SRC
|
||||
"csrc/libtorch_stable/cuda_view.cu"
|
||||
"csrc/libtorch_stable/cuda_utils_kernels.cu"
|
||||
@@ -929,7 +915,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
SRCS "${FP4_SM120_SRCS}"
|
||||
CUDA_ARCHS "${FP4_SM120_ARCHS}")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
|
||||
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
|
||||
message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}")
|
||||
@@ -962,7 +947,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
SRCS "${FP4_SM100_SRCS}"
|
||||
CUDA_ARCHS "${FP4_SM100_ARCHS}")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
|
||||
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
|
||||
message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}")
|
||||
|
||||
@@ -60,6 +60,7 @@ endif()
|
||||
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
|
||||
set(QUTLASS_SOURCES
|
||||
csrc/qutlass_registration.cpp
|
||||
${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp
|
||||
${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu
|
||||
${qutlass_SOURCE_DIR}/qutlass/csrc/gemm_ada.cu
|
||||
@@ -78,8 +79,19 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
|
||||
|
||||
if(CUTLASS_INCLUDE_DIR AND EXISTS "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h")
|
||||
list(APPEND QUTLASS_INCLUDES "${CUTLASS_INCLUDE_DIR}")
|
||||
if(CUTLASS_TOOLS_UTIL_INCLUDE_DIR AND
|
||||
EXISTS "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}/cutlass/util/packed_stride.hpp")
|
||||
list(APPEND QUTLASS_INCLUDES "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}")
|
||||
else()
|
||||
get_filename_component(_qutlass_cutlass_root "${CUTLASS_INCLUDE_DIR}" DIRECTORY)
|
||||
if(EXISTS "${_qutlass_cutlass_root}/tools/util/include/cutlass/util/packed_stride.hpp")
|
||||
list(APPEND QUTLASS_INCLUDES "${_qutlass_cutlass_root}/tools/util/include")
|
||||
endif()
|
||||
endif()
|
||||
elseif(EXISTS "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include/cutlass/cutlass.h")
|
||||
list(APPEND QUTLASS_INCLUDES "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include")
|
||||
list(APPEND QUTLASS_INCLUDES
|
||||
"${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include"
|
||||
"${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/tools/util/include")
|
||||
message(STATUS "[QUTLASS] Using QuTLASS vendored CUTLASS headers (no vLLM CUTLASS detected).")
|
||||
else()
|
||||
message(FATAL_ERROR "[QUTLASS] CUTLASS headers not found. "
|
||||
@@ -91,12 +103,23 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
|
||||
CUDA_ARCHS "${QUTLASS_ARCHS}"
|
||||
)
|
||||
|
||||
target_sources(_C PRIVATE ${QUTLASS_SOURCES})
|
||||
target_include_directories(_C PRIVATE ${QUTLASS_INCLUDES})
|
||||
target_compile_definitions(_C PRIVATE
|
||||
# QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
|
||||
# Keep it as its own extension (registers torch.ops._qutlass_C).
|
||||
define_extension_target(
|
||||
_qutlass_C
|
||||
DESTINATION vllm
|
||||
LANGUAGE ${VLLM_GPU_LANG}
|
||||
SOURCES ${QUTLASS_SOURCES}
|
||||
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
|
||||
ARCHITECTURES ${VLLM_GPU_ARCHES}
|
||||
INCLUDE_DIRECTORIES ${QUTLASS_INCLUDES}
|
||||
USE_SABI 3
|
||||
WITH_SOABI)
|
||||
|
||||
target_compile_definitions(_qutlass_C PRIVATE
|
||||
QUTLASS_DISABLE_PYBIND=1
|
||||
TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
|
||||
)
|
||||
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
|
||||
|
||||
set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
|
||||
@@ -111,4 +134,5 @@ else()
|
||||
"[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in "
|
||||
"CUDA_ARCHS='${CUDA_ARCHS}'.")
|
||||
endif()
|
||||
add_custom_target(_qutlass_C)
|
||||
endif()
|
||||
|
||||
@@ -4,6 +4,9 @@
|
||||
#ifdef CPU_CAPABILITY_AMXBF16
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_amx.hpp"
|
||||
#endif
|
||||
#if defined(__riscv_v)
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_rvv.hpp"
|
||||
#endif
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
|
||||
|
||||
#define VLLM_DISPATCH_CASE_16B_TYPES(...) \
|
||||
@@ -319,6 +322,8 @@ void cpu_gemm_wna16(
|
||||
return ISA::AMX;
|
||||
} else if (isa_hint == "vec") {
|
||||
return ISA::VEC;
|
||||
} else if (isa_hint == "rvv") {
|
||||
return ISA::RVV;
|
||||
} else {
|
||||
TORCH_CHECK(false, "unsupported isa hint: " + isa_hint);
|
||||
}
|
||||
@@ -397,6 +402,40 @@ void cpu_gemm_wna16(
|
||||
pack_factor);
|
||||
return;
|
||||
}
|
||||
} else if (isa == ISA::RVV) {
|
||||
using gemm_t = cpu_micro_gemm::MicroGemm<ISA::RVV, scalar_t>;
|
||||
if (has_zp) {
|
||||
using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, true, false>;
|
||||
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
|
||||
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
|
||||
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
|
||||
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
|
||||
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
|
||||
scales_group_stride, zeros_group_stride, group_num, group_size,
|
||||
pack_factor);
|
||||
return;
|
||||
}
|
||||
if (use_desc_act) {
|
||||
using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, true>;
|
||||
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
|
||||
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
|
||||
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
|
||||
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
|
||||
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
|
||||
scales_group_stride, zeros_group_stride, group_num, group_size,
|
||||
pack_factor);
|
||||
return;
|
||||
} else {
|
||||
using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, false>;
|
||||
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
|
||||
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
|
||||
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
|
||||
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
|
||||
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
|
||||
scales_group_stride, zeros_group_stride, group_num, group_size,
|
||||
pack_factor);
|
||||
return;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@@ -0,0 +1,228 @@
|
||||
#ifndef CPU_MICRO_GEMM_RVV_HPP
|
||||
#define CPU_MICRO_GEMM_RVV_HPP
|
||||
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
|
||||
|
||||
#if defined(__riscv_v)
|
||||
|
||||
namespace cpu_micro_gemm {
|
||||
namespace {
|
||||
|
||||
constexpr int32_t RVV_MGEMM_N8 = 8;
|
||||
constexpr int32_t RVV_MGEMM_B_GROUP_STRIDE = 16;
|
||||
|
||||
template <typename scalar_t>
|
||||
FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32(const scalar_t* ptr);
|
||||
|
||||
template <>
|
||||
FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32<float>(const float* ptr) {
|
||||
return RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, RVV_MGEMM_N8);
|
||||
}
|
||||
|
||||
template <>
|
||||
FORCE_INLINE fixed_fp32x8_t
|
||||
load_row8_b_as_f32<c10::Half>(const c10::Half* ptr) {
|
||||
#if defined(__riscv_zvfh)
|
||||
fixed_fp16x8_t vec = RVVI(__riscv_vle16_v_f16, LMUL_128)(
|
||||
reinterpret_cast<const _Float16*>(ptr), RVV_MGEMM_N8);
|
||||
return RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
|
||||
#else
|
||||
alignas(32) float values[RVV_MGEMM_N8];
|
||||
for (int32_t i = 0; i < RVV_MGEMM_N8; ++i) {
|
||||
values[i] = static_cast<float>(ptr[i]);
|
||||
}
|
||||
return RVVI(__riscv_vle32_v_f32, LMUL_256)(values, RVV_MGEMM_N8);
|
||||
#endif
|
||||
}
|
||||
|
||||
template <>
|
||||
FORCE_INLINE fixed_fp32x8_t
|
||||
load_row8_b_as_f32<c10::BFloat16>(const c10::BFloat16* ptr) {
|
||||
#if defined(__riscv_zvfbfmin)
|
||||
fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
|
||||
reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
|
||||
fixed_bf16x8_t vec =
|
||||
RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16, LMUL_128)(raw);
|
||||
return RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
|
||||
#else
|
||||
fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
|
||||
reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
|
||||
auto wide = RVVI(__riscv_vzext_vf2_u32, LMUL_256)(raw, RVV_MGEMM_N8);
|
||||
auto shifted = RVVI(__riscv_vsll_vx_u32, LMUL_256)(wide, 16, RVV_MGEMM_N8);
|
||||
return RVVI4(__riscv_vreinterpret_v_u32, LMUL_256, _f32, LMUL_256)(shifted);
|
||||
#endif
|
||||
}
|
||||
|
||||
// Mx8 RVV kernel. B points at one 8-channel half of a 16-channel packed group,
|
||||
// with rows separated by RVV_MGEMM_B_GROUP_STRIDE scalar elements.
|
||||
template <int32_t M, typename scalar_t>
|
||||
FORCE_INLINE void gemm_micro_rvv_fma_mx8_ku4(const scalar_t* __restrict__ a_ptr,
|
||||
const scalar_t* __restrict__ b_ptr,
|
||||
float* __restrict__ c_ptr,
|
||||
const int64_t lda,
|
||||
const int64_t ldc, const int32_t k,
|
||||
const bool accum_c) {
|
||||
static_assert(0 < M && M <= 8);
|
||||
|
||||
#define RVV_ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
|
||||
#define RVV_IF_M(i) if constexpr (M > (i))
|
||||
|
||||
#define RVV_DECL_A(i) const scalar_t* __restrict__ a##i = a_ptr + (i) * lda;
|
||||
RVV_ROWS_APPLY(RVV_DECL_A)
|
||||
#undef RVV_DECL_A
|
||||
|
||||
#define RVV_DECL_ACC(i) fixed_fp32x8_t acc##i;
|
||||
RVV_ROWS_APPLY(RVV_DECL_ACC)
|
||||
#undef RVV_DECL_ACC
|
||||
|
||||
#define RVV_INIT_ACC(i) \
|
||||
RVV_IF_M(i) { \
|
||||
if (accum_c) { \
|
||||
acc##i = RVVI(__riscv_vle32_v_f32, LMUL_256)(c_ptr + (i) * ldc, \
|
||||
RVV_MGEMM_N8); \
|
||||
} else { \
|
||||
acc##i = RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, RVV_MGEMM_N8); \
|
||||
} \
|
||||
}
|
||||
RVV_ROWS_APPLY(RVV_INIT_ACC)
|
||||
#undef RVV_INIT_ACC
|
||||
|
||||
int32_t k_idx = 0;
|
||||
for (; k_idx + 3 < k; k_idx += 4) {
|
||||
#define RVV_FMA_ROW(i, K_OFFSET) \
|
||||
RVV_IF_M(i) { \
|
||||
acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)( \
|
||||
acc##i, static_cast<float>(*(a##i + k_idx + (K_OFFSET))), b, \
|
||||
RVV_MGEMM_N8); \
|
||||
}
|
||||
|
||||
#define RVV_STEP_K(K_OFFSET) \
|
||||
{ \
|
||||
fixed_fp32x8_t b = load_row8_b_as_f32<scalar_t>( \
|
||||
b_ptr + (k_idx + (K_OFFSET)) * RVV_MGEMM_B_GROUP_STRIDE); \
|
||||
RVV_FMA_ROW(0, K_OFFSET) \
|
||||
RVV_FMA_ROW(1, K_OFFSET) \
|
||||
RVV_FMA_ROW(2, K_OFFSET) \
|
||||
RVV_FMA_ROW(3, K_OFFSET) \
|
||||
RVV_FMA_ROW(4, K_OFFSET) \
|
||||
RVV_FMA_ROW(5, K_OFFSET) \
|
||||
RVV_FMA_ROW(6, K_OFFSET) \
|
||||
RVV_FMA_ROW(7, K_OFFSET) \
|
||||
}
|
||||
|
||||
RVV_STEP_K(0)
|
||||
RVV_STEP_K(1)
|
||||
RVV_STEP_K(2)
|
||||
RVV_STEP_K(3)
|
||||
#undef RVV_STEP_K
|
||||
#undef RVV_FMA_ROW
|
||||
}
|
||||
|
||||
for (; k_idx < k; ++k_idx) {
|
||||
fixed_fp32x8_t b =
|
||||
load_row8_b_as_f32<scalar_t>(b_ptr + k_idx * RVV_MGEMM_B_GROUP_STRIDE);
|
||||
#define RVV_TAIL_ROW(i) \
|
||||
RVV_IF_M(i) { \
|
||||
acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)( \
|
||||
acc##i, static_cast<float>(*(a##i + k_idx)), b, RVV_MGEMM_N8); \
|
||||
}
|
||||
RVV_ROWS_APPLY(RVV_TAIL_ROW)
|
||||
#undef RVV_TAIL_ROW
|
||||
}
|
||||
|
||||
#define RVV_STORE_ROW(i) \
|
||||
RVV_IF_M(i) { \
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_256)(c_ptr + (i) * ldc, acc##i, \
|
||||
RVV_MGEMM_N8); \
|
||||
}
|
||||
RVV_ROWS_APPLY(RVV_STORE_ROW)
|
||||
#undef RVV_STORE_ROW
|
||||
|
||||
#undef RVV_ROWS_APPLY
|
||||
#undef RVV_IF_M
|
||||
}
|
||||
|
||||
template <int32_t M, typename scalar_t>
|
||||
FORCE_INLINE void gemm_micro_rvv_mx32_ku4(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
static_assert(0 < M && M <= 8);
|
||||
scalar_t* __restrict__ curr_b_0 = b_ptr;
|
||||
scalar_t* __restrict__ curr_b_1 = b_ptr + b_n_group_stride;
|
||||
|
||||
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0, c_ptr, lda, ldc, k, accum_c);
|
||||
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0 + RVV_MGEMM_N8,
|
||||
c_ptr + RVV_MGEMM_N8, lda, ldc, k, accum_c);
|
||||
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1, c_ptr + 16, lda, ldc, k,
|
||||
accum_c);
|
||||
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1 + RVV_MGEMM_N8, c_ptr + 24, lda,
|
||||
ldc, k, accum_c);
|
||||
}
|
||||
|
||||
class TileGemmRVV {
|
||||
public:
|
||||
template <typename scalar_t>
|
||||
FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
switch (m) {
|
||||
case 1:
|
||||
gemm_micro_rvv_mx32_ku4<1>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 2:
|
||||
gemm_micro_rvv_mx32_ku4<2>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 3:
|
||||
gemm_micro_rvv_mx32_ku4<3>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 4:
|
||||
gemm_micro_rvv_mx32_ku4<4>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 5:
|
||||
gemm_micro_rvv_mx32_ku4<5>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 6:
|
||||
gemm_micro_rvv_mx32_ku4<6>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 7:
|
||||
gemm_micro_rvv_mx32_ku4<7>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 8:
|
||||
gemm_micro_rvv_mx32_ku4<8>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
template <typename scalar_t>
|
||||
class MicroGemm<cpu_utils::ISA::RVV, scalar_t> {
|
||||
public:
|
||||
static constexpr int32_t MaxMSize = 8;
|
||||
static constexpr int32_t NSize = 32;
|
||||
|
||||
public:
|
||||
void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
TileGemmRVV::gemm<scalar_t>(CPU_MICRO_GEMM_PARAMS);
|
||||
}
|
||||
|
||||
static void pack_weight(const scalar_t* __restrict__ weight,
|
||||
scalar_t* __restrict__ packed_weight,
|
||||
const int32_t output_size, const int32_t input_size) {
|
||||
TORCH_CHECK_EQ(output_size % 16, 0);
|
||||
for (int32_t o_idx = 0; o_idx < output_size; ++o_idx) {
|
||||
const scalar_t* __restrict__ curr_weight = weight + o_idx * input_size;
|
||||
scalar_t* __restrict__ curr_packed_weight =
|
||||
packed_weight + (o_idx / 16) * (16 * input_size) + o_idx % 16;
|
||||
for (int32_t i_idx = 0; i_idx < input_size; ++i_idx) {
|
||||
*curr_packed_weight = *curr_weight;
|
||||
|
||||
curr_packed_weight += 16;
|
||||
++curr_weight;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace cpu_micro_gemm
|
||||
|
||||
#endif // defined(__riscv_v)
|
||||
|
||||
#endif // CPU_MICRO_GEMM_RVV_HPP
|
||||
+3
-1
@@ -8,13 +8,15 @@
|
||||
#include "cpu/cpu_types.hpp"
|
||||
|
||||
namespace cpu_utils {
|
||||
enum class ISA { AMX, VEC };
|
||||
enum class ISA { AMX, VEC, RVV };
|
||||
|
||||
inline ISA get_isa(const std::string& isa) {
|
||||
if (isa == "amx") {
|
||||
return ISA::AMX;
|
||||
} else if (isa == "vec") {
|
||||
return ISA::VEC;
|
||||
} else if (isa == "rvv") {
|
||||
return ISA::RVV;
|
||||
} else {
|
||||
TORCH_CHECK(false, "Invalid isa type: " + isa);
|
||||
}
|
||||
|
||||
@@ -136,8 +136,12 @@ typename T::Fmha::Arguments args_from_options(
|
||||
StrideQ stride_Q_pe = cute::make_tuple(
|
||||
static_cast<int64_t>(q_pe.stride(1)), _1{}, static_cast<int64_t>(q_pe.stride(0)));
|
||||
|
||||
// Read the token and page strides from the cache tensor instead of assuming
|
||||
// packed pages, so strided views (e.g. per-layer views into a cross-layer
|
||||
// block-major cache) are addressed correctly.
|
||||
StrideK stride_C = cute::make_tuple(
|
||||
static_cast<int64_t>(0 + D_latent + D_rope), _1{}, static_cast<int64_t>(page_size * (D_latent + D_rope)));
|
||||
static_cast<int64_t>(kv_c_and_k_pe_cache.stride(1)), _1{},
|
||||
static_cast<int64_t>(kv_c_and_k_pe_cache.stride(0)));
|
||||
StrideLSE stride_PT = cute::make_stride(_1{}, page_count_per_seq);
|
||||
StrideLSE stride_LSE = cute::make_tuple(_1{}, 0 + H);
|
||||
StrideO stride_O = cute::make_tuple(static_cast<int64_t>(0 + D_latent), _1{}, static_cast<int64_t>(0 + H * D_latent));
|
||||
@@ -268,9 +272,14 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba
|
||||
using TileShapeD = typename MlaSm100Type::TileShapeD;
|
||||
arguments.problem_shape =
|
||||
cute::make_tuple(TileShapeH{}, static_cast<int>(max_seq_len), TileShapeD{}, static_cast<int>(num_batches));
|
||||
// Assumes device 0 when getting sm_count.
|
||||
arguments.hw_info.sm_count =
|
||||
sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count;
|
||||
if (sm_count <= 0) {
|
||||
int current_device = 0;
|
||||
cudaGetDevice(¤t_device);
|
||||
arguments.hw_info.sm_count =
|
||||
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(current_device);
|
||||
} else {
|
||||
arguments.hw_info.sm_count = sm_count;
|
||||
}
|
||||
arguments.split_kv = static_cast<int>(num_kv_splits);
|
||||
MlaSm100Type::Fmha::set_split_kv(arguments);
|
||||
|
||||
|
||||
@@ -549,7 +549,7 @@ __global__ void indexer_k_quant_and_cache_kernel(
|
||||
const int head_dim, // dimension of each head
|
||||
const int quant_block_size, // quantization block size
|
||||
const int cache_block_size, // cache block size
|
||||
const int cache_stride, // stride for each token in kv_cache
|
||||
const int64_t cache_block_stride, // stride for each block in kv_cache
|
||||
|
||||
const bool use_ue8m0 // use ue8m0 scale format
|
||||
) {
|
||||
@@ -590,16 +590,15 @@ __global__ void indexer_k_quant_and_cache_kernel(
|
||||
scale = exp2f(ceilf(log2f(scale)));
|
||||
}
|
||||
|
||||
const int64_t dst_offset = block_idx * cache_block_size * cache_stride +
|
||||
block_offset * head_dim + head_dim_idx;
|
||||
const int64_t dst_offset =
|
||||
block_idx * cache_block_stride + block_offset * head_dim + head_dim_idx;
|
||||
for (int i = 0; i < VEC_SIZE; i++) {
|
||||
kv_cache[dst_offset + i] =
|
||||
fp8::scaled_convert<cache_t, scalar_t, kv_dt>(k_val_ptr[i], scale);
|
||||
}
|
||||
if (threadIdx.x == 0) {
|
||||
const int64_t dst_scale_idx =
|
||||
block_idx * cache_block_size * cache_stride +
|
||||
cache_block_size * head_dim +
|
||||
block_idx * cache_block_stride + cache_block_size * head_dim +
|
||||
(block_offset * head_dim + head_dim_idx) * 4 / quant_block_size;
|
||||
reinterpret_cast<float*>(kv_cache)[dst_scale_idx / 4] = scale;
|
||||
}
|
||||
@@ -1452,7 +1451,7 @@ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
reinterpret_cast<KV_T*>(k.data_ptr()), \
|
||||
reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
|
||||
slot_mapping.const_data_ptr<int64_t>(), head_dim, quant_block_size, \
|
||||
cache_block_size, cache_stride, use_ue8m0);
|
||||
cache_block_size, cache_block_stride, use_ue8m0);
|
||||
|
||||
void indexer_k_quant_and_cache(
|
||||
torch::stable::Tensor& k, // [num_tokens, head_dim]
|
||||
@@ -1463,7 +1462,7 @@ void indexer_k_quant_and_cache(
|
||||
int num_tokens = k.size(0);
|
||||
int head_dim = k.size(1);
|
||||
int cache_block_size = kv_cache.size(1);
|
||||
int cache_stride = kv_cache.size(2);
|
||||
int64_t cache_block_stride = kv_cache.stride(0);
|
||||
bool use_ue8m0 = scale_fmt == "ue8m0";
|
||||
|
||||
STD_TORCH_CHECK(k.device() == kv_cache.device(),
|
||||
|
||||
@@ -9,7 +9,7 @@
|
||||
#include <torch/headeronly/core/ScalarType.h>
|
||||
|
||||
#include "../../cuda_compat.h"
|
||||
#include "core/math.hpp"
|
||||
#include "libtorch_stable/core/math.hpp"
|
||||
#include "libtorch_stable/dispatch_utils.h"
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
|
||||
@@ -2,9 +2,25 @@
|
||||
|
||||
#include <torch/csrc/stable/library.h>
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include <torch/headeronly/util/Exception.h>
|
||||
|
||||
#include <optional>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include <torch/csrc/stable/ops.h>
|
||||
|
||||
inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) {
|
||||
// Ensure tensor is on CUDA
|
||||
STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device");
|
||||
|
||||
// Get the raw data pointer
|
||||
void* data_ptr = tensor.mutable_data_ptr();
|
||||
|
||||
/// Create a new tensor from the raw data pointer
|
||||
return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(),
|
||||
tensor.device(), tensor.scalar_type());
|
||||
}
|
||||
|
||||
void per_token_group_quant_fp8(const torch::stable::Tensor& input,
|
||||
torch::stable::Tensor& output_q,
|
||||
@@ -371,6 +387,18 @@ void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
|
||||
void silu_and_mul_clamp(torch::stable::Tensor& out,
|
||||
torch::stable::Tensor& input, double limit,
|
||||
double alpha = 1.0, double beta = 0.0);
|
||||
|
||||
void silu_and_mul_quant(torch::stable::Tensor& out,
|
||||
torch::stable::Tensor& input,
|
||||
torch::stable::Tensor& scale);
|
||||
|
||||
void persistent_masked_m_silu_mul_quant(
|
||||
const torch::stable::Tensor& input, // (E, T, 2*H)
|
||||
const torch::stable::Tensor& tokens_per_expert, // (E)
|
||||
torch::stable::Tensor& y_q, // (E, T, H) [OUT]
|
||||
torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT]
|
||||
bool use_ue8m0);
|
||||
|
||||
void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
|
||||
void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
|
||||
void gelu_tanh_and_mul(torch::stable::Tensor& out,
|
||||
|
||||
+63
-55
@@ -1,16 +1,12 @@
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <torch/all.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include <cmath>
|
||||
#include "core/math.hpp"
|
||||
#include "../cuda_compat.h"
|
||||
#include "dispatch_utils.h"
|
||||
|
||||
#include "libtorch_stable/core/math.hpp"
|
||||
#include "cuda_compat.h"
|
||||
#include "libtorch_stable/dispatch_utils.h"
|
||||
#include "quantization/w8a8/fp8/common.cuh"
|
||||
|
||||
#include <c10/util/Float8_e4m3fn.h>
|
||||
|
||||
#ifndef USE_ROCM
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
@@ -33,7 +29,6 @@ typedef __hip_fp8x4_e4m3_fnuz __nv_fp8x4_e4m3;
|
||||
#endif
|
||||
#endif
|
||||
|
||||
#include "core/registration.h"
|
||||
namespace vllm {
|
||||
|
||||
template <typename T>
|
||||
@@ -564,41 +559,47 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
|
||||
} // namespace vllm
|
||||
|
||||
// Launch activation, gating, and quantize kernel.
|
||||
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \
|
||||
int d = input.size(-1) / 2; \
|
||||
int64_t num_tokens = input.numel() / input.size(-1); \
|
||||
dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \
|
||||
dim3 block(std::min(d, 512)); \
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
|
||||
VLLM_DISPATCH_FLOATING_TYPES( \
|
||||
input.scalar_type(), "act_and_mul_kernel", [&] { \
|
||||
VLLM_DISPATCH_FP8_TYPES( \
|
||||
out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] { \
|
||||
vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>, \
|
||||
fp8_t> \
|
||||
<<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(), \
|
||||
input.data_ptr<scalar_t>(), \
|
||||
scale.data_ptr<float>(), d); \
|
||||
}); \
|
||||
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \
|
||||
int d = input.size(-1) / 2; \
|
||||
int64_t num_tokens = input.numel() / input.size(-1); \
|
||||
dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \
|
||||
dim3 block(std::min(d, 512)); \
|
||||
const torch::stable::accelerator::DeviceGuard device_guard( \
|
||||
input.get_device_index()); \
|
||||
const cudaStream_t stream = \
|
||||
get_current_cuda_stream(input.get_device_index()); \
|
||||
VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
|
||||
input.scalar_type(), "act_and_mul_kernel", [&] { \
|
||||
VLLM_STABLE_DISPATCH_FP8_TYPES( \
|
||||
out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] { \
|
||||
vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>, \
|
||||
fp8_t> \
|
||||
<<<grid, block, 0, stream>>>( \
|
||||
out.mutable_data_ptr<fp8_t>(), \
|
||||
input.const_data_ptr<scalar_t>(), \
|
||||
scale.const_data_ptr<float>(), d); \
|
||||
}); \
|
||||
});
|
||||
|
||||
void silu_and_mul_quant(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input, // [..., 2 * d]
|
||||
torch::Tensor& scale) {
|
||||
TORCH_CHECK(out.dtype() == torch::kFloat8_e4m3fn ||
|
||||
out.dtype() == torch::kFloat8_e4m3fnuz);
|
||||
TORCH_CHECK(input.dtype() == torch::kFloat16 ||
|
||||
input.dtype() == torch::kBFloat16);
|
||||
TORCH_CHECK(input.size(-1) % 2 == 0);
|
||||
void silu_and_mul_quant(torch::stable::Tensor& out, // [..., d]
|
||||
torch::stable::Tensor& input, // [..., 2 * d]
|
||||
torch::stable::Tensor& scale) {
|
||||
STD_TORCH_CHECK(
|
||||
out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
|
||||
out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
|
||||
STD_TORCH_CHECK(
|
||||
input.scalar_type() == torch::headeronly::ScalarType::Half ||
|
||||
input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
|
||||
"Input must be FP16 or BF16");
|
||||
STD_TORCH_CHECK(input.size(-1) % 2 == 0);
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel);
|
||||
}
|
||||
|
||||
void persistent_masked_m_silu_mul_quant(
|
||||
const at::Tensor& input, // (E, T, 2*H)
|
||||
const at::Tensor& tokens_per_expert, // (E)
|
||||
at::Tensor& y_q, // (E, T, H) [OUT]
|
||||
at::Tensor& y_s, // (E, T, H//group_size) [OUT]
|
||||
const torch::stable::Tensor& input, // (E, T, 2*H)
|
||||
const torch::stable::Tensor& tokens_per_expert, // (E)
|
||||
torch::stable::Tensor& y_q, // (E, T, H) [OUT]
|
||||
torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT]
|
||||
bool cast_scale_ue8m0) {
|
||||
#ifndef USE_ROCM
|
||||
|
||||
@@ -606,14 +607,18 @@ void persistent_masked_m_silu_mul_quant(
|
||||
// fixed GROUP_SIZE of 128.
|
||||
static constexpr int GROUP_SIZE = 128;
|
||||
|
||||
TORCH_CHECK(input.dtype() == torch::kBFloat16);
|
||||
TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn ||
|
||||
y_q.dtype() == torch::kFloat8_e4m3fnuz);
|
||||
TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
|
||||
STD_TORCH_CHECK(input.scalar_type() ==
|
||||
torch::headeronly::ScalarType::BFloat16);
|
||||
STD_TORCH_CHECK(
|
||||
y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
|
||||
y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
|
||||
STD_TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
|
||||
|
||||
bool const is_packed_ue8m0 =
|
||||
(y_s.dtype() == torch::kInt32 && cast_scale_ue8m0);
|
||||
TORCH_CHECK(y_s.dtype() == torch::kFloat32 || is_packed_ue8m0);
|
||||
(y_s.scalar_type() == torch::headeronly::ScalarType::Int &&
|
||||
cast_scale_ue8m0);
|
||||
STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Float ||
|
||||
is_packed_ue8m0);
|
||||
|
||||
using Idx_t = int64_t;
|
||||
|
||||
@@ -631,7 +636,7 @@ void persistent_masked_m_silu_mul_quant(
|
||||
|
||||
int const NUM_GROUPS = H / GROUP_SIZE;
|
||||
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
|
||||
|
||||
// TODO: Get this from cuda_arch ?
|
||||
static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32;
|
||||
@@ -643,18 +648,21 @@ void persistent_masked_m_silu_mul_quant(
|
||||
static constexpr int max_shared_mem_bytes = \
|
||||
GROUP_SIZE * 2 * STAGES * NUM_WARPS * 2; \
|
||||
dim3 grid(sms), block(THREAD_COUNT); \
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
|
||||
VLLM_DISPATCH_FP8_TYPES( \
|
||||
const torch::stable::accelerator::DeviceGuard device_guard( \
|
||||
input.get_device_index()); \
|
||||
VLLM_STABLE_DISPATCH_FP8_TYPES( \
|
||||
y_q.scalar_type(), "silu_mul_fp8_quant_deep_gemm_kernel", [&] { \
|
||||
vllm::silu_mul_fp8_quant_deep_gemm_kernel< \
|
||||
BLOCK_COUNT, max_shared_mem_bytes, fp8_t, scale_t, THREAD_COUNT, \
|
||||
Idx_t, CEIL_UE8M0, GROUP_SIZE, STAGES> \
|
||||
<<<grid, block, max_shared_mem_bytes + (E + 1) * 16, stream>>>( \
|
||||
reinterpret_cast<__nv_bfloat16*>(input.data_ptr()), \
|
||||
(fp8_t*)y_q.data_ptr(), \
|
||||
reinterpret_cast<scale_t*>(y_s.data_ptr()), \
|
||||
reinterpret_cast<int32_t*>(tokens_per_expert.data_ptr()), E, \
|
||||
T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \
|
||||
reinterpret_cast<const __nv_bfloat16*>( \
|
||||
input.const_data_ptr()), \
|
||||
y_q.mutable_data_ptr<fp8_t>(), \
|
||||
reinterpret_cast<scale_t*>(y_s.mutable_data_ptr()), \
|
||||
reinterpret_cast<const int32_t*>( \
|
||||
tokens_per_expert.const_data_ptr()), \
|
||||
E, T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \
|
||||
stride_yq_t, stride_yq_h, STRIDE_YS_E, STRIDE_YS_T, \
|
||||
STRIDE_YS_G, STRIDE_YS_P, stride_counts_e); \
|
||||
});
|
||||
@@ -679,7 +687,7 @@ void persistent_masked_m_silu_mul_quant(
|
||||
Idx_t stride_ys_g = y_s.stride(2);
|
||||
Idx_t stride_ys_p = 0;
|
||||
if (!cast_scale_ue8m0) {
|
||||
TORCH_CHECK(!is_packed_ue8m0);
|
||||
STD_TORCH_CHECK(!is_packed_ue8m0);
|
||||
LAUNCH_ON_H(float, stride_ys_e, stride_ys_t, stride_ys_g, stride_ys_p,
|
||||
false);
|
||||
return;
|
||||
@@ -692,8 +700,8 @@ void persistent_masked_m_silu_mul_quant(
|
||||
return;
|
||||
}
|
||||
|
||||
TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
|
||||
TORCH_CHECK(y_s.dtype() == torch::kInt32);
|
||||
STD_TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
|
||||
STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Int);
|
||||
|
||||
// Int32 packed ue8m0 scales tensor.
|
||||
// Let E, T, G be the number to experts, number of tokens and number of groups
|
||||
@@ -31,7 +31,7 @@
|
||||
|
||||
#include "cutlass/util/packed_stride.hpp"
|
||||
|
||||
#include "core/math.hpp"
|
||||
#include "libtorch_stable/core/math.hpp"
|
||||
#include "core/batch_invariant.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
@@ -31,7 +31,7 @@
|
||||
|
||||
#include "cutlass/util/packed_stride.hpp"
|
||||
|
||||
#include "core/math.hpp"
|
||||
#include "libtorch_stable/core/math.hpp"
|
||||
#include "core/batch_invariant.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
@@ -19,7 +19,7 @@
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
#include "cutlass/util/packed_stride.hpp"
|
||||
|
||||
#include "core/math.hpp"
|
||||
#include "libtorch_stable/core/math.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
// clang-format on
|
||||
|
||||
|
||||
@@ -14,7 +14,7 @@
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "core/math.hpp"
|
||||
#include "libtorch_stable/core/math.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
// clang-format on
|
||||
|
||||
|
||||
@@ -22,7 +22,7 @@
|
||||
#include "cutlass/epilogue/threadblock/fusion/visitors.hpp"
|
||||
#include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"
|
||||
|
||||
#include "core/math.hpp"
|
||||
#include "libtorch_stable/core/math.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/common.hpp"
|
||||
// clang-format on
|
||||
|
||||
|
||||
@@ -301,8 +301,9 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
|
||||
|
||||
const int sf_k_local = local_group_id % kGroupsPerBlockX;
|
||||
const int row_local = local_group_id / kGroupsPerBlockX;
|
||||
const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
|
||||
const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
|
||||
// Rows on grid.x: mn scales with tokens and can exceed the 65535 grid.y cap.
|
||||
const int sf_k_idx = blockIdx.y * kGroupsPerBlockX + sf_k_local;
|
||||
const int mn_idx = blockIdx.x * kRowsPerBlock + row_local;
|
||||
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
asm volatile("griddepcontrol.wait;");
|
||||
@@ -496,14 +497,15 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
|
||||
" is not a multiple of 4.");
|
||||
const int kx = GetGroupsPerBlockX(padded_groups_per_row);
|
||||
const int ry = 16 / kx;
|
||||
const int64_t blocks_x = padded_groups_per_row / kx;
|
||||
const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry;
|
||||
const int64_t row_blocks = (tma_aligned_mn + ry - 1) / ry;
|
||||
const int64_t sf_k_blocks = padded_groups_per_row / kx;
|
||||
const int num_threads = (kx * ry) * THREADS_PER_GROUP;
|
||||
// CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs.
|
||||
STD_TORCH_CHECK(blocks_x <= static_cast<int64_t>(INT32_MAX) &&
|
||||
blocks_y <= static_cast<int64_t>(INT32_MAX),
|
||||
// CUDA caps grid.x at 2^31 - 1 and grid.y at 2^16 - 1 (65535).
|
||||
constexpr int64_t kMaxGridDimYZ = 65535;
|
||||
STD_TORCH_CHECK(row_blocks <= static_cast<int64_t>(INT32_MAX) &&
|
||||
sf_k_blocks <= kMaxGridDimYZ,
|
||||
"per_token_group_quant_8bit_packed grid too large: (",
|
||||
blocks_x, ", ", blocks_y, ").");
|
||||
row_blocks, ", ", sf_k_blocks, ").");
|
||||
|
||||
auto dst_type = output_q.scalar_type();
|
||||
|
||||
@@ -513,8 +515,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
|
||||
#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
|
||||
do { \
|
||||
cudaLaunchConfig_t config = {}; \
|
||||
config.gridDim = dim3(static_cast<unsigned int>(blocks_x), \
|
||||
static_cast<unsigned int>(blocks_y)); \
|
||||
config.gridDim = dim3(static_cast<unsigned int>(row_blocks), \
|
||||
static_cast<unsigned int>(sf_k_blocks)); \
|
||||
config.blockDim = dim3(num_threads); \
|
||||
config.dynamicSmemBytes = 0; \
|
||||
config.stream = stream; \
|
||||
@@ -539,8 +541,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
|
||||
#else
|
||||
#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
|
||||
do { \
|
||||
dim3 grid(static_cast<unsigned int>(blocks_x), \
|
||||
static_cast<unsigned int>(blocks_y)); \
|
||||
dim3 grid(static_cast<unsigned int>(row_blocks), \
|
||||
static_cast<unsigned int>(sf_k_blocks)); \
|
||||
dim3 block(num_threads); \
|
||||
per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
|
||||
RY> \
|
||||
|
||||
@@ -34,6 +34,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
// TODO: Remove this once ROCm upgrade to torch 2.11.
|
||||
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
|
||||
|
||||
// Note about marlin kernel 'workspace' arguments:
|
||||
// Technically these should be mutable since they are modified by the kernel.
|
||||
// But since they are set back to zero once the kernel is finished we can
|
||||
// hand wave and say that they have no net effect.
|
||||
//
|
||||
// The reason to mark 'workspace' as immutable is so that they don't interfere
|
||||
// with using ScalarType arguments in the ops. If they are marked as mutable,
|
||||
// pytorch throws an assert in
|
||||
// 'torch._higher_order_ops._register_effectful_op' that prevents these
|
||||
// kernels from being torch.compile'd.
|
||||
// See the following document for more info on custom types and ops that use
|
||||
// custom types:
|
||||
// https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
|
||||
|
||||
// Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
|
||||
ops.def(
|
||||
"machete_supported_schedules("
|
||||
@@ -480,6 +494,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"Tensor workspace, int k, int max_seq_len) -> ()");
|
||||
|
||||
// Activation ops
|
||||
ops.def(
|
||||
"persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
|
||||
"y_q, Tensor! y_s, bool use_ue8m0) -> ()");
|
||||
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
|
||||
|
||||
// Activation function used in SwiGLU.
|
||||
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
|
||||
|
||||
@@ -492,6 +511,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
|
||||
"float alpha=1.0, float beta=0.0) -> ()");
|
||||
|
||||
// SwiGLU activation with FP8 quantization.
|
||||
ops.def(
|
||||
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
|
||||
|
||||
// Activation function used in GeGLU with `none` approximation.
|
||||
ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
|
||||
|
||||
@@ -690,6 +713,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
|
||||
|
||||
// Activation kernels (shared CUDA/ROCm)
|
||||
ops.impl("persistent_masked_m_silu_mul_quant",
|
||||
TORCH_BOX(&persistent_masked_m_silu_mul_quant));
|
||||
ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor));
|
||||
ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant));
|
||||
ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
|
||||
ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
|
||||
ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul));
|
||||
|
||||
-32
@@ -9,28 +9,6 @@
|
||||
|
||||
#include <vector>
|
||||
|
||||
torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
|
||||
// Ensure tensor is on CUDA
|
||||
if (!tensor.is_cuda()) {
|
||||
throw std::runtime_error("Tensor must be on CUDA device");
|
||||
}
|
||||
|
||||
// Get the raw data pointer
|
||||
void* data_ptr = tensor.data_ptr();
|
||||
|
||||
// Get tensor sizes and strides
|
||||
std::vector<int64_t> sizes = tensor.sizes().vec();
|
||||
std::vector<int64_t> strides = tensor.strides().vec();
|
||||
|
||||
// Get tensor options (dtype, device)
|
||||
auto options = tensor.options();
|
||||
|
||||
// Create a new tensor from the raw data pointer
|
||||
auto new_tensor = torch::from_blob(data_ptr, sizes, strides, options);
|
||||
|
||||
return new_tensor;
|
||||
}
|
||||
|
||||
// rms_norm and fused_add_rms_norm declarations also exist in
|
||||
// csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
|
||||
// because the CPU build still uses these torch::Tensor declarations.
|
||||
@@ -53,16 +31,6 @@ void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
|
||||
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit,
|
||||
double alpha = 1.0, double beta = 0.0);
|
||||
|
||||
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
|
||||
torch::Tensor& scale);
|
||||
|
||||
void persistent_masked_m_silu_mul_quant(
|
||||
const at::Tensor& input, // (E, T, 2*H)
|
||||
const at::Tensor& counts, // (E)
|
||||
at::Tensor& y_q, // (E, T, H) [OUT]
|
||||
at::Tensor& y_s, // (E, T, H//group_size) [OUT]
|
||||
bool use_ue8m0);
|
||||
|
||||
void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
|
||||
|
||||
void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
#include "core/registration.h"
|
||||
|
||||
// QuTLASS registers torch.ops._qutlass_C via TORCH_LIBRARY in bindings.cpp.
|
||||
// This stub lets Python import vllm._qutlass_C to trigger op registration.
|
||||
REGISTER_EXTENSION(_qutlass_C)
|
||||
@@ -20,17 +20,6 @@
|
||||
|
||||
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
// vLLM custom ops
|
||||
//
|
||||
|
||||
ops.def(
|
||||
"persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
|
||||
"y_q, Tensor! y_s,"
|
||||
"bool use_ue8m0) -> ()");
|
||||
ops.impl("persistent_masked_m_silu_mul_quant", torch::kCUDA,
|
||||
&persistent_masked_m_silu_mul_quant);
|
||||
|
||||
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
|
||||
ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
|
||||
|
||||
#ifdef USE_ROCM
|
||||
// TODO: Remove this once we upgrade to torch 2.11.
|
||||
@@ -39,35 +28,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
|
||||
ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
|
||||
&get_cuda_view_from_cpu_tensor);
|
||||
#endif
|
||||
|
||||
// Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
|
||||
ops.def(
|
||||
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
|
||||
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
|
||||
|
||||
// Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
|
||||
// GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
|
||||
// kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
|
||||
// full-cache bf16/fp8 variants).
|
||||
|
||||
// Quantization ops
|
||||
#ifndef USE_ROCM
|
||||
|
||||
// Note about marlin kernel 'workspace' arguments:
|
||||
// Technically these should be mutable since they are modified by the kernel.
|
||||
// But since they are set back to zero once the kernel is finished we can
|
||||
// hand wave and say that they have no net effect.
|
||||
//
|
||||
// The reason to mark 'workspace' as immutable is so that they don't interfere
|
||||
// with using ScalarType arguments in the ops. If they are marked as mutable,
|
||||
// pytorch throws an assert in
|
||||
// 'torch._higher_order_ops._register_effectful_op' that prevents these
|
||||
// kernels from being torch.compile'd.
|
||||
// See the following document for more info on custom types and ops that use
|
||||
// custom types:
|
||||
// https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
|
||||
|
||||
#endif
|
||||
}
|
||||
|
||||
|
||||
+10
-10
@@ -132,7 +132,7 @@ CMD ["/bin/bash"]
|
||||
FROM vllm-base AS ucx-nixl-build
|
||||
|
||||
ARG UCX_VERSION=v1.21.0-rc2
|
||||
ARG NIXL_VERSION=0.10.1
|
||||
ARG NIXL_VERSION=v1.2.0
|
||||
|
||||
# Build-time only: compiler, autotools, and verbs dev headers
|
||||
RUN apt-get update -y && apt-get install -y --no-install-recommends \
|
||||
@@ -149,25 +149,25 @@ RUN apt-get update -y && apt-get install -y --no-install-recommends \
|
||||
# patchelf (installed via uv) is used by the NIXL wheel build to rewrite
|
||||
# RPATH entries, making the wheel portable across stages.
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
git clone https://github.com/openucx/ucx /tmp/ucx_source && \
|
||||
cd /tmp/ucx_source && git checkout "${UCX_VERSION}" && \
|
||||
git clone --depth 1 --branch "${UCX_VERSION}" https://github.com/openucx/ucx /tmp/ucx_source && \
|
||||
cd /tmp/ucx_source && \
|
||||
bash autogen.sh && \
|
||||
./configure --prefix=/tmp/ucx_install --with-ze=yes --enable-examples --enable-mt && \
|
||||
make CFLAGS="-Wno-error=incompatible-pointer-types" -j8 && make install && \
|
||||
git clone https://github.com/ai-dynamo/nixl /tmp/nixl_source && \
|
||||
cd /tmp/nixl_source && git checkout "${NIXL_VERSION}" && \
|
||||
make CFLAGS="-Wno-error=incompatible-pointer-types" -j"$(nproc)" && make install && \
|
||||
git clone --depth 1 --branch "${NIXL_VERSION}" https://github.com/ai-dynamo/nixl /tmp/nixl_source && \
|
||||
cd /tmp/nixl_source && \
|
||||
uv pip install --upgrade meson pybind11 patchelf && \
|
||||
uv pip install -r requirements.txt && \
|
||||
PKG_CONFIG_PATH=/tmp/ucx_install/lib/pkgconfig \
|
||||
LD_LIBRARY_PATH=/tmp/ucx_install/lib \
|
||||
python -m pip wheel --no-deps . -w /tmp/nixl_wheels/ && \
|
||||
find /tmp/ucx_install -type f \( -name '*.a' -o -name '*.la' \) -delete && \
|
||||
rm -rf /tmp/ucx_install/include /tmp/ucx_install/share /tmp/ucx_install/etc /tmp/ucx_install/lib/cmake /tmp/ucx_install/bin && \
|
||||
rm -rf /tmp/ucx_source /tmp/nixl_source
|
||||
rm -rf /tmp/ucx_install/{include,share,etc,bin} /tmp/ucx_install/lib/cmake \
|
||||
/tmp/ucx_source /tmp/nixl_source
|
||||
|
||||
FROM vllm-base AS vllm-openai
|
||||
|
||||
ARG NIXL_VERSION=0.10.1
|
||||
ARG NIXL_VERSION=v1.2.0
|
||||
|
||||
# Copy compiled UCX runtime libraries and the pre-built NIXL wheel.
|
||||
# No compiler or autotools are installed in this stage.
|
||||
@@ -192,7 +192,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
librdmacm1t64 \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& uv pip install --no-deps /tmp/nixl_wheels/nixl*.whl \
|
||||
&& uv pip install nixl==${NIXL_VERSION} \
|
||||
&& uv pip install nixl==${NIXL_VERSION} && uv pip uninstall nixl-cu13 \
|
||||
&& rm -rf /tmp/nixl_wheels
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
|
||||
@@ -304,9 +304,15 @@ review process:
|
||||
resources. The reviewer will add `ready` label to the PR when the PR is
|
||||
ready to merge or a full CI run is needed.
|
||||
|
||||
### Escalating Stalled Contributions
|
||||
### Pull Request Limits and Escalation
|
||||
|
||||
If you have an important contribution that has not yet received maintainer attention, please email us at:
|
||||
vLLM uses GitHub's [pull request limit](https://github.blog/open-source/maintainers/how-pull-request-limits-are-cutting-down-the-noise/)
|
||||
for contributors without write access. The current cap is 6 open PRs. If this
|
||||
blocks well-intentioned critical work, contact a committer to request bypass
|
||||
list access.
|
||||
|
||||
If you need an expedited review for an important contribution, please email us
|
||||
at:
|
||||
|
||||
<pr-review-request@vllm.ai>
|
||||
|
||||
|
||||
@@ -133,10 +133,10 @@ The model should inherit protocol `IsAttentionFree` and also implement class met
|
||||
For the mamba layers themselves, please use the [`MambaMixer`](../../../vllm/model_executor/layers/mamba/mamba_mixer.py) (for Mamba-1) or [`MambaMixer2`](../../../vllm/model_executor/layers/mamba/mamba_mixer2.py) (for Mamba-2) classes.
|
||||
The model should also be added to the `MODELS_CONFIG_MAP` dictionary in [vllm/model_executor/models/config.py](../../../vllm/model_executor/models/config.py) to ensure that the runtime defaults are optimized.
|
||||
|
||||
For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`BambaForCausalLM`](../../../vllm/model_executor/models/bamba.py) (for an example of a model that uses Mamba-2 and attention together).
|
||||
For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`NemotronHForCausalLM`](../../../vllm/model_executor/models/nemotron_h.py) (for an example of a model that uses Mamba-2 and attention together).
|
||||
These models should follow the same instructions as case (1), but they should inherit protocol `IsHybrid` (instead of `IsAttentionFree`) and it is *not* necessary to add them to the `MODELS_CONFIG_MAP` (their runtime defaults will be inferred from the protocol).
|
||||
|
||||
For case (3), we recommend looking at the implementation of [`MiniMaxText01ForCausalLM`](../../../vllm/model_executor/models/minimax_text_01.py) or [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which use custom "mamba-like" layers `MiniMaxText01LinearAttention` and `ShortConv` respectively.
|
||||
For case (3), we recommend looking at the implementation of [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which uses a custom "mamba-like" layer `ShortConv`.
|
||||
Please follow the same guidelines as case (2) for implementing these models.
|
||||
We use "mamba-like" to refer to layers that possess a state that is updated in-place, rather than being appended-to (like KV cache for attention).
|
||||
For implementing new custom mamba-like layers, one should inherit from `MambaBase` and implement the methods `get_state_dtype`, `get_state_shape` to calculate the data types and state shapes at runtime, as well as `mamba_type` and `get_attn_backend`.
|
||||
@@ -144,5 +144,5 @@ It is also necessary to implement the "attention meta-data" class which handles
|
||||
Please see [`LinearAttentionMetadata`](../../../vllm/v1/attention/backends/linear_attn.py) or [`ShortConvAttentionMetadata`](../../../vllm/v1/attention/backends/short_conv_attn.py) for examples of this.
|
||||
It is also worth noting that we should update `MambaAttentionBackendEnum` in [`registry.py`](../../../vllm/v1/attention/backends/registry.py) when adding a new mamba backend.
|
||||
Finally, if one wants to support torch compile and CUDA graphs, it necessary to wrap the call to the mamba-like layer inside a custom op and register it.
|
||||
Please see the calls to `direct_register_custom_op` in [vllm/model_executor/models/minimax_text_01.py](../../../vllm/model_executor/models/minimax_text_01.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this.
|
||||
Please see the calls to `direct_register_custom_op` in [vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py](../../../vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this.
|
||||
The new custom op should then be added to the list `_attention_ops` in [vllm/config/compilation.py](../../../vllm/config/compilation.py) to ensure that piecewise CUDA graphs works as intended.
|
||||
|
||||
@@ -170,8 +170,8 @@ Priority is **1 = highest** (tried first).
|
||||
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
|
||||
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
|
||||
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
|
||||
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 7.x-9.x |
|
||||
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
|
||||
|
||||
@@ -74,6 +74,7 @@ vllm serve <model> \
|
||||
| `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
|
||||
| `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
|
||||
| `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
|
||||
| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
|
||||
| `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
|
||||
|
||||
## Secondary Tiers
|
||||
|
||||
@@ -321,15 +321,6 @@ For Qwen2.5, the chat template in tokenizer_config.json has already included sup
|
||||
|
||||
Flags: `--tool-call-parser hermes`
|
||||
|
||||
### MiniMax Models (`minimax_m1`)
|
||||
|
||||
Supported models:
|
||||
|
||||
* `MiniMaxAi/MiniMax-M1-40k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja))
|
||||
* `MiniMaxAi/MiniMax-M1-80k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja))
|
||||
|
||||
Flags: `--tool-call-parser minimax --chat-template examples/tool_chat_template_minimax_m1.jinja`
|
||||
|
||||
### DeepSeek-V3 Models (`deepseek_v3`)
|
||||
|
||||
Supported models:
|
||||
|
||||
@@ -61,7 +61,7 @@ Models of any architecture can be converted into embedding models using `--conve
|
||||
| `ColModernVBertForRetrieval` | ColModernVBERT | T / I | `ModernVBERT/colmodernvbert-merged` | | |
|
||||
| `ColPaliForRetrieval` | ColPali | T / I | `vidore/colpali-v1.3-hf` | | |
|
||||
| `ColQwen3` | Qwen3-VL | T / I | `TomoroAI/tomoro-colqwen3-embed-4b`, `TomoroAI/tomoro-colqwen3-embed-8b` | | |
|
||||
| `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3` | | |
|
||||
| `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3`, `vultr/VultronRetrieverPrime-Qwen3.5-8B` | | |
|
||||
| `OpsColQwen3Model` | Qwen3-VL | T / I | `OpenSearch-AI/Ops-Colqwen3-4B`, `OpenSearch-AI/Ops-Colqwen3-8B` | | |
|
||||
| `Qwen3VLNemotronEmbedModel` | Qwen3-VL | T / I | `nvidia/nemotron-colembed-vl-4b-v2`, `nvidia/nemotron-colembed-vl-8b-v2` | ✅︎ | ✅︎ |
|
||||
| `*ForConditionalGeneration`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | \* | N/A | \* | \* |
|
||||
|
||||
@@ -441,7 +441,6 @@ th {
|
||||
| `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
|
||||
| `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniMaxForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01-hf`, etc. | | |
|
||||
| `MiniMaxM2ForCausalLM` | MiniMax-M2, MiniMax-M2.1 | `MiniMaxAI/MiniMax-M2`, etc. | ✅︎ | ✅︎ |
|
||||
| `MistralForCausalLM` | Ministral-3, Mistral, Mistral-Instruct | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-7B-v0.1`, `mistralai/Mistral-7B-Instruct-v0.1`, etc. | ✅︎ | ✅︎ |
|
||||
| `MistralLarge3ForCausalLM` | Mistral-Large-3-675B-Base-2512, Mistral-Large-3-675B-Instruct-2512 | `mistralai/Mistral-Large-3-675B-Base-2512`, `mistralai/Mistral-Large-3-675B-Instruct-2512`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -487,8 +486,6 @@ th {
|
||||
| `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
|
||||
| `TeleChat3ForCausalLM` | TeleChat3 | `Tele-AI/TeleChat3-36B-Thinking`, `Tele-AI/TeleChat3-Coder-36B-Thinking`, etc. | ✅︎ | ✅︎ |
|
||||
| `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniMaxM1ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-M1-40k`, `MiniMaxAI/MiniMax-M1-80k`, etc. | | |
|
||||
| `MiniMaxText01ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01`, etc. | | |
|
||||
| `Zamba2ForCausalLM` | Zamba2 | `Zyphra/Zamba2-7B-instruct`, `Zyphra/Zamba2-2.7B-instruct`, `Zyphra/Zamba2-1.2B-instruct`, etc. | | |
|
||||
|
||||
!!! note
|
||||
@@ -595,6 +592,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
|
||||
| `MiniCPMO` | MiniCPM-O | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>E+</sup> | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniCPMV` | MiniCPM-V | T + I<sup>E+</sup> + V<sup>E+</sup> | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, etc. | ✅︎ | |
|
||||
| `MiniMaxM3SparseForConditionalGeneration` | MiniMax-M3 | T + I<sup>+</sup> + V<sup>+</sup> | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | |
|
||||
| `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
|
||||
| `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
|
||||
| `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -128,7 +128,7 @@ Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaFor
|
||||
Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., `BambaForCausalLM`,
|
||||
`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`, `Plamo2ForCausalLM`).
|
||||
|
||||
Hybrid models with mechanisms different to Mamba are also supported (e.g, `MiniMaxText01ForCausalLM`, `MiniMaxM1ForCausalLM`, `Lfm2ForCausalLM`).
|
||||
Hybrid models with mechanisms different to Mamba are also supported (e.g, `Lfm2ForCausalLM`).
|
||||
|
||||
Please note that prefix caching is not yet supported for any of the above models.
|
||||
|
||||
|
||||
@@ -1481,39 +1481,6 @@ def run_minicpmv(questions: list[str], modality: str) -> ModelRequestData:
|
||||
return run_minicpmv_base(questions, modality, "openbmb/MiniCPM-V-2_6")
|
||||
|
||||
|
||||
def run_minimax_vl_01(questions: list[str], modality: str) -> ModelRequestData:
|
||||
assert modality == "image"
|
||||
|
||||
model_name = "MiniMaxAI/MiniMax-VL-01"
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model=model_name,
|
||||
max_num_seqs=2,
|
||||
limit_mm_per_prompt={modality: 1},
|
||||
trust_remote_code=True,
|
||||
tensor_parallel_size=8,
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
messages = [
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "image"}, {"type": "text", "text": question}],
|
||||
}
|
||||
]
|
||||
for question in questions
|
||||
]
|
||||
prompts = tokenizer.apply_chat_template(
|
||||
messages, add_generation_prompt=True, tokenize=False
|
||||
)
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompts=prompts,
|
||||
)
|
||||
|
||||
|
||||
# Mistral-3 HF-format
|
||||
def run_mistral3(questions: list[str], modality: str) -> ModelRequestData:
|
||||
assert modality == "image"
|
||||
@@ -2485,7 +2452,6 @@ model_example_map = {
|
||||
"mantis": run_mantis,
|
||||
"minicpmo": run_minicpmo,
|
||||
"minicpmv": run_minicpmv,
|
||||
"minimax_vl_01": run_minimax_vl_01,
|
||||
"mistral3": run_mistral3,
|
||||
"molmo": run_molmo,
|
||||
"molmo2": run_molmo2,
|
||||
|
||||
@@ -7,11 +7,27 @@ ColQwen3.5 is a multi-modal ColBERT-style model based on Qwen3.5.
|
||||
It produces per-token embeddings and uses MaxSim scoring for retrieval
|
||||
and reranking. Supports both text and image inputs.
|
||||
|
||||
Works for any ColQwen3.5 checkpoint, e.g. `athrael-soju/colqwen3.5-4.5B-v3`
|
||||
or `vultr/VultronRetrieverPrime-Qwen3.5-8B`.
|
||||
|
||||
Start the server with:
|
||||
vllm serve athrael-soju/colqwen3.5-4.5B --max-model-len 4096
|
||||
vllm serve athrael-soju/colqwen3.5-4.5B-v3 --max-model-len 4096 \
|
||||
--mm-processor-kwargs '{"min_pixels": 65536, "max_pixels": 1835008}'
|
||||
|
||||
Then run this script:
|
||||
python colqwen3_5_rerank_online.py
|
||||
|
||||
Parity note (matching the native colpali ColQwen3_5Processor pipeline):
|
||||
- Visual-token budget: ColQwen3_5Processor uses max_num_visual_tokens=1792,
|
||||
i.e. max_pixels = 1792 * (patch_size*merge_size)^2 = 1792 * 32^2 = 1835008
|
||||
(with min_pixels = shortest_edge = 65536). Pass these via --mm-processor-kwargs
|
||||
as above; the default budget gives fewer visual tokens and lower retrieval ndcg.
|
||||
- When you build prompts yourself (token_embed), reproduce the processor exactly:
|
||||
image (document): wrap in the instruction template
|
||||
"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
|
||||
"Describe the image.<|im_end|><|endoftext|>"
|
||||
query: append the augmentation suffix <text> + "<|endoftext|>" * 10
|
||||
Omitting these reproduces a silent ~2.5 ndcg@10 drop vs the native pipeline.
|
||||
"""
|
||||
|
||||
import requests
|
||||
|
||||
@@ -1,91 +0,0 @@
|
||||
{{ '<begin_of_document>' -}}
|
||||
{%- if custom_tools is defined %}
|
||||
{%- set tools = custom_tools %}
|
||||
{%- endif %}
|
||||
{%- if not tools is defined %}
|
||||
{%- set tools = none %}
|
||||
{%- endif %}
|
||||
|
||||
{#- Extract system message #}
|
||||
{% set ns = namespace(system_prompt='') -%}
|
||||
{%- if messages[0]['role'] == 'system' %}
|
||||
{%- if messages[0]['content'] is string %}
|
||||
{%- set ns.system_prompt = messages[0]['content']|trim %}
|
||||
{%- else %}
|
||||
{%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %}
|
||||
{%- endif %}
|
||||
{%- set messages = messages[1:] %}
|
||||
{%- else %}
|
||||
{%- if tools is not none %}
|
||||
{%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
|
||||
{%- else %}
|
||||
{%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
|
||||
{#- System message #}
|
||||
{%- if ns.system_prompt != '' %}
|
||||
{{ '<beginning_of_sentence>system ai_setting=assistant\n' + ns.system_prompt + '<end_of_sentence>\n' -}}
|
||||
{%- endif %}
|
||||
|
||||
{#- Tools configuration #}
|
||||
{%- if tools is not none %}
|
||||
{{ '<beginning_of_sentence>system tool_setting=tools\nYou are provided with these tools:\n<tools>\n' -}}
|
||||
{%- for tool in tools %}
|
||||
{{ tool | tojson ~ '\n' -}}
|
||||
{%- endfor %}
|
||||
{{ '</tools>\n\nIf you need to call tools, please respond with <tool_calls></tool_calls> XML tags, and provide tool-name and json-object of arguments, following the format below:\n<tool_calls>\n{"name": <tool-name>, "arguments": <args-json-object>}\n...\n</tool_calls><end_of_sentence>\n' -}}
|
||||
{%- endif %}
|
||||
|
||||
{#- Process messages #}
|
||||
{%- for message in messages %}
|
||||
{%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
|
||||
{%- if message['role'] == 'user' %}
|
||||
{{ '<beginning_of_sentence>user name=user\n' -}}
|
||||
{%- if message['content'] is string %}
|
||||
{{ message['content']|trim -}}
|
||||
{%- else %}
|
||||
{%- for content in message['content'] %}
|
||||
{%- if content['type'] == 'text' %}
|
||||
{{ content['text']|trim -}}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{{ '<end_of_sentence>\n' -}}
|
||||
{%- elif message['role'] == 'assistant' %}
|
||||
{{ '<beginning_of_sentence>ai name=assistant\n' -}}
|
||||
{%- if message['content'] is string %}
|
||||
{{ message['content']|trim -}}
|
||||
{%- else %}
|
||||
{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %}
|
||||
{{ content['text']|trim -}}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{{ '<end_of_sentence>\n' -}}
|
||||
{%- endif %}
|
||||
{%- elif 'tool_calls' in message %}
|
||||
{{ '<beginning_of_sentence>ai name=assistant\n<tool_calls>\n' -}}
|
||||
{%- for tool_call in message.tool_calls %}
|
||||
{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}}
|
||||
{%- endfor %}
|
||||
{{ '</tool_calls><end_of_sentence>\n' -}}
|
||||
{%- elif message.role == "tool" or message.role == "ipython" %}
|
||||
{{ '<beginning_of_sentence>tool name=tools\n' -}}
|
||||
{%- if message.content is string %}
|
||||
{{ 'tool result: ' + message.content + '\n\n' -}}
|
||||
{%- else %}
|
||||
{%- for content in message['content'] %}
|
||||
{%- if content['type'] == 'text' %}
|
||||
{{ 'tool result: ' + content['text'] + '\n\n' -}}
|
||||
{%- elif content.get('name') %}
|
||||
{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{{ '<end_of_sentence>\n' -}}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
|
||||
{%- if add_generation_prompt %}
|
||||
{{ '<beginning_of_sentence>ai name=assistant\n' -}}
|
||||
{%- endif %}
|
||||
@@ -40,7 +40,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
|
||||
mteb[bm25s]>=2, <3 # required for mteb test
|
||||
transformers==5.5.3
|
||||
tokenizers==0.22.2
|
||||
schemathesis>=3.39.15 # Required for openai schema test.
|
||||
schemathesis>=4.0.0 # Required for openai schema test.
|
||||
# quantization
|
||||
bitsandbytes==0.49.2
|
||||
buildkite-test-collector==0.1.9
|
||||
|
||||
@@ -31,7 +31,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
|
||||
mteb[bm25s]>=2, <3 # required for mteb test
|
||||
transformers==5.5.3
|
||||
tokenizers==0.22.2
|
||||
schemathesis>=3.39.15 # Required for openai schema test.
|
||||
schemathesis>=4.0.0 # Required for openai schema test.
|
||||
# quantization
|
||||
bitsandbytes>=0.49.2
|
||||
buildkite-test-collector==0.1.9
|
||||
|
||||
@@ -39,7 +39,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
|
||||
mteb[bm25s]>=2, <3 # required for mteb test
|
||||
transformers==5.5.3
|
||||
tokenizers==0.22.2
|
||||
schemathesis>=3.39.15 # Required for openai schema test
|
||||
schemathesis>=4.0.0 # Required for openai schema test
|
||||
# quantization
|
||||
bitsandbytes==0.49.2
|
||||
buildkite-test-collector==0.1.9
|
||||
|
||||
@@ -386,7 +386,6 @@ mod tests {
|
||||
tool_chat_template_llama3.2_pythonic.jinja => String
|
||||
tool_chat_template_llama4_json.jinja => OpenAi
|
||||
tool_chat_template_llama4_pythonic.jinja => OpenAi
|
||||
tool_chat_template_minimax_m1.jinja => OpenAi
|
||||
tool_chat_template_mistral.jinja => String
|
||||
tool_chat_template_mistral3.jinja => OpenAi
|
||||
tool_chat_template_mistral_parallel.jinja => String
|
||||
|
||||
-91
@@ -1,91 +0,0 @@
|
||||
{{ '<begin_of_document>' -}}
|
||||
{%- if custom_tools is defined %}
|
||||
{%- set tools = custom_tools %}
|
||||
{%- endif %}
|
||||
{%- if not tools is defined %}
|
||||
{%- set tools = none %}
|
||||
{%- endif %}
|
||||
|
||||
{#- Extract system message #}
|
||||
{% set ns = namespace(system_prompt='') -%}
|
||||
{%- if messages[0]['role'] == 'system' %}
|
||||
{%- if messages[0]['content'] is string %}
|
||||
{%- set ns.system_prompt = messages[0]['content']|trim %}
|
||||
{%- else %}
|
||||
{%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %}
|
||||
{%- endif %}
|
||||
{%- set messages = messages[1:] %}
|
||||
{%- else %}
|
||||
{%- if tools is not none %}
|
||||
{%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
|
||||
{%- else %}
|
||||
{%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
|
||||
{%- endif %}
|
||||
{%- endif %}
|
||||
|
||||
{#- System message #}
|
||||
{%- if ns.system_prompt != '' %}
|
||||
{{ '<beginning_of_sentence>system ai_setting=assistant\n' + ns.system_prompt + '<end_of_sentence>\n' -}}
|
||||
{%- endif %}
|
||||
|
||||
{#- Tools configuration #}
|
||||
{%- if tools is not none %}
|
||||
{{ '<beginning_of_sentence>system tool_setting=tools\nYou are provided with these tools:\n<tools>\n' -}}
|
||||
{%- for tool in tools %}
|
||||
{{ tool | tojson ~ '\n' -}}
|
||||
{%- endfor %}
|
||||
{{ '</tools>\n\nIf you need to call tools, please respond with <tool_calls></tool_calls> XML tags, and provide tool-name and json-object of arguments, following the format below:\n<tool_calls>\n{"name": <tool-name>, "arguments": <args-json-object>}\n...\n</tool_calls><end_of_sentence>\n' -}}
|
||||
{%- endif %}
|
||||
|
||||
{#- Process messages #}
|
||||
{%- for message in messages %}
|
||||
{%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
|
||||
{%- if message['role'] == 'user' %}
|
||||
{{ '<beginning_of_sentence>user name=user\n' -}}
|
||||
{%- if message['content'] is string %}
|
||||
{{ message['content']|trim -}}
|
||||
{%- else %}
|
||||
{%- for content in message['content'] %}
|
||||
{%- if content['type'] == 'text' %}
|
||||
{{ content['text']|trim -}}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{{ '<end_of_sentence>\n' -}}
|
||||
{%- elif message['role'] == 'assistant' %}
|
||||
{{ '<beginning_of_sentence>ai name=assistant\n' -}}
|
||||
{%- if message['content'] is string %}
|
||||
{{ message['content']|trim -}}
|
||||
{%- else %}
|
||||
{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %}
|
||||
{{ content['text']|trim -}}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{{ '<end_of_sentence>\n' -}}
|
||||
{%- endif %}
|
||||
{%- elif 'tool_calls' in message %}
|
||||
{{ '<beginning_of_sentence>ai name=assistant\n<tool_calls>\n' -}}
|
||||
{%- for tool_call in message.tool_calls %}
|
||||
{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}}
|
||||
{%- endfor %}
|
||||
{{ '</tool_calls><end_of_sentence>\n' -}}
|
||||
{%- elif message.role == "tool" or message.role == "ipython" %}
|
||||
{{ '<beginning_of_sentence>tool name=tools\n' -}}
|
||||
{%- if message.content is string %}
|
||||
{{ 'tool result: ' + message.content + '\n\n' -}}
|
||||
{%- else %}
|
||||
{%- for content in message['content'] %}
|
||||
{%- if content['type'] == 'text' %}
|
||||
{{ 'tool result: ' + content['text'] + '\n\n' -}}
|
||||
{%- elif content.get('name') %}
|
||||
{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
{%- endif %}
|
||||
{{ '<end_of_sentence>\n' -}}
|
||||
{%- endif %}
|
||||
{%- endfor %}
|
||||
|
||||
{%- if add_generation_prompt %}
|
||||
{{ '<beginning_of_sentence>ai name=assistant\n' -}}
|
||||
{%- endif %}
|
||||
@@ -277,6 +277,12 @@ pub struct EngineCoreSamplingParams {
|
||||
pub max_tokens: u32,
|
||||
/// Minimum number of tokens to generate before EOS or stop-token handling.
|
||||
pub min_tokens: u32,
|
||||
/// Maximum number of reasoning ("thinking") tokens to emit before the
|
||||
/// reasoning section is force-closed. `None` means unlimited; the
|
||||
/// user-facing `-1` sentinel is normalized to `None` by the frontend before
|
||||
/// reaching this DTO, so only non-negative values are sent. Enforced
|
||||
/// engine-side (and only when a reasoning parser is configured).
|
||||
pub thinking_token_budget: Option<u64>,
|
||||
/// Number of log probabilities to return per generated token.
|
||||
///
|
||||
/// `None` disables sample logprobs. `-1` requests the full vocabulary.
|
||||
@@ -345,6 +351,7 @@ impl EngineCoreSamplingParams {
|
||||
seed: None,
|
||||
max_tokens: 65536,
|
||||
min_tokens: 0,
|
||||
thinking_token_budget: None,
|
||||
logprobs: None,
|
||||
prompt_logprobs: None,
|
||||
min_p: 0.0,
|
||||
|
||||
@@ -150,6 +150,7 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest {
|
||||
top_k: 8,
|
||||
max_tokens: 32,
|
||||
min_tokens: 1,
|
||||
thinking_token_budget: Some(256),
|
||||
stop_token_ids: vec![151643],
|
||||
eos_token_id: Some(151645),
|
||||
all_stop_token_ids: BTreeSet::from([151643, 151645]),
|
||||
@@ -2502,6 +2503,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
|
||||
seed: None,
|
||||
max_tokens: 16,
|
||||
min_tokens: 0,
|
||||
thinking_token_budget: None,
|
||||
logprobs: None,
|
||||
prompt_logprobs: None,
|
||||
min_p: 0.0,
|
||||
|
||||
@@ -39,6 +39,7 @@ class EngineCoreSamplingParams(msgspec.Struct, dict=True, omit_defaults=True):
|
||||
seed: int | None = None
|
||||
max_tokens: int = 16
|
||||
min_tokens: int = 0
|
||||
thinking_token_budget: int | None = None
|
||||
min_p: float = 0.0
|
||||
frequency_penalty: float = 0.0
|
||||
presence_penalty: float = 0.0
|
||||
@@ -122,6 +123,7 @@ request = EngineCoreRequest(
|
||||
seed=None,
|
||||
max_tokens=32,
|
||||
min_tokens=1,
|
||||
thinking_token_budget=256,
|
||||
min_p=0.0,
|
||||
frequency_penalty=0.0,
|
||||
presence_penalty=0.0,
|
||||
|
||||
@@ -103,6 +103,7 @@ fn is_request_validation_error(error: &vllm_text::Error) -> bool {
|
||||
| vllm_text::Error::EmptyPromptTokenIds { .. }
|
||||
| vllm_text::Error::Logprobs(_)
|
||||
| vllm_text::Error::OutOfVocab(_)
|
||||
| vllm_text::Error::InvalidThinkingTokenBudget
|
||||
// An empty tokenized prompt detected later, at request prepare
|
||||
// time, surfaces through the transparent Llm wrapper.
|
||||
| vllm_text::Error::Llm(vllm_llm::Error::EmptyPromptTokenIds { .. })
|
||||
@@ -127,6 +128,18 @@ mod tests {
|
||||
assert!(response.error.message.contains("9000"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_thinking_token_budget_maps_to_invalid_request() {
|
||||
let api_error = text_submit_error(
|
||||
"failed to submit completion request",
|
||||
vllm_text::Error::InvalidThinkingTokenBudget,
|
||||
);
|
||||
assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
|
||||
let response = api_error.to_error_response();
|
||||
assert_eq!(response.error.error_type, "invalid_request_error");
|
||||
assert!(response.error.message.contains("thinking_token_budget"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn chat_wrapped_prompt_too_long_maps_to_invalid_request() {
|
||||
let error = vllm_chat::Error::Text(vllm_text::Error::PromptTooLong {
|
||||
|
||||
@@ -150,6 +150,33 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_generate_request_forwards_thinking_token_budget() {
|
||||
let request: GenerateRequest = serde_json::from_value(json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"token_ids": [11, 22, 33],
|
||||
"sampling_params": {
|
||||
"thinking_token_budget": 64
|
||||
}
|
||||
}))
|
||||
.expect("parse request");
|
||||
|
||||
let prepared = prepare_generate_request(
|
||||
request,
|
||||
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
|
||||
ResolvedRequestContext::default(),
|
||||
)
|
||||
.expect("prepare");
|
||||
|
||||
// The raw inference route shares `vllm_text::SamplingParams`, so the
|
||||
// field is carried through to lowering exactly like the OpenAI routes
|
||||
// (normalization/validation then happens in `lower_sampling_params`).
|
||||
assert_eq!(
|
||||
prepared.text_request.sampling_params.thinking_token_budget,
|
||||
Some(64)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_generate_request_gates_continuous_usage_on_include_usage() {
|
||||
let request: GenerateRequest = serde_json::from_value(json!({
|
||||
|
||||
@@ -115,6 +115,7 @@ pub(super) fn prepare_chat_request(
|
||||
seed: request.seed,
|
||||
max_tokens: request.max_completion_tokens,
|
||||
min_tokens: request.min_tokens,
|
||||
thinking_token_budget: request.thinking_token_budget,
|
||||
logprobs: request.logprobs.then_some(top_logprobs),
|
||||
prompt_logprobs,
|
||||
min_p: request.min_p,
|
||||
@@ -613,6 +614,31 @@ mod tests {
|
||||
assert_eq!(prepared.chat_request.sampling_params, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_chat_request_passes_through_thinking_token_budget() {
|
||||
let prepare = |budget: Option<i64>| {
|
||||
prepare_chat_request(
|
||||
ChatCompletionRequest {
|
||||
thinking_token_budget: budget,
|
||||
..base_request()
|
||||
},
|
||||
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
|
||||
ResolvedRequestContext::default(),
|
||||
)
|
||||
.expect("request is valid")
|
||||
.chat_request
|
||||
.sampling_params
|
||||
.thinking_token_budget
|
||||
};
|
||||
|
||||
// The convert layer forwards the raw value verbatim (including the `-1`
|
||||
// "unlimited" sentinel); normalization/validation happens during
|
||||
// lowering (see `vllm_text::lower`).
|
||||
assert_eq!(prepare(Some(64)), Some(64));
|
||||
assert_eq!(prepare(Some(-1)), Some(-1));
|
||||
assert_eq!(prepare(None), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_chat_request_accepts_developer_messages() {
|
||||
let request = ChatCompletionRequest {
|
||||
|
||||
@@ -165,8 +165,10 @@ pub struct ChatCompletionRequest {
|
||||
pub bad_words: Option<Vec<String>>,
|
||||
|
||||
// -------- Extra vLLM Parameters --------
|
||||
/// Token budget for reasoning/thinking
|
||||
pub thinking_token_budget: Option<u32>,
|
||||
/// Token budget for reasoning/thinking. Accepts a non-negative integer, or
|
||||
/// `-1` for unlimited (mirroring the Python frontend, which normalizes `-1`
|
||||
/// to "no budget").
|
||||
pub thinking_token_budget: Option<i64>,
|
||||
|
||||
/// Whether to include reasoning content in the response
|
||||
#[serde(default = "default_true")]
|
||||
|
||||
@@ -108,11 +108,6 @@ pub(super) fn validate_request_compat(
|
||||
"truncate_prompt_tokens",
|
||||
"truncate_prompt_tokens is not supported.",
|
||||
)?;
|
||||
reject_non_default(
|
||||
request.thinking_token_budget.as_ref(),
|
||||
"thinking_token_budget",
|
||||
"thinking_token_budget is not supported.",
|
||||
)?;
|
||||
reject_non_default(
|
||||
request.media_io_kwargs.as_ref(),
|
||||
"media_io_kwargs",
|
||||
|
||||
@@ -108,6 +108,7 @@ pub(super) fn prepare_completion_request(
|
||||
seed: request.seed,
|
||||
max_tokens,
|
||||
min_tokens: request.min_tokens,
|
||||
thinking_token_budget: request.thinking_token_budget,
|
||||
logprobs,
|
||||
prompt_logprobs,
|
||||
min_p: request.min_p,
|
||||
@@ -266,6 +267,34 @@ mod tests {
|
||||
assert!(!prepared.text_request.decode_options.skip_special_tokens);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_completion_request_passes_through_thinking_token_budget() {
|
||||
let prepare = |budget: serde_json::Value| {
|
||||
let request: CompletionRequest = serde_json::from_value(json!({
|
||||
"model": "Qwen/Qwen1.5-0.5B-Chat",
|
||||
"prompt": "hello",
|
||||
"thinking_token_budget": budget,
|
||||
}))
|
||||
.expect("parse request");
|
||||
prepare_completion_request(
|
||||
request,
|
||||
&served(&["Qwen/Qwen1.5-0.5B-Chat"]),
|
||||
ResolvedRequestContext::default(),
|
||||
)
|
||||
.expect("prepare")
|
||||
.text_request
|
||||
.sampling_params
|
||||
.thinking_token_budget
|
||||
};
|
||||
|
||||
// The convert layer forwards the raw value verbatim (including the `-1`
|
||||
// "unlimited" sentinel); normalization/validation happens during
|
||||
// lowering (see `vllm_text::lower`).
|
||||
assert_eq!(prepare(json!(64)), Some(64));
|
||||
assert_eq!(prepare(json!(-1)), Some(-1));
|
||||
assert_eq!(prepare(json!(null)), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prepare_completion_request_maps_stream_usage_and_token_format_options() {
|
||||
let request: CompletionRequest = serde_json::from_value(json!({
|
||||
|
||||
@@ -146,6 +146,11 @@ pub struct CompletionRequest {
|
||||
/// Additional kwargs for structured outputs
|
||||
pub structured_outputs: Option<Value>,
|
||||
|
||||
/// Token budget for reasoning/thinking. Accepts a non-negative integer, or
|
||||
/// `-1` for unlimited (mirroring the Python frontend, which normalizes `-1`
|
||||
/// to "no budget").
|
||||
pub thinking_token_budget: Option<i64>,
|
||||
|
||||
/// Request scheduling priority (lower means earlier; default 0)
|
||||
pub priority: Option<i32>,
|
||||
|
||||
|
||||
@@ -20,6 +20,8 @@ pub enum Error {
|
||||
Logprobs(#[from] LogprobsError),
|
||||
#[error(transparent)]
|
||||
OutOfVocab(#[from] OutOfVocabError),
|
||||
#[error("`thinking_token_budget` must be a non-negative integer or -1 for unlimited.")]
|
||||
InvalidThinkingTokenBudget,
|
||||
#[error("text request stream `{request_id}` closed before terminal output")]
|
||||
StreamClosedBeforeTerminalOutput { request_id: String },
|
||||
#[error(transparent)]
|
||||
|
||||
@@ -87,6 +87,7 @@ pub fn lower_sampling_params(
|
||||
seed,
|
||||
max_tokens,
|
||||
min_tokens,
|
||||
thinking_token_budget,
|
||||
logprobs,
|
||||
prompt_logprobs,
|
||||
min_p,
|
||||
@@ -128,6 +129,7 @@ pub fn lower_sampling_params(
|
||||
prompt_len,
|
||||
)?;
|
||||
let min_tokens = min_tokens.unwrap_or(0);
|
||||
let thinking_token_budget = normalize_thinking_token_budget(thinking_token_budget)?;
|
||||
let frequency_penalty = frequency_penalty.unwrap_or(0.0);
|
||||
let presence_penalty = presence_penalty.unwrap_or(0.0);
|
||||
|
||||
@@ -149,6 +151,7 @@ pub fn lower_sampling_params(
|
||||
seed,
|
||||
max_tokens,
|
||||
min_tokens,
|
||||
thinking_token_budget,
|
||||
logprobs,
|
||||
prompt_logprobs,
|
||||
min_p,
|
||||
@@ -170,6 +173,21 @@ pub fn lower_sampling_params(
|
||||
Ok(params)
|
||||
}
|
||||
|
||||
/// Normalize the user-facing `thinking_token_budget` into the engine value.
|
||||
///
|
||||
/// Mirrors Python's `validate_thinking_token_budget`
|
||||
/// (<https://github.com/vllm-project/vllm/blob/ecf9d83520eb217401b47d8a5451a27c5231b8c2/vllm/sampling_params.py#L35-L55>):
|
||||
/// `None` and the `-1` "unlimited" sentinel both map to `None`; any other
|
||||
/// negative value is rejected; non-negative values pass through unchanged. Like
|
||||
/// Python's `int`, no upper bound is imposed.
|
||||
fn normalize_thinking_token_budget(value: Option<i64>) -> Result<Option<u64>> {
|
||||
match value {
|
||||
None | Some(-1) => Ok(None),
|
||||
Some(budget) if budget >= 0 => Ok(Some(budget as u64)),
|
||||
Some(_) => Err(Error::InvalidThinkingTokenBudget),
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert bad-word strings into token-ID sequences, following the Python vLLM
|
||||
/// logic in `SamplingParams.update_from_tokenizer()`.
|
||||
///
|
||||
@@ -366,6 +384,36 @@ mod tests {
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lower_sampling_params_normalizes_thinking_token_budget() {
|
||||
let lower = |budget: Option<i64>| {
|
||||
lower_sampling_params_with_limits(
|
||||
SamplingParams {
|
||||
thinking_token_budget: budget,
|
||||
..SamplingParams::default()
|
||||
},
|
||||
sample_sampling_limits(),
|
||||
)
|
||||
};
|
||||
|
||||
// Non-negative budgets (including 0) pass through unchanged.
|
||||
assert_eq!(lower(Some(256)).unwrap().thinking_token_budget, Some(256));
|
||||
assert_eq!(lower(Some(0)).unwrap().thinking_token_budget, Some(0));
|
||||
// `None` and the `-1` "unlimited" sentinel both disable the budget.
|
||||
assert_eq!(lower(None).unwrap().thinking_token_budget, None);
|
||||
assert_eq!(lower(Some(-1)).unwrap().thinking_token_budget, None);
|
||||
// No upper bound is imposed, matching Python's `int`.
|
||||
assert_eq!(
|
||||
lower(Some(i64::from(u32::MAX) + 1)).unwrap().thinking_token_budget,
|
||||
Some(u64::from(u32::MAX) + 1)
|
||||
);
|
||||
// Other negatives are rejected.
|
||||
assert!(matches!(
|
||||
lower(Some(-2)),
|
||||
Err(Error::InvalidThinkingTokenBudget)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lower_text_request_applies_python_style_eos_hints() {
|
||||
let prepared = lower_text_request(
|
||||
@@ -386,6 +434,7 @@ mod tests {
|
||||
seed: None,
|
||||
max_tokens: 999997,
|
||||
min_tokens: 0,
|
||||
thinking_token_budget: None,
|
||||
logprobs: None,
|
||||
prompt_logprobs: None,
|
||||
min_p: 0.0,
|
||||
@@ -437,6 +486,7 @@ mod tests {
|
||||
seed: None,
|
||||
max_tokens: 999997,
|
||||
min_tokens: 0,
|
||||
thinking_token_budget: None,
|
||||
logprobs: None,
|
||||
prompt_logprobs: None,
|
||||
min_p: 0.0,
|
||||
@@ -567,6 +617,7 @@ mod tests {
|
||||
seed: None,
|
||||
max_tokens: 40957,
|
||||
min_tokens: 0,
|
||||
thinking_token_budget: None,
|
||||
logprobs: None,
|
||||
prompt_logprobs: None,
|
||||
min_p: 0.0,
|
||||
@@ -628,6 +679,7 @@ mod tests {
|
||||
seed: None,
|
||||
max_tokens: 999997,
|
||||
min_tokens: 0,
|
||||
thinking_token_budget: None,
|
||||
logprobs: None,
|
||||
prompt_logprobs: None,
|
||||
min_p: 0.0,
|
||||
@@ -697,6 +749,7 @@ mod tests {
|
||||
seed: None,
|
||||
max_tokens: 32,
|
||||
min_tokens: 2,
|
||||
thinking_token_budget: None,
|
||||
logprobs: None,
|
||||
prompt_logprobs: None,
|
||||
min_p: 0.1,
|
||||
@@ -929,6 +982,7 @@ mod tests {
|
||||
seed: None,
|
||||
max_tokens: 128,
|
||||
min_tokens: 0,
|
||||
thinking_token_budget: None,
|
||||
logprobs: None,
|
||||
prompt_logprobs: None,
|
||||
min_p: 0.1,
|
||||
|
||||
@@ -309,7 +309,7 @@ fn matches_stop_string(stops: &[String], output: &str, new_bytes: usize) -> Opti
|
||||
.find_map(|(ss_idx, (ss, len, start_off))| {
|
||||
output[start_off..]
|
||||
.windows(len)
|
||||
.rposition(|w| w == ss)
|
||||
.position(|w| w == ss)
|
||||
.map(|pos| (ss_idx, start_off + pos))
|
||||
})
|
||||
}
|
||||
@@ -562,6 +562,13 @@ mod tests {
|
||||
assert_eq!(result, Some((0, 4)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stop_string_matches_leftmost_with_multiple_new_bytes() {
|
||||
let stops = vec!["\n".to_string()];
|
||||
let result = matches_stop_string(&stops, "Answer\n\n", 2);
|
||||
assert_eq!(result, Some((0, 6)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stop_string_matches_at_beginning() {
|
||||
let stops = vec!["say".to_string()];
|
||||
|
||||
@@ -56,6 +56,12 @@ pub struct SamplingParams {
|
||||
pub max_tokens: Option<u32>,
|
||||
/// Minimum number of tokens to generate before EOS or stop-token handling.
|
||||
pub min_tokens: Option<u32>,
|
||||
/// Maximum number of reasoning ("thinking") tokens to emit before the
|
||||
/// reasoning section is force-closed. `None` or the user-facing `-1`
|
||||
/// "unlimited" sentinel both disable the budget. The raw value is carried
|
||||
/// here; `-1` is normalized to `None` (and other negatives rejected) during
|
||||
/// lowering (see `lower_sampling_params`).
|
||||
pub thinking_token_budget: Option<i64>,
|
||||
/// Number of log probabilities to return per generated token.
|
||||
///
|
||||
/// `None` disables sample logprobs. `-1` requests the full vocabulary.
|
||||
@@ -116,6 +122,7 @@ impl Default for SamplingParams {
|
||||
seed: None,
|
||||
max_tokens: None,
|
||||
min_tokens: None,
|
||||
thinking_token_budget: None,
|
||||
logprobs: None,
|
||||
prompt_logprobs: None,
|
||||
min_p: None,
|
||||
|
||||
@@ -769,6 +769,7 @@ class precompiled_wheel_utils:
|
||||
"vllm/_C.abi3.so",
|
||||
"vllm/_C_stable_libtorch.abi3.so",
|
||||
"vllm/_moe_C_stable_libtorch.abi3.so",
|
||||
"vllm/_qutlass_C.abi3.so",
|
||||
"vllm/_flashmla_C.abi3.so",
|
||||
"vllm/_flashmla_extension_C.abi3.so",
|
||||
"vllm/_sparse_flashmla_C.abi3.so",
|
||||
@@ -1135,6 +1136,7 @@ if _is_cuda():
|
||||
# DeepGEMM requires CUDA 12.3+ (SM90/SM100)
|
||||
# Optional since it won't build on unsupported architectures
|
||||
ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", optional=True))
|
||||
ext_modules.append(CMakeExtension(name="vllm._qutlass_C", optional=True))
|
||||
# fmha_sm100 is a Python/CuTe-DSL package installed into vllm.third_party.
|
||||
ext_modules.append(CMakeExtension(name="vllm.fmha_sm100", optional=True))
|
||||
|
||||
@@ -1149,7 +1151,8 @@ if _is_cpu():
|
||||
ext_modules.append(CMakeExtension(name="vllm._C"))
|
||||
|
||||
if _build_custom_ops():
|
||||
ext_modules.append(CMakeExtension(name="vllm._C"))
|
||||
if _is_hip():
|
||||
ext_modules.append(CMakeExtension(name="vllm._C"))
|
||||
if _is_cuda() or _is_hip():
|
||||
ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch"))
|
||||
ext_modules.append(CMakeExtension(name="vllm._moe_C_stable_libtorch"))
|
||||
|
||||
@@ -22,7 +22,7 @@ import torch
|
||||
|
||||
import vllm.config
|
||||
from tests.compile.backend import TestBackend
|
||||
from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
|
||||
from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
|
||||
from vllm.config import (
|
||||
@@ -83,9 +83,8 @@ class _ViewDoubleQuantModel(torch.nn.Module):
|
||||
[_NoViewDoubleQuantModel, _ViewDoubleQuantModel],
|
||||
ids=["no_view", "with_view"],
|
||||
)
|
||||
@pytest.mark.skipif(
|
||||
not is_aiter_found_and_supported(),
|
||||
reason="Only test on ROCm with AITER installed and supported",
|
||||
@pytest.mark.skip(
|
||||
reason="Skipping for now because pytorch compiler removes one the two quant ops"
|
||||
)
|
||||
def test_double_aiter_rms_fp8_group_quant_fusion(
|
||||
model_cls: type[torch.nn.Module],
|
||||
|
||||
@@ -175,7 +175,7 @@ MULTIMODAL_MODELS = {
|
||||
"facebook/chameleon-7b": PPTestSettings.fast(),
|
||||
"adept/fuyu-8b": PPTestSettings.fast(),
|
||||
"zai-org/glm-4v-9b": PPTestSettings.fast(),
|
||||
"OpenGVLab/InternVL2-1B": PPTestSettings.fast(),
|
||||
"OpenGVLab/InternVL3-1B": PPTestSettings.fast(),
|
||||
"llava-hf/llava-1.5-7b-hf": PPTestSettings.fast(),
|
||||
"llava-hf/llava-v1.6-mistral-7b-hf": PPTestSettings.fast(),
|
||||
"llava-hf/LLaVA-NeXT-Video-7B-hf": PPTestSettings.fast(),
|
||||
@@ -203,7 +203,7 @@ TEST_MODELS = [
|
||||
"intfloat/e5-mistral-7b-instruct",
|
||||
"BAAI/bge-multilingual-gemma2",
|
||||
# [MULTIMODAL GENERATION]
|
||||
"OpenGVLab/InternVL2-1B",
|
||||
"OpenGVLab/InternVL3-1B",
|
||||
"microsoft/Phi-3.5-vision-instruct",
|
||||
"fixie-ai/ultravox-v0_5-llama-3_2-1b",
|
||||
# [LANGUAGE GENERATION - HYBRID ARCH]
|
||||
|
||||
@@ -649,3 +649,196 @@ def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch):
|
||||
args = EngineArgs(model="s3://bucket/model", tokenizer="s3://bucket/tokenizer")
|
||||
assert args.model == "s3://bucket/model"
|
||||
assert args.tokenizer == "s3://bucket/tokenizer"
|
||||
|
||||
|
||||
class TestDeviceIds:
|
||||
def test_device_ids_with_cvd_out_of_range(self, monkeypatch):
|
||||
"""--device-ids index beyond the CVD set raises ValueError."""
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
key = current_platform.device_control_env_var
|
||||
monkeypatch.setenv(key, "4,5")
|
||||
args = EngineArgs(model="m", device_ids=[0, 2])
|
||||
with pytest.raises(ValueError, match="out of range"):
|
||||
args._resolve_device_ids()
|
||||
|
||||
def test_device_ids_with_cvd_resolve_to_physical_ids(self, monkeypatch):
|
||||
"""--device-ids are CVD-local indices resolved to physical ids."""
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
key = current_platform.device_control_env_var
|
||||
monkeypatch.setenv(key, "4,5")
|
||||
args = EngineArgs(model="m", device_ids=[0, 1])
|
||||
assert args._resolve_device_ids() == [4, 5]
|
||||
|
||||
def test_device_ids_with_uuid_cvd_resolve_to_physical_ids(self, monkeypatch):
|
||||
"""--device-ids support UUID CVD values resolved by the platform."""
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
key = current_platform.device_control_env_var
|
||||
monkeypatch.setenv(key, "GPU-abcd1234,GPU-ef567890")
|
||||
monkeypatch.setattr(
|
||||
type(current_platform),
|
||||
"device_control_id_to_physical_device_id",
|
||||
classmethod(
|
||||
lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
|
||||
),
|
||||
)
|
||||
|
||||
args = EngineArgs(model="m", device_ids=[0, 1])
|
||||
assert args._resolve_device_ids() == [4, 5]
|
||||
|
||||
def test_device_ids_with_uuid_args_resolve_to_physical_ids(self, monkeypatch):
|
||||
"""UUID --device-ids are resolved to physical IDs immediately."""
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
monkeypatch.setattr(
|
||||
type(current_platform),
|
||||
"device_control_id_to_physical_device_id",
|
||||
classmethod(lambda cls, device_id: {"GPU-abcd1234": 4}[device_id]),
|
||||
)
|
||||
|
||||
args = EngineArgs(model="m", device_ids=["GPU-abcd1234"])
|
||||
assert args._resolve_device_ids() == [4]
|
||||
|
||||
def test_device_ids_reject_mixed_integer_and_uuid_args(self):
|
||||
"""--device-ids must not mix CVD indices and UUIDs."""
|
||||
args = EngineArgs(model="m", device_ids=[0, "GPU-abcd1234"])
|
||||
with pytest.raises(ValueError, match="must not mix"):
|
||||
args._resolve_device_ids()
|
||||
|
||||
def test_no_device_ids(self):
|
||||
"""No --device-ids returns None."""
|
||||
args = EngineArgs(model="m")
|
||||
assert args._resolve_device_ids() is None
|
||||
|
||||
def test_cli_parsing(self):
|
||||
"""--device-ids parses comma-separated string from CLI."""
|
||||
parser = FlexibleArgumentParser()
|
||||
EngineArgs.add_cli_args(parser)
|
||||
parsed = parser.parse_args(["--model", "m", "--device-ids", "0,2,4"])
|
||||
assert parsed.device_ids == [0, 2, 4]
|
||||
|
||||
def test_cli_parsing_uuid(self):
|
||||
"""--device-ids parses comma-separated UUID strings from CLI."""
|
||||
parser = FlexibleArgumentParser()
|
||||
EngineArgs.add_cli_args(parser)
|
||||
parsed = parser.parse_args(
|
||||
["--model", "m", "--device-ids", "GPU-abcd1234,GPU-ef567890"]
|
||||
)
|
||||
assert parsed.device_ids == ["GPU-abcd1234", "GPU-ef567890"]
|
||||
|
||||
def test_assigned_physical_gpu_ids_are_physical_with_cvd(self, monkeypatch):
|
||||
"""assigned_physical_gpu_ids are already physical and not composed with CVD."""
|
||||
import vllm.platforms.interface as platform_interface
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [4, 5])
|
||||
monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
|
||||
|
||||
assert current_platform.device_id_to_physical_device_id(0) == 4
|
||||
assert current_platform.device_id_to_physical_device_id(1) == 5
|
||||
assert current_platform.logical_device_id_to_visible_device_id(0) == 0
|
||||
assert current_platform.logical_device_id_to_visible_device_id(1) == 1
|
||||
|
||||
def test_assigned_physical_gpu_ids_map_to_visible_uuid_cvd(self, monkeypatch):
|
||||
"""Physical IDs map back to visible ordinals when CVD uses UUIDs."""
|
||||
import vllm.platforms.interface as platform_interface
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [5])
|
||||
monkeypatch.setenv(
|
||||
current_platform.device_control_env_var,
|
||||
"GPU-abcd1234,GPU-ef567890",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
type(current_platform),
|
||||
"device_control_id_to_physical_device_id",
|
||||
classmethod(
|
||||
lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
|
||||
),
|
||||
)
|
||||
|
||||
assert current_platform.logical_device_id_to_visible_device_id(0) == 1
|
||||
|
||||
def test_device_ids_reject_duplicates(self):
|
||||
"""--device-ids must not contain duplicate entries."""
|
||||
args = EngineArgs(model="m", device_ids=[2, 2])
|
||||
with pytest.raises(ValueError, match="duplicates"):
|
||||
args._resolve_device_ids()
|
||||
|
||||
def test_cli_parsing_strips_whitespace(self):
|
||||
"""--device-ids tolerates whitespace around commas."""
|
||||
parser = FlexibleArgumentParser()
|
||||
EngineArgs.add_cli_args(parser)
|
||||
parsed = parser.parse_args(["--model", "m", "--device-ids", "0, 2, 4"])
|
||||
assert parsed.device_ids == [0, 2, 4]
|
||||
|
||||
def test_visible_ordinal_to_physical_ignores_assigned_ids(self, monkeypatch):
|
||||
"""visible_device_id_to_physical_device_id maps torch device ordinals,
|
||||
independent of the logical-to-physical mapping.
|
||||
|
||||
Regression test: CustomAllreduce passes device.index (a visible
|
||||
ordinal) and must not index into assigned_physical_gpu_ids, which
|
||||
raised IndexError for non-identity --device-ids like [2, 3].
|
||||
"""
|
||||
import vllm.platforms.interface as platform_interface
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [2, 3])
|
||||
monkeypatch.delenv(current_platform.device_control_env_var, raising=False)
|
||||
|
||||
# CVD unset: visible ordinal == physical ID, even beyond the
|
||||
# assigned list's length.
|
||||
assert current_platform.visible_device_id_to_physical_device_id(2) == 2
|
||||
assert current_platform.visible_device_id_to_physical_device_id(3) == 3
|
||||
|
||||
monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
|
||||
assert current_platform.visible_device_id_to_physical_device_id(1) == 5
|
||||
with pytest.raises(IndexError, match="out of range"):
|
||||
current_platform.visible_device_id_to_physical_device_id(2)
|
||||
|
||||
|
||||
class TestDpDeviceIdSharding:
|
||||
def test_dp_supervisor_device_ids_stay_env_relative(self):
|
||||
"""Regression test: the DP supervisor must pass env-relative indices,
|
||||
not physical IDs, because each child re-resolves --device-ids
|
||||
against its inherited device-control env var."""
|
||||
import argparse
|
||||
|
||||
from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
|
||||
|
||||
args = argparse.Namespace(
|
||||
tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=None
|
||||
)
|
||||
assert _build_device_ids(args, local_rank=0) == [0, 1]
|
||||
assert _build_device_ids(args, local_rank=1) == [2, 3]
|
||||
|
||||
def test_dp_supervisor_shards_user_device_ids(self):
|
||||
"""User-provided --device-ids are sharded across DP children."""
|
||||
import argparse
|
||||
|
||||
from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
|
||||
|
||||
args = argparse.Namespace(
|
||||
tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=[4, 5, 6, 7]
|
||||
)
|
||||
assert _build_device_ids(args, local_rank=0) == [4, 5]
|
||||
assert _build_device_ids(args, local_rank=1) == [6, 7]
|
||||
with pytest.raises(ValueError, match="needs devices"):
|
||||
_build_device_ids(args, local_rank=2)
|
||||
|
||||
def test_dp_rank_shards_user_assigned_gpu_ids(self):
|
||||
"""get_physical_gpu_ids_for_local_dp_rank slices the user-provided
|
||||
--device-ids list instead of recomputing from the env var."""
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank
|
||||
|
||||
evar = current_platform.device_control_env_var
|
||||
assert get_physical_gpu_ids_for_local_dp_rank(
|
||||
evar, local_dp_rank=1, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
|
||||
) == [6, 7]
|
||||
with pytest.raises(ValueError, match="needs devices"):
|
||||
get_physical_gpu_ids_for_local_dp_rank(
|
||||
evar, local_dp_rank=2, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
|
||||
)
|
||||
|
||||
@@ -8,6 +8,8 @@ AnthropicServingMessages._convert_anthropic_to_openai_request().
|
||||
Also covers extended-thinking edge cases such as ``redacted_thinking``
|
||||
blocks echoed back by Anthropic clients, and streaming conversion in
|
||||
``message_stream_converter``.
|
||||
|
||||
Also covers cache usage computation in ``_build_anthropic_usage``.
|
||||
"""
|
||||
|
||||
import json
|
||||
@@ -18,7 +20,11 @@ import pytest
|
||||
from vllm.entrypoints.anthropic.protocol import (
|
||||
AnthropicMessagesRequest,
|
||||
)
|
||||
from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
|
||||
from vllm.entrypoints.anthropic.serving import (
|
||||
AnthropicServingMessages,
|
||||
_build_anthropic_usage,
|
||||
_get_cached_tokens,
|
||||
)
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import (
|
||||
ChatCompletionResponseStreamChoice,
|
||||
ChatCompletionStreamResponse,
|
||||
@@ -27,6 +33,7 @@ from vllm.entrypoints.openai.engine.protocol import (
|
||||
DeltaFunctionCall,
|
||||
DeltaMessage,
|
||||
DeltaToolCall,
|
||||
PromptTokenUsageInfo,
|
||||
UsageInfo,
|
||||
)
|
||||
|
||||
@@ -653,6 +660,108 @@ class TestThinkingBlockConversion:
|
||||
assert asst.get("content") == "Hi!"
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# Cache usage computation
|
||||
# ======================================================================
|
||||
|
||||
|
||||
class TestGetCachedTokens:
|
||||
"""Tests for _get_cached_tokens helper."""
|
||||
|
||||
def test_none_usage(self):
|
||||
assert _get_cached_tokens(None) is None
|
||||
|
||||
def test_no_prompt_tokens_details(self):
|
||||
usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
|
||||
assert _get_cached_tokens(usage) is None
|
||||
|
||||
def test_cached_tokens_present(self):
|
||||
usage = UsageInfo(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=10,
|
||||
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
|
||||
)
|
||||
assert _get_cached_tokens(usage) == 80
|
||||
|
||||
def test_cached_tokens_zero(self):
|
||||
"""Zero cached tokens should return 0, not None."""
|
||||
usage = UsageInfo(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=10,
|
||||
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
|
||||
)
|
||||
assert _get_cached_tokens(usage) == 0
|
||||
|
||||
def test_cached_tokens_none_in_details(self):
|
||||
usage = UsageInfo(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=10,
|
||||
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=None),
|
||||
)
|
||||
assert _get_cached_tokens(usage) is None
|
||||
|
||||
|
||||
class TestBuildAnthropicUsage:
|
||||
"""Tests for _build_anthropic_usage helper.
|
||||
|
||||
Anthropic defines: total_input = input_tokens + cache_read + cache_creation
|
||||
vLLM's prompt_tokens is the total.
|
||||
"""
|
||||
|
||||
def test_no_cache_info(self):
|
||||
"""When cache info is unavailable, return raw prompt_tokens."""
|
||||
result = _build_anthropic_usage(100, 10, None)
|
||||
assert result.input_tokens == 100
|
||||
assert result.output_tokens == 10
|
||||
assert result.cache_read_input_tokens is None
|
||||
assert result.cache_creation_input_tokens is None
|
||||
|
||||
def test_cache_hit(self):
|
||||
"""When cache is hit, input_tokens excludes cached tokens."""
|
||||
usage = UsageInfo(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=10,
|
||||
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
|
||||
)
|
||||
result = _build_anthropic_usage(100, 10, usage)
|
||||
assert result.input_tokens == 20 # 100 - 80
|
||||
assert result.output_tokens == 10
|
||||
assert result.cache_read_input_tokens == 80
|
||||
assert result.cache_creation_input_tokens == 0
|
||||
|
||||
def test_zero_cached_tokens(self):
|
||||
"""Zero cached tokens should still set cache_creation to 0."""
|
||||
usage = UsageInfo(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=10,
|
||||
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
|
||||
)
|
||||
result = _build_anthropic_usage(100, 10, usage)
|
||||
assert result.input_tokens == 100 # 100 - 0
|
||||
assert result.cache_read_input_tokens == 0
|
||||
assert result.cache_creation_input_tokens == 0
|
||||
|
||||
def test_all_tokens_cached(self):
|
||||
"""When all tokens are cached, input_tokens should be 0."""
|
||||
usage = UsageInfo(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=10,
|
||||
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=100),
|
||||
)
|
||||
result = _build_anthropic_usage(100, 10, usage)
|
||||
assert result.input_tokens == 0
|
||||
assert result.cache_read_input_tokens == 100
|
||||
assert result.cache_creation_input_tokens == 0
|
||||
|
||||
def test_no_prompt_tokens_details(self):
|
||||
"""UsageInfo without prompt_tokens_details returns no cache info."""
|
||||
usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
|
||||
result = _build_anthropic_usage(100, 10, usage)
|
||||
assert result.input_tokens == 100
|
||||
assert result.cache_read_input_tokens is None
|
||||
assert result.cache_creation_input_tokens is None
|
||||
|
||||
|
||||
class TestInlineSystemMessageInMessagesArray:
|
||||
"""Verify that ``role: system`` messages embedded inside the ``messages``
|
||||
array are preserved in their original position.
|
||||
@@ -1098,6 +1207,135 @@ class TestMessageStartIncludesTypeAndRole:
|
||||
assert message["role"] == "assistant"
|
||||
|
||||
|
||||
class TestStreamingCacheUsageSemantics:
|
||||
"""Locks in the documented streaming behavior of cache usage fields.
|
||||
|
||||
vLLM's OpenAI chat completion streaming only attaches
|
||||
``prompt_tokens_details`` to the terminal usage chunk. The Anthropic layer
|
||||
mirrors that contract: cache fields are omitted on ``message_start`` (key
|
||||
absence signals "unknown") and populated on ``message_delta`` (the final
|
||||
cumulative count). This is intentionally consistent with vLLM's OpenAI
|
||||
behavior, even though Anthropic's upstream API populates cache fields on
|
||||
``message_start``; closing that gap requires plumbing cache info into the
|
||||
first chunk at the OpenAI layer, which is out of scope here.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_cache_fields_absent_then_populated(self):
|
||||
"""First chunk lacks prompt_tokens_details (vLLM contract);
|
||||
message_start omits cache fields. The final chunk carries
|
||||
prompt_tokens_details, so message_delta carries resolved values."""
|
||||
|
||||
async def sse_input():
|
||||
yield _make_stream_chunk(
|
||||
delta=DeltaMessage(role="assistant", content="hi"),
|
||||
usage=UsageInfo(prompt_tokens=100, total_tokens=100),
|
||||
)
|
||||
yield _make_stream_chunk(finish_reason="stop")
|
||||
yield _make_stream_chunk(
|
||||
choices=[],
|
||||
usage=UsageInfo(
|
||||
prompt_tokens=100,
|
||||
completion_tokens=5,
|
||||
total_tokens=105,
|
||||
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
|
||||
),
|
||||
)
|
||||
yield "data: [DONE]"
|
||||
|
||||
converter = _make_stream_converter()
|
||||
output = []
|
||||
async for event in converter.message_stream_converter(sse_input()):
|
||||
output.append(event)
|
||||
events = _parse_sse_events(output)
|
||||
|
||||
# message_start: cache fields unknown → omitted from JSON entirely.
|
||||
start_usage = events[0][1]["message"]["usage"]
|
||||
assert events[0][0] == "message_start"
|
||||
assert start_usage["input_tokens"] == 100
|
||||
assert "cache_read_input_tokens" not in start_usage
|
||||
assert "cache_creation_input_tokens" not in start_usage
|
||||
|
||||
# message_delta: authoritative usage with cache fields populated.
|
||||
delta_usage = next(
|
||||
data["usage"] for ev, data in events if ev == "message_delta"
|
||||
)
|
||||
assert delta_usage["input_tokens"] == 20 # 100 - 80
|
||||
assert delta_usage["cache_read_input_tokens"] == 80
|
||||
assert delta_usage["cache_creation_input_tokens"] == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_no_cache_hit(self):
|
||||
"""When the final chunk reports cached_tokens=0, message_delta carries
|
||||
cache fields = 0 (cache miss); message_start still omits them."""
|
||||
|
||||
async def sse_input():
|
||||
yield _make_stream_chunk(
|
||||
delta=DeltaMessage(role="assistant"),
|
||||
usage=UsageInfo(prompt_tokens=50, total_tokens=50),
|
||||
)
|
||||
yield _make_stream_chunk(finish_reason="stop")
|
||||
yield _make_stream_chunk(
|
||||
choices=[],
|
||||
usage=UsageInfo(
|
||||
prompt_tokens=50,
|
||||
completion_tokens=5,
|
||||
total_tokens=55,
|
||||
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
|
||||
),
|
||||
)
|
||||
yield "data: [DONE]"
|
||||
|
||||
converter = _make_stream_converter()
|
||||
output = []
|
||||
async for event in converter.message_stream_converter(sse_input()):
|
||||
output.append(event)
|
||||
events = _parse_sse_events(output)
|
||||
|
||||
start_usage = events[0][1]["message"]["usage"]
|
||||
delta_usage = next(
|
||||
data["usage"] for ev, data in events if ev == "message_delta"
|
||||
)
|
||||
assert start_usage["input_tokens"] == 50
|
||||
assert "cache_read_input_tokens" not in start_usage
|
||||
assert "cache_creation_input_tokens" not in start_usage
|
||||
assert delta_usage["input_tokens"] == 50 # 50 - 0
|
||||
assert delta_usage["cache_read_input_tokens"] == 0
|
||||
assert delta_usage["cache_creation_input_tokens"] == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_no_prompt_tokens_details_at_all(self):
|
||||
"""If --enable-prompt-tokens-details is off, no chunk carries cache
|
||||
info; both message_start and message_delta omit cache fields."""
|
||||
|
||||
async def sse_input():
|
||||
yield _make_stream_chunk(
|
||||
delta=DeltaMessage(role="assistant"),
|
||||
usage=UsageInfo(prompt_tokens=30, total_tokens=30),
|
||||
)
|
||||
yield _make_stream_chunk(finish_reason="stop")
|
||||
yield _make_stream_chunk(
|
||||
choices=[],
|
||||
usage=UsageInfo(prompt_tokens=30, completion_tokens=2, total_tokens=32),
|
||||
)
|
||||
yield "data: [DONE]"
|
||||
|
||||
converter = _make_stream_converter()
|
||||
output = []
|
||||
async for event in converter.message_stream_converter(sse_input()):
|
||||
output.append(event)
|
||||
events = _parse_sse_events(output)
|
||||
|
||||
start_usage = events[0][1]["message"]["usage"]
|
||||
delta_usage = next(
|
||||
data["usage"] for ev, data in events if ev == "message_delta"
|
||||
)
|
||||
assert "cache_read_input_tokens" not in start_usage
|
||||
assert "cache_creation_input_tokens" not in start_usage
|
||||
assert "cache_read_input_tokens" not in delta_usage
|
||||
assert "cache_creation_input_tokens" not in delta_usage
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# Auto-detection of system-first template requirement
|
||||
# ======================================================================
|
||||
|
||||
@@ -364,7 +364,7 @@ class MockVLLMServer:
|
||||
await self._serve_task
|
||||
|
||||
|
||||
def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]):
|
||||
def launch_mock_vllm(child_args: argparse.Namespace):
|
||||
logger.info("Launching mock vLLM on port %s", child_args.port)
|
||||
mock_vllm = MockVLLMServer(
|
||||
port=child_args.port,
|
||||
@@ -375,7 +375,7 @@ def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]
|
||||
|
||||
|
||||
def launch_mock_vllm_with_drain(
|
||||
child_args: argparse.Namespace, env_updates: dict[str, str]
|
||||
child_args: argparse.Namespace,
|
||||
):
|
||||
logger.info("Launching mock vLLM with 15s drain on port %s", child_args.port)
|
||||
mock_vllm = MockVLLMServer(
|
||||
|
||||
@@ -6,15 +6,22 @@ from typing import Final
|
||||
import pytest
|
||||
import schemathesis
|
||||
from hypothesis import HealthCheck, settings
|
||||
from schemathesis import GenerationConfig
|
||||
from schemathesis.models import Case
|
||||
from schemathesis import GenerationMode
|
||||
from schemathesis.config import (
|
||||
ChecksConfig,
|
||||
CoveragePhaseConfig,
|
||||
GenerationConfig,
|
||||
PhasesConfig,
|
||||
PositiveDataAcceptanceConfig,
|
||||
ProjectConfig,
|
||||
ProjectsConfig,
|
||||
SchemathesisConfig,
|
||||
)
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ...utils import RemoteOpenAIServer
|
||||
|
||||
schemathesis.experimental.OPEN_API_3_1.enable()
|
||||
|
||||
MODEL_NAME = "HuggingFaceTB/SmolVLM-256M-Instruct"
|
||||
MAXIMUM_IMAGES = 2
|
||||
_ROCM_TIMEOUT_MULTIPLIER = 3 if current_platform.is_rocm() else 1
|
||||
@@ -44,21 +51,38 @@ def server():
|
||||
@pytest.fixture(scope="module")
|
||||
def get_schema(server):
|
||||
# avoid generating null (\x00) bytes in strings during test case generation
|
||||
return schemathesis.openapi.from_uri(
|
||||
return schemathesis.openapi.from_url(
|
||||
f"{server.url_root}/openapi.json",
|
||||
generation_config=GenerationConfig(allow_x00=False),
|
||||
config=SchemathesisConfig(
|
||||
projects=ProjectsConfig(
|
||||
default=ProjectConfig(
|
||||
generation=GenerationConfig(
|
||||
allow_x00=False,
|
||||
modes=[GenerationMode.POSITIVE],
|
||||
),
|
||||
checks=ChecksConfig(
|
||||
positive_data_acceptance=PositiveDataAcceptanceConfig(
|
||||
enabled=False,
|
||||
),
|
||||
),
|
||||
phases=PhasesConfig(
|
||||
coverage=CoveragePhaseConfig(enabled=False),
|
||||
),
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
schema = schemathesis.from_pytest_fixture("get_schema")
|
||||
schema = schemathesis.pytest.from_fixture("get_schema")
|
||||
|
||||
|
||||
@schemathesis.hook
|
||||
def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
|
||||
def before_generate_case(context: schemathesis.HookContext, strategy):
|
||||
op = context.operation
|
||||
assert op is not None
|
||||
|
||||
def no_invalid_types(case: schemathesis.models.Case):
|
||||
def no_invalid_types(case: schemathesis.Case):
|
||||
"""
|
||||
Skips tool_calls with `"type": "custom"` which schemathesis incorrectly
|
||||
generates instead of the valid `"type": "function"`.
|
||||
@@ -68,39 +92,25 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
|
||||
-d '{"messages": [{"role": "assistant", "tool_calls": [{"custom": {"input": "", "name": ""}, "id": "", "type": "custom"}]}]}' \
|
||||
http://localhost:8000/v1/chat/completions
|
||||
""" # noqa: E501
|
||||
if hasattr(case, "body") and isinstance(case.body, dict):
|
||||
if (
|
||||
"messages" in case.body
|
||||
and isinstance(case.body["messages"], list)
|
||||
and len(case.body["messages"]) > 0
|
||||
):
|
||||
for message in case.body["messages"]:
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
if (
|
||||
hasattr(case, "body")
|
||||
and isinstance(case.body, dict)
|
||||
and "messages" in case.body
|
||||
and isinstance(case.body["messages"], list)
|
||||
and len(case.body["messages"]) > 0
|
||||
):
|
||||
for message in case.body["messages"]:
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
|
||||
tool_calls = message.get("tool_calls", [])
|
||||
if isinstance(tool_calls, list):
|
||||
for tool_call in tool_calls:
|
||||
if isinstance(tool_call, dict):
|
||||
if tool_call.get("type") != "function":
|
||||
return False
|
||||
if "custom" in tool_call:
|
||||
return False
|
||||
|
||||
# Sometimes structured_outputs.grammar is generated to be empty
|
||||
# Causing a server error in EBNF grammar parsing
|
||||
# https://github.com/vllm-project/vllm/pull/22587#issuecomment-3195253421
|
||||
structured_outputs = case.body.get("structured_outputs", {})
|
||||
grammar = (
|
||||
structured_outputs.get("grammar")
|
||||
if isinstance(structured_outputs, dict)
|
||||
else None
|
||||
)
|
||||
|
||||
if grammar == "":
|
||||
# Allow None (will be handled as no grammar)
|
||||
# But skip empty strings
|
||||
return False
|
||||
tool_calls = message.get("tool_calls", [])
|
||||
if isinstance(tool_calls, list):
|
||||
for tool_call in tool_calls:
|
||||
if isinstance(tool_call, dict):
|
||||
if tool_call.get("type") != "function":
|
||||
return False
|
||||
if "custom" in tool_call:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
@@ -108,7 +118,6 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
|
||||
|
||||
|
||||
@schema.parametrize()
|
||||
@schema.override(headers={"Content-Type": "application/json"})
|
||||
@settings(
|
||||
deadline=LONG_TIMEOUT_SECONDS * 1000,
|
||||
max_examples=50,
|
||||
@@ -122,7 +131,7 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
|
||||
# generating large-but-valid request bodies before vLLM is called.
|
||||
suppress_health_check=[HealthCheck.filter_too_much, HealthCheck.data_too_large],
|
||||
)
|
||||
def test_openapi_stateless(case: Case):
|
||||
def test_openapi_stateless(case: schemathesis.Case):
|
||||
key = (
|
||||
case.operation.method.upper(),
|
||||
case.operation.path,
|
||||
@@ -151,4 +160,8 @@ def test_openapi_stateless(case: Case):
|
||||
}.get(key, DEFAULT_TIMEOUT_SECONDS)
|
||||
|
||||
# No need to verify SSL certificate for localhost
|
||||
case.call_and_validate(verify=False, timeout=timeout)
|
||||
case.call_and_validate(
|
||||
verify=False,
|
||||
timeout=timeout,
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
|
||||
@@ -25,7 +25,7 @@ def server():
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--max-model-len",
|
||||
"5000",
|
||||
"16384",
|
||||
"--enforce-eager",
|
||||
"--limit-mm-per-prompt",
|
||||
json.dumps({"video": MAXIMUM_VIDEOS}),
|
||||
@@ -143,4 +143,4 @@ def test_chat_video_url_request(server: RemoteOpenAIServer, model_name: str):
|
||||
assert output.model == model_name
|
||||
assert len(output.data) == 1
|
||||
assert len(output.data[0].probs) == 2
|
||||
assert output.usage.prompt_tokens == 4807
|
||||
assert output.usage.prompt_tokens == 8993
|
||||
|
||||
@@ -8,6 +8,7 @@ import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from tests.utils import RemoteLaunchRenderServer
|
||||
from vllm.tokenizers import get_tokenizer
|
||||
|
||||
MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
|
||||
|
||||
@@ -486,3 +487,438 @@ async def test_derender_completion_kv_transfer_params_passthrough(client):
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["kv_transfer_params"] == kv
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# E2E: render -> derender roundtrip with parser (reasoning + tool calls)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
PARSER_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
|
||||
|
||||
_E2E_TOOLS = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather for a city",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def parser_server():
|
||||
args = [
|
||||
"--enable-auto-tool-choice",
|
||||
"--tool-call-parser",
|
||||
"hermes",
|
||||
"--reasoning-parser",
|
||||
"deepseek_r1",
|
||||
]
|
||||
with RemoteLaunchRenderServer(PARSER_MODEL, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def parser_client(parser_server):
|
||||
async with httpx.AsyncClient(
|
||||
base_url=parser_server.url_for(""), timeout=60.0
|
||||
) as http_client:
|
||||
yield http_client
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def parser_tokenizer():
|
||||
return get_tokenizer(PARSER_MODEL)
|
||||
|
||||
|
||||
def _encode(tokenizer, text: str) -> list[int]:
|
||||
return tokenizer.encode(text, add_special_tokens=False)
|
||||
|
||||
|
||||
def _decoded(tokenizer, token_ids: list[int]) -> str:
|
||||
return tokenizer.decode(token_ids, skip_special_tokens=True)
|
||||
|
||||
|
||||
def _require_markers_survive(tokenizer, text: str, *markers: str) -> list[int]:
|
||||
"""Encode text and skip the test if any marker is lost in roundtrip."""
|
||||
ids = _encode(tokenizer, text)
|
||||
decoded = tokenizer.decode(ids, skip_special_tokens=False)
|
||||
for m in markers:
|
||||
if m not in decoded:
|
||||
pytest.skip(f"Marker {m!r} lost in encode->decode roundtrip")
|
||||
return ids
|
||||
|
||||
|
||||
async def _e2e_render_chat(
|
||||
client: httpx.AsyncClient,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
) -> dict:
|
||||
resp = await client.post(
|
||||
"/v1/chat/completions/render",
|
||||
json={"model": model, "messages": messages},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
return resp.json()
|
||||
|
||||
|
||||
def _e2e_generate_response(
|
||||
token_ids: list[int],
|
||||
request_id: str = "chatcmpl-e2e-test",
|
||||
) -> dict:
|
||||
return {
|
||||
"request_id": request_id,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"token_ids": token_ids,
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_e2e_plain_roundtrip(parser_client, parser_tokenizer):
|
||||
"""Plain text without reasoning markers roundtrips correctly."""
|
||||
messages = [{"role": "user", "content": "What is 2+2?"}]
|
||||
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
|
||||
|
||||
answer = "The answer is four."
|
||||
output_ids = _encode(parser_tokenizer, answer)
|
||||
expected = _decoded(parser_tokenizer, output_ids)
|
||||
|
||||
resp = await parser_client.post(
|
||||
"/v1/chat/completions/derender",
|
||||
json={
|
||||
"model": PARSER_MODEL,
|
||||
"generate_response": _e2e_generate_response(output_ids),
|
||||
"prompt_tokens": len(gen_req["token_ids"]),
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
content = resp.json()["choices"][0]["message"]["content"]
|
||||
assert content == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_e2e_token_identity(parser_client, parser_tokenizer):
|
||||
"""encode(derender(token_ids)) == token_ids (RL invariant)."""
|
||||
messages = [{"role": "user", "content": "Hi"}]
|
||||
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
|
||||
|
||||
answer = "Hello! How can I help?"
|
||||
output_ids = _encode(parser_tokenizer, answer)
|
||||
|
||||
resp = await parser_client.post(
|
||||
"/v1/chat/completions/derender",
|
||||
json={
|
||||
"model": PARSER_MODEL,
|
||||
"generate_response": _e2e_generate_response(output_ids),
|
||||
"prompt_tokens": len(gen_req["token_ids"]),
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
content = resp.json()["choices"][0]["message"]["content"]
|
||||
re_encoded = _encode(parser_tokenizer, content)
|
||||
assert output_ids == re_encoded
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_e2e_non_ascii_roundtrip(parser_client, parser_tokenizer):
|
||||
"""CJK + emoji roundtrip without U+FFFD."""
|
||||
messages = [{"role": "user", "content": "Reply in Chinese"}]
|
||||
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
|
||||
|
||||
answer = "你好世界 😀"
|
||||
output_ids = _encode(parser_tokenizer, answer)
|
||||
|
||||
resp = await parser_client.post(
|
||||
"/v1/chat/completions/derender",
|
||||
json={
|
||||
"model": PARSER_MODEL,
|
||||
"generate_response": _e2e_generate_response(output_ids),
|
||||
"prompt_tokens": len(gen_req["token_ids"]),
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
content = resp.json()["choices"][0]["message"]["content"]
|
||||
assert "�" not in content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_e2e_parsed_reasoning(parser_client, parser_tokenizer):
|
||||
"""<think>...</think> splits into reasoning + content."""
|
||||
messages = [{"role": "user", "content": "What is 2+3?"}]
|
||||
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
|
||||
|
||||
reasoning_text = "The user wants 2 plus 3. That is 5."
|
||||
answer_text = "The answer is 5."
|
||||
output_text = f"<think>{reasoning_text}</think>{answer_text}"
|
||||
output_ids = _require_markers_survive(parser_tokenizer, output_text, "</think>")
|
||||
|
||||
resp = await parser_client.post(
|
||||
"/v1/chat/completions/derender",
|
||||
json={
|
||||
"model": PARSER_MODEL,
|
||||
"generate_response": _e2e_generate_response(output_ids),
|
||||
"prompt_tokens": len(gen_req["token_ids"]),
|
||||
"chat_request": {
|
||||
"model": PARSER_MODEL,
|
||||
"messages": messages,
|
||||
"include_reasoning": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
msg = resp.json()["choices"][0]["message"]
|
||||
assert msg["reasoning"] is not None
|
||||
assert reasoning_text in msg["reasoning"]
|
||||
assert answer_text in msg["content"]
|
||||
assert "<think>" not in msg["content"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_e2e_parsed_tool_call(parser_client, parser_tokenizer):
|
||||
"""<tool_call> extracted into tool_calls field."""
|
||||
messages = [{"role": "user", "content": "Weather in Paris?"}]
|
||||
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
|
||||
|
||||
output_text = (
|
||||
"<think>Let me check the weather.</think>"
|
||||
'<tool_call>\n{"name": "get_weather", '
|
||||
'"arguments": {"city": "Paris"}}\n</tool_call>'
|
||||
)
|
||||
output_ids = _require_markers_survive(
|
||||
parser_tokenizer,
|
||||
output_text,
|
||||
"</think>",
|
||||
"<tool_call>",
|
||||
"</tool_call>",
|
||||
)
|
||||
|
||||
resp = await parser_client.post(
|
||||
"/v1/chat/completions/derender",
|
||||
json={
|
||||
"model": PARSER_MODEL,
|
||||
"generate_response": _e2e_generate_response(output_ids),
|
||||
"prompt_tokens": len(gen_req["token_ids"]),
|
||||
"chat_request": {
|
||||
"model": PARSER_MODEL,
|
||||
"messages": messages,
|
||||
"tools": _E2E_TOOLS,
|
||||
"tool_choice": "auto",
|
||||
},
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
choice = resp.json()["choices"][0]
|
||||
assert choice["message"]["tool_calls"]
|
||||
assert choice["message"]["tool_calls"][0]["function"]["name"] == "get_weather"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_e2e_parsed_reasoning_and_tool_call(parser_client, parser_tokenizer):
|
||||
"""Reasoning + tool call in the same output."""
|
||||
messages = [{"role": "user", "content": "Weather in Paris?"}]
|
||||
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
|
||||
|
||||
reasoning_text = "I should look up the weather."
|
||||
tool_text = (
|
||||
'<tool_call>\n{"name": "get_weather", '
|
||||
'"arguments": {"city": "Paris"}}\n</tool_call>'
|
||||
)
|
||||
output_text = f"<think>{reasoning_text}</think>{tool_text}"
|
||||
output_ids = _require_markers_survive(
|
||||
parser_tokenizer, output_text, "</think>", "<tool_call>"
|
||||
)
|
||||
|
||||
resp = await parser_client.post(
|
||||
"/v1/chat/completions/derender",
|
||||
json={
|
||||
"model": PARSER_MODEL,
|
||||
"generate_response": _e2e_generate_response(output_ids),
|
||||
"prompt_tokens": len(gen_req["token_ids"]),
|
||||
"chat_request": {
|
||||
"model": PARSER_MODEL,
|
||||
"messages": messages,
|
||||
"tools": _E2E_TOOLS,
|
||||
"tool_choice": "auto",
|
||||
"include_reasoning": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
choice = resp.json()["choices"][0]
|
||||
assert choice["message"]["reasoning"] is not None
|
||||
assert reasoning_text in choice["message"]["reasoning"]
|
||||
assert choice["message"]["tool_calls"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_e2e_no_chat_request_fallback(parser_client, parser_tokenizer):
|
||||
"""Without chat_request, derender falls back to plain detokenization."""
|
||||
messages = [{"role": "user", "content": "Hello"}]
|
||||
gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
|
||||
|
||||
answer = "Hi there!"
|
||||
output_ids = _encode(parser_tokenizer, answer)
|
||||
|
||||
resp = await parser_client.post(
|
||||
"/v1/chat/completions/derender",
|
||||
json={
|
||||
"model": PARSER_MODEL,
|
||||
"generate_response": _e2e_generate_response(output_ids),
|
||||
"prompt_tokens": len(gen_req["token_ids"]),
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
content = resp.json()["choices"][0]["message"]["content"]
|
||||
assert "Hi" in content
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# E2E: HarmonyParser + GPT-OSS
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
HARMONY_MODEL = "openai/gpt-oss-20b"
|
||||
|
||||
|
||||
def _ensure_harmony_vocab():
|
||||
"""Pre-cache the o200k_base BPE file needed by openai-harmony.
|
||||
|
||||
The Rust tiktoken-rs backend downloads from Azure Blob Storage, which
|
||||
may be unreachable in some environments. When the cache is cold we
|
||||
fetch the file ourselves and place it in ``/tmp/tiktoken-rs-cache/``
|
||||
using the SHA-1(URL) filename that tiktoken-rs expects.
|
||||
"""
|
||||
import hashlib
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
url = "https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken"
|
||||
cache_dir = Path("/tmp/tiktoken-rs-cache")
|
||||
cache_key = hashlib.sha1(url.encode()).hexdigest()
|
||||
cache_file = cache_dir / cache_key
|
||||
if not cache_file.exists():
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
urllib.request.urlretrieve(url, cache_file)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def harmony_server():
|
||||
_ensure_harmony_vocab()
|
||||
args = [
|
||||
"--trust-remote-code",
|
||||
"--enable-auto-tool-choice",
|
||||
"--tool-call-parser",
|
||||
"openai",
|
||||
"--reasoning-parser",
|
||||
"openai_gptoss",
|
||||
]
|
||||
with RemoteLaunchRenderServer(HARMONY_MODEL, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def harmony_client(harmony_server):
|
||||
async with httpx.AsyncClient(
|
||||
base_url=harmony_server.url_for(""), timeout=60.0
|
||||
) as http_client:
|
||||
yield http_client
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def harmony_tokenizer():
|
||||
return get_tokenizer(HARMONY_MODEL, trust_remote_code=True)
|
||||
|
||||
|
||||
def _harmony_extract_assistant_ids(
|
||||
tokenizer, assistant_msg: dict, user_content: str = "test"
|
||||
) -> list[int]:
|
||||
"""Extract assistant token IDs via apply_chat_template diff."""
|
||||
prompt = [{"role": "user", "content": user_content}]
|
||||
full = prompt + [assistant_msg]
|
||||
text_prompt = tokenizer.apply_chat_template(
|
||||
prompt, add_generation_prompt=True, tokenize=False
|
||||
)
|
||||
text_full = tokenizer.apply_chat_template(
|
||||
full, add_generation_prompt=False, tokenize=False
|
||||
)
|
||||
prompt_ids = tokenizer.encode(text_prompt)
|
||||
full_ids = tokenizer.encode(text_full)
|
||||
assistant_ids = list(full_ids[len(prompt_ids) :])
|
||||
if not assistant_ids:
|
||||
pytest.skip("Could not extract assistant tokens for Harmony")
|
||||
return assistant_ids
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_e2e_harmony_plain_roundtrip(harmony_client, harmony_tokenizer):
|
||||
"""GPT-OSS content-only roundtrip."""
|
||||
messages = [{"role": "user", "content": "What is 2+2?"}]
|
||||
gen_req = await _e2e_render_chat(harmony_client, HARMONY_MODEL, messages)
|
||||
|
||||
assistant_msg = {"role": "assistant", "content": "Four."}
|
||||
output_ids = _harmony_extract_assistant_ids(harmony_tokenizer, assistant_msg)
|
||||
|
||||
resp = await harmony_client.post(
|
||||
"/v1/chat/completions/derender",
|
||||
json={
|
||||
"model": HARMONY_MODEL,
|
||||
"generate_response": _e2e_generate_response(output_ids),
|
||||
"prompt_tokens": len(gen_req["token_ids"]),
|
||||
"chat_request": {
|
||||
"model": HARMONY_MODEL,
|
||||
"messages": messages,
|
||||
},
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
content = resp.json()["choices"][0]["message"]["content"]
|
||||
assert content is not None and len(content) > 0
|
||||
assert "Four" in content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_e2e_harmony_reasoning(harmony_client, harmony_tokenizer):
|
||||
"""GPT-OSS reasoning: analysis channel extracted."""
|
||||
messages = [{"role": "user", "content": "Add 2 and 3."}]
|
||||
gen_req = await _e2e_render_chat(harmony_client, HARMONY_MODEL, messages)
|
||||
|
||||
reasoning_text = "The user wants 2 plus 3."
|
||||
answer_text = "The answer is 5."
|
||||
assistant_msg = {
|
||||
"role": "assistant",
|
||||
"thinking": reasoning_text,
|
||||
"content": answer_text,
|
||||
}
|
||||
output_ids = _harmony_extract_assistant_ids(harmony_tokenizer, assistant_msg)
|
||||
|
||||
decoded = harmony_tokenizer.decode(output_ids)
|
||||
if reasoning_text not in decoded:
|
||||
pytest.skip("Harmony template did not render thinking")
|
||||
|
||||
resp = await harmony_client.post(
|
||||
"/v1/chat/completions/derender",
|
||||
json={
|
||||
"model": HARMONY_MODEL,
|
||||
"generate_response": _e2e_generate_response(output_ids),
|
||||
"prompt_tokens": len(gen_req["token_ids"]),
|
||||
"chat_request": {
|
||||
"model": HARMONY_MODEL,
|
||||
"messages": messages,
|
||||
"include_reasoning": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
msg = resp.json()["choices"][0]["message"]
|
||||
assert msg["reasoning"] is not None
|
||||
assert reasoning_text in msg["reasoning"]
|
||||
assert answer_text in (msg["content"] or "")
|
||||
|
||||
@@ -15,7 +15,7 @@ from vllm.entrypoints.serve.tokenize.protocol import (
|
||||
TokenizeChatRequest,
|
||||
TokenizeCompletionRequest,
|
||||
)
|
||||
from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
|
||||
from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
|
||||
MODEL_NAME = "openai-community/gpt2"
|
||||
@@ -58,7 +58,7 @@ class MockModelConfig:
|
||||
return self.diff_sampling_param or {}
|
||||
|
||||
|
||||
def _build_serving_tokenization(engine: AsyncLLM) -> OpenAIServingTokenization:
|
||||
def _build_serving_tokenization(engine: AsyncLLM) -> ServingTokenization:
|
||||
models = OpenAIServingModels(
|
||||
engine_client=engine,
|
||||
base_model_paths=BASE_MODEL_PATHS,
|
||||
@@ -71,8 +71,7 @@ def _build_serving_tokenization(engine: AsyncLLM) -> OpenAIServingTokenization:
|
||||
chat_template=None,
|
||||
chat_template_content_format="auto",
|
||||
)
|
||||
return OpenAIServingTokenization(
|
||||
engine,
|
||||
return ServingTokenization(
|
||||
models,
|
||||
openai_serving_render=serving_render,
|
||||
request_logger=None,
|
||||
|
||||
@@ -78,7 +78,16 @@ def test_gsm8k_correctness(config_filename):
|
||||
"Skipping DeepSeek-V3.2 and DeepSeek-R1 on ROCm platforms "
|
||||
"due to agent pool disk space issues and pod evictions."
|
||||
)
|
||||
if current_platform.is_rocm() and (
|
||||
"Qwen3.5-35B-A3B-MXFP4-AITER-TP2" in config_filename.name
|
||||
):
|
||||
from vllm.platforms.rocm import on_gfx950
|
||||
|
||||
if not on_gfx950():
|
||||
pytest.skip(
|
||||
"Skipping Qwen3.5-35B-A3B-MXFP4-AITER-TP2 on non-GFX950 platforms. "
|
||||
"The quantization scheme is not supported on non-GFX950 platforms."
|
||||
)
|
||||
# Parse server arguments from config (use shlex to handle quoted strings)
|
||||
server_args_str = eval_config.get("server_args", "")
|
||||
server_args = shlex.split(server_args_str) if server_args_str else []
|
||||
|
||||
@@ -15,16 +15,14 @@ from vllm.config import (
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.platforms.cpu import CpuPlatform
|
||||
|
||||
# CudaPlatform and RocmPlatform import their respective compiled C extensions
|
||||
# at module level, raising ModuleNotFoundError on incompatible builds.
|
||||
try:
|
||||
if current_platform.is_cuda():
|
||||
from vllm.platforms.cuda import CudaPlatform
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
else:
|
||||
CudaPlatform = None
|
||||
|
||||
try:
|
||||
if current_platform.is_rocm():
|
||||
from vllm.platforms.rocm import RocmPlatform
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
else:
|
||||
RocmPlatform = None
|
||||
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
@@ -434,9 +432,15 @@ def test_per_head_quant_scales_backend_selection(
|
||||
[
|
||||
("FLASH_ATTN", True, True), # FlashAttn supports non-causal
|
||||
("FLASH_ATTN", False, True), # FlashAttn also works with causal
|
||||
("FLASHINFER", True, False), # FlashInfer does not support non-causal
|
||||
("FLASHINFER", False, True), # FlashInfer works with causal
|
||||
],
|
||||
]
|
||||
+ (
|
||||
[
|
||||
("FLASHINFER", True, True), # FlashInfer supports non-causal
|
||||
("FLASHINFER", False, True), # FlashInfer works with causal
|
||||
]
|
||||
if CudaPlatform is not None
|
||||
else []
|
||||
),
|
||||
)
|
||||
def test_non_causal_backend_selection(
|
||||
backend_name: str, use_non_causal: bool, should_succeed: bool
|
||||
@@ -459,11 +463,12 @@ def test_non_causal_backend_selection(
|
||||
attention_config=attention_config, cache_config=cache_config
|
||||
)
|
||||
|
||||
if CudaPlatform is None:
|
||||
pytest.skip("CudaPlatform not available")
|
||||
platform = CudaPlatform or RocmPlatform
|
||||
if platform is None:
|
||||
pytest.skip("CudaPlatform and RocmPlatform are not available")
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch("vllm.platforms.current_platform", CudaPlatform()),
|
||||
patch("vllm.platforms.current_platform", platform()),
|
||||
):
|
||||
if should_succeed:
|
||||
backend = get_attn_backend(
|
||||
|
||||
@@ -212,3 +212,69 @@ def test_cutlass_mla_decode(
|
||||
print(
|
||||
f"{t:.3f} ms, {FLOPS / 10**9 / t:.0f} TFLOPS,", f"{bytes / 10**6 / t:.0f} GB/s"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.has_device_capability(100),
|
||||
reason=CUTLASS_MLA_UNSUPPORTED_REASON,
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_cutlass_mla_decode_cross_layer_view():
|
||||
"""The kernel must read the cache's page-dim stride instead of assuming
|
||||
pages are packed back-to-back. A per-layer view into a cross-layer
|
||||
(block-major) cache has stride(0) inflated by num_layers; outputs must
|
||||
match a contiguous cache holding the same data exactly."""
|
||||
device = torch.device("cuda:0")
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
torch.set_default_device(device)
|
||||
torch.manual_seed(42)
|
||||
|
||||
b, mean_sk, d, dv, block_size = 4, 512, 576, 512, 64
|
||||
num_layers, layer_idx = 3, 1
|
||||
scale = math.sqrt(d) ** (-1)
|
||||
|
||||
num_pages = b * (mean_sk // block_size)
|
||||
cache_seqlens = torch.full((b,), mean_sk, dtype=torch.int32)
|
||||
block_table = torch.arange(num_pages, dtype=torch.int32).view(
|
||||
b, mean_sk // block_size
|
||||
)
|
||||
|
||||
kv_contig = torch.randn(num_pages, block_size, d)
|
||||
# Neighbor layers hold random data so packed-pages addressing reads
|
||||
# garbage rather than zeros.
|
||||
kv_cross_layer = torch.randn(num_pages, num_layers, block_size, d)
|
||||
kv_view = kv_cross_layer[:, layer_idx]
|
||||
kv_view.copy_(kv_contig)
|
||||
assert kv_view.stride(0) == num_layers * block_size * d
|
||||
|
||||
q_nope = torch.randn(b, 128, dv)
|
||||
q_pe = torch.randn(b, 128, d - dv)
|
||||
sm_count = num_compute_units(device.index)
|
||||
workspace_size = ops.sm100_cutlass_mla_get_workspace_size(
|
||||
mean_sk, b, sm_count, num_kv_splits=1
|
||||
)
|
||||
workspace = torch.empty(workspace_size, dtype=torch.uint8)
|
||||
|
||||
def run(cache):
|
||||
out = torch.empty(b, 128, dv)
|
||||
lse = torch.empty(b, 128, dtype=torch.float32)
|
||||
ops.sm100_cutlass_mla_decode(
|
||||
out,
|
||||
lse,
|
||||
q_nope,
|
||||
q_pe,
|
||||
cache,
|
||||
cache_seqlens,
|
||||
block_table,
|
||||
workspace,
|
||||
scale,
|
||||
1,
|
||||
)
|
||||
return out, lse
|
||||
|
||||
out_contig, lse_contig = run(kv_contig)
|
||||
out_view, lse_view = run(kv_view)
|
||||
|
||||
# Same data and same compute order; only addressing differs.
|
||||
assert torch.equal(out_contig, out_view)
|
||||
assert torch.equal(lse_contig, lse_view)
|
||||
|
||||
@@ -0,0 +1,566 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Bit-exact kernel equivalence for MLA decode/write kernels on the
|
||||
cross-layer (block-major) KV cache layout.
|
||||
|
||||
The cross-layer layout carves each layer's per-block page out of a single
|
||||
unified slot, so the per-layer view has an inflated ``stride(0)`` (the full
|
||||
unified slot) and a non-zero storage offset. These tests confirm the MLA
|
||||
kernels behind the backends that opt in to the layout (FlashMLA dense,
|
||||
FlashInfer MLA dense, FlashMLA fp8 sparse, plus the ``concat_and_cache_mla``
|
||||
write) honor that strided view bit-identically to a contiguous per-layer
|
||||
cache, and that writes do not bleed into neighbouring layers' segments.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not torch.cuda.is_available(), reason="MLA cache kernels require CUDA"
|
||||
)
|
||||
|
||||
|
||||
def test_concat_and_cache_mla_into_unified_slot_view():
|
||||
"""concat_and_cache_mla must write correctly into a per-layer view whose
|
||||
block stride is the full unified slot (block-major), with zero bleed into
|
||||
the other layers' segments of the same slot."""
|
||||
from vllm import _custom_ops as ops
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
kv_lora_rank = 512
|
||||
pe = 64
|
||||
entry = kv_lora_rank + pe
|
||||
page = 64
|
||||
num_blocks = 32
|
||||
ntok = 200
|
||||
|
||||
kv_c = torch.randn(ntok, kv_lora_rank, device=dev, dtype=torch.bfloat16)
|
||||
k_pe = torch.randn(ntok, pe, device=dev, dtype=torch.bfloat16)
|
||||
slot = torch.randperm(num_blocks * page, device=dev, dtype=torch.int64)[:ntok]
|
||||
scale = torch.tensor(1.0, device=dev)
|
||||
|
||||
def write(cache):
|
||||
ops.concat_and_cache_mla(kv_c, k_pe, cache, slot, "auto", scale)
|
||||
|
||||
# Contiguous per-layer reference: (num_blocks, page, entry).
|
||||
ref = torch.zeros(num_blocks, page, entry, device=dev, dtype=torch.bfloat16)
|
||||
write(ref)
|
||||
|
||||
# Unified slot holding three layer pages per block. Carve the middle
|
||||
# layer's view (non-zero offset, block stride == full unified slot).
|
||||
layer_page_elems = page * entry
|
||||
n_layers = 3
|
||||
unified_slot_elems = n_layers * layer_page_elems
|
||||
big = torch.zeros(num_blocks, unified_slot_elems, device=dev, dtype=torch.bfloat16)
|
||||
flat = big.view(-1)
|
||||
offset = layer_page_elems # middle layer
|
||||
view = torch.as_strided(
|
||||
flat,
|
||||
size=(num_blocks, page, entry),
|
||||
stride=(unified_slot_elems, entry, 1),
|
||||
storage_offset=offset,
|
||||
)
|
||||
assert not view.is_contiguous()
|
||||
assert view.stride(0) == unified_slot_elems
|
||||
write(view)
|
||||
|
||||
# Bit-exact equivalence and zero bleed into the neighbour segments.
|
||||
max_diff = (ref.float() - view.float()).abs().max().item()
|
||||
assert max_diff == 0.0, f"max|Δ| = {max_diff}"
|
||||
|
||||
neighbour_lo = torch.as_strided(
|
||||
flat, (num_blocks, layer_page_elems), (unified_slot_elems, 1), 0
|
||||
)
|
||||
neighbour_hi = torch.as_strided(
|
||||
flat,
|
||||
(num_blocks, layer_page_elems),
|
||||
(unified_slot_elems, 1),
|
||||
2 * layer_page_elems,
|
||||
)
|
||||
assert neighbour_lo.abs().max().item() == 0.0
|
||||
assert neighbour_hi.abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_flashmla_dense_decode_unified_slot_view():
|
||||
"""FlashMLA dense decode (FLASHMLA backend, e.g. Kimi-K2-style dense MLA
|
||||
on Hopper) must read a unified-slot block-major view bit-identically to a
|
||||
contiguous per-layer cache."""
|
||||
import vllm.v1.attention.ops.flashmla as fm
|
||||
|
||||
ok, reason = fm.is_flashmla_dense_supported()
|
||||
if not ok:
|
||||
pytest.skip(reason)
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
dt = torch.bfloat16
|
||||
head_dim = 576
|
||||
hdv = 512
|
||||
h_q = 128
|
||||
page = 64
|
||||
num_blocks = 64
|
||||
bs = 4
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
|
||||
q = torch.randn(bs, 1, h_q, head_dim, device=dev, dtype=dt) * 0.1
|
||||
kv_data = torch.randn(num_blocks, page, 1, head_dim, device=dev, dtype=dt) * 0.1
|
||||
|
||||
# (A) contiguous per-layer reference.
|
||||
cache_contiguous = kv_data.clone().contiguous()
|
||||
|
||||
# (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
|
||||
unified = (
|
||||
torch.randn(num_blocks, n_layers, page, 1, head_dim, device=dev, dtype=dt) * 0.1
|
||||
)
|
||||
unified[:, layer].copy_(kv_data)
|
||||
cache_view = unified[:, layer]
|
||||
assert not cache_view.is_contiguous()
|
||||
assert cache_view.stride(0) == n_layers * page * 1 * head_dim
|
||||
|
||||
max_blk = num_blocks // bs
|
||||
block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
|
||||
bs, max_blk
|
||||
)
|
||||
cache_seqlens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
|
||||
|
||||
def run(kc):
|
||||
meta, num_splits = fm.get_mla_metadata()
|
||||
out, _ = fm.flash_mla_with_kvcache(
|
||||
q=q,
|
||||
k_cache=kc,
|
||||
block_table=block_table,
|
||||
cache_seqlens=cache_seqlens,
|
||||
head_dim_v=hdv,
|
||||
tile_scheduler_metadata=meta,
|
||||
num_splits=num_splits,
|
||||
softmax_scale=head_dim**-0.5,
|
||||
causal=True,
|
||||
)
|
||||
return out.clone().float()
|
||||
|
||||
out_ref = run(cache_contiguous)
|
||||
out_view = run(cache_view)
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert out_ref.abs().max().item() > 0.0
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_flashinfer_mla_dense_decode_unified_slot_view():
|
||||
"""FlashInfer MLA dense decode must read a unified-slot block-major view
|
||||
(inflated stride(0), non-zero storage offset) bit-identically to a
|
||||
contiguous per-layer cache."""
|
||||
try:
|
||||
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
|
||||
except ImportError:
|
||||
pytest.skip("flashinfer is not available")
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.is_device_capability_family(100):
|
||||
pytest.skip("FlashInfer trtllm-gen MLA requires sm100")
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
dt = torch.bfloat16
|
||||
kv_lora_rank = 512
|
||||
qk_rope_head_dim = 64
|
||||
qk_nope_head_dim = 128
|
||||
head_dim = kv_lora_rank + qk_rope_head_dim # 576
|
||||
num_qo_heads = 128
|
||||
page = 64
|
||||
num_blocks = 64
|
||||
bs = 4
|
||||
n_layers = 3 # >1 so the per-layer view's block stride is inflated.
|
||||
layer = 1
|
||||
|
||||
q = torch.randn(bs, 1, num_qo_heads, head_dim, device=dev, dtype=dt)
|
||||
kv_data = torch.randn(num_blocks, 1, page, head_dim, device=dev, dtype=dt)
|
||||
|
||||
# (A) contiguous per-layer reference.
|
||||
kv_contiguous = kv_data.clone().contiguous()
|
||||
|
||||
# (B) unified slot: block b of every layer packed together; view one layer
|
||||
# -> stride(0) is n_layers x larger and storage offset is non-zero.
|
||||
unified = torch.randn(num_blocks, n_layers, 1, page, head_dim, device=dev, dtype=dt)
|
||||
unified[:, layer].copy_(kv_data)
|
||||
kv_view = unified[:, layer]
|
||||
assert not kv_view.is_contiguous()
|
||||
assert kv_view.stride(0) == n_layers * 1 * page * head_dim
|
||||
|
||||
max_blk = num_blocks // bs
|
||||
block_tables = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
|
||||
bs, max_blk
|
||||
)
|
||||
seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
|
||||
ws = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=dev)
|
||||
scale = head_dim**-0.5
|
||||
|
||||
def run(kv):
|
||||
return trtllm_batch_decode_with_kv_cache_mla(
|
||||
query=q,
|
||||
kv_cache=kv,
|
||||
workspace_buffer=ws,
|
||||
qk_nope_head_dim=qk_nope_head_dim,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
block_tables=block_tables,
|
||||
seq_lens=seq_lens,
|
||||
max_seq_len=int(seq_lens.max().item()),
|
||||
bmm1_scale=scale,
|
||||
bmm2_scale=1.0,
|
||||
).clone()
|
||||
|
||||
out_ref = run(kv_contiguous).float()
|
||||
out_view = run(kv_view).float()
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_flashmla_fp8_sparse_decode_unified_slot_view():
|
||||
"""FlashMLA fp8 sparse decode (DeepSeek V3.2/V4 DSA path) must read a
|
||||
unified-slot block-major view bit-identically to a contiguous fp8_ds_mla
|
||||
cache, with finite nonzero output."""
|
||||
import vllm.v1.attention.ops.flashmla as fm
|
||||
|
||||
ok, reason = fm.is_flashmla_sparse_supported()
|
||||
if not ok:
|
||||
pytest.skip(reason)
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
entry = 656 # fp8_ds_mla bytes per token
|
||||
page = 64
|
||||
num_blocks = 32
|
||||
h_q = 128
|
||||
head_dim = 576
|
||||
hdv = 512
|
||||
batch = 2
|
||||
topk = 128
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
|
||||
q = torch.randn(batch, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1
|
||||
|
||||
# Structurally valid fp8 ds_mla payload: 512B fp8 + 16B f32 scales + 128B
|
||||
# bf16 rope (random bytes corrupt the scale region and yield NaNs).
|
||||
nope = (torch.randn(num_blocks, page, 1, 512, device=dev) * 0.1).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
scales = torch.ones(num_blocks, page, 1, 4, device=dev, dtype=torch.float32)
|
||||
rope = (torch.randn(num_blocks, page, 1, 64, device=dev) * 0.1).to(torch.bfloat16)
|
||||
payload = torch.cat(
|
||||
[
|
||||
nope.view(torch.uint8).view(num_blocks, page, 1, 512),
|
||||
scales.view(torch.uint8).view(num_blocks, page, 1, 16),
|
||||
rope.view(torch.uint8).view(num_blocks, page, 1, 128),
|
||||
],
|
||||
dim=-1,
|
||||
).contiguous()
|
||||
assert payload.shape[-1] == entry and payload.dtype == torch.uint8
|
||||
|
||||
# (A) contiguous reference.
|
||||
cache_contiguous = payload.clone().contiguous()
|
||||
|
||||
# (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
|
||||
unified = torch.randint(
|
||||
0, 256, (num_blocks, n_layers, page, 1, entry), device=dev, dtype=torch.uint8
|
||||
)
|
||||
unified[:, layer].copy_(payload)
|
||||
cache_view = unified[:, layer]
|
||||
assert not cache_view.is_contiguous()
|
||||
assert cache_view.stride(0) == n_layers * page * 1 * entry
|
||||
|
||||
# Sparse indices: each batch uses its own disjoint blocks.
|
||||
blocks_per_batch = num_blocks // batch
|
||||
idx = torch.full((batch, 1, topk), -1, device=dev, dtype=torch.int32)
|
||||
for b in range(batch):
|
||||
slots: list[int] = []
|
||||
for blk in range(b * blocks_per_batch, (b + 1) * blocks_per_batch):
|
||||
slots.extend(blk * page + off for off in range(page))
|
||||
slots_t = torch.tensor(slots[:topk], device=dev, dtype=torch.int32)
|
||||
idx[b, 0, : slots_t.numel()] = slots_t
|
||||
|
||||
def run(kc):
|
||||
meta, num_splits = fm.get_mla_metadata()
|
||||
out, _ = fm.flash_mla_with_kvcache(
|
||||
q=q,
|
||||
k_cache=kc,
|
||||
block_table=None,
|
||||
cache_seqlens=None,
|
||||
head_dim_v=hdv,
|
||||
tile_scheduler_metadata=meta,
|
||||
is_fp8_kvcache=True,
|
||||
indices=idx,
|
||||
softmax_scale=head_dim**-0.5,
|
||||
)
|
||||
return out.clone().float()
|
||||
|
||||
out_ref = run(cache_contiguous)
|
||||
out_view = run(cache_view)
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert out_ref.abs().max().item() > 0.0
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_indexer_k_quant_and_cache_into_unified_slot_view():
|
||||
"""indexer_k_quant_and_cache (DeepSeek V3.2/V4 DSA indexer K write) must
|
||||
write correctly into a per-layer view whose block stride is the full
|
||||
unified slot, with zero bleed into the other layers' segments."""
|
||||
from vllm import _custom_ops as ops
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
head_dim = 128
|
||||
quant_block_size = 128
|
||||
block_size = 64
|
||||
num_blocks = 16
|
||||
ntok = 100
|
||||
# Indexer cache layout per token: head_dim fp8 bytes followed by
|
||||
# head_dim * 4 / quant_block_size scale bytes.
|
||||
cache_stride = head_dim + head_dim * 4 // quant_block_size
|
||||
|
||||
k = torch.randn(ntok, head_dim, device=dev, dtype=torch.bfloat16)
|
||||
slot = torch.randperm(num_blocks * block_size, device=dev, dtype=torch.int64)[:ntok]
|
||||
|
||||
def write(cache):
|
||||
ops.indexer_k_quant_and_cache(k, cache, slot, quant_block_size, "ue8m0")
|
||||
|
||||
# Contiguous per-layer reference.
|
||||
ref = torch.zeros(
|
||||
num_blocks, block_size, cache_stride, device=dev, dtype=torch.uint8
|
||||
)
|
||||
write(ref)
|
||||
|
||||
# Unified slot holding three layer pages per block; carve the middle one.
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
unified = torch.zeros(
|
||||
num_blocks, n_layers, block_size, cache_stride, device=dev, dtype=torch.uint8
|
||||
)
|
||||
view = unified[:, layer]
|
||||
assert not view.is_contiguous()
|
||||
assert view.stride(0) == n_layers * block_size * cache_stride
|
||||
write(view)
|
||||
|
||||
assert torch.equal(ref, view.contiguous())
|
||||
# Zero bleed into the neighbour layers' segments.
|
||||
assert unified[:, 0].abs().max().item() == 0
|
||||
assert unified[:, 2].abs().max().item() == 0
|
||||
|
||||
|
||||
def test_flashattn_mla_dense_decode_unified_slot_view():
|
||||
"""FA3 decode (FLASH_ATTN_MLA backend) must read a unified-slot
|
||||
block-major view bit-identically to a contiguous per-layer cache."""
|
||||
try:
|
||||
from vllm.vllm_flash_attn import flash_attn_varlen_func
|
||||
except ImportError:
|
||||
pytest.skip("vllm_flash_attn is not available")
|
||||
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
|
||||
|
||||
if not flash_attn_supports_mla():
|
||||
pytest.skip("FA3 MLA requires a Hopper device")
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
dt = torch.bfloat16
|
||||
kv_lora_rank = 512
|
||||
rope_dim = 64
|
||||
entry = kv_lora_rank + rope_dim # 576
|
||||
h_q = 16
|
||||
page = 64
|
||||
num_blocks = 64
|
||||
bs = 4
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
|
||||
q_pe = torch.randn(bs, h_q, rope_dim, device=dev, dtype=dt) * 0.1
|
||||
q_nope = torch.randn(bs, h_q, kv_lora_rank, device=dev, dtype=dt) * 0.1
|
||||
kv_data = torch.randn(num_blocks, page, entry, device=dev, dtype=dt) * 0.1
|
||||
|
||||
# (A) contiguous per-layer reference.
|
||||
cache_contiguous = kv_data.clone().contiguous()
|
||||
|
||||
# (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
|
||||
unified = torch.randn(num_blocks, n_layers, page, entry, device=dev, dtype=dt) * 0.1
|
||||
unified[:, layer].copy_(kv_data)
|
||||
cache_view = unified[:, layer]
|
||||
assert not cache_view.is_contiguous()
|
||||
assert cache_view.stride(0) == n_layers * page * entry
|
||||
|
||||
max_blk = num_blocks // bs
|
||||
block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
|
||||
bs, max_blk
|
||||
)
|
||||
seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
|
||||
cu_seqlens_q = torch.arange(bs + 1, device=dev, dtype=torch.int32)
|
||||
|
||||
def run(cache):
|
||||
kv_c_cache = cache[..., :kv_lora_rank]
|
||||
k_pe_cache = cache[..., kv_lora_rank:]
|
||||
out = flash_attn_varlen_func(
|
||||
q=q_pe,
|
||||
k=k_pe_cache.unsqueeze(-2), # Add head dim of 1
|
||||
v=kv_c_cache.unsqueeze(-2), # Add head dim of 1
|
||||
q_v=q_nope,
|
||||
max_seqlen_q=1,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
max_seqlen_k=int(seq_lens.max().item()),
|
||||
seqused_k=seq_lens,
|
||||
block_table=block_table,
|
||||
softmax_scale=entry**-0.5,
|
||||
causal=True,
|
||||
fa_version=3,
|
||||
)
|
||||
return out.clone().float()
|
||||
|
||||
out_ref = run(cache_contiguous)
|
||||
out_view = run(cache_view)
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert out_ref.abs().max().item() > 0.0
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_flashmla_dense_fp8_decode_unified_slot_view():
|
||||
"""FlashMLA dense fp8 decode (FLASHMLA backend with quantized KV cache)
|
||||
must read a unified-slot block-major view bit-identically to a contiguous
|
||||
per-layer fp8 cache."""
|
||||
import vllm.v1.attention.ops.flashmla as fm
|
||||
|
||||
ok, reason = fm.is_flashmla_dense_supported()
|
||||
if not ok:
|
||||
pytest.skip(reason)
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
head_dim = 576
|
||||
hdv = 512
|
||||
h_q = 128
|
||||
page = 64
|
||||
num_blocks = 64
|
||||
bs = 4
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
|
||||
q = torch.randn(bs, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1
|
||||
kv_data = (torch.randn(num_blocks, page, head_dim, device=dev) * 0.1).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
|
||||
# (A) contiguous per-layer reference.
|
||||
cache_contiguous = kv_data.clone().contiguous()
|
||||
|
||||
# (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
|
||||
unified = (torch.randn(num_blocks, n_layers, page, head_dim, device=dev) * 0.1).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
unified[:, layer].copy_(kv_data)
|
||||
cache_view = unified[:, layer]
|
||||
assert not cache_view.is_contiguous()
|
||||
assert cache_view.stride(0) == n_layers * page * head_dim
|
||||
|
||||
max_blk = num_blocks // bs
|
||||
block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
|
||||
bs, max_blk
|
||||
)
|
||||
cache_seqlens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
|
||||
descale = torch.ones(1, device=dev, dtype=torch.float32)
|
||||
|
||||
def run(kc):
|
||||
tile_md, num_splits = fm.get_mla_metadata_dense_fp8(cache_seqlens, h_q, 1)
|
||||
out, _ = fm.flash_mla_with_kvcache_fp8(
|
||||
q=q,
|
||||
k_cache=kc.unsqueeze(-2), # Add head dim of 1
|
||||
block_table=block_table,
|
||||
cache_seqlens=cache_seqlens,
|
||||
head_dim_v=hdv,
|
||||
tile_scheduler_metadata=tile_md,
|
||||
num_splits=num_splits,
|
||||
softmax_scale=head_dim**-0.5,
|
||||
causal=True,
|
||||
descale_q=descale,
|
||||
descale_k=descale,
|
||||
)
|
||||
return out.clone().float()
|
||||
|
||||
out_ref = run(cache_contiguous)
|
||||
out_view = run(cache_view)
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert out_ref.abs().max().item() > 0.0
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_flashinfer_mla_dense_fp8_decode_unified_slot_view():
|
||||
"""FlashInfer MLA dense decode with an fp8 KV cache must read a
|
||||
unified-slot block-major view bit-identically to a contiguous per-layer
|
||||
cache."""
|
||||
try:
|
||||
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
|
||||
except ImportError:
|
||||
pytest.skip("flashinfer is not available")
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.is_device_capability_family(100):
|
||||
pytest.skip("FlashInfer trtllm-gen MLA requires sm100")
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
kv_lora_rank = 512
|
||||
qk_rope_head_dim = 64
|
||||
qk_nope_head_dim = 128
|
||||
head_dim = kv_lora_rank + qk_rope_head_dim # 576
|
||||
num_qo_heads = 128
|
||||
page = 64
|
||||
num_blocks = 64
|
||||
bs = 4
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
|
||||
# With a quantized KV cache the decode query is quantized to fp8 as well
|
||||
# (trtllm-gen has no bf16-query x fp8-cache decode kernel).
|
||||
q = (torch.randn(bs, 1, num_qo_heads, head_dim, device=dev) * 0.1).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
kv_data = (torch.randn(num_blocks, 1, page, head_dim, device=dev) * 0.1).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
|
||||
# (A) contiguous per-layer reference.
|
||||
kv_contiguous = kv_data.clone().contiguous()
|
||||
|
||||
# (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
|
||||
unified = (
|
||||
torch.randn(num_blocks, n_layers, 1, page, head_dim, device=dev) * 0.1
|
||||
).to(torch.float8_e4m3fn)
|
||||
unified[:, layer].copy_(kv_data)
|
||||
kv_view = unified[:, layer]
|
||||
assert not kv_view.is_contiguous()
|
||||
assert kv_view.stride(0) == n_layers * 1 * page * head_dim
|
||||
|
||||
max_blk = num_blocks // bs
|
||||
block_tables = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
|
||||
bs, max_blk
|
||||
)
|
||||
seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
|
||||
ws = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=dev)
|
||||
scale = head_dim**-0.5
|
||||
|
||||
def run(kv):
|
||||
return trtllm_batch_decode_with_kv_cache_mla(
|
||||
query=q,
|
||||
kv_cache=kv,
|
||||
workspace_buffer=ws,
|
||||
qk_nope_head_dim=qk_nope_head_dim,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
block_tables=block_tables,
|
||||
seq_lens=seq_lens,
|
||||
max_seq_len=int(seq_lens.max().item()),
|
||||
bmm1_scale=scale,
|
||||
bmm2_scale=1.0,
|
||||
).clone()
|
||||
|
||||
out_ref = run(kv_contiguous).float()
|
||||
out_view = run(kv_view).float()
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
@@ -5,10 +5,12 @@ import math
|
||||
import random
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from contextlib import nullcontext
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.attention import SDPBackend, sdpa_kernel
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
|
||||
@@ -557,15 +559,21 @@ def test_contexted_kv_attention_alibi(
|
||||
query_len, seq_len, alibi_slopes, device, dtype
|
||||
)
|
||||
|
||||
# Compute attention
|
||||
out = F.scaled_dot_product_attention(
|
||||
q_sdpa,
|
||||
k_sdpa,
|
||||
v_sdpa,
|
||||
attn_mask=alibi_mask,
|
||||
dropout_p=0.0,
|
||||
scale=scale,
|
||||
)
|
||||
# Compute attention. On ROCm we force use of the Math SDPA backend rather than
|
||||
# the Flash or Mem-Efficient backends for increased numerical accuracy
|
||||
if current_platform.is_rocm():
|
||||
sdpa_context = sdpa_kernel(SDPBackend.MATH)
|
||||
else:
|
||||
sdpa_context = nullcontext()
|
||||
with sdpa_context:
|
||||
out = F.scaled_dot_product_attention(
|
||||
q_sdpa,
|
||||
k_sdpa,
|
||||
v_sdpa,
|
||||
attn_mask=alibi_mask,
|
||||
dropout_p=0.0,
|
||||
scale=scale,
|
||||
)
|
||||
|
||||
# Reshape output back to [query_len, num_heads, head_size]
|
||||
out = out.view(num_heads, query_len, head_size).permute(1, 0, 2)
|
||||
|
||||
@@ -90,7 +90,9 @@ def _ref_sparse_prefill_ragged(
|
||||
return out.to(torch.bfloat16)
|
||||
|
||||
|
||||
def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
|
||||
def _pack_fp8_ds_mla_cache(
|
||||
kv: torch.Tensor, block_size: int, is_extra: bool = False
|
||||
) -> torch.Tensor:
|
||||
assert kv.shape[-1] == HEAD_DIM
|
||||
num_tokens = kv.shape[0]
|
||||
num_blocks = (num_tokens + block_size - 1) // block_size
|
||||
@@ -101,7 +103,9 @@ def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
|
||||
)
|
||||
cache_flat = cache.view(torch.uint8).flatten()
|
||||
kv_nope_fp8 = (
|
||||
kv[:, :NOPE_HEAD_DIM].to(current_platform.fp8_dtype()).view(torch.uint8)
|
||||
kv[:, :NOPE_HEAD_DIM]
|
||||
.to(torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype())
|
||||
.view(torch.uint8)
|
||||
)
|
||||
kv_rope_u8 = kv[:, NOPE_HEAD_DIM:].contiguous().view(torch.uint8)
|
||||
|
||||
@@ -120,7 +124,7 @@ def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
|
||||
|
||||
|
||||
def _read_fp8_ds_mla_cache(
|
||||
cache: torch.Tensor, slot: int, block_size: int
|
||||
cache: torch.Tensor, slot: int, block_size: int, is_extra: bool = False
|
||||
) -> torch.Tensor:
|
||||
cache_flat = cache.view(torch.uint8).flatten()
|
||||
block_idx = slot // block_size
|
||||
@@ -129,7 +133,9 @@ def _read_fp8_ds_mla_cache(
|
||||
token_base = block_base + pos * 576
|
||||
|
||||
nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM]
|
||||
nope = nope_u8.view(current_platform.fp8_dtype()).to(torch.float32)
|
||||
nope = nope_u8.view(
|
||||
torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype()
|
||||
).to(torch.float32)
|
||||
rope_u8 = cache_flat[
|
||||
token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
|
||||
]
|
||||
@@ -157,7 +163,9 @@ def _ref_sparse_decode_ragged(
|
||||
]
|
||||
if extra_cache is not None and extra_rows is not None:
|
||||
row_kv.extend(
|
||||
_read_fp8_ds_mla_cache(extra_cache, int(slot), block_size)
|
||||
_read_fp8_ds_mla_cache(
|
||||
extra_cache, int(slot), block_size, is_extra=True
|
||||
)
|
||||
for slot in extra_rows[query_idx]
|
||||
)
|
||||
|
||||
@@ -326,7 +334,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None:
|
||||
main_kv = torch.randn(6, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
|
||||
extra_kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
|
||||
main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size)
|
||||
extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
|
||||
extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
|
||||
main_indices = torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device)
|
||||
main_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device=device)
|
||||
extra_indices = torch.tensor([1, 3, 0], dtype=torch.int32, device=device)
|
||||
@@ -477,7 +485,7 @@ def test_sparse_attn_decode_split_k_kernel(
|
||||
rows = [[1, 3, 0, 5, 2, 4], [3, 0, 6]]
|
||||
extra_kv = torch.randn(7, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
|
||||
extra_rows = rows
|
||||
extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
|
||||
extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
|
||||
extra_indices, extra_indptr = _ragged_from_rows(rows, device)
|
||||
|
||||
attn_sink = (
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user