forked from Karylab-cklius/vllm
Compare commits
69
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
71475ffb12 | ||
|
|
17bcc8ee3c | ||
|
|
607f3a50a3 | ||
|
|
9167ef8dd7 | ||
|
|
d7045619c1 | ||
|
|
206aee7dd9 | ||
|
|
674e6ffdb6 | ||
|
|
8b49cf3a37 | ||
|
|
2ae73c758c | ||
|
|
d95d03c719 | ||
|
|
803b9d7881 | ||
|
|
1312f07531 | ||
|
|
fa1b9840f6 | ||
|
|
916e56c05c | ||
|
|
a085b5257d | ||
|
|
7fd05e05ae | ||
|
|
99255f3cb5 | ||
|
|
75a7cf2c10 | ||
|
|
4b95e9cec4 | ||
|
|
856b15c62c | ||
|
|
6fb3f7b46b | ||
|
|
d109eacd05 | ||
|
|
e68fa1b90a | ||
|
|
f05f3664c3 | ||
|
|
e9f8f31e9a | ||
|
|
de3fe8dc62 | ||
|
|
0899f436aa | ||
|
|
358a755e43 | ||
|
|
a60883644b | ||
|
|
5aa371dc8e | ||
|
|
de3da0b97c | ||
|
|
9e92de51c6 | ||
|
|
bde0efdbb7 | ||
|
|
ea74f701db | ||
|
|
a8208e6a81 | ||
|
|
76c9cccc36 | ||
|
|
ed57f77192 | ||
|
|
7a1eb8ac2e | ||
|
|
c2e88a281c | ||
|
|
fd74c90d9c | ||
|
|
146f44b77d | ||
|
|
0d4f714208 | ||
|
|
03aeed802f | ||
|
|
2c8b76c5cb | ||
|
|
407b34be26 | ||
|
|
4c7c69b4e0 | ||
|
|
5e2c37facd | ||
|
|
c8bbe05189 | ||
|
|
6232fb4b66 | ||
|
|
2c06cf3486 | ||
|
|
e6f710a87f | ||
|
|
c245d35ff4 | ||
|
|
f8ac0c7cf0 | ||
|
|
ebf862c351 | ||
|
|
8d8062d0a7 | ||
|
|
985961345a | ||
|
|
706a04d34b | ||
|
|
22631f80a0 | ||
|
|
2cc008e7b4 | ||
|
|
5d5c776444 | ||
|
|
592ae6805c | ||
|
|
7b1bc0a3eb | ||
|
|
c0879d9483 | ||
|
|
f5f9878514 | ||
|
|
2ce95a761b | ||
|
+8 |
4d51588e23 | ||
|
|
32e45636e3 | ||
|
|
b39c266dae | ||
|
|
9558f43903 |
@@ -126,5 +126,4 @@ steps:
|
||||
'cd tests &&
|
||||
pytest -v -s lora/test_default_mm_loras.py &&
|
||||
(pytest -v -s lora/test_qwen3_unembed.py || true) &&
|
||||
(pytest -v -s lora/test_qwenvl.py || true) &&
|
||||
pytest -v -s lora/test_whisper.py'
|
||||
|
||||
+31
-31
@@ -388,18 +388,18 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# Multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- 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
|
||||
# Pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# Features demo
|
||||
- python3 offline_inference/prefix_caching.py
|
||||
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
|
||||
- python3 offline_inference/llm_engine_example.py
|
||||
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
|
||||
- python3 offline_inference/spec_decode.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
|
||||
- python3 offline_inference/spec_decode.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
|
||||
- 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
|
||||
- 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
|
||||
|
||||
#---------------------------------------------------------- mi250 · kernels ----------------------------------------------------------#
|
||||
|
||||
@@ -1168,13 +1168,13 @@ steps:
|
||||
- vllm/v1/attention/backends/
|
||||
- vllm/v1/attention/selector.py
|
||||
- tests/distributed/test_context_parallel.py
|
||||
- examples/offline_inference/data_parallel.py
|
||||
- examples/features/data_parallel/data_parallel_offline.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export TORCH_NCCL_BLOCKING_WAIT=1
|
||||
- pytest -v -s tests/distributed/test_context_parallel.py
|
||||
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization
|
||||
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization
|
||||
|
||||
- label: Distributed Tests (4xA100-4xMI300) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -1203,7 +1203,7 @@ steps:
|
||||
- tests/distributed/test_torchrun_example.py
|
||||
- tests/distributed/test_torchrun_example_moe.py
|
||||
- examples/rl/
|
||||
- tests/examples/offline_inference/data_parallel.py
|
||||
- tests/examples/features/data_parallel/data_parallel_offline.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export TORCH_NCCL_BLOCKING_WAIT=1
|
||||
@@ -1213,7 +1213,7 @@ steps:
|
||||
- PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
|
||||
- DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
|
||||
- TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
|
||||
- python3 ../examples/offline_inference/data_parallel.py --enforce-eager
|
||||
- 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
|
||||
@@ -1266,7 +1266,7 @@ steps:
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- examples/offline_inference/torchrun_dp_example.py
|
||||
- examples/features/torchrun/torchrun_dp_example_offline.py
|
||||
- vllm/config/parallel.py
|
||||
- vllm/distributed/
|
||||
- vllm/v1/engine/llm_engine.py
|
||||
@@ -1275,7 +1275,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export TORCH_NCCL_BLOCKING_WAIT=1
|
||||
- torchrun --nproc-per-node=8 ../examples/offline_inference/torchrun_dp_example.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-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
|
||||
|
||||
#-------------------------------------------------------- mi300 · entrypoints --------------------------------------------------------#
|
||||
|
||||
@@ -1647,18 +1647,18 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# Multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- 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
|
||||
# Pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# Features demo
|
||||
- python3 offline_inference/prefix_caching.py
|
||||
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
|
||||
- python3 offline_inference/llm_engine_example.py
|
||||
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
|
||||
- python3 offline_inference/spec_decode.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
|
||||
- python3 offline_inference/spec_decode.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
|
||||
- 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
|
||||
- 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
|
||||
|
||||
#---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------#
|
||||
|
||||
@@ -1951,8 +1951,8 @@ steps:
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
|
||||
#------------------------------------------------------- mi300 · quantization --------------------------------------------------------#
|
||||
|
||||
@@ -2302,7 +2302,7 @@ steps:
|
||||
commands:
|
||||
- export TORCH_NCCL_BLOCKING_WAIT=1
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
|
||||
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
|
||||
- pytest -v -s tests/distributed/test_packed_tensor.py
|
||||
@@ -2713,7 +2713,7 @@ steps:
|
||||
- vllm/v1/attention/selector.py
|
||||
- tests/distributed/test_context_parallel.py
|
||||
- tests/v1/distributed/test_dbo.py
|
||||
- examples/offline_inference/data_parallel.py
|
||||
- examples/features/data_parallel/data_parallel_offline.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
@@ -2930,18 +2930,18 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# Multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- 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
|
||||
# Pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# Features demo
|
||||
- python3 offline_inference/prefix_caching.py
|
||||
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
|
||||
- python3 offline_inference/llm_engine_example.py
|
||||
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
|
||||
- python3 offline_inference/spec_decode.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
|
||||
- python3 offline_inference/spec_decode.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
|
||||
- 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
|
||||
- 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
|
||||
|
||||
#---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------#
|
||||
|
||||
|
||||
@@ -88,9 +88,8 @@ steps:
|
||||
- vllm/distributed/
|
||||
- tests/distributed/test_torchrun_example.py
|
||||
- tests/distributed/test_torchrun_example_moe.py
|
||||
- examples/offline_inference/rlhf_colocate.py
|
||||
- examples/rl/
|
||||
- tests/examples/offline_inference/data_parallel.py
|
||||
- tests/examples/features/data_parallel/data_parallel_offline.py
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
@@ -107,7 +106,7 @@ steps:
|
||||
# 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
|
||||
# test with internal dp
|
||||
- python3 examples/offline_inference/data_parallel.py --enforce-eager
|
||||
- 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
|
||||
@@ -159,7 +158,7 @@ steps:
|
||||
num_devices: 8
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- examples/offline_inference/torchrun_dp_example.py
|
||||
- examples/features/torchrun/torchrun_dp_example_offline.py
|
||||
- vllm/config/parallel.py
|
||||
- vllm/distributed/
|
||||
- vllm/v1/engine/llm_engine.py
|
||||
@@ -169,7 +168,7 @@ steps:
|
||||
# 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/offline_inference/torchrun_dp_example.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-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
|
||||
|
||||
- label: Distributed Tests (4 GPUs)(A100)
|
||||
device: a100
|
||||
@@ -194,7 +193,7 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s tests/distributed/test_context_parallel.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
|
||||
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
|
||||
- pytest -v -s tests/distributed/test_packed_tensor.py
|
||||
@@ -222,9 +221,9 @@ steps:
|
||||
- vllm/executor/
|
||||
- vllm/model_executor/models/
|
||||
- tests/distributed/
|
||||
- tests/examples/offline_inference/data_parallel.py
|
||||
- tests/examples/features/data_parallel/data_parallel_offline.py
|
||||
commands:
|
||||
- ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"
|
||||
- ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"
|
||||
|
||||
- label: Pipeline + Context Parallelism (4 GPUs)
|
||||
timeout_in_minutes: 60
|
||||
|
||||
@@ -95,11 +95,13 @@ steps:
|
||||
- tests/kernels/moe/test_deepgemm.py
|
||||
- tests/kernels/moe/test_batched_deepgemm.py
|
||||
- tests/kernels/attention/test_deepgemm_attention.py
|
||||
- tests/quantization/test_cutlass_w4a16.py
|
||||
commands:
|
||||
- pytest -v -s kernels/quantization/test_block_fp8.py
|
||||
- pytest -v -s kernels/moe/test_deepgemm.py
|
||||
- pytest -v -s kernels/moe/test_batched_deepgemm.py
|
||||
- pytest -v -s kernels/attention/test_deepgemm_attention.py
|
||||
- pytest -v -s quantization/test_cutlass_w4a16.py
|
||||
|
||||
- label: Kernels (B200)
|
||||
timeout_in_minutes: 30
|
||||
|
||||
@@ -113,19 +113,19 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# for multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- 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
|
||||
# for pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# for features demo
|
||||
- python3 offline_inference/prefix_caching.py
|
||||
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
|
||||
- python3 offline_inference/llm_engine_example.py
|
||||
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
|
||||
- python3 offline_inference/spec_decode.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
|
||||
- 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 offline_inference/spec_decode.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
|
||||
- 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: Metrics, Tracing (2 GPUs)
|
||||
timeout_in_minutes: 20
|
||||
|
||||
@@ -31,8 +31,9 @@ steps:
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/v1/core/sched/
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
- examples/offline_inference/
|
||||
- examples/basic/offline_inference/
|
||||
- examples/generate/multimodal/
|
||||
- examples/features/
|
||||
- examples/pooling/embed/vision_embedding_offline.py
|
||||
- examples/others/tensorize_vllm_model.py
|
||||
commands:
|
||||
@@ -44,19 +45,19 @@ steps:
|
||||
#- 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
|
||||
# for multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- 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
|
||||
# for pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# for features demo
|
||||
- python3 offline_inference/prefix_caching.py
|
||||
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
|
||||
- python3 offline_inference/llm_engine_example.py
|
||||
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
|
||||
- python3 offline_inference/spec_decode.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
|
||||
- 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 offline_inference/spec_decode.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
|
||||
- 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)
|
||||
timeout_in_minutes: 45
|
||||
|
||||
@@ -69,9 +69,9 @@ steps:
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
# Whisper needs spawn method to avoid deadlock
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
|
||||
- label: Transformers Backward Compatibility Models Test
|
||||
working_dir: "/vllm-workspace/"
|
||||
@@ -83,7 +83,7 @@ steps:
|
||||
- pytest -v -s tests/models/test_transformers.py
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/offline_inference/basic/chat.py
|
||||
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
# Whisper needs spawn method to avoid deadlock
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
|
||||
+3
-8
@@ -308,8 +308,7 @@ pull_request_rules:
|
||||
- files=benchmarks/benchmark_serving_structured_output.py
|
||||
- files=benchmarks/run_structured_output_benchmark.sh
|
||||
- files=docs/features/structured_outputs.md
|
||||
- files=examples/offline_inference/structured_outputs.py
|
||||
- files=examples/online_serving/structured_outputs/structured_outputs.py
|
||||
- files=^examples/features/structured_outputs/
|
||||
- files~=^tests/v1/structured_output/
|
||||
- files=tests/entrypoints/llm/test_struct_output_generate.py
|
||||
- files~=^vllm/v1/structured_output/
|
||||
@@ -325,7 +324,7 @@ pull_request_rules:
|
||||
- or:
|
||||
- files~=^vllm/v1/spec_decode/
|
||||
- files~=^tests/v1/spec_decode/
|
||||
- files~=^examples/.*(spec_decode|mlpspeculator|eagle|speculation).*\.py
|
||||
- files=^examples/features/speculative_decoding/
|
||||
- files~=^vllm/model_executor/models/.*eagle.*\.py
|
||||
- files=vllm/model_executor/models/mlp_speculator.py
|
||||
- files~=^vllm/transformers_utils/configs/(eagle|medusa|mlp_speculator)\.py
|
||||
@@ -389,11 +388,7 @@ pull_request_rules:
|
||||
- files~=^tests/entrypoints/anthropic/.*tool.*
|
||||
- files~=^vllm/tool_parsers/
|
||||
- files=docs/features/tool_calling.md
|
||||
- files~=^examples/tool_chat_*
|
||||
- files=examples/offline_inference/chat_with_tools.py
|
||||
- files=examples/online_serving/openai_chat_completion_client_with_tools_required.py
|
||||
- files=examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py
|
||||
- files=examples/online_serving/openai_chat_completion_client_with_tools.py
|
||||
- files~=^examples/tool_calling/
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
|
||||
@@ -564,27 +564,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
"in CUDA target architectures.")
|
||||
endif()
|
||||
|
||||
# DeepSeek V4 indexer top-k. Needs thread-block clusters + TMA + PDL, so
|
||||
# builds for Hopper (sm_90a) and Blackwell datacenter (sm_100/sm_103). Not
|
||||
# supported on sm_120 (consumer Blackwell, no clusters). Requires CUDA >=
|
||||
# 12.4 for the cuda::ptx mbarrier wrappers. Ported from sglang.
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(DSV4_TOPK_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(DSV4_TOPK_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.4 AND DSV4_TOPK_ARCHS)
|
||||
set(DSV4_TOPK_SRC "csrc/deepseek_v4/fast_topk_v2.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${DSV4_TOPK_SRC}"
|
||||
CUDA_ARCHS "${DSV4_TOPK_ARCHS}")
|
||||
list(APPEND VLLM_EXT_SRC ${DSV4_TOPK_SRC})
|
||||
message(STATUS "Building deepseek_v4 fast_topk_v2 for archs: ${DSV4_TOPK_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building deepseek_v4 fast_topk_v2 (needs CUDA >= 12.4 "
|
||||
"and a compatible Hopper+ arch).")
|
||||
endif()
|
||||
|
||||
#
|
||||
# Machete kernels
|
||||
|
||||
|
||||
@@ -1,183 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Microbench: fast_topk_v2 vs persistent_topk for k in {512, 1024}.
|
||||
|
||||
Both ops select the top-k entries per row of a `[B, L]` float32 score
|
||||
tensor. vLLM's `persistent_topk` is the existing path used by the indexer;
|
||||
`fast_topk_v2` is the sm_90+ port from sglang that adds Hopper thread-block
|
||||
clusters.
|
||||
|
||||
V4-Flash uses `index_topk = 512`; V4-Pro uses `index_topk = 1024`. We bench
|
||||
both Ks at the realistic shape regimes (small-B, L up to 256K compressed).
|
||||
|
||||
Timing uses **CUDA graph replay** to amortize launch overhead (~3-5 µs on
|
||||
Blackwell). We capture N invocations of the same kernel, replay the graph
|
||||
many times, divide.
|
||||
|
||||
Run::
|
||||
|
||||
.venv/bin/python benchmarks/kernels/benchmark_fast_topk_v2.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import statistics
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
import vllm._C # noqa: F401 ensures schemas are registered
|
||||
from vllm.v1.attention.ops.deepseek_v4_ops.fast_topk import (
|
||||
fast_topk_v2_raw,
|
||||
plan_topk_v2,
|
||||
workspace_ints_per_batch,
|
||||
)
|
||||
|
||||
RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 # bytes; matches sparse_attn_indexer.py
|
||||
|
||||
|
||||
def _capture_graph(callable_fn, *, calls_per_graph: int) -> torch.cuda.CUDAGraph:
|
||||
for _ in range(3):
|
||||
callable_fn()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
g = torch.cuda.CUDAGraph()
|
||||
s = torch.cuda.Stream()
|
||||
s.wait_stream(torch.cuda.current_stream())
|
||||
with torch.cuda.stream(s):
|
||||
with torch.cuda.graph(g, stream=s):
|
||||
for _ in range(calls_per_graph):
|
||||
callable_fn()
|
||||
torch.cuda.current_stream().wait_stream(s)
|
||||
return g
|
||||
|
||||
|
||||
def time_graph_us(graph: torch.cuda.CUDAGraph, *, calls_per_graph: int,
|
||||
warmup: int = 5, replays: int = 30) -> float:
|
||||
for _ in range(warmup):
|
||||
graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
samples = []
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
for _ in range(replays):
|
||||
start.record()
|
||||
graph.replay()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
samples.append(start.elapsed_time(end) * 1000.0 / calls_per_graph)
|
||||
return statistics.median(samples)
|
||||
|
||||
|
||||
def make_inputs(batch_size: int, seq_len: int, *, seed: int = 0):
|
||||
device = torch.device("cuda")
|
||||
g = torch.Generator(device=device).manual_seed(seed)
|
||||
L = (seq_len + 3) & ~3
|
||||
scores = torch.randn(batch_size, L, generator=g, dtype=torch.float32,
|
||||
device=device)
|
||||
seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32,
|
||||
device=device)
|
||||
return scores, seq_lens, L
|
||||
|
||||
|
||||
def bench_persistent_topk(scores, seq_lens, k, *, calls_per_graph: int) -> float:
|
||||
B = scores.shape[0]
|
||||
output = scores.new_empty((B, k), dtype=torch.int32)
|
||||
workspace = scores.new_empty((RADIX_TOPK_WORKSPACE_SIZE,), dtype=torch.uint8)
|
||||
max_seq_len = scores.shape[1]
|
||||
|
||||
def run():
|
||||
torch.ops._C.persistent_topk(
|
||||
scores, seq_lens, output, workspace, k, max_seq_len)
|
||||
|
||||
graph = _capture_graph(run, calls_per_graph=calls_per_graph)
|
||||
return time_graph_us(graph, calls_per_graph=calls_per_graph)
|
||||
|
||||
|
||||
def bench_fast_topk_v2(scores, seq_lens, k, *,
|
||||
calls_per_graph: int) -> float:
|
||||
B = scores.shape[0]
|
||||
metadata = plan_topk_v2(seq_lens)
|
||||
workspace = scores.new_empty((B, workspace_ints_per_batch()),
|
||||
dtype=torch.int32)
|
||||
topk_indices = scores.new_empty((B, k), dtype=torch.int32)
|
||||
|
||||
def run():
|
||||
fast_topk_v2_raw(scores, seq_lens, topk=k,
|
||||
metadata=metadata, workspace=workspace,
|
||||
topk_indices=topk_indices)
|
||||
|
||||
graph = _capture_graph(run, calls_per_graph=calls_per_graph)
|
||||
return time_graph_us(graph, calls_per_graph=calls_per_graph)
|
||||
|
||||
|
||||
def fmt(us: float) -> str:
|
||||
return f"{us:8.2f}"
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--batch-sizes", type=int, nargs="+",
|
||||
default=[1, 4, 16, 32, 64, 128, 256])
|
||||
parser.add_argument("--seq-lens", type=int, nargs="+",
|
||||
default=[1024, 4096, 16384, 32768, 65536, 131072])
|
||||
parser.add_argument("--ks", type=int, nargs="+",
|
||||
default=[512, 1024])
|
||||
parser.add_argument("--calls-per-graph", type=int, default=64)
|
||||
parser.add_argument("--replays", type=int, default=30)
|
||||
args = parser.parse_args()
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA is required for this benchmark.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
print(f"GPU: {torch.cuda.get_device_name(0)} "
|
||||
f"(SM {torch.cuda.get_device_capability(0)})")
|
||||
print(f"calls_per_graph={args.calls_per_graph}, replays={args.replays}")
|
||||
print("Per-call medians via CUDA graph replay (host launch overhead "
|
||||
"amortized).\n")
|
||||
|
||||
for k in args.ks:
|
||||
print(f"=== k = {k} ===")
|
||||
print(f"{'B':>4} {'L':>7} | {'persistent_topk':>17} | "
|
||||
f"{'fast_topk_v2':>14} | {'speedup':>8} | {'path':<14}")
|
||||
print("-" * 80)
|
||||
for B in args.batch_sizes:
|
||||
for L in args.seq_lens:
|
||||
# Skip seq_lens beyond persistent_topk's k-dependent useful
|
||||
# range. Both kernels handle up to 256K with k=1024.
|
||||
try:
|
||||
scores, seq_lens, _ = make_inputs(B, L, seed=B * L * k)
|
||||
p_us = bench_persistent_topk(
|
||||
scores, seq_lens, k,
|
||||
calls_per_graph=args.calls_per_graph)
|
||||
f_us = bench_fast_topk_v2(
|
||||
scores, seq_lens, k,
|
||||
calls_per_graph=args.calls_per_graph)
|
||||
speedup = p_us / f_us if f_us > 0 else float("inf")
|
||||
|
||||
if L <= k:
|
||||
path = "trivial"
|
||||
elif L <= 4 * 4 * 1024:
|
||||
path = "register-1p"
|
||||
elif L <= 32768:
|
||||
path = "register-2p"
|
||||
elif B <= 15:
|
||||
path = "cluster-fused"
|
||||
else:
|
||||
path = "cluster-2stg"
|
||||
|
||||
print(
|
||||
f"{B:>4} {L:>7} | "
|
||||
f"{fmt(p_us):>14} us | "
|
||||
f"{fmt(f_us):>11} us | "
|
||||
f"{speedup:>5.2f}x | {path}"
|
||||
)
|
||||
except RuntimeError as e:
|
||||
print(f"{B:>4} {L:>7} | ERROR: {e}")
|
||||
print()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,324 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Benchmarks FP8 vs BF16 ViT attention via FlashInfer cuDNN backend.
|
||||
#
|
||||
# == Usage Examples ==
|
||||
#
|
||||
# Benchmark mode (default, FlashInfer CUDAGraph Bench)
|
||||
# python3 benchmark_vit_fp8_attn.py
|
||||
#
|
||||
# Profile mode (PyTorch profiler, saves TensorBoard traces):
|
||||
# python3 benchmark_vit_fp8_attn.py --profile
|
||||
# python3 benchmark_vit_fp8_attn.py --profile --profile-output-dir ./profile_traces
|
||||
#
|
||||
# Custom seq_lens:
|
||||
# python3 benchmark_vit_fp8_attn.py --seq-lens 4096 8192 16384
|
||||
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.profiler import ProfilerActivity, profile, record_function
|
||||
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
# Qwen3-VL defaults
|
||||
NUM_HEADS = 16
|
||||
HEAD_DIM = 72
|
||||
DEFAULT_SEQ_LENS = [2304, 4096, 8192, 16384]
|
||||
|
||||
|
||||
def _setup_fp8_attention(num_heads: int, head_dim: int) -> tuple:
|
||||
"""Create FP8 and BF16 attention modules + workspace."""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
_get_flashinfer_workspace_buffer,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
old_dtype = torch.get_default_dtype()
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
|
||||
backend_patch = patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
)
|
||||
|
||||
# FP8 attention
|
||||
mm_config_fp8 = MultiModalConfig(mm_encoder_attn_dtype="fp8")
|
||||
vllm_config_fp8 = VllmConfig()
|
||||
vllm_config_fp8.model_config = SimpleNamespace(multimodal_config=mm_config_fp8)
|
||||
with set_current_vllm_config(vllm_config_fp8), backend_patch:
|
||||
attn_fp8 = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
# BF16 attention (no FP8)
|
||||
with set_current_vllm_config(VllmConfig()), backend_patch:
|
||||
attn_bf16 = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
torch.set_default_dtype(old_dtype)
|
||||
|
||||
workspace = _get_flashinfer_workspace_buffer()
|
||||
return attn_fp8, attn_bf16, workspace
|
||||
|
||||
|
||||
def _build_meta(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
fp8: bool,
|
||||
):
|
||||
"""Build cu_seqlens, max_seqlen, sequence_lengths."""
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.utils.math_utils import round_up
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
cu_np = np.array([0, seq_len], dtype=np.int32)
|
||||
fp8_padded = num_heads * round_up(head_dim, 16) if fp8 else None
|
||||
|
||||
seq_lengths = MMEncoderAttention.maybe_compute_seq_lens(
|
||||
AttentionBackendEnum.FLASHINFER, cu_np, torch.device("cuda")
|
||||
)
|
||||
max_seqlen = torch.tensor(
|
||||
MMEncoderAttention.compute_max_seqlen(AttentionBackendEnum.FLASHINFER, cu_np),
|
||||
dtype=torch.int32,
|
||||
)
|
||||
cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
cu_np,
|
||||
num_heads * head_dim,
|
||||
1,
|
||||
torch.device("cuda"),
|
||||
fp8_padded_hidden_size=fp8_padded,
|
||||
)
|
||||
return cu_seqlens, max_seqlen, seq_lengths
|
||||
|
||||
|
||||
def run_benchmark(
|
||||
seq_lens: list[int],
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
method: str,
|
||||
):
|
||||
"""Benchmark FP8 vs BF16 attention across seq_lens.
|
||||
|
||||
Uses FlashInfer GPU-level timing to measure pure kernel time,
|
||||
excluding CPU launch overhead.
|
||||
"""
|
||||
if method == "cupti":
|
||||
from flashinfer.testing import bench_gpu_time_with_cupti as bench_fn
|
||||
|
||||
bench_fn = partial(bench_fn, use_cuda_graph=True, cold_l2_cache=False)
|
||||
elif method == "cudagraph":
|
||||
from flashinfer.testing import (
|
||||
bench_gpu_time_with_cudagraph as bench_fn,
|
||||
)
|
||||
|
||||
bench_fn = partial(bench_fn, cold_l2_cache=False)
|
||||
else:
|
||||
raise ValueError(f"Invalid method: {method}")
|
||||
|
||||
attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
|
||||
|
||||
print(f"Timing method: {method}")
|
||||
print(f"{'seq_len':>8} {'BF16 (us)':>12} {'FP8 (us)':>12} {'Speedup':>10}")
|
||||
print("-" * 46)
|
||||
|
||||
for seq_len in seq_lens:
|
||||
torch.manual_seed(42)
|
||||
|
||||
q = torch.randn(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
|
||||
# we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
|
||||
cu_bf16 = cu_fp8.clone()
|
||||
|
||||
def bf16_fn(q=q, k=k, v=v, cu=cu_bf16, ms=max_s, sl=seq_l):
|
||||
attn_bf16._forward_flashinfer(q, k, v, cu, ms, sl)
|
||||
|
||||
def fp8_fn(q=q, k=k, v=v, cu=cu_fp8, ms=max_s, sl=seq_l):
|
||||
attn_fp8._forward_flashinfer(q, k, v, cu, ms, sl)
|
||||
|
||||
# bench_fn returns List[float] of per-iteration times in ms
|
||||
bf16_times = bench_fn(bf16_fn)
|
||||
fp8_times = bench_fn(fp8_fn)
|
||||
|
||||
bf16_us = np.median(bf16_times) * 1e3 # ms -> us
|
||||
fp8_us = np.median(fp8_times) * 1e3
|
||||
speedup = bf16_us / fp8_us if fp8_us > 0 else float("inf")
|
||||
|
||||
print(f"{seq_len:>8} {bf16_us:>12.1f} {fp8_us:>12.1f} {speedup:>9.2f}x")
|
||||
|
||||
|
||||
def _make_trace_handler(output_dir: str, worker_name: str, label: str):
|
||||
"""Create a trace handler that saves to TensorBoard and prints summary."""
|
||||
|
||||
def handler(prof):
|
||||
torch.profiler.tensorboard_trace_handler(output_dir, worker_name)(prof)
|
||||
print(f"\n{'=' * 80}")
|
||||
print(label)
|
||||
print(f"{'=' * 80}")
|
||||
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))
|
||||
|
||||
return handler
|
||||
|
||||
|
||||
def run_profile(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
warmup: int,
|
||||
output_dir: str,
|
||||
):
|
||||
"""Profile FP8 vs BF16 attention with PyTorch profiler."""
|
||||
attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
|
||||
|
||||
torch.manual_seed(42)
|
||||
q = torch.randn(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
|
||||
# we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
|
||||
cu_bf16 = cu_fp8.clone()
|
||||
|
||||
sched = torch.profiler.schedule(wait=0, warmup=warmup, active=1)
|
||||
|
||||
# Profile BF16 (warmup handled by profiler schedule)
|
||||
with profile(
|
||||
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
|
||||
schedule=sched,
|
||||
on_trace_ready=_make_trace_handler(
|
||||
output_dir,
|
||||
f"bf16_h{head_dim}_s{seq_len}",
|
||||
f"BF16 Attention (seq_len={seq_len}, heads={num_heads}, "
|
||||
f"head_dim={head_dim})",
|
||||
),
|
||||
) as prof_bf16:
|
||||
for _ in range(warmup + 1):
|
||||
with record_function("bf16_attention"):
|
||||
attn_bf16._forward_flashinfer(
|
||||
q.clone(), k.clone(), v.clone(), cu_bf16, max_s, seq_l
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
prof_bf16.step()
|
||||
|
||||
# Profile FP8 (warmup handled by profiler schedule)
|
||||
with profile(
|
||||
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
|
||||
schedule=sched,
|
||||
on_trace_ready=_make_trace_handler(
|
||||
output_dir,
|
||||
f"fp8_h{head_dim}_s{seq_len}",
|
||||
f"FP8 Attention (seq_len={seq_len}, heads={num_heads}, "
|
||||
f"head_dim={head_dim})",
|
||||
),
|
||||
) as prof_fp8:
|
||||
for _ in range(warmup + 1):
|
||||
with record_function("fp8_attention"):
|
||||
attn_fp8._forward_flashinfer(
|
||||
q.clone(), k.clone(), v.clone(), cu_fp8, max_s, seq_l
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
prof_fp8.step()
|
||||
|
||||
print(f"\nTensorBoard traces saved to: {output_dir}")
|
||||
print(f"View with: tensorboard --logdir={output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = FlexibleArgumentParser(description="Benchmark FP8 vs BF16 ViT attention.")
|
||||
parser.add_argument(
|
||||
"--seq-lens",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=DEFAULT_SEQ_LENS,
|
||||
help="Sequence lengths to benchmark",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-heads",
|
||||
type=int,
|
||||
default=NUM_HEADS,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--head-dim",
|
||||
type=int,
|
||||
default=HEAD_DIM,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--method",
|
||||
choices=["cupti", "cudagraph"],
|
||||
default="cudagraph",
|
||||
help="GPU timing method: cupti (CUPTI kernel timing) or "
|
||||
"cudagraph (CUDA graph capture/replay). Default: cudagraph",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Warmup iterations (profile mode only)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile",
|
||||
action="store_true",
|
||||
help="Run PyTorch profiler instead of benchmark",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-seq-len",
|
||||
type=int,
|
||||
default=8192,
|
||||
help="Sequence length for profiling (default: 8192)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-output-dir",
|
||||
type=str,
|
||||
default="./profile_traces",
|
||||
help="Output directory for TensorBoard traces (default: ./profile_traces)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.profile:
|
||||
run_profile(
|
||||
args.profile_seq_len,
|
||||
args.num_heads,
|
||||
args.head_dim,
|
||||
args.warmup,
|
||||
args.profile_output_dir,
|
||||
)
|
||||
else:
|
||||
run_benchmark(
|
||||
args.seq_lens,
|
||||
args.num_heads,
|
||||
args.head_dim,
|
||||
args.method,
|
||||
)
|
||||
+82
-25
@@ -11,29 +11,74 @@
|
||||
namespace vllm {
|
||||
|
||||
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
|
||||
bool act_first>
|
||||
bool act_first, bool HAS_CLAMP>
|
||||
__device__ __forceinline__ scalar_t compute(const scalar_t& x,
|
||||
const scalar_t& y) {
|
||||
return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
|
||||
const scalar_t& y,
|
||||
const float limit) {
|
||||
if constexpr (act_first) {
|
||||
scalar_t gate = x;
|
||||
scalar_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
gate = (scalar_t)fminf((float)gate, limit);
|
||||
up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
|
||||
}
|
||||
return ACT_FN(gate) * up;
|
||||
} else {
|
||||
scalar_t gate = x;
|
||||
scalar_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
|
||||
up = (scalar_t)fminf((float)up, limit);
|
||||
}
|
||||
return gate * ACT_FN(up);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
|
||||
bool act_first>
|
||||
bool act_first, bool HAS_CLAMP>
|
||||
__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
|
||||
const packed_t& y) {
|
||||
return act_first ? packed_mul(PACKED_ACT_FN(x), y)
|
||||
: packed_mul(x, PACKED_ACT_FN(y));
|
||||
const packed_t& y,
|
||||
const float limit) {
|
||||
if constexpr (act_first) {
|
||||
packed_t gate = x;
|
||||
packed_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
float2 g = cast_to_float2(gate);
|
||||
float2 u = cast_to_float2(up);
|
||||
g.x = fminf(g.x, limit);
|
||||
g.y = fminf(g.y, limit);
|
||||
u.x = fmaxf(fminf(u.x, limit), -limit);
|
||||
u.y = fmaxf(fminf(u.y, limit), -limit);
|
||||
gate = cast_to_packed<packed_t>(g);
|
||||
up = cast_to_packed<packed_t>(u);
|
||||
}
|
||||
return packed_mul(PACKED_ACT_FN(gate), up);
|
||||
} else {
|
||||
packed_t gate = x;
|
||||
packed_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
float2 g = cast_to_float2(gate);
|
||||
float2 u = cast_to_float2(up);
|
||||
g.x = fmaxf(fminf(g.x, limit), -limit);
|
||||
g.y = fmaxf(fminf(g.y, limit), -limit);
|
||||
u.x = fminf(u.x, limit);
|
||||
u.y = fminf(u.y, limit);
|
||||
gate = cast_to_packed<packed_t>(g);
|
||||
up = cast_to_packed<packed_t>(u);
|
||||
}
|
||||
return packed_mul(gate, PACKED_ACT_FN(up));
|
||||
}
|
||||
}
|
||||
|
||||
// Activation and gating kernel template.
|
||||
template <typename scalar_t, typename packed_t,
|
||||
scalar_t (*ACT_FN)(const scalar_t&),
|
||||
packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
|
||||
bool use_vec, bool use_256b = false>
|
||||
bool use_vec, bool HAS_CLAMP, bool use_256b = false>
|
||||
__global__ void act_and_mul_kernel(
|
||||
scalar_t* __restrict__ out, // [..., d]
|
||||
const scalar_t* __restrict__ input, // [..., 2, d]
|
||||
const int d) {
|
||||
const int d, const float limit) {
|
||||
const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
|
||||
const scalar_t* y_ptr = x_ptr + d;
|
||||
scalar_t* out_ptr = out + blockIdx.x * d;
|
||||
@@ -58,8 +103,9 @@ __global__ void act_and_mul_kernel(
|
||||
}
|
||||
#pragma unroll
|
||||
for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
|
||||
x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
|
||||
x.elts[j], y.elts[j]);
|
||||
x.elts[j] =
|
||||
packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
|
||||
x.elts[j], y.elts[j], limit);
|
||||
}
|
||||
if constexpr (use_256b) {
|
||||
st256(x, &out_vec[i]);
|
||||
@@ -72,7 +118,8 @@ __global__ void act_and_mul_kernel(
|
||||
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
|
||||
const scalar_t x = VLLM_LDG(&x_ptr[idx]);
|
||||
const scalar_t y = VLLM_LDG(&y_ptr[idx]);
|
||||
out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
|
||||
out_ptr[idx] =
|
||||
compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -151,8 +198,11 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
|
||||
// Launch activation and gating kernel.
|
||||
// Use ACT_FIRST (bool) indicating whether to apply the activation function
|
||||
// first.
|
||||
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST) \
|
||||
// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
|
||||
// clamped (max only) and up input is clamped (both sides) before the
|
||||
// activation function is applied.
|
||||
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST, \
|
||||
HAS_CLAMP, LIMIT) \
|
||||
auto dtype = input.scalar_type(); \
|
||||
int d = input.size(-1) / 2; \
|
||||
int64_t num_tokens = input.numel() / input.size(-1); \
|
||||
@@ -177,8 +227,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
|
||||
KERNEL<scalar_t>, \
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, true, true><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
|
||||
ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
|
||||
}); \
|
||||
} else { \
|
||||
VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
|
||||
@@ -186,8 +236,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
|
||||
KERNEL<scalar_t>, \
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, true, false><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
|
||||
ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
|
||||
}); \
|
||||
} \
|
||||
} else { \
|
||||
@@ -197,8 +247,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
|
||||
KERNEL<scalar_t>, \
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, false><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
|
||||
ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
|
||||
}); \
|
||||
}
|
||||
|
||||
@@ -206,7 +256,14 @@ void silu_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
true);
|
||||
true, false, 0.0f);
|
||||
}
|
||||
|
||||
void silu_and_mul_clamp(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input, // [..., 2 * d]
|
||||
double limit) {
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
true, true, (float)limit);
|
||||
}
|
||||
|
||||
void mul_and_silu(torch::Tensor& out, // [..., d]
|
||||
@@ -215,21 +272,21 @@ void mul_and_silu(torch::Tensor& out, // [..., d]
|
||||
// The difference between mul_and_silu and silu_and_mul is that mul_and_silu
|
||||
// applies the silu to the latter half of the input.
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
false);
|
||||
false, false, 0.0f);
|
||||
}
|
||||
|
||||
void gelu_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
|
||||
true);
|
||||
true, false, 0.0f);
|
||||
}
|
||||
|
||||
void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
|
||||
vllm::packed_gelu_tanh_kernel, true);
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(
|
||||
vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
|
||||
}
|
||||
|
||||
namespace vllm {
|
||||
|
||||
@@ -1,693 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// DeepSeek V4 indexer top-k (k = 512 for Flash, k = 1024 for Pro). Ported
|
||||
// from sglang's jit_kernel/csrc/deepseek_v4/topk_v2.cuh.
|
||||
//
|
||||
// Combines three strategies (Register / Streaming / Cluster) dispatched per
|
||||
// row by a separate plan kernel that decides a `cluster_threshold` from the
|
||||
// observed seq_lens distribution. The host side picks one of three launch
|
||||
// shapes:
|
||||
// 1. all rows fit in the small (register) path -> single short kernel
|
||||
// 2. small batch (<= kNumClusters) with some long rows -> fused cluster
|
||||
// kernel (stage 1 + tie-break in one launch)
|
||||
// 3. larger batch -> persistent cluster stage 1 + non-cluster stage 2
|
||||
//
|
||||
// Architecture support: Hopper (sm_90a) and Blackwell datacenter (sm_100/
|
||||
// sm_103). Requires thread-block clusters, TMA bulk async copy, mbarrier,
|
||||
// and Programmatic Dependent Launch — sm_120 (consumer Blackwell) lacks
|
||||
// clusters and is not supported. The heuristic constants in `topk_plan`
|
||||
// were tuned on B200 (sglang upstream); they are functionally correct on
|
||||
// H100/H200 too but may be suboptimal until retuned.
|
||||
|
||||
#include "topk/cluster.cuh"
|
||||
#include "topk/common.cuh"
|
||||
#include "topk/register.cuh"
|
||||
#include "topk/streaming.cuh"
|
||||
#include "topk/utils.cuh"
|
||||
|
||||
#include "core/registration.h"
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/util/Exception.h>
|
||||
#include <cooperative_groups.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <torch/all.h>
|
||||
#include <torch/library.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
|
||||
namespace vllm::dsv4_topk {
|
||||
|
||||
// All K-dependent type and constant lookups go through these aliases / vars
|
||||
// so the kernels can be templated on K. Kernel and Smem sizes happen to be
|
||||
// K-independent (e.g., kMaxTies, kMax2PassLength, kHistBins are all set in
|
||||
// terms of kBlockSize/kHistBits, not K), so we don't pay extra smem for the
|
||||
// 1024 instantiation.
|
||||
template <uint32_t K> using Large = ClusterTopK<K>;
|
||||
template <uint32_t K> using Medium = StreamingTopK<K>;
|
||||
template <uint32_t K> using Small = RegisterTopK<K>;
|
||||
|
||||
// Metadata struct layout is K-independent — pick any K to grab the type.
|
||||
using Metadata = Large<512>::Metadata;
|
||||
constexpr uint32_t kNumClusters = 15; // hardware-capped persistent count
|
||||
constexpr uint32_t kClusterSize = Large<512>::kClusterSize;
|
||||
constexpr uint32_t kMax2PassLength = Small<512>::kMax2PassLength;
|
||||
constexpr uint32_t kMaxSupportedLength = Large<512>::kMaxLength;
|
||||
|
||||
// Row 0 of the metadata tensor stores GlobalMetadata; rows [1..N+1) hold the
|
||||
// per-item Metadata entries that the persistent stage-1 consumes.
|
||||
struct alignas(16) GlobalMetadata {
|
||||
uint32_t cluster_threshold;
|
||||
uint32_t num_cluster_items;
|
||||
uint32_t reserved[2];
|
||||
};
|
||||
static_assert(sizeof(GlobalMetadata) == sizeof(Metadata),
|
||||
"metadata row 0 layout must match Metadata stride");
|
||||
|
||||
#define VLLM_SMALL_TOPK_KERNEL __global__ __launch_bounds__(kBlockSize, 2)
|
||||
#define VLLM_LARGE_CLUSTER __cluster_dims__(1, kClusterSize, 1)
|
||||
// Stage 1 is persistent + cluster -> high smem -> occupancy 1.
|
||||
#define VLLM_LARGE_TOPK_STAGE_1 \
|
||||
__global__ __launch_bounds__(kBlockSize, 1) VLLM_LARGE_CLUSTER
|
||||
// Stage 2 is non-cluster + small smem -> occupancy 2.
|
||||
#define VLLM_LARGE_TOPK_STAGE_2 __global__ __launch_bounds__(kBlockSize, 2)
|
||||
#define VLLM_FUSED_COMBINE_KERNEL \
|
||||
__global__ __launch_bounds__(kBlockSize, 1) VLLM_LARGE_CLUSTER
|
||||
#define VLLM_PLAN_KERNEL __global__ __launch_bounds__(kBlockSize, 1)
|
||||
|
||||
struct TopKParams {
|
||||
const uint32_t* __restrict__ seq_lens;
|
||||
const float* __restrict__ scores;
|
||||
const int32_t* __restrict__ page_table;
|
||||
int32_t* __restrict__ page_indices;
|
||||
int64_t score_stride;
|
||||
int64_t page_table_stride;
|
||||
uint8_t* __restrict__ workspace;
|
||||
const Metadata* __restrict__ metadata = nullptr;
|
||||
int64_t workspace_stride; // bytes per batch
|
||||
uint32_t batch_size;
|
||||
uint32_t page_bits;
|
||||
|
||||
VLLM_DSV4_DEVICE const float* get_scores(uint32_t batch_id) const {
|
||||
return scores + batch_id * score_stride;
|
||||
}
|
||||
template <uint32_t K, bool kRawOutput>
|
||||
VLLM_DSV4_DEVICE TransformParamsT<kRawOutput> get_transform(
|
||||
uint32_t batch_id, int32_t* indices) const {
|
||||
return {
|
||||
.page_table = page_table + batch_id * page_table_stride,
|
||||
.indices_in = indices,
|
||||
.indices_out = page_indices + batch_id * K,
|
||||
.page_bits = page_bits,
|
||||
};
|
||||
}
|
||||
VLLM_DSV4_DEVICE const GlobalMetadata& get_global_metadata() const {
|
||||
return *reinterpret_cast<const GlobalMetadata*>(metadata);
|
||||
}
|
||||
VLLM_DSV4_DEVICE const Metadata& get_item_metadata(uint32_t work_id) const {
|
||||
return metadata[1 + work_id]; // skip the GlobalMetadata row
|
||||
}
|
||||
};
|
||||
|
||||
VLLM_DSV4_DEVICE uint2 partition_work(uint32_t length, uint32_t rank) {
|
||||
constexpr uint32_t kTMAAlign = 4;
|
||||
const auto total_units = (length + kTMAAlign - 1) / kTMAAlign;
|
||||
const auto base = total_units / kClusterSize;
|
||||
const auto extra = total_units % kClusterSize;
|
||||
const auto local_units = base + (rank < extra ? 1u : 0u);
|
||||
const auto offset_units = rank * base + min(rank, extra);
|
||||
const auto offset = offset_units * kTMAAlign;
|
||||
const auto finish = min(offset + local_units * kTMAAlign, length);
|
||||
return {offset, finish - offset};
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Plan kernel: decides cluster_threshold from the observed seq_lens
|
||||
// distribution and compacts items with seq_len > threshold into metadata[1..].
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
VLLM_PLAN_KERNEL void topk_plan(const uint32_t* __restrict__ seq_lens,
|
||||
Metadata* __restrict__ metadata,
|
||||
uint32_t batch_size,
|
||||
uint32_t static_cluster_threshold) {
|
||||
// (threshold, max_batch_size_for_that_threshold). Tuned on B200 by sglang.
|
||||
struct Pair {
|
||||
uint32_t threshold;
|
||||
uint32_t max_batch_size;
|
||||
};
|
||||
constexpr Pair kCandidates[] = {
|
||||
{32768, 30}, {40960, 45}, {49152, 45}, {65536, 60},
|
||||
{98304, 60}, {131072, 75}, {196608, 90}, {262144, 105},
|
||||
};
|
||||
constexpr uint32_t kNumCandidates =
|
||||
sizeof(kCandidates) / sizeof(kCandidates[0]);
|
||||
constexpr uint32_t kMinBatchSize = kCandidates[0].max_batch_size;
|
||||
static_assert(kCandidates[0].threshold == kMax2PassLength);
|
||||
static_assert(kCandidates[kNumCandidates - 1].threshold ==
|
||||
kMaxSupportedLength);
|
||||
|
||||
__shared__ uint32_t s_count;
|
||||
__shared__ uint32_t s_counts[kNumCandidates];
|
||||
__shared__ uint32_t s_threshold;
|
||||
|
||||
const auto tx = threadIdx.x;
|
||||
if (tx == 0) s_count = 0;
|
||||
if (tx < kNumCandidates) s_counts[tx] = 0;
|
||||
__syncthreads();
|
||||
|
||||
if (static_cluster_threshold > 0) {
|
||||
if (tx == 0) s_threshold = static_cluster_threshold;
|
||||
} else if (batch_size <= kMinBatchSize) {
|
||||
if (tx == 0) s_threshold = kMax2PassLength;
|
||||
} else {
|
||||
for (uint32_t i = tx; i < batch_size; i += kBlockSize) {
|
||||
const uint32_t sl = seq_lens[i];
|
||||
assert(sl <= kMaxSupportedLength);
|
||||
uint32_t count = 0;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < kNumCandidates; ++j) {
|
||||
count += (sl > kCandidates[j].threshold ? 1 : 0);
|
||||
}
|
||||
if (count > 0) {
|
||||
atomicAdd(&s_counts[count - 1], 1);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
if (tx == 0) {
|
||||
uint32_t accum = 0;
|
||||
uint32_t chosen = kMaxSupportedLength;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumCandidates; ++i) {
|
||||
const auto j = kNumCandidates - 1 - i;
|
||||
accum += s_counts[j];
|
||||
if (accum > kCandidates[j].max_batch_size) break;
|
||||
chosen = kCandidates[j].threshold;
|
||||
}
|
||||
s_threshold = chosen;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
const auto cluster_threshold = max(s_threshold, kMax2PassLength);
|
||||
|
||||
// Compact items with seq_len > cluster_threshold into metadata[1..N+1).
|
||||
for (uint32_t i = tx; i < batch_size; i += kBlockSize) {
|
||||
const uint32_t sl = seq_lens[i];
|
||||
if (sl > cluster_threshold) {
|
||||
const auto pos = atomicAdd(&s_count, 1);
|
||||
metadata[1 + pos] = {i, sl, false};
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
const auto N = s_count;
|
||||
|
||||
// has_next chain for the persistent consumer + sentinel slots.
|
||||
for (uint32_t i = tx; i < N; i += kBlockSize) {
|
||||
if (i + kNumClusters < N) metadata[1 + i].has_next = true;
|
||||
}
|
||||
if (tx < kNumClusters && tx >= N) metadata[1 + tx] = {0, 0, false};
|
||||
if (tx == 0) {
|
||||
auto* g = reinterpret_cast<GlobalMetadata*>(metadata);
|
||||
*g = {
|
||||
.cluster_threshold = cluster_threshold,
|
||||
.num_cluster_items = N,
|
||||
.reserved = {0, 0},
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Short kernel: all rows fit in the register path (max_seq_len <=
|
||||
// Small::kMax1PassLength).
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
template <uint32_t K, bool kRawOutput>
|
||||
VLLM_SMALL_TOPK_KERNEL void topk_short_transform(
|
||||
const __grid_constant__ TopKParams params) {
|
||||
alignas(128) extern __shared__ uint8_t smem[];
|
||||
__shared__ int32_t s_topk_indices[K];
|
||||
const auto batch_id = blockIdx.x;
|
||||
const auto seq_len = params.seq_lens[batch_id];
|
||||
const auto transform =
|
||||
params.template get_transform<K, kRawOutput>(batch_id, s_topk_indices);
|
||||
if (seq_len <= K) {
|
||||
trivial_transform(transform, seq_len, K);
|
||||
} else {
|
||||
Small<K>::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem,
|
||||
/*use_pdl=*/true);
|
||||
pdl_trigger_secondary<true>();
|
||||
Small<K>::transform(transform);
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Persistent stage 1 (cluster). One CTA per cluster; the persistent block
|
||||
// walks `metadata[1..N]` round-robin and runs Large::stage1 per item.
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
template <uint32_t K, bool kRawOutput>
|
||||
VLLM_LARGE_TOPK_STAGE_1 void topk_combine_preprocess(
|
||||
const __grid_constant__ TopKParams params) {
|
||||
alignas(128) extern __shared__ uint8_t smem[];
|
||||
__shared__ int32_t s_topk_indices[K];
|
||||
uint32_t work_id = blockIdx.x;
|
||||
uint32_t batch_id = 0, seq_len = 0, length = 0, offset = 0;
|
||||
bool has_next = false;
|
||||
const auto cluster_rank = blockIdx.y;
|
||||
|
||||
const auto prefetch_metadata = [&] {
|
||||
const auto m = params.get_item_metadata(work_id);
|
||||
batch_id = m.batch_id;
|
||||
seq_len = m.seq_len;
|
||||
has_next = m.has_next;
|
||||
work_id += kNumClusters;
|
||||
};
|
||||
const auto launch_prologue = [&] {
|
||||
const auto partition = partition_work(seq_len, cluster_rank);
|
||||
offset = partition.x;
|
||||
length = partition.y;
|
||||
Large<K>::stage1_prologue(params.get_scores(batch_id) + offset, length,
|
||||
smem);
|
||||
};
|
||||
|
||||
pdl_wait_primary<true>();
|
||||
pdl_trigger_secondary<true>();
|
||||
|
||||
prefetch_metadata();
|
||||
if (seq_len == 0) return;
|
||||
Large<K>::stage1_init(smem);
|
||||
launch_prologue();
|
||||
while (true) {
|
||||
const auto this_length = length;
|
||||
const auto this_offset = offset;
|
||||
const auto need_prefetch = has_next;
|
||||
const auto transform =
|
||||
params.template get_transform<K, kRawOutput>(batch_id, s_topk_indices);
|
||||
const auto ws = params.workspace + batch_id * params.workspace_stride;
|
||||
if (need_prefetch) prefetch_metadata();
|
||||
Large<K>::stage1(s_topk_indices, this_length, smem, /*reuse=*/true);
|
||||
if (need_prefetch) launch_prologue();
|
||||
Large<K>::stage1_epilogue(transform, this_offset, ws, smem);
|
||||
if (!need_prefetch) break;
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Stage 2 (non-cluster). Per-row dispatch: trivial / Small / Medium / Large.
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
template <uint32_t K, bool kRawOutput>
|
||||
VLLM_LARGE_TOPK_STAGE_2 void topk_combine_transform(
|
||||
const __grid_constant__ TopKParams params) {
|
||||
alignas(128) extern __shared__ uint8_t smem[];
|
||||
__shared__ int32_t s_topk_indices[K];
|
||||
const auto batch_id = blockIdx.x;
|
||||
const auto seq_len = params.seq_lens[batch_id];
|
||||
const auto cluster_threshold = params.get_global_metadata().cluster_threshold;
|
||||
const auto transform =
|
||||
params.template get_transform<K, kRawOutput>(batch_id, s_topk_indices);
|
||||
if (seq_len <= K) {
|
||||
trivial_transform(transform, seq_len, K);
|
||||
} else if (seq_len <= kMax2PassLength) {
|
||||
if (seq_len <= Small<K>::kMax1PassLength) {
|
||||
Small<K>::run(params.get_scores(batch_id), s_topk_indices, seq_len,
|
||||
smem);
|
||||
} else {
|
||||
__syncwarp();
|
||||
Small<K>::template run<true>(params.get_scores(batch_id),
|
||||
s_topk_indices, seq_len, smem);
|
||||
}
|
||||
Small<K>::transform(transform);
|
||||
} else if (seq_len <= cluster_threshold) {
|
||||
Medium<K>::run(params.get_scores(batch_id), seq_len, s_topk_indices, smem);
|
||||
Medium<K>::transform(transform, smem);
|
||||
} else {
|
||||
const auto ws = params.workspace + batch_id * params.workspace_stride;
|
||||
pdl_wait_primary<true>();
|
||||
Large<K>::transform(transform, ws, smem);
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Fused kernel for small batches. Both stage 1 and the tie-break run inside
|
||||
// the same launch; cluster rank 0 finishes the row.
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
template <uint32_t K, bool kRawOutput>
|
||||
VLLM_FUSED_COMBINE_KERNEL void topk_fused_transform(
|
||||
const __grid_constant__ TopKParams params) {
|
||||
alignas(128) extern __shared__ uint8_t smem[];
|
||||
__shared__ int32_t s_topk_indices[K];
|
||||
const auto batch_id = blockIdx.x;
|
||||
const auto cluster_rank = blockIdx.y;
|
||||
const auto seq_len = params.seq_lens[batch_id];
|
||||
const auto transform =
|
||||
params.template get_transform<K, kRawOutput>(batch_id, s_topk_indices);
|
||||
if (seq_len <= K) {
|
||||
if (cluster_rank != 0) return;
|
||||
trivial_transform(transform, seq_len, K);
|
||||
} else if (seq_len <= Small<K>::kMax1PassLength) {
|
||||
if (cluster_rank != 0) return;
|
||||
Small<K>::run(params.get_scores(batch_id), s_topk_indices, seq_len, smem,
|
||||
/*use_pdl=*/true);
|
||||
Small<K>::transform(transform);
|
||||
} else {
|
||||
const auto partition = partition_work(seq_len, cluster_rank);
|
||||
const auto offset = partition.x;
|
||||
const auto length = partition.y;
|
||||
const auto ws = params.workspace + batch_id * params.workspace_stride;
|
||||
Large<K>::stage1_init(smem);
|
||||
pdl_wait_primary<true>();
|
||||
Large<K>::stage1_prologue(params.get_scores(batch_id) + offset, length,
|
||||
smem);
|
||||
Large<K>::stage1(s_topk_indices, length, smem);
|
||||
Large<K>::stage1_epilogue(transform, offset, ws, smem);
|
||||
cooperative_groups::this_cluster().sync();
|
||||
if (cluster_rank != 0) return;
|
||||
Large<K>::transform(transform, ws, smem);
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t K> constexpr size_t kStage1SMEM = sizeof(typename Large<K>::Smem) + 128;
|
||||
template <uint32_t K> constexpr size_t kStage2SMEM =
|
||||
(sizeof(typename Small<K>::Smem) > sizeof(typename Medium<K>::Smem)
|
||||
? sizeof(typename Small<K>::Smem)
|
||||
: sizeof(typename Medium<K>::Smem)) +
|
||||
128;
|
||||
|
||||
// Per-(kernel, smem) memoization: each instantiation has its own static. This
|
||||
// matters because cudaFuncSetAttribute is per-function and we want it to fire
|
||||
// exactly once per kernel symbol.
|
||||
template <auto* f, size_t kSmem>
|
||||
void setup_kernel_smem_once() {
|
||||
[[maybe_unused]] static const auto result = [] {
|
||||
return cudaFuncSetAttribute(reinterpret_cast<const void*>(f),
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
static_cast<int>(kSmem));
|
||||
}();
|
||||
TORCH_CHECK(result == cudaSuccess,
|
||||
"fast_topk_v2: cudaFuncSetAttribute failed: ",
|
||||
cudaGetErrorString(result));
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Host-side launchers
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
|
||||
#define CHECK_DTYPE(x, t) \
|
||||
TORCH_CHECK(x.scalar_type() == (t), #x " must be ", #t)
|
||||
#define CHECK_CONTIG(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
|
||||
|
||||
} // namespace vllm::dsv4_topk
|
||||
|
||||
void fast_topk_v2_plan(const torch::Tensor& seq_lens, torch::Tensor& metadata,
|
||||
int64_t static_cluster_threshold) {
|
||||
using namespace vllm::dsv4_topk;
|
||||
CHECK_CUDA(seq_lens);
|
||||
CHECK_CUDA(metadata);
|
||||
CHECK_DTYPE(seq_lens, torch::kInt32);
|
||||
CHECK_DTYPE(metadata, torch::kInt32);
|
||||
TORCH_CHECK(seq_lens.dim() == 1);
|
||||
TORCH_CHECK(metadata.dim() == 2 && metadata.size(1) == 4);
|
||||
TORCH_CHECK(metadata.size(0) == seq_lens.size(0) + 1,
|
||||
"metadata must be (batch_size + 1, 4)");
|
||||
CHECK_CONTIG(seq_lens);
|
||||
CHECK_CONTIG(metadata);
|
||||
|
||||
const auto batch_size = static_cast<uint32_t>(seq_lens.size(0));
|
||||
if (batch_size <= kNumClusters) return; // metadata unused in fused path
|
||||
|
||||
const auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
cudaLaunchConfig_t cfg{};
|
||||
cfg.gridDim = dim3(1);
|
||||
cfg.blockDim = dim3(kBlockSize);
|
||||
cfg.dynamicSmemBytes = 0;
|
||||
cfg.stream = stream;
|
||||
cfg.numAttrs = 0;
|
||||
TORCH_CHECK(cudaLaunchKernelEx(
|
||||
&cfg, &topk_plan,
|
||||
reinterpret_cast<const uint32_t*>(seq_lens.data_ptr<int32_t>()),
|
||||
reinterpret_cast<Metadata*>(metadata.data_ptr<int32_t>()),
|
||||
batch_size,
|
||||
static_cast<uint32_t>(static_cluster_threshold)) == cudaSuccess,
|
||||
"fast_topk_v2_plan launch failed: ",
|
||||
cudaGetErrorString(cudaGetLastError()));
|
||||
}
|
||||
|
||||
namespace vllm::dsv4_topk {
|
||||
|
||||
// Shared dispatch path for fast_topk_v2 and fast_topk_v2_raw. Templated on
|
||||
// (K, kRawOutput). K is the top-k value (512 for V4-Flash, 1024 for V4-Pro);
|
||||
// kRawOutput=false folds the page-table gather, kRawOutput=true emits raw
|
||||
// row-local indices. The set of input tensors is the same modulo
|
||||
// (page_table, page_size), which the caller has already validated.
|
||||
template <uint32_t K, bool kRawOutput>
|
||||
static void launch_dispatch(const TopKParams& params, uint32_t batch_size,
|
||||
uint32_t max_seq_len, cudaStream_t stream) {
|
||||
// Helper: build a cudaLaunchConfig with optional PDL + cluster attributes.
|
||||
// The attribute storage must outlive cudaLaunchKernelEx (cfg.attrs points
|
||||
// into it), so it lives in each call site below as a stack local.
|
||||
auto make_cfg = [&](dim3 grid, dim3 block, size_t smem,
|
||||
cudaLaunchAttribute* attrs, bool enable_cluster,
|
||||
bool enable_pdl) {
|
||||
cudaLaunchConfig_t cfg{};
|
||||
cfg.gridDim = grid;
|
||||
cfg.blockDim = block;
|
||||
cfg.dynamicSmemBytes = static_cast<unsigned>(smem);
|
||||
cfg.stream = stream;
|
||||
int n = 0;
|
||||
if (enable_pdl) {
|
||||
attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
attrs[n].val.programmaticStreamSerializationAllowed = 1;
|
||||
++n;
|
||||
}
|
||||
if (enable_cluster) {
|
||||
attrs[n].id = cudaLaunchAttributeClusterDimension;
|
||||
attrs[n].val.clusterDim = {1, kClusterSize, 1};
|
||||
++n;
|
||||
}
|
||||
cfg.numAttrs = n;
|
||||
cfg.attrs = n ? attrs : nullptr;
|
||||
return cfg;
|
||||
};
|
||||
|
||||
auto check_launch = [](cudaError_t err) {
|
||||
TORCH_CHECK(err == cudaSuccess,
|
||||
"fast_topk_v2 launch failed: ", cudaGetErrorString(err));
|
||||
};
|
||||
|
||||
constexpr size_t kS1 = kStage1SMEM<K>;
|
||||
constexpr size_t kS2 = kStage2SMEM<K>;
|
||||
if (max_seq_len <= Small<K>::kMax1PassLength) {
|
||||
setup_kernel_smem_once<&topk_short_transform<K, kRawOutput>, kS2>();
|
||||
cudaLaunchAttribute attrs[2];
|
||||
auto cfg = make_cfg(dim3(batch_size), dim3(kBlockSize), kS2, attrs,
|
||||
/*cluster=*/false, /*pdl=*/true);
|
||||
check_launch(cudaLaunchKernelEx(
|
||||
&cfg, topk_short_transform<K, kRawOutput>, params));
|
||||
} else if (batch_size <= kNumClusters) {
|
||||
constexpr size_t kFusedSMEM = kS1 > kS2 ? kS1 : kS2;
|
||||
setup_kernel_smem_once<&topk_fused_transform<K, kRawOutput>, kFusedSMEM>();
|
||||
cudaLaunchAttribute attrs[2];
|
||||
auto cfg = make_cfg(dim3(batch_size, kClusterSize), dim3(kBlockSize),
|
||||
kFusedSMEM, attrs, /*cluster=*/true, /*pdl=*/true);
|
||||
check_launch(cudaLaunchKernelEx(
|
||||
&cfg, topk_fused_transform<K, kRawOutput>, params));
|
||||
} else {
|
||||
const auto num_clusters = std::min<uint32_t>(batch_size, kNumClusters);
|
||||
setup_kernel_smem_once<&topk_combine_preprocess<K, kRawOutput>, kS1>();
|
||||
cudaLaunchAttribute attrs1[2];
|
||||
auto cfg1 = make_cfg(dim3(num_clusters, kClusterSize), dim3(kBlockSize),
|
||||
kS1, attrs1, /*cluster=*/true, /*pdl=*/true);
|
||||
check_launch(cudaLaunchKernelEx(
|
||||
&cfg1, topk_combine_preprocess<K, kRawOutput>, params));
|
||||
|
||||
setup_kernel_smem_once<&topk_combine_transform<K, kRawOutput>, kS2>();
|
||||
cudaLaunchAttribute attrs2[2];
|
||||
auto cfg2 = make_cfg(dim3(batch_size), dim3(kBlockSize), kS2, attrs2,
|
||||
/*cluster=*/false, /*pdl=*/true);
|
||||
check_launch(cudaLaunchKernelEx(
|
||||
&cfg2, topk_combine_transform<K, kRawOutput>, params));
|
||||
}
|
||||
}
|
||||
|
||||
// Top-level K dispatcher: validate the runtime topk argument and route to
|
||||
// the right template instantiation.
|
||||
template <bool kRawOutput>
|
||||
static void launch_dispatch_k(int64_t topk, const TopKParams& params,
|
||||
uint32_t batch_size, uint32_t max_seq_len,
|
||||
cudaStream_t stream) {
|
||||
if (topk == 512) {
|
||||
launch_dispatch<512, kRawOutput>(params, batch_size, max_seq_len, stream);
|
||||
} else if (topk == 1024) {
|
||||
launch_dispatch<1024, kRawOutput>(params, batch_size, max_seq_len, stream);
|
||||
} else {
|
||||
TORCH_CHECK(false,
|
||||
"fast_topk_v2 supports topk in {512, 1024}, got ", topk);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace vllm::dsv4_topk
|
||||
|
||||
void fast_topk_v2(const torch::Tensor& scores, const torch::Tensor& seq_lens,
|
||||
const torch::Tensor& page_table, torch::Tensor& page_indices,
|
||||
int64_t page_size, const torch::Tensor& workspace,
|
||||
const torch::Tensor& metadata, int64_t topk) {
|
||||
using namespace vllm::dsv4_topk;
|
||||
CHECK_CUDA(scores);
|
||||
CHECK_CUDA(seq_lens);
|
||||
CHECK_CUDA(page_table);
|
||||
CHECK_CUDA(page_indices);
|
||||
CHECK_CUDA(workspace);
|
||||
CHECK_CUDA(metadata);
|
||||
CHECK_DTYPE(scores, torch::kFloat32);
|
||||
CHECK_DTYPE(seq_lens, torch::kInt32);
|
||||
CHECK_DTYPE(page_table, torch::kInt32);
|
||||
CHECK_DTYPE(page_indices, torch::kInt32);
|
||||
CHECK_DTYPE(workspace, torch::kInt32);
|
||||
CHECK_DTYPE(metadata, torch::kInt32);
|
||||
|
||||
TORCH_CHECK(scores.dim() == 2 && scores.stride(1) == 1,
|
||||
"scores must be 2D with last stride 1");
|
||||
TORCH_CHECK(seq_lens.dim() == 1 && seq_lens.is_contiguous());
|
||||
TORCH_CHECK(page_table.dim() == 2 && page_table.stride(1) == 1,
|
||||
"page_table must be 2D with last stride 1");
|
||||
TORCH_CHECK(page_indices.dim() == 2 && page_indices.is_contiguous() &&
|
||||
page_indices.size(1) == topk,
|
||||
"page_indices must be (B, topk) contiguous");
|
||||
// workspace size is K-independent (it stages cluster-path ties whose
|
||||
// count is bounded by kMaxTies, not K), so this check uses any K.
|
||||
TORCH_CHECK(workspace.dim() == 2 && workspace.stride(1) == 1 &&
|
||||
workspace.size(1) == Large<512>::kWorkspaceInts,
|
||||
"workspace must be (B, kWorkspaceInts) with last stride 1");
|
||||
TORCH_CHECK(metadata.dim() == 2 && metadata.size(1) == 4 &&
|
||||
metadata.is_contiguous(),
|
||||
"metadata must be (B + 1, 4) contiguous");
|
||||
|
||||
const auto batch_size = static_cast<uint32_t>(scores.size(0));
|
||||
TORCH_CHECK(seq_lens.size(0) == batch_size);
|
||||
TORCH_CHECK(page_table.size(0) == batch_size);
|
||||
TORCH_CHECK(page_indices.size(0) == batch_size);
|
||||
TORCH_CHECK(workspace.size(0) == batch_size);
|
||||
TORCH_CHECK(metadata.size(0) == batch_size + 1);
|
||||
|
||||
const auto max_seq_len = static_cast<uint32_t>(scores.size(1));
|
||||
TORCH_CHECK(page_size > 0 && (page_size & (page_size - 1)) == 0,
|
||||
"page_size must be a positive power of 2");
|
||||
TORCH_CHECK(scores.stride(0) % 4 == 0,
|
||||
"score stride must be a multiple of 4 (TMA 16-byte alignment)");
|
||||
|
||||
// page_bits = log2(page_size). __builtin_ctzll is a host-side compiler
|
||||
// builtin available under C++17 (vLLM compiles host code with C++17).
|
||||
const auto page_bits = static_cast<uint32_t>(
|
||||
__builtin_ctzll(static_cast<unsigned long long>(page_size)));
|
||||
TopKParams params{
|
||||
.seq_lens =
|
||||
reinterpret_cast<const uint32_t*>(seq_lens.data_ptr<int32_t>()),
|
||||
.scores = scores.data_ptr<float>(),
|
||||
.page_table = page_table.data_ptr<int32_t>(),
|
||||
.page_indices = page_indices.data_ptr<int32_t>(),
|
||||
.score_stride = scores.stride(0),
|
||||
.page_table_stride = page_table.stride(0),
|
||||
.workspace = reinterpret_cast<uint8_t*>(workspace.data_ptr<int32_t>()),
|
||||
.metadata =
|
||||
reinterpret_cast<const Metadata*>(metadata.data_ptr<int32_t>()),
|
||||
.workspace_stride =
|
||||
workspace.stride(0) * static_cast<int64_t>(sizeof(int32_t)),
|
||||
.batch_size = batch_size,
|
||||
.page_bits = page_bits,
|
||||
};
|
||||
|
||||
launch_dispatch_k<false>(topk, params, batch_size, max_seq_len,
|
||||
at::cuda::getCurrentCUDAStream().stream());
|
||||
}
|
||||
|
||||
// Top-k only: skip the page-table gather and emit raw row-local indices.
|
||||
// Same selection algorithm as fast_topk_v2; just doesn't touch a page
|
||||
// table. Output semantics match torch.ops._C.persistent_topk and the V4
|
||||
// indexer's existing topk_indices_buffer contract.
|
||||
void fast_topk_v2_raw(const torch::Tensor& scores,
|
||||
const torch::Tensor& seq_lens,
|
||||
torch::Tensor& topk_indices,
|
||||
const torch::Tensor& workspace,
|
||||
const torch::Tensor& metadata,
|
||||
int64_t topk) {
|
||||
using namespace vllm::dsv4_topk;
|
||||
CHECK_CUDA(scores);
|
||||
CHECK_CUDA(seq_lens);
|
||||
CHECK_CUDA(topk_indices);
|
||||
CHECK_CUDA(workspace);
|
||||
CHECK_CUDA(metadata);
|
||||
CHECK_DTYPE(scores, torch::kFloat32);
|
||||
CHECK_DTYPE(seq_lens, torch::kInt32);
|
||||
CHECK_DTYPE(topk_indices, torch::kInt32);
|
||||
CHECK_DTYPE(workspace, torch::kInt32);
|
||||
CHECK_DTYPE(metadata, torch::kInt32);
|
||||
|
||||
TORCH_CHECK(scores.dim() == 2 && scores.stride(1) == 1,
|
||||
"scores must be 2D with last stride 1");
|
||||
TORCH_CHECK(seq_lens.dim() == 1 && seq_lens.is_contiguous());
|
||||
TORCH_CHECK(topk_indices.dim() == 2 && topk_indices.is_contiguous() &&
|
||||
topk_indices.size(1) == topk,
|
||||
"topk_indices must be (B, topk) contiguous");
|
||||
TORCH_CHECK(workspace.dim() == 2 && workspace.stride(1) == 1 &&
|
||||
workspace.size(1) == Large<512>::kWorkspaceInts,
|
||||
"workspace must be (B, kWorkspaceInts) with last stride 1");
|
||||
TORCH_CHECK(metadata.dim() == 2 && metadata.size(1) == 4 &&
|
||||
metadata.is_contiguous(),
|
||||
"metadata must be (B + 1, 4) contiguous");
|
||||
|
||||
const auto batch_size = static_cast<uint32_t>(scores.size(0));
|
||||
TORCH_CHECK(seq_lens.size(0) == batch_size);
|
||||
TORCH_CHECK(topk_indices.size(0) == batch_size);
|
||||
TORCH_CHECK(workspace.size(0) == batch_size);
|
||||
TORCH_CHECK(metadata.size(0) == batch_size + 1);
|
||||
|
||||
const auto max_seq_len = static_cast<uint32_t>(scores.size(1));
|
||||
TORCH_CHECK(scores.stride(0) % 4 == 0,
|
||||
"score stride must be a multiple of 4 (TMA 16-byte alignment)");
|
||||
|
||||
// page_table / page_bits are unused on the raw path; passing nullptr/0 is
|
||||
// safe because every kernel call site is gated by `if constexpr
|
||||
// (kRawOutput)` so the page-table loads are eliminated at compile time.
|
||||
TopKParams params{
|
||||
.seq_lens =
|
||||
reinterpret_cast<const uint32_t*>(seq_lens.data_ptr<int32_t>()),
|
||||
.scores = scores.data_ptr<float>(),
|
||||
.page_table = nullptr,
|
||||
.page_indices = topk_indices.data_ptr<int32_t>(),
|
||||
.score_stride = scores.stride(0),
|
||||
.page_table_stride = 0,
|
||||
.workspace = reinterpret_cast<uint8_t*>(workspace.data_ptr<int32_t>()),
|
||||
.metadata =
|
||||
reinterpret_cast<const Metadata*>(metadata.data_ptr<int32_t>()),
|
||||
.workspace_stride =
|
||||
workspace.stride(0) * static_cast<int64_t>(sizeof(int32_t)),
|
||||
.batch_size = batch_size,
|
||||
.page_bits = 0,
|
||||
};
|
||||
|
||||
launch_dispatch_k<true>(topk, params, batch_size, max_seq_len,
|
||||
at::cuda::getCurrentCUDAStream().stream());
|
||||
}
|
||||
|
||||
int64_t fast_topk_v2_workspace_ints() {
|
||||
// Workspace size is K-independent (kMaxTies, not K, drives it).
|
||||
return static_cast<int64_t>(vllm::dsv4_topk::Large<512>::kWorkspaceInts);
|
||||
}
|
||||
|
||||
// Register impls here (instead of in torch_bindings.cpp) so they only exist
|
||||
// when CMake compiles this source — i.e., when the target build has a
|
||||
// compatible Hopper / Blackwell-datacenter arch. On other configs the schema
|
||||
// remains defined but a call surfaces a clear "no impl" runtime error.
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
|
||||
m.impl("fast_topk_v2_plan", &fast_topk_v2_plan);
|
||||
m.impl("fast_topk_v2", &fast_topk_v2);
|
||||
m.impl("fast_topk_v2_raw", &fast_topk_v2_raw);
|
||||
}
|
||||
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CompositeExplicitAutograd, m) {
|
||||
m.impl("fast_topk_v2_workspace_ints", &fast_topk_v2_workspace_ints);
|
||||
}
|
||||
@@ -1,266 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Cluster top-k strategy for very large N. Uses Hopper thread-block clusters
|
||||
// (cooperative_groups::this_cluster) to parallelize histogram + scatter across
|
||||
// up to ``kClusterSize`` blocks per row. Each row is processed in two stages:
|
||||
// stage 1: per-block histogram, all-reduce across the cluster, threshold
|
||||
// scatter, and an epilogue that page-translates strictly-above
|
||||
// entries to global memory and stages ties into a per-row workspace.
|
||||
// stage 2: tie-break across the cluster's combined ties (run by cluster
|
||||
// rank 0 in the fused kernel, or as a separate launch otherwise).
|
||||
// Ported from
|
||||
// jit_kernel/include/sgl_kernel/deepseek_v4/topk/cluster.cuh.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "common.cuh"
|
||||
#include "ptx.cuh"
|
||||
#include "utils.cuh"
|
||||
|
||||
#include <cooperative_groups.h>
|
||||
#include <cstdint>
|
||||
|
||||
namespace vllm::dsv4_topk {
|
||||
|
||||
template <uint32_t K>
|
||||
struct ClusterTopK {
|
||||
static constexpr uint32_t kClusterSize = 8;
|
||||
static constexpr uint32_t kHistBits = 10;
|
||||
static constexpr uint32_t kHistBins = 1 << kHistBits;
|
||||
static constexpr uint32_t kElemPerStage = 8;
|
||||
static constexpr uint32_t kSizePerStage = kElemPerStage * kBlockSize;
|
||||
static constexpr uint32_t kNumStages = 4;
|
||||
static constexpr uint32_t kMaxLength = kClusterSize * kNumStages * kSizePerStage;
|
||||
static constexpr uint32_t kAboveBits = 11;
|
||||
|
||||
struct Smem {
|
||||
uint64_t barrier[kNumStages];
|
||||
uint32_t local_above_equal[kClusterSize];
|
||||
uint32_t prefix_above_equal;
|
||||
alignas(128) uint32_t counter_gt;
|
||||
alignas(128) uint32_t counter_eq;
|
||||
alignas(128) MatchBin match;
|
||||
alignas(128) uint32_t warp_sum[kNumWarps];
|
||||
uint32_t histogram[kHistBins];
|
||||
alignas(128) float score_buffer[kNumStages][kSizePerStage];
|
||||
Tie tie_buffer[kMaxTies];
|
||||
};
|
||||
|
||||
// Per-row metadata produced by the plan kernel and consumed by the fused /
|
||||
// stage-1 kernels. {batch_id, seq_len, has_next} arranged in an int4-sized
|
||||
// 16-byte struct so the planner can do contiguous int32x4 stores.
|
||||
struct alignas(16) Metadata {
|
||||
uint32_t batch_id;
|
||||
uint32_t seq_len;
|
||||
bool has_next;
|
||||
};
|
||||
|
||||
// Per-row workspace storing {(num_above, num_ties)} + the gathered ties.
|
||||
struct WorkSpace {
|
||||
uint2 metadata;
|
||||
Tie ties[kMaxTies];
|
||||
};
|
||||
|
||||
static constexpr uint32_t kWorkspaceInts = sizeof(WorkSpace) / sizeof(uint32_t);
|
||||
|
||||
VLLM_DSV4_DEVICE static void stage1_init(void* _smem) {
|
||||
const auto tx = threadIdx.x;
|
||||
__builtin_assume(tx < kBlockSize);
|
||||
const auto smem = static_cast<Smem*>(_smem);
|
||||
if (tx < kHistBins) smem->histogram[tx] = 0;
|
||||
if (tx < kNumStages) ptx::mbarrier_init(&smem->barrier[tx], 1);
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE static void stage1_prologue(const float* scores,
|
||||
uint32_t length, void* _smem) {
|
||||
if (threadIdx.x == 0) {
|
||||
const auto smem = static_cast<Smem*>(_smem);
|
||||
const auto num_stages = (length + kSizePerStage - 1) / kSizePerStage;
|
||||
const auto length_aligned = (length + 3u) & ~3u;
|
||||
#pragma unroll
|
||||
for (uint32_t stage = 0; stage < kNumStages; stage++) {
|
||||
if (stage >= num_stages) break;
|
||||
const auto offset = stage * kSizePerStage;
|
||||
const auto size = min(kSizePerStage, length_aligned - offset);
|
||||
const auto size_bytes = size * sizeof(float);
|
||||
const auto bar = &smem->barrier[stage];
|
||||
ptx::tma_load(smem->score_buffer[stage], scores + offset, size_bytes,
|
||||
bar);
|
||||
ptx::mbarrier_arrive_expect_tx(bar, size_bytes);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE static void stage1(int32_t* indices, uint32_t length,
|
||||
void* _smem, bool reuse = false) {
|
||||
const auto smem = static_cast<Smem*>(_smem);
|
||||
const auto tx = threadIdx.x;
|
||||
__builtin_assume(tx < kBlockSize);
|
||||
const auto lane_id = tx % kWarpThreads;
|
||||
const auto warp_id = tx / kWarpThreads;
|
||||
|
||||
// Local histogram.
|
||||
#pragma unroll
|
||||
for (uint32_t stage = 0; stage < kNumStages; stage++) {
|
||||
const auto offset = stage * kSizePerStage;
|
||||
if (offset >= length) break;
|
||||
const auto size = min(kSizePerStage, length - offset);
|
||||
if (lane_id == 0) ptx::mbarrier_wait(&smem->barrier[stage], 0);
|
||||
__syncwarp();
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kElemPerStage; ++i) {
|
||||
const auto idx = tx + i * kBlockSize;
|
||||
if (idx >= size) break;
|
||||
const auto score = smem->score_buffer[stage][idx];
|
||||
const auto bin = extract_coarse_bin<kHistBits>(score);
|
||||
atomicAdd(&smem->histogram[bin], 1);
|
||||
}
|
||||
}
|
||||
|
||||
static_assert(kHistBins <= kBlockSize);
|
||||
|
||||
// Two-shot all-reduce across the cluster.
|
||||
{
|
||||
auto cluster = cooperative_groups::this_cluster();
|
||||
cluster.sync();
|
||||
const auto cluster_rank = blockIdx.y;
|
||||
const auto kLocalSize = kHistBins / kClusterSize;
|
||||
const auto offset = kLocalSize * cluster_rank;
|
||||
|
||||
const auto src_tx = tx / kClusterSize;
|
||||
const auto src_rank = tx % kClusterSize;
|
||||
|
||||
if (tx < kHistBins) {
|
||||
const auto addr = &smem->histogram[offset + src_tx];
|
||||
const auto src_addr = cluster.map_shared_rank(addr, src_rank);
|
||||
*src_addr = warp_reduce_sum<kClusterSize>(*src_addr);
|
||||
}
|
||||
cluster.sync();
|
||||
}
|
||||
|
||||
// Each block now holds the full cluster histogram. Find the threshold.
|
||||
{
|
||||
const auto value = tx < kHistBins ? smem->histogram[tx] : 0;
|
||||
const auto warp_inc = warp_inclusive_sum(lane_id, value);
|
||||
if (lane_id == kWarpThreads - 1) {
|
||||
smem->warp_sum[warp_id] = warp_inc;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
const auto tmp = smem->warp_sum[lane_id];
|
||||
const auto total_length = warp_reduce_sum(tmp);
|
||||
uint32_t prefix_sum = warp_reduce_sum(lane_id < warp_id ? tmp : 0);
|
||||
prefix_sum += warp_inc;
|
||||
const auto above = total_length - prefix_sum;
|
||||
if (tx < kHistBins && above < K && above + value >= K) {
|
||||
smem->counter_gt = smem->counter_eq = 0;
|
||||
smem->match = {
|
||||
.bin = tx,
|
||||
.above_count = above,
|
||||
.equal_count = value,
|
||||
};
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
const auto thr_bin = smem->match.bin;
|
||||
|
||||
// Scatter strictly-above entries to `indices`, stash ties in tie_buffer.
|
||||
#pragma unroll
|
||||
for (uint32_t stage = 0; stage < kNumStages; stage++) {
|
||||
const auto offset = stage * kSizePerStage;
|
||||
if (offset >= length) break;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kElemPerStage; ++i) {
|
||||
const auto buf_idx = tx + i * kBlockSize;
|
||||
const auto global_idx = offset + buf_idx;
|
||||
if (global_idx >= length) break;
|
||||
const auto score = smem->score_buffer[stage][buf_idx];
|
||||
const auto bin = extract_coarse_bin<kHistBits>(score);
|
||||
if (bin > thr_bin) {
|
||||
indices[atomicAdd(&smem->counter_gt, 1)] = global_idx;
|
||||
} else if (bin == thr_bin) {
|
||||
const auto pos = atomicAdd(&smem->counter_eq, 1);
|
||||
if (pos < kMaxTies) smem->tie_buffer[pos] = {global_idx, score};
|
||||
}
|
||||
}
|
||||
}
|
||||
if (reuse) {
|
||||
const auto num_stages = (length + kSizePerStage - 1) / kSizePerStage;
|
||||
if (tx < kHistBins) smem->histogram[tx] = 0;
|
||||
if (tx < num_stages) ptx::mbarrier_arrive(&smem->barrier[tx]);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
template <typename TParams>
|
||||
VLLM_DSV4_DEVICE static void stage1_epilogue(TParams params,
|
||||
uint32_t offset, void* _ws,
|
||||
void* _smem) {
|
||||
auto cluster = cooperative_groups::this_cluster();
|
||||
const auto smem = static_cast<Smem*>(_smem);
|
||||
const auto tx = threadIdx.x;
|
||||
const auto local_above = smem->counter_gt;
|
||||
const auto local_equal = smem->counter_eq;
|
||||
const auto cluster_rank = blockIdx.y;
|
||||
|
||||
constexpr uint32_t kAboveMask = (1 << kAboveBits) - 1;
|
||||
static_assert(kAboveMask >= K);
|
||||
|
||||
static_assert(kMaxTies <= kBlockSize);
|
||||
const auto idx_above = tx < local_above ? params.indices_in[tx] : 0;
|
||||
const auto tie_value = tx < local_equal ? smem->tie_buffer[tx] : Tie{0, 0.0f};
|
||||
|
||||
// Push counts to remote shared memory to reduce inter-block latency.
|
||||
if (tx < kClusterSize) {
|
||||
const auto value = (local_equal << kAboveBits) | local_above;
|
||||
const auto dst_addr = cluster.map_shared_rank(smem->local_above_equal, tx);
|
||||
dst_addr[cluster_rank] = value;
|
||||
}
|
||||
// After this final sync, every block can read only its own smem (peer
|
||||
// ranks may have already exited), so we don't touch remote smem again.
|
||||
cluster.sync();
|
||||
if (tx < kClusterSize) {
|
||||
const auto value = tx < cluster_rank ? smem->local_above_equal[tx] : 0;
|
||||
const auto kActiveMask = (1u << kClusterSize) - 1;
|
||||
smem->prefix_above_equal = warp_reduce_sum<kClusterSize>(value, kActiveMask);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto prefix_packed = smem->prefix_above_equal;
|
||||
const auto prefix_above = prefix_packed & kAboveMask;
|
||||
const auto prefix_equal = prefix_packed >> kAboveBits;
|
||||
|
||||
// Page-translate strictly-above entries.
|
||||
if (tx < local_above) {
|
||||
params.write(tx + prefix_above, idx_above + offset);
|
||||
}
|
||||
// Stage ties into the per-row workspace (regular global writes).
|
||||
const auto ws = static_cast<WorkSpace*>(_ws);
|
||||
if (tx < local_equal && tx + prefix_equal < kMaxTies) {
|
||||
ws->ties[tx + prefix_equal] = {tie_value.idx + offset, tie_value.score};
|
||||
}
|
||||
// Last cluster rank publishes the sums into ws->metadata.
|
||||
if (cluster_rank == kClusterSize - 1 && tx == 0) {
|
||||
const auto sum_above = prefix_above + local_above;
|
||||
const auto sum_equal = prefix_equal + local_equal;
|
||||
ws->metadata = make_uint2(sum_above, sum_equal);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename TParams>
|
||||
VLLM_DSV4_DEVICE static void transform(TParams params, const void* _ws,
|
||||
void* _smem) {
|
||||
const auto ws = static_cast<const WorkSpace*>(_ws);
|
||||
const auto meta = &ws->metadata;
|
||||
const auto num_above = meta->x;
|
||||
const auto num_equal = meta->y;
|
||||
if (num_above >= K || num_equal == 0) return;
|
||||
const auto clamped_ties = min(num_equal, kMaxTies);
|
||||
tie_handle_transform(ws->ties, clamped_ties, num_above, K, params, _smem);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace vllm::dsv4_topk
|
||||
@@ -1,219 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Shared types/utilities for the three DeepSeek V4 top-k strategies
|
||||
// (Register / Streaming / Cluster). Ported from sglang's
|
||||
// jit_kernel/include/sgl_kernel/deepseek_v4/topk/common.cuh.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "utils.cuh"
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
#include <cstdint>
|
||||
|
||||
namespace vllm::dsv4_topk {
|
||||
|
||||
inline constexpr uint32_t kMaxTopK = 1024;
|
||||
inline constexpr uint32_t kBlockSize = 1024;
|
||||
inline constexpr uint32_t kNumWarps = kBlockSize / kWarpThreads;
|
||||
// 1 element per thread in the tie-breaking pass.
|
||||
inline constexpr uint32_t kMaxTies = 1024;
|
||||
inline constexpr uint32_t kRadixBins = 256;
|
||||
static_assert(kMaxTopK <= kBlockSize && kMaxTies <= kBlockSize);
|
||||
|
||||
// Always vectorize global loads as float4.
|
||||
using Vec4 = AlignedVector<float, 4>;
|
||||
|
||||
// page_to_indices: convert a flat compressed-token index into a (block * page_size + offset)
|
||||
// page-table-resolved index. page_size must be a power of 2; page_bits = log2(page_size).
|
||||
VLLM_DSV4_DEVICE int32_t page_to_indices(const int32_t* __restrict__ page_table,
|
||||
uint32_t i, uint32_t page_bits) {
|
||||
const uint32_t mask = (1u << page_bits) - 1u;
|
||||
return (page_table[i >> page_bits] << page_bits) | (i & mask);
|
||||
}
|
||||
|
||||
// Output-side description of how each strategy commits its top-k output.
|
||||
//
|
||||
// Two modes, picked at compile time via ``kRawOutput``:
|
||||
// - kRawOutput=false (paged): fold the page-table gather into the output
|
||||
// store. ``write(dst, src)`` emits ``page_to_indices(table, src, bits)``;
|
||||
// ``transform(idx)`` reads ``indices_in[idx]`` and re-emits via the
|
||||
// page lookup. This is the original kernel behavior.
|
||||
// - kRawOutput=true (raw): skip the page lookup entirely. The kernel
|
||||
// just writes row-local raw indices, matching ``persistent_topk``'s
|
||||
// output contract. ``page_table`` and ``page_bits`` are unused; the
|
||||
// compiler eliminates the dead loads via ``if constexpr``.
|
||||
template <bool kRawOutput>
|
||||
struct TransformParamsT {
|
||||
const int32_t* __restrict__ page_table;
|
||||
const int32_t* __restrict__ indices_in;
|
||||
int32_t* __restrict__ indices_out;
|
||||
uint32_t page_bits;
|
||||
|
||||
VLLM_DSV4_DEVICE void transform(uint32_t idx) const {
|
||||
if constexpr (kRawOutput) {
|
||||
indices_out[idx] = static_cast<int32_t>(indices_in[idx]);
|
||||
} else {
|
||||
indices_out[idx] =
|
||||
page_to_indices(page_table, indices_in[idx], page_bits);
|
||||
}
|
||||
}
|
||||
VLLM_DSV4_DEVICE void write(uint32_t dst, uint32_t src) const {
|
||||
if constexpr (kRawOutput) {
|
||||
indices_out[dst] = static_cast<int32_t>(src);
|
||||
} else {
|
||||
indices_out[dst] = page_to_indices(page_table, src, page_bits);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Back-compat alias. The four kernels in fast_topk_v2.cu instantiate both
|
||||
// variants explicitly via templates.
|
||||
using TransformParams = TransformParamsT<false>;
|
||||
|
||||
struct alignas(16) MatchBin {
|
||||
uint32_t bin;
|
||||
uint32_t above_count;
|
||||
uint32_t equal_count;
|
||||
};
|
||||
|
||||
struct alignas(8) Tie {
|
||||
uint32_t idx;
|
||||
float score;
|
||||
};
|
||||
|
||||
// Shared-memory layout for the final tie-breaking radix pass. Reused by both
|
||||
// the streaming kernel (overlapping `score_buffer`) and the cluster kernel.
|
||||
struct TieHandleSmem {
|
||||
alignas(128) uint32_t counter;
|
||||
alignas(128) MatchBin match;
|
||||
uint32_t histogram[kRadixBins];
|
||||
uint32_t warp_sum[kNumWarps];
|
||||
};
|
||||
|
||||
// Order-preserving fp32 -> uint key, truncated to the top kBits. Used for the
|
||||
// coarse histogram pass.
|
||||
template <uint32_t kBits>
|
||||
VLLM_DSV4_DEVICE uint32_t extract_coarse_bin(float x) {
|
||||
static_assert(0 < kBits && kBits < 15);
|
||||
__half h = __float2half_rn(x);
|
||||
uint16_t bits = __half_as_ushort(h);
|
||||
uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits)
|
||||
: static_cast<uint16_t>(bits | 0x8000);
|
||||
return key >> (16 - kBits);
|
||||
}
|
||||
|
||||
// Full 32-bit order-preserving key, used in tie-breaking.
|
||||
VLLM_DSV4_DEVICE uint32_t extract_exact_bin(float x) {
|
||||
uint32_t bits = __float_as_uint(x);
|
||||
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) {
|
||||
static_assert(kWarpThreads == 32);
|
||||
#pragma unroll
|
||||
for (uint32_t offset = 1; offset < 32; offset *= 2) {
|
||||
uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset);
|
||||
if (lane_id >= offset) val += n;
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
// Fast path when seq_len <= K: identity mapping, padded to K with -1.
|
||||
template <typename TParams>
|
||||
VLLM_DSV4_DEVICE void trivial_transform(const TParams& params, uint32_t length,
|
||||
uint32_t K) {
|
||||
const auto tx = threadIdx.x;
|
||||
if (tx < length) {
|
||||
params.write(tx, tx);
|
||||
} else if (tx < K) {
|
||||
params.indices_out[tx] = -1;
|
||||
}
|
||||
}
|
||||
|
||||
// Tie-break the threshold-bin candidates that didn't fit in the strict-above
|
||||
// region. One block-wide radix pass over the full 32-bit key (fp32 bit
|
||||
// pattern, with idx as a secondary key). Writes at most `K - num_above`
|
||||
// entries via params.write(...).
|
||||
template <typename TParams>
|
||||
VLLM_DSV4_DEVICE void tie_handle_transform(const Tie* __restrict__ ties,
|
||||
uint32_t num_ties, uint32_t num_above,
|
||||
uint32_t K, TParams params,
|
||||
void* _smem) {
|
||||
auto* smem = static_cast<TieHandleSmem*>(_smem);
|
||||
const auto tx = threadIdx.x;
|
||||
const auto lane_id = tx % kWarpThreads;
|
||||
const auto warp_id = tx / kWarpThreads;
|
||||
|
||||
const bool has_elem = tx < num_ties;
|
||||
const auto tie = has_elem ? ties[tx] : Tie{0, 0.0f};
|
||||
const uint32_t key = extract_exact_bin(tie.score);
|
||||
const uint32_t idx = tie.idx;
|
||||
bool active = has_elem;
|
||||
uint32_t topk_remain = K - num_above;
|
||||
uint32_t write_pos = K;
|
||||
|
||||
smem->counter = 0;
|
||||
__syncthreads();
|
||||
|
||||
// 256 bins / 32 lanes = 8 warps span the histogram inter-warp prefix.
|
||||
constexpr uint32_t kRadixWarps = kRadixBins / kWarpThreads;
|
||||
|
||||
#pragma unroll
|
||||
for (int round = 0; round < 4; round++) {
|
||||
const uint32_t shift = 24 - round * 8;
|
||||
const uint32_t bin = (key >> shift) & 0xFFu;
|
||||
|
||||
// 1. Histogram.
|
||||
if (tx < kRadixBins) smem->histogram[tx] = 0;
|
||||
__syncthreads();
|
||||
if (active) atomicAdd(&smem->histogram[bin], 1);
|
||||
__syncthreads();
|
||||
|
||||
// 2. Two-pass prefix sum across the 256 bins.
|
||||
uint32_t hist_val = 0;
|
||||
uint32_t warp_inc = 0;
|
||||
if (tx < kRadixBins) {
|
||||
hist_val = smem->histogram[tx];
|
||||
warp_inc = warp_inclusive_sum(lane_id, hist_val);
|
||||
if (lane_id == kWarpThreads - 1) smem->warp_sum[warp_id] = warp_inc;
|
||||
}
|
||||
__syncthreads();
|
||||
if (tx < kRadixBins) {
|
||||
const auto tmp = (lane_id < kRadixWarps) ? smem->warp_sum[lane_id] : 0;
|
||||
const auto total = warp_reduce_sum(tmp);
|
||||
const auto inter = warp_reduce_sum(lane_id < warp_id ? tmp : 0);
|
||||
const auto prefix = inter + warp_inc;
|
||||
const auto above = total - prefix;
|
||||
// 3. Find threshold bin.
|
||||
if (above < topk_remain && above + hist_val >= topk_remain) {
|
||||
smem->match = {tx, above, topk_remain - above};
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto thr = smem->match.bin;
|
||||
const auto n_above = smem->match.above_count;
|
||||
|
||||
// 4. Scatter.
|
||||
if (active) {
|
||||
if (bin > thr) {
|
||||
write_pos = num_above + atomicAdd(&smem->counter, 1);
|
||||
active = false;
|
||||
} else if (bin < thr) {
|
||||
active = false;
|
||||
} else if (round == 3) {
|
||||
write_pos = K - atomicAdd(&smem->match.equal_count, -1u);
|
||||
}
|
||||
// bin == thr && round < 3: stay active for the next radix round.
|
||||
}
|
||||
|
||||
topk_remain -= n_above;
|
||||
if (topk_remain == 0) break;
|
||||
}
|
||||
|
||||
if (write_pos < K) params.write(write_pos, idx);
|
||||
}
|
||||
|
||||
} // namespace vllm::dsv4_topk
|
||||
@@ -1,66 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Thin wrappers around the CUDA PTX intrinsics used by the top-k pipeline.
|
||||
// All of these require sm_90+. Ported from sglang's
|
||||
// jit_kernel/include/sgl_kernel/deepseek_v4/topk/ptx.cuh.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "utils.cuh"
|
||||
|
||||
#include <cuda/ptx>
|
||||
#include <cstdint>
|
||||
|
||||
namespace vllm::dsv4_topk::ptx {
|
||||
|
||||
VLLM_DSV4_DEVICE void mbarrier_init(uint64_t* addr, uint32_t arrives) {
|
||||
cuda::ptx::mbarrier_init(addr, arrives);
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE void mbarrier_arrive(uint64_t* addr) {
|
||||
cuda::ptx::mbarrier_arrive(cuda::ptx::sem_relaxed, cuda::ptx::scope_cta,
|
||||
cuda::ptx::space_shared, addr);
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE void mbarrier_arrive_expect_tx(uint64_t* addr, uint32_t tx) {
|
||||
cuda::ptx::mbarrier_arrive_expect_tx(cuda::ptx::sem_relaxed,
|
||||
cuda::ptx::scope_cta,
|
||||
cuda::ptx::space_shared, addr, tx);
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE void mbarrier_wait(uint64_t* addr, uint32_t phase) {
|
||||
while (!cuda::ptx::mbarrier_try_wait_parity(cuda::ptx::sem_relaxed,
|
||||
cuda::ptx::scope_cta, addr,
|
||||
phase))
|
||||
;
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE void tma_load(void* dst, const void* src, uint32_t num_bytes,
|
||||
uint64_t* mbar) {
|
||||
cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, cuda::ptx::space_global,
|
||||
dst, src, num_bytes, mbar);
|
||||
}
|
||||
|
||||
// elect.sync: pick a single arbitrary thread out of an active mask. Used to
|
||||
// fire a single TMA load per warp without the full ``if (tx == 0)`` cost.
|
||||
VLLM_DSV4_DEVICE uint32_t elect_sync() {
|
||||
uint32_t pred = 0;
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
".reg .pred %%px;\n\t"
|
||||
"elect.sync _|%%px, %1;\n\t"
|
||||
"@%%px mov.s32 %0, 1;\n\t"
|
||||
"}"
|
||||
: "+r"(pred)
|
||||
: "r"(0xFFFFFFFF));
|
||||
return pred;
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE bool elect_sync_cta(uint32_t tx) {
|
||||
const auto warp_id = tx / 32;
|
||||
const auto uniform_warp_id = __shfl_sync(0xFFFFFFFF, warp_id, 0);
|
||||
return (uniform_warp_id == 0 && elect_sync());
|
||||
}
|
||||
|
||||
} // namespace vllm::dsv4_topk::ptx
|
||||
@@ -1,314 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Register-resident top-k strategy for the DeepSeek V4 indexer (small N
|
||||
// fast path). One block per row; up to ``kMax2PassLength`` scores per row
|
||||
// streamed through registers, with a single 12-bit-coarse radix pass and
|
||||
// a final tie-break round. Ported from
|
||||
// jit_kernel/include/sgl_kernel/deepseek_v4/topk/register.cuh.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "common.cuh"
|
||||
#include "ptx.cuh"
|
||||
#include "utils.cuh"
|
||||
|
||||
#include <cfloat>
|
||||
#include <cstdint>
|
||||
|
||||
namespace vllm::dsv4_topk {
|
||||
|
||||
template <uint32_t K>
|
||||
struct RegisterTopK {
|
||||
static constexpr uint32_t kHistBits = 12;
|
||||
static constexpr uint32_t kHistBins = 1 << kHistBits;
|
||||
static constexpr uint32_t kVecsPerThread = 4;
|
||||
static constexpr uint32_t kMaxTolerance = 0;
|
||||
// Length covered by registers in a single pass.
|
||||
static constexpr uint32_t kMax1PassLength = kVecsPerThread * 4 * kBlockSize;
|
||||
// Extra length staged through shared memory in the 2-pass path.
|
||||
static constexpr uint32_t kMaxExtraLength = kMax1PassLength;
|
||||
static constexpr uint32_t kMax2PassLength = kMax1PassLength + kMaxExtraLength;
|
||||
|
||||
struct Smem {
|
||||
using HistVec = AlignedVector<uint32_t, kHistBins / kBlockSize>;
|
||||
alignas(128) uint32_t counter_gt;
|
||||
alignas(128) uint32_t counter_eq;
|
||||
uint64_t mbarrier; // for the cp.async.bulk in the 2-pass path
|
||||
MatchBin match;
|
||||
uint32_t warp_sum[kNumWarps];
|
||||
union {
|
||||
uint32_t histogram[kHistBins];
|
||||
HistVec histogram_vec[kBlockSize];
|
||||
Tie tie_buffer[kMaxTies];
|
||||
};
|
||||
alignas(16) float score_buffer[kMaxExtraLength];
|
||||
};
|
||||
|
||||
template <bool kIs2Pass = false>
|
||||
VLLM_DSV4_DEVICE static void run(const float* scores, int32_t* indices,
|
||||
uint32_t length, void* _smem,
|
||||
bool use_pdl = false) {
|
||||
const auto smem = static_cast<Smem*>(_smem);
|
||||
const auto tx = threadIdx.x;
|
||||
const auto lane_id = tx % kWarpThreads;
|
||||
const auto warp_id = tx / kWarpThreads;
|
||||
|
||||
// Init histogram + counters.
|
||||
{
|
||||
typename Smem::HistVec hist_vec;
|
||||
hist_vec.fill(0);
|
||||
smem->histogram_vec[tx] = hist_vec;
|
||||
if (tx == 0) {
|
||||
smem->counter_gt = smem->counter_eq = 0;
|
||||
if constexpr (kIs2Pass) {
|
||||
ptx::mbarrier_init(&smem->mbarrier, 1);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (use_pdl) pdl_wait_primary<true>();
|
||||
|
||||
// Stream the first `kMax1PassLength` scores into registers.
|
||||
Vec4 local[kVecsPerThread];
|
||||
#pragma unroll
|
||||
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
|
||||
const uint32_t base = (tx + v * kBlockSize) * 4;
|
||||
if (base >= length) break;
|
||||
local[v].load(scores, tx + v * kBlockSize);
|
||||
}
|
||||
|
||||
// Issue the 2-pass TMA prefetch (next chunk of scores into smem).
|
||||
if constexpr (kIs2Pass) {
|
||||
if (ptx::elect_sync_cta(tx)) {
|
||||
const auto length_aligned = (length + 3u - kMax1PassLength) & ~3u;
|
||||
const auto size_bytes = length_aligned * sizeof(float);
|
||||
ptx::tma_load(smem->score_buffer, scores + kMax1PassLength, size_bytes,
|
||||
&smem->mbarrier);
|
||||
ptx::mbarrier_arrive_expect_tx(&smem->mbarrier, size_bytes);
|
||||
}
|
||||
__syncwarp();
|
||||
}
|
||||
|
||||
// Phase 1: histogram via shared-memory atomics.
|
||||
#pragma unroll
|
||||
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
|
||||
#pragma unroll
|
||||
for (uint32_t e = 0; e < 4; ++e) {
|
||||
if constexpr (!kIs2Pass) {
|
||||
const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
|
||||
if (idx >= length) goto LABEL_ACC_FINISH;
|
||||
}
|
||||
atomicAdd(&smem->histogram[extract_coarse_bin<kHistBits>(local[v][e])],
|
||||
1);
|
||||
}
|
||||
}
|
||||
if constexpr (kIs2Pass) {
|
||||
if (lane_id == 0) ptx::mbarrier_wait(&smem->mbarrier, 0);
|
||||
__syncwarp();
|
||||
for (uint32_t i = tx; i + kMax1PassLength < length; i += kBlockSize) {
|
||||
const auto val = smem->score_buffer[i];
|
||||
atomicAdd(&smem->histogram[extract_coarse_bin<kHistBits>(val)], 1);
|
||||
}
|
||||
}
|
||||
[[maybe_unused]] LABEL_ACC_FINISH:
|
||||
__syncthreads();
|
||||
|
||||
// Phase 2: prefix scan over the histogram, locate the threshold bin.
|
||||
{
|
||||
constexpr uint32_t kItems = kHistBins / kBlockSize;
|
||||
uint32_t orig[kItems];
|
||||
const auto hist_vec = smem->histogram_vec[tx];
|
||||
uint32_t tmp_local_sum = 0;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kItems; ++i) {
|
||||
orig[i] = hist_vec[i];
|
||||
tmp_local_sum += orig[i];
|
||||
}
|
||||
|
||||
const auto warp_inc = warp_inclusive_sum(lane_id, tmp_local_sum);
|
||||
const auto warp_exc = warp_inc - tmp_local_sum;
|
||||
if (lane_id == kWarpThreads - 1) {
|
||||
smem->warp_sum[warp_id] = warp_inc;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
const auto tmp = smem->warp_sum[lane_id];
|
||||
// Exactly one bin satisfies above < K && above + count >= K.
|
||||
uint32_t prefix_sum = warp_reduce_sum(lane_id < warp_id ? tmp : 0);
|
||||
prefix_sum += warp_exc;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kItems; ++i) {
|
||||
prefix_sum += orig[i];
|
||||
const auto above = length - prefix_sum;
|
||||
if (above < K && above + orig[i] >= K) {
|
||||
smem->match = {
|
||||
.bin = tx * kItems + i,
|
||||
.above_count = above,
|
||||
.equal_count = orig[i],
|
||||
};
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
const auto thr_bin = smem->match.bin;
|
||||
const auto num_above = smem->match.above_count;
|
||||
const auto num_equal = smem->match.equal_count;
|
||||
|
||||
// Phase 3: Scatter.
|
||||
// - bin > thr -> write directly to output (strictly above).
|
||||
// - bin == thr -> when no tie-break is needed, admit first-come;
|
||||
// otherwise stash into tie_buffer for phase 4.
|
||||
const bool need_tiebreak = (num_equal + num_above > K + kMaxTolerance);
|
||||
const auto topk_indices = indices;
|
||||
const auto tie_buffer = smem->tie_buffer;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
|
||||
#pragma unroll
|
||||
for (uint32_t e = 0; e < 4; ++e) {
|
||||
const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
|
||||
if constexpr (!kIs2Pass) {
|
||||
if (idx >= length) goto LABEL_SCATTER_DONE;
|
||||
}
|
||||
const uint32_t bin = extract_coarse_bin<kHistBits>(local[v][e]);
|
||||
if (bin > thr_bin) {
|
||||
topk_indices[atomicAdd(&smem->counter_gt, 1)] = idx;
|
||||
} else if (bin == thr_bin) {
|
||||
const auto pos = atomicAdd(&smem->counter_eq, 1);
|
||||
if (need_tiebreak) {
|
||||
if (pos < kMaxTies) {
|
||||
tie_buffer[pos] = {.idx = idx, .score = local[v][e]};
|
||||
}
|
||||
} else {
|
||||
if (const auto which = pos + num_above; which < K) {
|
||||
topk_indices[which] = idx;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// 2-pass: pull the next chunk in from the staged smem buffer.
|
||||
if constexpr (kIs2Pass) {
|
||||
local[v].load(smem->score_buffer, tx + v * kBlockSize);
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (kIs2Pass) {
|
||||
#pragma unroll
|
||||
for (uint32_t v = 0; v < kVecsPerThread; ++v) {
|
||||
#pragma unroll
|
||||
for (uint32_t e = 0; e < 4; ++e) {
|
||||
const uint32_t idx =
|
||||
(tx + v * kBlockSize) * 4 + e + kMax1PassLength;
|
||||
if (idx >= length) goto LABEL_SCATTER_DONE;
|
||||
const uint32_t bin = extract_coarse_bin<kHistBits>(local[v][e]);
|
||||
if (bin > thr_bin) {
|
||||
topk_indices[atomicAdd(&smem->counter_gt, 1)] = idx;
|
||||
} else if (bin == thr_bin) {
|
||||
const auto pos = atomicAdd(&smem->counter_eq, 1);
|
||||
if (need_tiebreak) {
|
||||
if (pos < kMaxTies) {
|
||||
tie_buffer[pos] = {.idx = idx, .score = local[v][e]};
|
||||
}
|
||||
} else {
|
||||
if (const auto which = pos + num_above; which < K) {
|
||||
topk_indices[which] = idx;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
[[maybe_unused]] LABEL_SCATTER_DONE:
|
||||
if (!need_tiebreak) return;
|
||||
|
||||
// Phase 4: tie-break within the threshold bin. We assume num_ties <=
|
||||
// kBlockSize (one block of ties), so each thread takes one tied element,
|
||||
// counts the number of tied elements with strictly higher (score, -idx),
|
||||
// and writes to output if its rank is below the remaining quota.
|
||||
__syncthreads();
|
||||
static_assert(kMaxTies <= kBlockSize);
|
||||
|
||||
const uint32_t num_ties = min(num_equal, kMaxTies);
|
||||
const uint32_t topk_remain = K - num_above;
|
||||
|
||||
const auto is_greater = [](const Tie& a, const Tie& b) {
|
||||
return (a.score > b.score) || (a.score == b.score && a.idx < b.idx);
|
||||
};
|
||||
|
||||
if (num_ties <= kWarpThreads) {
|
||||
static_assert(kWarpThreads <= kNumWarps);
|
||||
if (lane_id >= num_ties || warp_id >= num_ties) return;
|
||||
const uint32_t mask = (1ull << num_ties) - 1u;
|
||||
const auto tie = tie_buffer[lane_id];
|
||||
const auto target_tie = tie_buffer[warp_id];
|
||||
const bool pred = is_greater(tie, target_tie);
|
||||
const auto rank =
|
||||
static_cast<uint32_t>(__popc(__ballot_sync(mask, pred)));
|
||||
if (lane_id == 0 && rank < topk_remain) {
|
||||
topk_indices[num_above + rank] = target_tie.idx;
|
||||
}
|
||||
} else if (num_ties <= kWarpThreads * 2) {
|
||||
// 64x64 case: each thread takes 2 elements.
|
||||
const auto lane_id_1 = lane_id + kWarpThreads;
|
||||
const auto warp_id_1 = warp_id + kWarpThreads;
|
||||
const auto invalid = Tie{.idx = 0xFFFFFFFFu, .score = -FLT_MAX};
|
||||
const auto tie_0 = tie_buffer[lane_id];
|
||||
const auto tie_1 = lane_id_1 < num_ties ? tie_buffer[lane_id_1] : invalid;
|
||||
{
|
||||
const auto target = tie_buffer[warp_id];
|
||||
const bool pred_0 = is_greater(tie_0, target);
|
||||
const bool pred_1 = is_greater(tie_1, target);
|
||||
const auto rank_0 =
|
||||
static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_0)));
|
||||
const auto rank_1 =
|
||||
static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_1)));
|
||||
const auto rank = rank_0 + rank_1;
|
||||
if (lane_id == 0 && rank < topk_remain) {
|
||||
topk_indices[num_above + rank] = target.idx;
|
||||
}
|
||||
}
|
||||
if (warp_id_1 < num_ties) {
|
||||
const auto target = tie_buffer[warp_id_1];
|
||||
const bool pred_0 = is_greater(tie_0, target);
|
||||
const bool pred_1 = is_greater(tie_1, target);
|
||||
const auto rank_0 =
|
||||
static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_0)));
|
||||
const auto rank_1 =
|
||||
static_cast<uint32_t>(__popc(__ballot_sync(0xFFFFFFFF, pred_1)));
|
||||
const auto rank = rank_0 + rank_1;
|
||||
if (lane_id == 0 && rank < topk_remain) {
|
||||
topk_indices[num_above + rank] = target.idx;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
[[unlikely]];
|
||||
// Block-wide fallback. Rarely reached.
|
||||
for (auto i = warp_id; i < num_ties; i += kNumWarps) {
|
||||
const auto target_tie = tie_buffer[i];
|
||||
uint32_t local_rank = 0;
|
||||
for (auto j = lane_id; j < num_ties; j += kWarpThreads) {
|
||||
const auto tie = tie_buffer[j];
|
||||
if (is_greater(tie, target_tie)) local_rank++;
|
||||
}
|
||||
const auto rank = warp_reduce_sum(local_rank);
|
||||
if (lane_id == 0 && rank < topk_remain) {
|
||||
topk_indices[num_above + rank] = target_tie.idx;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename TParams>
|
||||
VLLM_DSV4_DEVICE static void transform(TParams params) {
|
||||
__syncthreads();
|
||||
if (const auto tx = threadIdx.x; tx < K) params.transform(tx);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace vllm::dsv4_topk
|
||||
@@ -1,209 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Streaming top-k strategy for medium N. Uses a TMA-driven double-buffered
|
||||
// histogram pass + scatter pass over chunks of `kSizePerStage` floats.
|
||||
// Ported from
|
||||
// jit_kernel/include/sgl_kernel/deepseek_v4/topk/streaming.cuh.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "common.cuh"
|
||||
#include "ptx.cuh"
|
||||
#include "utils.cuh"
|
||||
|
||||
#include <cfloat>
|
||||
#include <cstdint>
|
||||
|
||||
namespace vllm::dsv4_topk {
|
||||
|
||||
template <uint32_t K>
|
||||
struct StreamingTopK {
|
||||
static constexpr uint32_t kHistBits = 12;
|
||||
static constexpr uint32_t kHistBins = 1 << kHistBits;
|
||||
static constexpr uint32_t kElemPerStage = 8;
|
||||
static constexpr uint32_t kSizePerStage = kElemPerStage * kBlockSize;
|
||||
static constexpr uint32_t kNumStages = 2; // double buffer
|
||||
|
||||
static constexpr uint32_t kHistItems = kHistBins / kBlockSize; // 4
|
||||
static_assert(kHistItems * kBlockSize == kHistBins);
|
||||
using HistVec = AlignedVector<uint32_t, kHistItems>;
|
||||
|
||||
struct Smem {
|
||||
// [phase = 0 (histogram) | 1 (scatter)] x [buffer = 0 | 1]
|
||||
uint64_t barrier[2][kNumStages];
|
||||
alignas(128) uint32_t counter_gt;
|
||||
alignas(128) uint32_t counter_eq;
|
||||
alignas(128) MatchBin match;
|
||||
alignas(128) uint32_t warp_sum[kNumWarps];
|
||||
union {
|
||||
uint32_t histogram[kHistBins];
|
||||
HistVec histogram_vec[kBlockSize];
|
||||
Tie tie_buffer[kMaxTies];
|
||||
};
|
||||
union {
|
||||
float score_buffer[kNumStages][kSizePerStage];
|
||||
TieHandleSmem stage2; // reused for the tie-handling phase
|
||||
};
|
||||
};
|
||||
|
||||
// length must be 4-aligned (caller rounds up); TMA wants 16-byte alignment.
|
||||
template <bool kIsScatter>
|
||||
VLLM_DSV4_DEVICE static void issue_tma(const float* scores, uint32_t stage,
|
||||
uint32_t length, Smem* smem) {
|
||||
const auto buf_idx = stage % kNumStages;
|
||||
const auto offset = stage * kSizePerStage;
|
||||
const auto size = min(kSizePerStage, length - offset);
|
||||
const auto size_bytes = size * sizeof(float);
|
||||
const auto bar = &smem->barrier[kIsScatter][buf_idx];
|
||||
ptx::tma_load(smem->score_buffer[buf_idx], scores + offset, size_bytes,
|
||||
bar);
|
||||
ptx::mbarrier_arrive_expect_tx(bar, size_bytes);
|
||||
}
|
||||
|
||||
// Unified streaming pass. kIsScatter=false: build histogram (phase A).
|
||||
// kIsScatter=true: scatter using the threshold bin (phase C). Each barrier
|
||||
// is reused across iterations via the reuse-arrive pattern.
|
||||
template <bool kIsScatter>
|
||||
VLLM_DSV4_DEVICE static void stream_pass(const float* scores, uint32_t length,
|
||||
uint32_t thr_bin,
|
||||
int32_t* s_topk_indices,
|
||||
Smem* smem) {
|
||||
const auto tx = threadIdx.x;
|
||||
const auto num_iters = (length + kSizePerStage - 1) / kSizePerStage;
|
||||
const auto lane_id = tx % kWarpThreads;
|
||||
|
||||
const auto length_aligned = (length + 3u) & ~3u;
|
||||
if (tx == 0) {
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumStages; i++) {
|
||||
if (i >= num_iters) break;
|
||||
issue_tma<kIsScatter>(scores, i, length_aligned, smem);
|
||||
}
|
||||
}
|
||||
|
||||
for (uint32_t iter = 0; iter < num_iters; iter++) {
|
||||
const auto buf_idx = iter % kNumStages;
|
||||
const auto offset = iter * kSizePerStage;
|
||||
const auto this_size = min(kSizePerStage, length - offset);
|
||||
|
||||
if (lane_id == 1) {
|
||||
const auto phase_bit = (iter / kNumStages) & 1;
|
||||
ptx::mbarrier_wait(&smem->barrier[kIsScatter][buf_idx], phase_bit);
|
||||
}
|
||||
__syncwarp();
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kElemPerStage; i++) {
|
||||
const auto local_idx = tx + i * kBlockSize;
|
||||
if (local_idx >= this_size) break;
|
||||
const auto score = smem->score_buffer[buf_idx][local_idx];
|
||||
const auto bin = extract_coarse_bin<kHistBits>(score);
|
||||
if constexpr (kIsScatter) {
|
||||
const auto global_idx = offset + local_idx;
|
||||
if (bin > thr_bin) {
|
||||
const auto pos = atomicAdd(&smem->counter_gt, 1);
|
||||
if (pos < K) s_topk_indices[pos] = global_idx;
|
||||
} else if (bin == thr_bin) {
|
||||
const auto pos = atomicAdd(&smem->counter_eq, 1);
|
||||
if (pos < kMaxTies) smem->tie_buffer[pos] = {global_idx, score};
|
||||
}
|
||||
} else {
|
||||
atomicAdd(&smem->histogram[bin], 1);
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
if (tx == 0) {
|
||||
if (const auto next_iter = iter + kNumStages; next_iter < num_iters) {
|
||||
issue_tma<kIsScatter>(scores, next_iter, length_aligned, smem);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Phase B: locate threshold bin via warp-level prefix scan.
|
||||
VLLM_DSV4_DEVICE static void find_threshold(uint32_t length, Smem* smem) {
|
||||
const auto tx = threadIdx.x;
|
||||
const auto lane_id = tx % kWarpThreads;
|
||||
const auto warp_id = tx / kWarpThreads;
|
||||
|
||||
uint32_t orig[kHistItems];
|
||||
const auto hist_vec = smem->histogram_vec[tx];
|
||||
uint32_t local_sum = 0;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kHistItems; ++i) {
|
||||
orig[i] = hist_vec[i];
|
||||
local_sum += orig[i];
|
||||
}
|
||||
|
||||
const auto warp_inc = warp_inclusive_sum(lane_id, local_sum);
|
||||
const auto warp_exc = warp_inc - local_sum;
|
||||
if (lane_id == kWarpThreads - 1) smem->warp_sum[warp_id] = warp_inc;
|
||||
__syncthreads();
|
||||
|
||||
const auto tmp = smem->warp_sum[lane_id];
|
||||
uint32_t prefix_sum = warp_reduce_sum(lane_id < warp_id ? tmp : 0);
|
||||
prefix_sum += warp_exc;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kHistItems; ++i) {
|
||||
prefix_sum += orig[i];
|
||||
const auto above = length - prefix_sum;
|
||||
if (above < K && above + orig[i] >= K) {
|
||||
smem->match = {
|
||||
.bin = tx * kHistItems + i,
|
||||
.above_count = above,
|
||||
.equal_count = orig[i],
|
||||
};
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
VLLM_DSV4_DEVICE static void run(const float* scores, uint32_t length,
|
||||
int32_t* topk_indices, void* _smem) {
|
||||
const auto smem = static_cast<Smem*>(_smem);
|
||||
const auto tx = threadIdx.x;
|
||||
__builtin_assume(tx < kBlockSize);
|
||||
|
||||
{
|
||||
HistVec zero;
|
||||
zero.fill(0);
|
||||
smem->histogram_vec[tx] = zero;
|
||||
if (tx < 2 * kNumStages) {
|
||||
const auto base_barrier = &smem->barrier[0][0];
|
||||
ptx::mbarrier_init(&base_barrier[tx], 1);
|
||||
}
|
||||
if (tx == 0) {
|
||||
smem->counter_gt = 0;
|
||||
smem->counter_eq = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Phase A: histogram.
|
||||
stream_pass<false>(scores, length, 0, nullptr, smem);
|
||||
|
||||
// Phase B: threshold bin.
|
||||
find_threshold(length, smem);
|
||||
|
||||
// Phase C: scatter.
|
||||
stream_pass<true>(scores, length, smem->match.bin, topk_indices, smem);
|
||||
}
|
||||
|
||||
template <typename TParams>
|
||||
VLLM_DSV4_DEVICE static void transform(TParams params, void* _smem) {
|
||||
// Phase D: page-translate above entries, then refine ties.
|
||||
const auto smem = static_cast<Smem*>(_smem);
|
||||
const auto tx = threadIdx.x;
|
||||
const auto num_above = smem->match.above_count;
|
||||
if (tx < num_above) params.transform(tx);
|
||||
const auto num_equal = smem->counter_eq;
|
||||
if (num_above >= K || num_equal == 0) return;
|
||||
const auto clamped_ties = min(num_equal, kMaxTies);
|
||||
tie_handle_transform(smem->tie_buffer, clamped_ties, num_above, K, params,
|
||||
&smem->stage2);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace vllm::dsv4_topk
|
||||
@@ -1,75 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Minimal device-side utilities used by the DeepSeek V4 indexer top-k port.
|
||||
// Replaces sgl_kernel/{utils,warp,vec,type}.cuh — we only need the bits the
|
||||
// top-k kernels actually touch.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
namespace vllm::dsv4_topk {
|
||||
|
||||
#define VLLM_DSV4_DEVICE __forceinline__ __device__
|
||||
|
||||
inline constexpr uint32_t kWarpThreads = 32u;
|
||||
inline constexpr uint32_t kFullMask = 0xffffffffu;
|
||||
|
||||
// Programmatic Dependent Launch (sm_90+). When enabled, the kernel waits for
|
||||
// the predecessor on the same stream to advance past its dependents-launch
|
||||
// trigger before doing anything memory-dependent. Used to overlap the
|
||||
// fp8_paged_mqa_logits epilogue with the first stage of top-k.
|
||||
template <bool kUsePDL>
|
||||
VLLM_DSV4_DEVICE void pdl_wait_primary() {
|
||||
if constexpr (kUsePDL) {
|
||||
asm volatile("griddepcontrol.wait;" ::: "memory");
|
||||
}
|
||||
}
|
||||
|
||||
template <bool kUsePDL>
|
||||
VLLM_DSV4_DEVICE void pdl_trigger_secondary() {
|
||||
if constexpr (kUsePDL) {
|
||||
asm volatile("griddepcontrol.launch_dependents;" :::);
|
||||
}
|
||||
}
|
||||
|
||||
// Warp-level XOR-shuffle reduce. kThreads must be a power of 2 and <= 32.
|
||||
template <uint32_t kThreads = kWarpThreads, typename T>
|
||||
VLLM_DSV4_DEVICE T warp_reduce_sum(T value, uint32_t active_mask = kFullMask) {
|
||||
#pragma unroll
|
||||
for (auto offset = kThreads >> 1; offset > 0; offset >>= 1) {
|
||||
value = value + __shfl_xor_sync(active_mask, value, offset, 32);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
// 128-bit-aligned vector of N elements of T (N must be a power of 2, total
|
||||
// size <= 16 bytes). Used for vectorized loads/stores into shared memory.
|
||||
template <typename T, std::size_t N>
|
||||
struct alignas(sizeof(T) * N) AlignedVector {
|
||||
static_assert(N > 0 && (N & (N - 1)) == 0, "N must be a power of two");
|
||||
static_assert(sizeof(T) * N <= 16,
|
||||
"AlignedVector exceeds the 128-bit CUDA vector limit");
|
||||
|
||||
T data[N];
|
||||
|
||||
VLLM_DSV4_DEVICE void load(const void* ptr, std::size_t offset = 0) {
|
||||
*reinterpret_cast<AlignedVector*>(this) =
|
||||
reinterpret_cast<const AlignedVector*>(ptr)[offset];
|
||||
}
|
||||
VLLM_DSV4_DEVICE void store(void* ptr, std::size_t offset = 0) const {
|
||||
reinterpret_cast<AlignedVector*>(ptr)[offset] = *this;
|
||||
}
|
||||
VLLM_DSV4_DEVICE void fill(T value) {
|
||||
#pragma unroll
|
||||
for (std::size_t i = 0; i < N; ++i) data[i] = value;
|
||||
}
|
||||
VLLM_DSV4_DEVICE T& operator[](std::size_t i) { return data[i]; }
|
||||
VLLM_DSV4_DEVICE const T& operator[](std::size_t i) const { return data[i]; }
|
||||
};
|
||||
|
||||
} // namespace vllm::dsv4_topk
|
||||
@@ -137,15 +137,18 @@ fused_add_rms_norm_static_fp8_quant_kernel(
|
||||
_f16Vec<scalar_t, width> res = residual_v[id];
|
||||
_f16Vec<scalar_t, width> w = weight_v[idx];
|
||||
using Converter = _typeConvert<scalar_t>;
|
||||
using HipT = typename Converter::hip_type;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < width; ++i) {
|
||||
float x = Converter::convert(res.data[i]);
|
||||
float wf = Converter::convert(w.data[i]);
|
||||
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
|
||||
// to match the unfused composite path at FP8 boundaries.
|
||||
scalar_t out_norm = Converter::convert(x * s_variance * wf);
|
||||
// to match the unfused composite path at FP8 boundaries. We use the
|
||||
// backend's hip_type for the intermediate since c10::Half/BFloat16 has
|
||||
// ambiguous conversions on CUDA and no implicit conversion on ROCm.
|
||||
HipT out_norm_h = Converter::convert(x * s_variance * wf);
|
||||
out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
|
||||
static_cast<float>(out_norm), scale_inv);
|
||||
Converter::convert(out_norm_h), scale_inv);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+2
-34
@@ -125,40 +125,6 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
|
||||
torch::Tensor& output, torch::Tensor& workspace, int64_t k,
|
||||
int64_t max_seq_len);
|
||||
|
||||
// DeepSeek V4 indexer top-k (k = 512). Hopper (sm_90a) and Blackwell
|
||||
// datacenter (sm_100/sm_103) — needs thread-block clusters, TMA, and PDL.
|
||||
// Two-step API:
|
||||
// 1. fast_topk_v2_plan inspects the seq_lens distribution and writes a
|
||||
// cluster_threshold + per-row Metadata into a (B+1, 4) int32 tensor. The
|
||||
// plan is amortized when cudagraph-captured: once per shape, reused across
|
||||
// layers.
|
||||
// 2. fast_topk_v2 selects the top-512 indices per row, folds the page-table
|
||||
// gather into the radix store, and writes (B, 512) int32 page indices.
|
||||
// Dispatches per row to one of three strategies (Register / Streaming /
|
||||
// Cluster) using the planned threshold.
|
||||
//
|
||||
// Returns the size in int32s of the per-row workspace required by
|
||||
// fast_topk_v2 (allocate `(B, fast_topk_v2_workspace_ints())` int32 contig).
|
||||
void fast_topk_v2_plan(const torch::Tensor& seq_lens, torch::Tensor& metadata,
|
||||
int64_t static_cluster_threshold);
|
||||
|
||||
void fast_topk_v2(const torch::Tensor& scores, const torch::Tensor& seq_lens,
|
||||
const torch::Tensor& page_table, torch::Tensor& page_indices,
|
||||
int64_t page_size, const torch::Tensor& workspace,
|
||||
const torch::Tensor& metadata, int64_t topk);
|
||||
|
||||
// Top-k only, no page-table fold-in. Same selection as fast_topk_v2 but
|
||||
// emits raw row-local indices into ``topk_indices`` (drop-in for
|
||||
// persistent_topk's output contract). topk must be one of {512, 1024}.
|
||||
void fast_topk_v2_raw(const torch::Tensor& scores,
|
||||
const torch::Tensor& seq_lens,
|
||||
torch::Tensor& topk_indices,
|
||||
const torch::Tensor& workspace,
|
||||
const torch::Tensor& metadata,
|
||||
int64_t topk);
|
||||
|
||||
int64_t fast_topk_v2_workspace_ints();
|
||||
|
||||
void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input,
|
||||
torch::Tensor& weight, torch::Tensor& scale,
|
||||
double epsilon);
|
||||
@@ -197,6 +163,8 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
|
||||
|
||||
void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
|
||||
|
||||
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
|
||||
|
||||
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
|
||||
torch::Tensor& scale);
|
||||
|
||||
|
||||
+6
-20
@@ -106,6 +106,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
|
||||
ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
|
||||
|
||||
// SwiGLU activation with input clamping.
|
||||
ops.def(
|
||||
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
|
||||
"-> ()");
|
||||
ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
|
||||
|
||||
ops.def(
|
||||
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
|
||||
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
|
||||
@@ -215,26 +221,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"Tensor workspace, int k, int max_seq_len) -> ()");
|
||||
ops.impl("persistent_topk", torch::kCUDA, &persistent_topk);
|
||||
|
||||
// DeepSeek V4 indexer top-k (k=512), ported from sglang's topk_v2 family.
|
||||
// Built for sm_90a (Hopper) + sm_100a/sm_103 (Blackwell datacenter).
|
||||
// Schema only here; impl is registered in csrc/deepseek_v4/fast_topk_v2.cu
|
||||
// so it's only present when CMake compiles the source for a supported arch.
|
||||
ops.def(
|
||||
"fast_topk_v2_plan(Tensor seq_lens, Tensor! metadata, "
|
||||
"int static_cluster_threshold) -> ()");
|
||||
|
||||
ops.def(
|
||||
"fast_topk_v2(Tensor scores, Tensor seq_lens, Tensor page_table, "
|
||||
"Tensor! page_indices, int page_size, Tensor workspace, "
|
||||
"Tensor metadata, int topk) -> ()");
|
||||
|
||||
ops.def(
|
||||
"fast_topk_v2_raw(Tensor scores, Tensor seq_lens, "
|
||||
"Tensor! topk_indices, Tensor workspace, Tensor metadata, int topk)"
|
||||
" -> ()");
|
||||
|
||||
ops.def("fast_topk_v2_workspace_ints() -> int");
|
||||
|
||||
// Layernorm-quant
|
||||
// Apply Root Mean Square (RMS) Normalization to the input tensor.
|
||||
ops.def(
|
||||
|
||||
+12
-5
@@ -538,9 +538,11 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
|
||||
cuda-nvrtc-${CUDA_VERSION_DASH} \
|
||||
cuda-cuobjdump-${CUDA_VERSION_DASH} \
|
||||
libcurand-dev-${CUDA_VERSION_DASH} \
|
||||
libcublas-${CUDA_VERSION_DASH} \
|
||||
libcublas-dev-${CUDA_VERSION_DASH} \
|
||||
# Required by fastsafetensors (fixes #20384)
|
||||
libnuma-dev && \
|
||||
libnuma-dev \
|
||||
# numactl CLI for NUMA binding at runtime
|
||||
numactl && \
|
||||
# Fixes nccl_allocator requiring nccl.h at runtime
|
||||
# https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
|
||||
# NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
|
||||
@@ -583,9 +585,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
ARG FLASHINFER_VERSION=0.6.8.post1
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
|
||||
--extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
|
||||
&& flashinfer show-config \
|
||||
&& flashinfer download-cubin
|
||||
--extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
|
||||
|
||||
# ============================================================
|
||||
# OPENAI API SERVER DEPENDENCIES
|
||||
@@ -667,6 +667,13 @@ RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm
|
||||
uv pip install --system ep_kernels/dist/*.whl --verbose \
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
|
||||
|
||||
# Download FlashInfer precompiled cubins AFTER all pip installs are done.
|
||||
# This must run after the vLLM wheel and EP kernels installs above, because
|
||||
# those can reinstall/touch flashinfer packages. Downloading cubins earlier
|
||||
# (in the flashinfer-jit-cache layer) causes ~2.5 GB of layer duplication
|
||||
# when a later pip install overwrites flashinfer package files.
|
||||
RUN flashinfer show-config && flashinfer download-cubin
|
||||
|
||||
# CUDA image changed from /usr/local/nvidia to /usr/local/cuda in 12.8 but will
|
||||
# return to /usr/local/nvidia in 13.0 to allow container providers to mount drivers
|
||||
# consistently from the host (see https://github.com/vllm-project/vllm/issues/18859).
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 156 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 182 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 193 KiB |
+2
-2
@@ -163,7 +163,7 @@ Running with a local file:
|
||||
|
||||
```bash
|
||||
vllm run-batch \
|
||||
-i offline_inference/openai_batch/openai_example_batch.jsonl \
|
||||
-i features/openai_batch/openai_example_batch.jsonl \
|
||||
-o results.jsonl \
|
||||
--model meta-llama/Meta-Llama-3-8B-Instruct
|
||||
```
|
||||
@@ -172,7 +172,7 @@ Using remote file:
|
||||
|
||||
```bash
|
||||
vllm run-batch \
|
||||
-i https://raw.githubusercontent.com/vllm-project/vllm/main/examples/offline_inference/openai_batch/openai_example_batch.jsonl \
|
||||
-i https://raw.githubusercontent.com/vllm-project/vllm/main/examples/features/openai_batch/openai_example_batch.jsonl \
|
||||
-o results.jsonl \
|
||||
--model meta-llama/Meta-Llama-3-8B-Instruct
|
||||
```
|
||||
|
||||
@@ -23,7 +23,7 @@ llm = LLM(model="ibm-granite/granite-3.1-8b-instruct", tensor_parallel_size=2)
|
||||
!!! note
|
||||
With tensor parallelism enabled, each process will read the whole model and split it into chunks, which makes the disk reading time even longer (proportional to the size of tensor parallelism).
|
||||
|
||||
You can convert the model checkpoint to a sharded checkpoint using [examples/offline_inference/save_sharded_state.py](../../examples/offline_inference/save_sharded_state.py). The conversion process might take some time, but later you can load the sharded checkpoint much faster. The model loading time should remain constant regardless of the size of tensor parallelism.
|
||||
You can convert the model checkpoint to a sharded checkpoint using [examples/features/sharded_state/load_sharded_state_offline.py](../../examples/features/sharded_state/load_sharded_state_offline.py). The conversion process might take some time, but later you can load the sharded checkpoint much faster. The model loading time should remain constant regardless of the size of tensor parallelism.
|
||||
|
||||
## Quantization
|
||||
|
||||
|
||||
@@ -42,7 +42,7 @@ Traces can be visualized using <https://ui.perfetto.dev/>.
|
||||
|
||||
#### Offline Inference
|
||||
|
||||
Refer to [examples/offline_inference/simple_profiling.py](../../examples/offline_inference/simple_profiling.py) for an example.
|
||||
Refer to [examples/features/profiling/simple_profiling_offline.py](../../examples/features/profiling/simple_profiling_offline.py) for an example.
|
||||
|
||||
#### OpenAI Server
|
||||
|
||||
|
||||
@@ -36,7 +36,7 @@ th {
|
||||
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
|
||||
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
|
||||
|
||||
!!! info "Table key"
|
||||
1. All types: mxfp4, nvfp4, int4, int8, fp8
|
||||
|
||||
@@ -11,7 +11,7 @@ Automatic Prefix Caching (APC in short) caches the KV cache of existing queries,
|
||||
|
||||
Set `enable_prefix_caching=True` in vLLM engine to enable APC. Here is an example:
|
||||
|
||||
[examples/offline_inference/automatic_prefix_caching.py](../../examples/offline_inference/automatic_prefix_caching.py)
|
||||
[examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py](../../examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py)
|
||||
|
||||
## Example workloads
|
||||
|
||||
|
||||
@@ -6,12 +6,12 @@ This directory contains examples for extending the context length of models usin
|
||||
|
||||
## Offline Inference Example
|
||||
|
||||
The [`context_extension.py`](../../examples/offline_inference/context_extension) script demonstrates how to extend the context length of a Qwen model using the YARN method (rope_parameters) and run a simple chat example.
|
||||
The [`context_extension.py`](../../examples/features/context_extension/context_extension_offline.py) script demonstrates how to extend the context length of a Qwen model using the YARN method (rope_parameters) and run a simple chat example.
|
||||
|
||||
### Usage
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/context_extension.py
|
||||
python examples/features/context_extension/context_extension_offline.py
|
||||
```
|
||||
|
||||
## OpenAI Online Method
|
||||
|
||||
@@ -47,7 +47,7 @@ the third parameter is the path to the LoRA adapter.
|
||||
)
|
||||
```
|
||||
|
||||
Check out [examples/offline_inference/multilora_inference.py](../../examples/offline_inference/multilora_inference.py) for an example of how to use LoRA adapters with the async engine and how to use more advanced configuration options.
|
||||
Check out [examples/features/lora/multilora_offline.py](../../examples/features/lora/multilora_offline.py) for an example of how to use LoRA adapters with the async engine and how to use more advanced configuration options.
|
||||
|
||||
## Serving LoRA Adapters
|
||||
|
||||
|
||||
@@ -68,7 +68,7 @@ You can pass a single image to the `'image'` field of the multi-modal dictionary
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
Full example: [examples/offline_inference/vision_language.py](../../examples/offline_inference/vision_language.py)
|
||||
Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
|
||||
|
||||
To substitute multiple images inside the same text prompt, you can pass in a list of images instead:
|
||||
|
||||
@@ -101,7 +101,7 @@ To substitute multiple images inside the same text prompt, you can pass in a lis
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
Full example: [examples/offline_inference/vision_language_multi_image.py](../../examples/offline_inference/vision_language_multi_image.py)
|
||||
Full example: [examples/generate/multimodal/vision_language_multi_image_offline.py](../../examples/generate/multimodal/vision_language_multi_image_offline.py)
|
||||
|
||||
If using the [LLM.chat](../models/generative_models.md#llmchat) method, you can pass images directly in the message content using various formats: image URLs, PIL Image objects, or pre-computed embeddings:
|
||||
|
||||
@@ -287,13 +287,13 @@ Instead of NumPy arrays, you can also pass `'torch.Tensor'` instances, as shown
|
||||
!!! note
|
||||
'process_vision_info' is only applicable to Qwen2.5-VL and similar models.
|
||||
|
||||
Full example: [examples/offline_inference/vision_language.py](../../examples/offline_inference/vision_language.py)
|
||||
Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
|
||||
|
||||
### Audio Inputs
|
||||
|
||||
You can pass a tuple `(array, sampling_rate)` to the `'audio'` field of the multi-modal dictionary.
|
||||
|
||||
Full example: [examples/offline_inference/audio_language.py](../../examples/offline_inference/audio_language.py)
|
||||
Full example: [examples/generate/multimodal/audio_language_offline.py](../../examples/generate/multimodal/audio_language_offline.py)
|
||||
|
||||
#### Chunking Long Audio for Transcription
|
||||
|
||||
@@ -674,7 +674,7 @@ Then, you can use the OpenAI client as follows:
|
||||
print("Chat completion output:", chat_response.choices[0].message.content)
|
||||
```
|
||||
|
||||
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
|
||||
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
|
||||
|
||||
!!! tip
|
||||
Loading from local file paths is also supported on vLLM: You can specify the allowed local media path via `--allowed-local-media-path` when launching the API server/engine,
|
||||
@@ -745,7 +745,7 @@ Then, you can use the OpenAI client as follows:
|
||||
print("Chat completion output from image url:", result)
|
||||
```
|
||||
|
||||
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
|
||||
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
|
||||
|
||||
!!! note
|
||||
By default, the timeout for fetching videos through HTTP URL is `30` seconds.
|
||||
@@ -958,7 +958,7 @@ Alternatively, you can pass `audio_url`, which is the audio counterpart of `imag
|
||||
print("Chat completion output from audio url:", result)
|
||||
```
|
||||
|
||||
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
|
||||
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
|
||||
|
||||
!!! note
|
||||
By default, the timeout for fetching audios through HTTP URL is `10` seconds.
|
||||
|
||||
@@ -16,7 +16,7 @@ To input multi-modal data, follow this schema in [vllm.inputs.EmbedsPrompt][]:
|
||||
|
||||
You can pass prompt embeddings from Hugging Face Transformers models to the `'prompt_embeds'` field of the prompt embedding dictionary, as shown in the following examples:
|
||||
|
||||
[examples/offline_inference/prompt_embed_inference.py](../../examples/offline_inference/prompt_embed_inference.py)
|
||||
[examples/features/prompt_embed/prompt_embed_offline.py](../../examples/features/prompt_embed/prompt_embed_offline.py)
|
||||
|
||||
## Online Serving
|
||||
|
||||
@@ -41,4 +41,4 @@ vllm serve meta-llama/Llama-3.2-1B-Instruct --runner generate \
|
||||
|
||||
Then, you can use the OpenAI client as follows:
|
||||
|
||||
[examples/online_serving/prompt_embed_inference_with_openai_client.py](../../examples/online_serving/prompt_embed_inference_with_openai_client.py)
|
||||
[examples/features/prompt_embed/prompt_embed_inference_with_openai_client.py](../../examples/features/prompt_embed/prompt_embed_inference_with_openai_client.py)
|
||||
|
||||
@@ -20,6 +20,7 @@ The following are the supported quantization formats for vLLM:
|
||||
- [AMD Quark](quark.md)
|
||||
- [Quantized KV Cache](quantized_kvcache.md)
|
||||
- [TorchAO](torchao.md)
|
||||
- [FP8 ViT Encoder Attention](fp8_vit_attn.md)
|
||||
|
||||
## Supported Hardware
|
||||
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
# FP8 ViT Encoder Attention
|
||||
|
||||
For visual understanding workloads with large images (e.g. QHD, 4K) and relatively
|
||||
short text prompts/generation, the ViT encoder attention can become a significant
|
||||
bottleneck, especially when the text model is quantized (e.g. NVFP4). vLLM
|
||||
supports optional FP8 quantization for the ViT encoder attention via the
|
||||
FlashInfer cuDNN backend. Q/K/V are quantized on-the-fly to FP8 before the
|
||||
cuDNN attention call.
|
||||
|
||||
!!! note
|
||||
- Currently supports Qwen3-VL family models only (`qwen3_vl`, `qwen3_vl_moe`,
|
||||
`qwen3_5`, `qwen3_5_moe`, and other models using Qwen3 ViT).
|
||||
- Dynamic scaling is not compatible with ViT full CUDA graphs.
|
||||
- Performance gains are mostly visible at QHD/4K resolutions or multi-image
|
||||
requests. Smaller images may see no speedup due to quantization overhead
|
||||
(3 quantization kernel launches + un-padding).
|
||||
- FP8 tensor-core speedup is more pronounced on GB300 than GB200.
|
||||
|
||||
## Requirements
|
||||
|
||||
- FlashInfer cuDNN backend with cuDNN >= 9.17.1.
|
||||
|
||||
## Usage
|
||||
|
||||
Enable FP8 ViT attention by passing `--mm-encoder-attn-dtype fp8` together
|
||||
with `--mm-encoder-attn-backend FLASHINFER`:
|
||||
|
||||
```bash
|
||||
vllm serve $MODEL \
|
||||
--mm-encoder-attn-backend FLASHINFER \
|
||||
--mm-encoder-attn-dtype fp8
|
||||
```
|
||||
|
||||
By default (no scale file), **dynamic scaling** is used: a 16-entry circular
|
||||
buffer of observed Q/K/V amax values drives per-forward scale updates. This
|
||||
matches BF16 accuracy without any calibration but adds a small per-forward
|
||||
overhead.
|
||||
|
||||
## Calibrate-Once, Reuse Workflow (Recommended)
|
||||
|
||||
For production, calibrate static scales on a representative dataset once and
|
||||
reuse them to avoid the dynamic overhead:
|
||||
|
||||
```bash
|
||||
# Step 1: calibrate and save scales (runs dynamic scaling for 16 passes,
|
||||
# then dumps the learned scales to JSON).
|
||||
vllm bench mm-processor \
|
||||
--model $MODEL --mm-encoder-attn-backend FLASHINFER \
|
||||
--mm-encoder-attn-dtype fp8 \
|
||||
--mm-encoder-fp8-scale-save-path /path/to/scales.json \
|
||||
--dataset-name hf --dataset-path lmarena-ai/VisionArena-Chat \
|
||||
--num-prompts 100
|
||||
|
||||
# Step 2: serve with static scales (no dynamic overhead).
|
||||
vllm serve $MODEL \
|
||||
--mm-encoder-attn-backend FLASHINFER \
|
||||
--mm-encoder-attn-dtype fp8 \
|
||||
--mm-encoder-fp8-scale-path /path/to/scales.json
|
||||
```
|
||||
|
||||
Saved scales are multiplied by `--mm-encoder-fp8-scale-save-margin` (default
|
||||
`1.5`) to leave headroom against activation outliers not present in the
|
||||
calibration set. The default has been validated to generalize across datasets
|
||||
(e.g. VisionArena-Chat calibration maintains BF16 accuracy on ChartQA).
|
||||
|
||||
## Scale File Format
|
||||
|
||||
```json
|
||||
{
|
||||
"visual.blocks.0.attn.attn": {"q": 224.0, "k": 198.0, "v": 210.0},
|
||||
"visual.blocks.1.attn.attn": {"q": 218.0, "k": 195.0, "v": 207.0}
|
||||
}
|
||||
```
|
||||
|
||||
Keys `q_scale` / `k_scale` / `v_scale` are accepted as aliases.
|
||||
|
||||
## Performance
|
||||
|
||||
**Core cuDNN attention kernel** (PyTorch profiler, `cudnn_generated_fort_native_sdpa_sm100_flash_fprop`, head_dim=128, seq_len=8192):
|
||||
|
||||
| Hardware | BF16 | FP8 | Speedup |
|
||||
| -------- | ---- | ---- | ------- |
|
||||
| GB200 | 350 us | 312 us | **1.12x** |
|
||||
| GB300 | 300 us | 211 us | **1.42x** |
|
||||
|
||||
**End-to-end encoder forward time** (Qwen3-VL-30B-A3B-Instruct on GB200, 3 images/request):
|
||||
|
||||
| Resolution | BF16 median | FP8 median | Speedup |
|
||||
| ---------- | ----------- | ---------- | ------- |
|
||||
| HD (720x1280) | 31.77 ms | 36.39 ms | 0.87x |
|
||||
| FullHD (1080x1920) | 57.99 ms | 58.73 ms | ~same |
|
||||
| QHD (1440x2560) | 131.83 ms | 122.30 ms | **1.08x** |
|
||||
| 4K (2160x3840) | 543.44 ms | 460.31 ms | **1.18x** |
|
||||
|
||||
Crossover is around FullHD with 3 images/request. At QHD and above, FP8 wins.
|
||||
|
||||
## Accuracy
|
||||
|
||||
ChartQA, Qwen3-VL-8B-Instruct, 500 samples. FP8 static uses scales calibrated
|
||||
on VisionArena-Chat (with default 1.5x margin):
|
||||
|
||||
| Metric | BF16 | FP8 dynamic | FP8 static |
|
||||
| ------ | ---- | ----------- | ---------- |
|
||||
| relaxed_accuracy | 0.780 | 0.776 | 0.780 |
|
||||
| anywhere_accuracy | 0.806 | 0.816 | 0.814 |
|
||||
| exact_match | 0.584 | 0.582 | 0.578 |
|
||||
|
||||
All three configurations match within statistical noise, confirming that
|
||||
static scales calibrated on one dataset generalize to another.
|
||||
@@ -13,6 +13,7 @@ vLLM currently supports the following reasoning models:
|
||||
|
||||
| Model Series | Parser Name | Structured Output Support | Tool Calling |
|
||||
| ------------ | ----------- | ---------------- | ----------- |
|
||||
| [Cohere Command A Reasoning](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) | `cohere_command3` | `json`, `regex` | ✅ |
|
||||
| [DeepSeek R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d) | `deepseek_r1` | `json`, `regex` | ❌ |
|
||||
| [DeepSeek-V3.1](https://huggingface.co/collections/deepseek-ai/deepseek-v31-68a491bed32bd77e7fca048f) | `deepseek_v3` | `json`, `regex` | ❌ |
|
||||
| [ERNIE-4.5-VL series](https://huggingface.co/baidu/ERNIE-4.5-VL-28B-A3B-PT) | `ernie45` | `json`, `regex` | ❌ |
|
||||
@@ -202,7 +203,7 @@ The reasoning content is also available when both tool calling and the reasoning
|
||||
print(f"Arguments: {tool_call.arguments}")
|
||||
```
|
||||
|
||||
For more examples, please refer to [examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py](../../examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py).
|
||||
For more examples, please refer to [examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py](../../examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py).
|
||||
|
||||
## Server-Level Default Chat Template Kwargs
|
||||
|
||||
|
||||
@@ -32,7 +32,7 @@ depend on your model family, traffic pattern, hardware, and sampling settings.
|
||||
| Suffix decoding | Low to medium gain | Medium gain | No extra draft model; dynamic speculation depth. |
|
||||
|
||||
For reproducible measurements in your environment, use
|
||||
[`examples/offline_inference/spec_decode.py`](../../../examples/offline_inference/spec_decode.py)
|
||||
[`examples/features/speculative_decoding/spec_decode_offline.py`](../../../examples/features/speculative_decoding/spec_decode_offline.py)
|
||||
or the [benchmark CLI guide](../../benchmarking/cli.md).
|
||||
|
||||
## `--speculative-config` schema
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# EAGLE Draft Models
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by an [EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)](https://arxiv.org/pdf/2401.15077) based draft model. A more detailed example for offline mode, including how to extract request level acceptance rate, can be found in [examples/offline_inference/spec_decode.py](../../../examples/offline_inference/spec_decode.py)
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by an [EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)](https://arxiv.org/pdf/2401.15077) based draft model. A more detailed example for offline mode, including how to extract request level acceptance rate, can be found in [examples/features/speculative_decoding/spec_decode_offline.py](../../../examples/features/speculative_decoding/spec_decode_offline.py)
|
||||
|
||||
## Eagle Drafter Example
|
||||
|
||||
|
||||
@@ -165,7 +165,7 @@ As an example, we can use to define a specific format of simplified SQL queries:
|
||||
print(completion.choices[0].message.content)
|
||||
```
|
||||
|
||||
See also: [full example](../examples/online_serving/structured_outputs.md)
|
||||
See also: [full example](../../examples/features/structured_outputs/README.md)
|
||||
|
||||
## Reasoning Outputs
|
||||
|
||||
@@ -208,7 +208,7 @@ Note that you can use reasoning with any provided structured outputs feature. Th
|
||||
print("content: ", completion.choices[0].message.content)
|
||||
```
|
||||
|
||||
See also: [full example](../examples/online_serving/structured_outputs.md)
|
||||
See also: [full example](../../examples/features/structured_outputs/README.md)
|
||||
|
||||
!!! note
|
||||
When using Qwen3 Coder models with reasoning enabled, structured outputs might become disabled if the reasoning content does not get parsed into the `reasoning` field separately (v0.11.2+).
|
||||
@@ -304,7 +304,7 @@ Step #2: explanation="Next, let's isolate 'x' by dividing both sides of the equa
|
||||
Answer: x = -29/8
|
||||
```
|
||||
|
||||
An example of using `structural_tag` can be found here: [examples/online_serving/structured_outputs](../../examples/online_serving/structured_outputs)
|
||||
An example of using `structural_tag` can be found here: [examples/features/structured_outputs](../../examples/features/structured_outputs/README.md)
|
||||
|
||||
## Offline Inference
|
||||
|
||||
@@ -339,4 +339,4 @@ shown below:
|
||||
print(outputs[0].outputs[0].text)
|
||||
```
|
||||
|
||||
See also: [full example](../examples/online_serving/structured_outputs.md)
|
||||
See also: [full example](../../examples/features/structured_outputs/structured_outputs_offline.py)
|
||||
|
||||
@@ -369,6 +369,16 @@ Flags:
|
||||
* For non-reasoning: `--tool-call-parser hunyuan_a13b`
|
||||
* For reasoning: `--tool-call-parser hunyuan_a13b --reasoning-parser hunyuan_a13b`
|
||||
|
||||
### Cohere Command A Reasoning (`cohere_command3`)
|
||||
|
||||
Supported models:
|
||||
|
||||
* [`CohereLabs/command-a-reasoning-08-2025`](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025)
|
||||
|
||||
Flags: `--tool-call-parser cohere_command3 --reasoning-parser cohere_command3`
|
||||
|
||||
Note: the Cohere tool parser requires the `cohere_melody` package, which is not installed by default. Before using this parser please install the [cohere_melody](https://pypi.org/project/cohere-melody/) package.
|
||||
|
||||
### LongCat-Flash-Chat Models (`longcat`)
|
||||
|
||||
Supported models:
|
||||
|
||||
@@ -101,7 +101,7 @@ vllm serve /path/to/sharded/model \
|
||||
--model-loader-extra-config '{"pattern":"custom-model-rank-{rank}-part-{part}.safetensors"}'
|
||||
```
|
||||
|
||||
To create sharded model files, you can use the script provided in [examples/offline_inference/save_sharded_state.py](../../../examples/offline_inference/save_sharded_state.py). This script demonstrates how to save a model in the sharded format that is compatible with the Run:ai Model Streamer sharded loader.
|
||||
To create sharded model files, you can use the script provided in [examples/features/sharded_state/save_sharded_state_offline.py](../../../examples/features/sharded_state/save_sharded_state_offline.py). This script demonstrates how to save a model in the sharded format that is compatible with the Run:ai Model Streamer sharded loader.
|
||||
|
||||
The sharded loader supports all the same tunable parameters as the regular Run:ai Model Streamer, including `concurrency` and `memory_limit`. These can be configured in the same way:
|
||||
|
||||
|
||||
@@ -439,6 +439,7 @@ th {
|
||||
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
|
||||
| `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
|
||||
| `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. | | |
|
||||
@@ -590,6 +591,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | | ✅︎ |
|
||||
| `LlavaOnevisionForConditionalGeneration` | LLaVA-Onevision | T + I<sup>+</sup> + V<sup>+</sup> | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc. | | ✅︎ |
|
||||
| `MiDashengLMModel` | MiDashengLM | T + A<sup>+</sup> | `mispeech/midashenglm-7b` | | ✅︎ |
|
||||
| `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. | ✅︎ | |
|
||||
| `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
|
||||
|
||||
@@ -16,7 +16,7 @@ For MoE models, when any requests are in progress in any rank, we must ensure th
|
||||
|
||||
In all cases, it is beneficial to load-balance requests between DP ranks. For online deployments, this balancing can be optimized by taking into account the state of each DP engine - in particular its currently scheduled and waiting (queued) requests, and KV cache state. Each DP engine has an independent KV cache, and the benefit of prefix caching can be maximized by directing prompts intelligently.
|
||||
|
||||
This document focuses on online deployments (with the API server). DP + EP is also supported for offline usage (via the LLM class), for an example see [examples/offline_inference/data_parallel.py](../../examples/offline_inference/data_parallel.py).
|
||||
This document focuses on online deployments (with the API server). DP + EP is also supported for offline usage (via the LLM class), for an example see [examples/features/data_parallel/data_parallel_offline.py](../../examples/features/data_parallel/data_parallel_offline.py).
|
||||
|
||||
There are two distinct modes supported for online deployments - self-contained with internal load balancing, or externally per-rank process deployment and load balancing.
|
||||
|
||||
|
||||
@@ -251,7 +251,7 @@ The following extra parameters are supported:
|
||||
Our Responses API is compatible with [OpenAI's Responses API](https://platform.openai.com/docs/api-reference/responses);
|
||||
you can use the [official OpenAI Python client](https://github.com/openai/openai-python) to interact with it.
|
||||
|
||||
Code example: [examples/online_serving/openai_responses_client_with_tools.py](../../examples/online_serving/openai_responses_client_with_tools.py)
|
||||
Code example: [examples/online_serving/openai_responses_client_with_tools.py](../../examples/tool_calling/openai_responses_client_with_tools.py)
|
||||
|
||||
#### Extra parameters
|
||||
|
||||
@@ -279,7 +279,7 @@ you can use the [official OpenAI Python client](https://github.com/openai/openai
|
||||
!!! note
|
||||
To use the Transcriptions API, please install with extra audio dependencies using `pip install vllm[audio]`.
|
||||
|
||||
Code example: [examples/online_serving/openai_transcription_client.py](../../examples/online_serving/openai_transcription_client.py)
|
||||
Code example: [examples/speech_to_text/openai/openai_transcription_client.py](../../examples/speech_to_text/openai/openai_transcription_client.py)
|
||||
|
||||
NOTE: beam search is currently supported in the transcriptions endpoint for encoder-decoder multimodal models, e.g., whisper, but highly inefficient as work for handling the encoder/decoder cache is actively ongoing. This is an active point of ongoing optimization and will be handled properly in the very near future.
|
||||
|
||||
@@ -397,7 +397,7 @@ Please mind that the popular `openai/whisper-large-v3-turbo` model does not supp
|
||||
!!! note
|
||||
To use the Translation API, please install with extra audio dependencies using `pip install vllm[audio]`.
|
||||
|
||||
Code example: [examples/online_serving/openai_translation_client.py](../../examples/online_serving/openai_translation_client.py)
|
||||
Code example: [examples/speech_to_text/openai/openai_translation_client.py](../../examples/speech_to_text/openai/openai_translation_client.py)
|
||||
|
||||
#### Extra Parameters
|
||||
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
# What is Layerwise (Re)loading?
|
||||
|
||||
Layerwise reloading is the system used to handle the loading of new weight data into existing weight data destinations without triggering recompilation of the cuda graph and other runtime artifacts. This system is used to enable [QeRL](https://arxiv.org/pdf/2510.11696)-style post training flows, where full-precision trainer weights are quantized and loaded into a target vLLM instance for fast, high-exploration rollouts. The core implementation can be found in [layerwise.py](../../vllm/model_executor/model_loader/reload/layerwise.py).
|
||||
|
||||

|
||||
|
||||
## Layerwise Reloading for QeRL
|
||||
|
||||
In order to load new weights into existing weight data destinations, a weight must undergo the following operations:
|
||||
|
||||
- Transfer: weights must be transferred from trainer model to target node/device
|
||||
- Fuse: weight partitions must be fused, for example qkv/gate_up
|
||||
- Process: this typically means online quantization and kernel-specific padding or striding
|
||||
- Shard: weights must be sharded according to the selected parallelism strategy
|
||||
- Copy: weights must be copied into the existing weight data destinations
|
||||
|
||||
Layerwise reloading achieves this using the following steps:
|
||||
|
||||
1. Weights are **transferred** from the trainer to the target (see [weight_transfer](weight_transfer/README.md))
|
||||
2. Weights loaded via `model.load_weights`, during which they are **sharded** and **fused**
|
||||
3. Weights are **processed** in an online fashion as soon as all of a layer's weights are loaded
|
||||
4. Weights are **copied** into the existing weight data destinations
|
||||
|
||||
For more information on implementation, see [Low Level `layerwise` API](#low-level-layerwise-api).
|
||||
|
||||
## Layerwise Loading with Online Quantization
|
||||
|
||||
Online quantization refers to when a user provides full precision weights and those weights are quantized on-the-fly as they are loaded into the model. The layerwise reloading system handles this by treating online quantization as a **processing** step, which is then handled in an online way both during first-time load and during reload. A typical online quantization method implementation should look like this:
|
||||
|
||||
```python
|
||||
class Fp8OnlineLinearMethod(Fp8LinearMethod):
|
||||
"""Online version of Fp8LinearMethod which loads a full precision checkpoint
|
||||
and quantizes weights during loading."""
|
||||
|
||||
uses_meta_device: bool = True
|
||||
|
||||
def create_weights(self, layer: torch.nn.Module, ...):
|
||||
# weight is materialized and processed during loading
|
||||
layer.weight = ModelWeightParameter(
|
||||
data=torch.empty(..., device="meta"),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
# set up online processing
|
||||
initialize_online_processing(layer)
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||||
return
|
||||
|
||||
layer.weight, layer.weight_scale = ops.scaled_fp8_quant(layer.weight)
|
||||
|
||||
# Prevent duplicate processing (e.g., during weight reload)
|
||||
layer._already_called_process_weights_after_loading = True
|
||||
```
|
||||
|
||||
## Example Usages
|
||||
|
||||
### High Level Weight Transfer API
|
||||
|
||||
The layerwise reloading system is integrated with the post-training weight transfer system. To use layerwise reloading in conjunction to the weight transfer system, follow the examples found [here](../../examples/rl/). Layerwise reloading is controlled by the `WeightTransferUpdateInfo.is_checkpoint_format` flag and is set to `True` by default.
|
||||
|
||||
### Mid Level `reload_weights` API
|
||||
|
||||
Layerwise reloading is also exposed via the `reload_weights` API. This interface can be called using the following code:
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
|
||||
llm = LLM("Qwen/Qwen3-0.6B")
|
||||
llm.collective_rpc("reload_weights")
|
||||
```
|
||||
|
||||
This interface also allows specifying a `weights_path` which can be used to select a checkpoint path to load from:
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
|
||||
# fine tuned model checkpoints for testing
|
||||
mul_path = "inference-optimization/Qwen3-0.6B-debug-multiply"
|
||||
add_path = "inference-optimization/Qwen3-0.6B-debug-add"
|
||||
|
||||
llm = LLM("Qwen/Qwen3-0.6B")
|
||||
llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_path})
|
||||
llm.generate("3 4 = ") # 12
|
||||
|
||||
llm.collective_rpc("reload_weights", kwargs={"weights_path": add_path})
|
||||
llm.generate("3 4 = ") # 7
|
||||
```
|
||||
|
||||
Finally, a `weights_iterator` can be provided directly. This iterator can be lazy or eagerly defined.
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
|
||||
weights_iterator = [("q_proj", ...), ("k_proj", ...), ...]
|
||||
|
||||
llm = LLM("Qwen/Qwen3-0.6B")
|
||||
llm.collective_rpc("reload_weights", kwargs={"weights_iterator": weights_iterator})
|
||||
```
|
||||
|
||||
### Low Level `layerwise` API
|
||||
|
||||
[layerwise.py](../../vllm/model_executor/model_loader/reload/layerwise.py) Implements the following functions to execute its lifecycle:
|
||||
|
||||
| Function | Purpose | Quantized Reload | Online Quantization |
|
||||
| - | - | - | - |
|
||||
| `record_metadata_for_reloading` | Record tensor metadata so that layers can be restored on the meta device | Called by `BaseModelLoader` | Called by `BaseModelLoader` |
|
||||
| `restore_layer_on_meta` | Restore layer to model format at start of reload | Called by `initialize_layerwise_reload` | Not called. Online quantized weights already start on meta device via `...OnlineLinearMethod.create_weights` |
|
||||
| `initialize_online_processing` | Wrap weight loaders with the `online_process_loader` wrapper, which buffers weights until all layer weights have been loaded | Called by `initialize_layerwise_reload` | Called by `...OnlineLinearMethod.create_weights` |
|
||||
| `_layerwise_process` | Process layer once all weights are loaded | Called by `online_process_loader` during loading | Called by `online_process_loader` during loading |
|
||||
| `_copy_and_restore_kernel_tensors` | Copy processed weights into original tensor locations to affect compiled cuda graphs, etc. | Called by `_layerwise_process` after `process_weights_after_loading` | Not called. There is no compiled cuda graph yet |
|
||||
| `finalize_layerwise_processing` | Catch any layers which did not load all weights (for example attention weights or weights with padding) | Called by `BaseModelLoader` | Called by `BaseModelLoader` |
|
||||
|
||||
You can plug into this lifecycle directly by calling the `initialize_layerwise_reload`, loading weights, then calling `finalize_layerwise_processing`:
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
from vllm.model_executor.model_loader.reload import initialize_layerwise_reload, finalize_layerwise_processing
|
||||
|
||||
llm = LLM("Qwen/Qwen3-0.6B")
|
||||
|
||||
# this model path requires `VLLM_ENABLE_V1_MULTIPROCESSING=0` and is not stable
|
||||
model = llm.llm_engine.engine_core.engine_core.model_executor.driver_worker.worker.get_model()
|
||||
|
||||
# layerwise reload
|
||||
initialize_layerwise_reload(model)
|
||||
model.load_weights(...)
|
||||
finalize_layerwise_processing(model, llm.model_config)
|
||||
```
|
||||
|
||||
## Troubleshooting Excessive Memory Usage
|
||||
|
||||
Layerwise reloading allows users to incrementally load and process weights as they are loaded into the model. This system relies on buffering layer weights on device until all weights of a layer have been loaded. However, without offloading, this approach necessarily causes excessive buffering if weights are loaded out of order.
|
||||
|
||||
For this reason, users must take care as to the order of weights when they are reloading into the model. Weight should be loaded "in order", meaning that each layer's weights are fully loaded before beginning to load the next layer's weights. "Out of order" loading can cause layer weights to stay buffered while other layer weights are loading, leading to excessive memory usage. In the example below, q_proj, k_proj, v_proj, and up_proj are all buffered at the same time, using more memory than if up_proj was loaded after q_proj, k_proj and v_proj.
|
||||
|
||||
| Correct Loading | Incorrect Loading |
|
||||
| - | - |
|
||||
|  |  |
|
||||
|
||||
Users will see a warning like the one below if weights are loaded out-of-order.
|
||||
|
||||
```console
|
||||
WARNING [layerwise.py:198] Allocating 28.5 MB of device memory to buffers to load ["QKVParallelLinear", "MergedColumnParallelLinear"] layers. This extra memory usage can be avoided by ordering weights by their parent layer when reloading.
|
||||
```
|
||||
@@ -7,7 +7,7 @@ reproducible results:
|
||||
or enable [batch invariance](../features/batch_invariance.md) to make the outputs insensitive to scheduling.
|
||||
- In online mode, you can only enable [batch invariance](../features/batch_invariance.md).
|
||||
|
||||
Example: [examples/offline_inference/reproducibility.py](../../examples/offline_inference/reproducibility.py)
|
||||
Example: [examples/features/batch_invariance/reproducibility_offline.py](../../examples/features/batch_invariance/reproducibility_offline.py)
|
||||
|
||||
!!! warning
|
||||
|
||||
|
||||
+26
-5
@@ -138,14 +138,22 @@ When `--api-key` is configured, the following `/v1` endpoints require Bearer tok
|
||||
|
||||
- `/v1/models` - List available models
|
||||
- `/v1/chat/completions` - Chat completions
|
||||
- `/v1/chat/completions/batch` - Batch chat completions
|
||||
- `/v1/chat/completions/render` - Render chat completion requests
|
||||
- `/v1/completions` - Text completions
|
||||
- `/v1/completions/render` - Render completion requests
|
||||
- `/v1/embeddings` - Generate embeddings
|
||||
- `/v1/audio/transcriptions` - Audio transcription
|
||||
- `/v1/audio/translations` - Audio translation
|
||||
- `/v1/messages` - Anthropic-compatible messages API
|
||||
- `/v1/responses` - Response management
|
||||
- `/v1/messages/count_tokens` - Count tokens for Anthropic messages
|
||||
- `/v1/responses` - Create a response
|
||||
- `/v1/responses/{response_id}` - Retrieve a response
|
||||
- `/v1/responses/{response_id}/cancel` - Cancel a response
|
||||
- `/v1/score` - Scoring API
|
||||
- `/v1/rerank` - Reranking API
|
||||
- `/v1/load_lora_adapter` - Load a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`)
|
||||
- `/v1/unload_lora_adapter` - Unload a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`)
|
||||
|
||||
### Unprotected Endpoints (No API Key Required)
|
||||
|
||||
@@ -155,16 +163,23 @@ The following endpoints **do not require authentication** even when `--api-key`
|
||||
|
||||
- `/invocations` - SageMaker-compatible endpoint (routes to the same inference functions as `/v1` endpoints)
|
||||
- `/inference/v1/generate` - Generate completions
|
||||
- `/generative_scoring` - Generative scoring API
|
||||
- `/pooling` - Pooling API
|
||||
- `/classify` - Classification API
|
||||
- `/score` - Scoring API (non-`/v1` variant)
|
||||
- `/rerank` - Reranking API (non-`/v1` variant)
|
||||
|
||||
**Operational control endpoints (always enabled):**
|
||||
**Operational control endpoints (only when `"generate"` task is supported):**
|
||||
|
||||
- `/pause` - Pause generation (causes denial of service)
|
||||
- `/resume` - Resume generation
|
||||
- `/is_paused` - Check if generation is paused
|
||||
- `/scale_elastic_ep` - Trigger scaling operations
|
||||
- `/is_scaling_elastic_ep` - Check if scaling is in progress
|
||||
- `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF
|
||||
- `/update_weights` - Update model weights (can alter model behavior)
|
||||
- `/get_world_size` - Get distributed world size
|
||||
- `/abort_requests` - Abort in-flight requests (only when `--tokens-only` is also set)
|
||||
|
||||
**Utility endpoints:**
|
||||
|
||||
@@ -207,9 +222,9 @@ These endpoints are only available when profiling is enabled and should only be
|
||||
|
||||
An attacker who can reach the vLLM HTTP server can:
|
||||
|
||||
1. **Bypass authentication** by using non-`/v1` endpoints like `/invocations`, `/inference/v1/generate`, `/pooling`, `/classify`, `/score`, or `/rerank` to run arbitrary inference without credentials
|
||||
2. **Cause denial of service** by calling `/pause` or `/scale_elastic_ep` without a token
|
||||
3. **Access operational controls** to manipulate server state (e.g., pausing generation)
|
||||
1. **Bypass authentication** by using non-`/v1` endpoints like `/invocations`, `/inference/v1/generate`, `/generative_scoring`, `/pooling`, `/classify`, `/score`, or `/rerank` to run arbitrary inference without credentials
|
||||
2. **Cause denial of service** by calling `/pause`, `/scale_elastic_ep`, or `/abort_requests` without a token
|
||||
3. **Access operational controls** to manipulate server state (e.g., pausing generation, updating model weights via `/update_weights`)
|
||||
4. **If `--enable-tokenizer-info-endpoint` is set:** Access sensitive tokenizer configuration including chat templates, which may reveal prompt engineering strategies or other implementation details
|
||||
5. **If `VLLM_SERVER_DEV_MODE=1` is set:** Execute arbitrary RPC commands via `/collective_rpc`, reset caches, put the engine to sleep, and access detailed server configuration
|
||||
|
||||
@@ -288,6 +303,12 @@ To disable the Python code interpreter specifically, omit `code_interpreter` fro
|
||||
|
||||
**Consider a custom implementation**: The GPT-OSS Python tool is a reference implementation. For production deployments, consider implementing a custom code execution sandbox with stricter isolation guarantees. See the [GPT-OSS documentation](https://github.com/openai/gpt-oss?tab=readme-ov-file#python) for guidance.
|
||||
|
||||
## Dynamic LoRA Loading
|
||||
|
||||
vLLM supports dynamically loading and unloading LoRA adapters at runtime via the `/v1/load_lora_adapter` and `/v1/unload_lora_adapter` API endpoints. This functionality is **not enabled by default** — it requires both `--enable-lora` and the environment variable `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True` to be set.
|
||||
|
||||
**Warning:** Dynamic LoRA loading is not a secure operation and should not be enabled in deployments exposed to untrusted clients. If you must enable dynamic LoRA loading, restrict access to the `/v1/load_lora_adapter` and `/v1/unload_lora_adapter` endpoints to trusted administrators only, using a reverse proxy or network-level access controls. Do not expose these endpoints to end users. For details on configuring LoRA adapters, see the [LoRA Adapters documentation](../features/lora.md).
|
||||
|
||||
## Reporting Security Vulnerabilities
|
||||
|
||||
If you believe you have found a security vulnerability in vLLM, please report it following the project's security policy. For more information on how to report security issues and the project's security policy, please see the [vLLM Security Policy](https://github.com/vllm-project/vllm/blob/main/SECURITY.md).
|
||||
|
||||
+1
-1
@@ -15,7 +15,7 @@ compares the generation time for two queries that share the same prefix
|
||||
but ask different questions.
|
||||
|
||||
Run:
|
||||
python examples/offline_inference/automatic_prefix_caching.py
|
||||
python examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py
|
||||
"""
|
||||
|
||||
import time
|
||||
+1
-1
@@ -6,7 +6,7 @@ of a Qwen model using the YARN method (rope_parameters)
|
||||
and run a simple chat example.
|
||||
|
||||
Usage:
|
||||
python examples/offline_inference/context_extension.py
|
||||
python examples/features/context_extension/context_extension_offline.py
|
||||
"""
|
||||
|
||||
from vllm import LLM, RequestOutput, SamplingParams
|
||||
+3
-3
@@ -3,14 +3,14 @@
|
||||
"""
|
||||
Usage:
|
||||
Single node:
|
||||
python examples/offline_inference/data_parallel.py \
|
||||
python examples/features/data_parallel/data_parallel_offline.py \
|
||||
--model="ibm-research/PowerMoE-3b" \
|
||||
-dp=2 \
|
||||
-tp=2
|
||||
|
||||
Multi-node:
|
||||
Node 0 (assume the node has ip of 10.99.48.128):
|
||||
python examples/offline_inference/data_parallel.py \
|
||||
python examples/features/data_parallel/data_parallel_offline.py \
|
||||
--model="ibm-research/PowerMoE-3b" \
|
||||
-dp=2 \
|
||||
-tp=2 \
|
||||
@@ -19,7 +19,7 @@ Multi-node:
|
||||
--dp-master-addr=10.99.48.128 \
|
||||
--dp-master-port=13345
|
||||
Node 1:
|
||||
python examples/offline_inference/data_parallel.py \
|
||||
python examples/features/data_parallel/data_parallel_offline.py \
|
||||
--model="ibm-research/PowerMoE-3b" \
|
||||
-dp=2 \
|
||||
-tp=2 \
|
||||
+1
-1
@@ -12,7 +12,7 @@ from vllm.v1.metrics.loggers import AggregatedLoggingStatLogger
|
||||
"""
|
||||
To run this example, run the following commands simultaneously with
|
||||
different CUDA_VISIBLE_DEVICES:
|
||||
python examples/online_serving/multi_instance_data_parallel.py
|
||||
python examples/features/data_parallel/multi_instance_data_parallel.py
|
||||
|
||||
vllm serve ibm-research/PowerMoE-3b -dp 2 -dpr 1 \
|
||||
--data-parallel-address 127.0.0.1 --data-parallel-rpc-port 62300 \
|
||||
+3
-3
@@ -9,7 +9,7 @@ This directory contains examples demonstrating how to use custom logits processo
|
||||
Demonstrates how to instantiate vLLM with a custom logits processor class that operates at the batch level. The example uses a `DummyLogitsProcessor` that masks out all tokens except a specified `target_token` when passed via `SamplingParams.extra_args`.
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/logits_processor/custom.py
|
||||
python examples/features/logits_processor/custom.py
|
||||
```
|
||||
|
||||
### `custom_req.py` — Request-level logits processor wrapper
|
||||
@@ -17,7 +17,7 @@ python examples/offline_inference/logits_processor/custom.py
|
||||
Shows how to wrap a request-level logits processor (which operates on individual requests) to be compatible with vLLM's batch-level logits processing interface.
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/logits_processor/custom_req.py
|
||||
python examples/features/logits_processor/custom_req.py
|
||||
```
|
||||
|
||||
### `custom_req_init.py` — Request-level processor with engine config
|
||||
@@ -25,7 +25,7 @@ python examples/offline_inference/logits_processor/custom_req.py
|
||||
A special case of wrapping a request-level logits processor where the processor needs access to engine configuration or model metadata during initialization (e.g., vocabulary size, tokenizer info).
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/logits_processor/custom_req_init.py
|
||||
python examples/features/logits_processor/custom_req_init.py
|
||||
```
|
||||
|
||||
## Key Concepts
|
||||
+11
-11
@@ -8,7 +8,7 @@ This is a guide to performing batch inference using the OpenAI batch file format
|
||||
|
||||
The OpenAI batch file format consists of a series of json objects on new lines.
|
||||
|
||||
[See here for an example file.](https://github.com/vllm-project/vllm/blob/main/examples/offline_inference/openai_batch/openai_example_batch.jsonl)
|
||||
[See here for an example file.](https://github.com/vllm-project/vllm/blob/main/examples/features/openai_batch/openai_example_batch.jsonl)
|
||||
|
||||
Each line represents a separate request. See the [OpenAI package reference](https://platform.openai.com/docs/api-reference/batch/requestInput) for more details.
|
||||
|
||||
@@ -30,13 +30,13 @@ We currently support `/v1/chat/completions`, `/v1/embeddings`, and `/v1/score` e
|
||||
To follow along with this example, you can download the example batch, or create your own batch file in your working directory.
|
||||
|
||||
```bash
|
||||
wget https://raw.githubusercontent.com/vllm-project/vllm/main/examples/offline_inference/openai_batch/openai_example_batch.jsonl
|
||||
wget https://raw.githubusercontent.com/vllm-project/vllm/main/examples/features/openai_batch/openai_example_batch.jsonl
|
||||
```
|
||||
|
||||
Once you've created your batch file it should look like this
|
||||
|
||||
```bash
|
||||
cat offline_inference/openai_batch/openai_example_batch.jsonl
|
||||
cat features/openai_batch/openai_example_batch.jsonl
|
||||
{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "meta-llama/Meta-Llama-3-8B-Instruct", "messages": [{"role": "system", "content": "You are a helpful assistant."},{"role": "user", "content": "Hello world!"}],"max_completion_tokens": 1000}}
|
||||
{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "meta-llama/Meta-Llama-3-8B-Instruct", "messages": [{"role": "system", "content": "You are an unhelpful assistant."},{"role": "user", "content": "Hello world!"}],"max_completion_tokens": 1000}}
|
||||
```
|
||||
@@ -49,7 +49,7 @@ You can run the batch with the following command, which will write its results t
|
||||
|
||||
```bash
|
||||
python -m vllm.entrypoints.openai.run_batch \
|
||||
-i offline_inference/openai_batch/openai_example_batch.jsonl \
|
||||
-i features/openai_batch/openai_example_batch.jsonl \
|
||||
-o results.jsonl \
|
||||
--model meta-llama/Meta-Llama-3-8B-Instruct
|
||||
```
|
||||
@@ -58,7 +58,7 @@ or use command-line:
|
||||
|
||||
```bash
|
||||
vllm run-batch \
|
||||
-i offline_inference/openai_batch/openai_example_batch.jsonl \
|
||||
-i features/openai_batch/openai_example_batch.jsonl \
|
||||
-o results.jsonl \
|
||||
--model meta-llama/Meta-Llama-3-8B-Instruct
|
||||
```
|
||||
@@ -77,11 +77,11 @@ cat results.jsonl
|
||||
|
||||
The batch runner supports remote input and output urls that are accessible via http/https.
|
||||
|
||||
For example, to run against our example input file located at `https://raw.githubusercontent.com/vllm-project/vllm/main/examples/offline_inference/openai_batch/openai_example_batch.jsonl`, you can run
|
||||
For example, to run against our example input file located at `https://raw.githubusercontent.com/vllm-project/vllm/main/examples/features/openai_batch/openai_example_batch.jsonl`, you can run
|
||||
|
||||
```bash
|
||||
python -m vllm.entrypoints.openai.run_batch \
|
||||
-i https://raw.githubusercontent.com/vllm-project/vllm/main/examples/offline_inference/openai_batch/openai_example_batch.jsonl \
|
||||
-i https://raw.githubusercontent.com/vllm-project/vllm/main/examples/features/openai_batch/openai_example_batch.jsonl \
|
||||
-o results.jsonl \
|
||||
--model meta-llama/Meta-Llama-3-8B-Instruct
|
||||
```
|
||||
@@ -90,7 +90,7 @@ or use command-line:
|
||||
|
||||
```bash
|
||||
vllm run-batch \
|
||||
-i https://raw.githubusercontent.com/vllm-project/vllm/main/examples/offline_inference/openai_batch/openai_example_batch.jsonl \
|
||||
-i https://raw.githubusercontent.com/vllm-project/vllm/main/examples/features/openai_batch/openai_example_batch.jsonl \
|
||||
-o results.jsonl \
|
||||
--model meta-llama/Meta-Llama-3-8B-Instruct
|
||||
```
|
||||
@@ -113,13 +113,13 @@ To integrate with cloud blob storage, we recommend using presigned urls.
|
||||
To follow along with this example, you can download the example batch, or create your own batch file in your working directory.
|
||||
|
||||
```bash
|
||||
wget https://raw.githubusercontent.com/vllm-project/vllm/main/examples/offline_inference/openai_batch/openai_example_batch.jsonl
|
||||
wget https://raw.githubusercontent.com/vllm-project/vllm/main/examples/features/openai_batch/openai_example_batch.jsonl
|
||||
```
|
||||
|
||||
Once you've created your batch file it should look like this
|
||||
|
||||
```bash
|
||||
cat offline_inference/openai_batch/openai_example_batch.jsonl
|
||||
cat features/openai_batch/openai_example_batch.jsonl
|
||||
{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "meta-llama/Meta-Llama-3-8B-Instruct", "messages": [{"role": "system", "content": "You are a helpful assistant."},{"role": "user", "content": "Hello world!"}],"max_completion_tokens": 1000}}
|
||||
{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "meta-llama/Meta-Llama-3-8B-Instruct", "messages": [{"role": "system", "content": "You are an unhelpful assistant."},{"role": "user", "content": "Hello world!"}],"max_completion_tokens": 1000}}
|
||||
```
|
||||
@@ -127,7 +127,7 @@ cat offline_inference/openai_batch/openai_example_batch.jsonl
|
||||
Now upload your batch file to your S3 bucket.
|
||||
|
||||
```bash
|
||||
aws s3 cp offline_inference/openai_batch/openai_example_batch.jsonl s3://MY_BUCKET/MY_INPUT_FILE.jsonl
|
||||
aws s3 cp features/openai_batch/openai_example_batch.jsonl s3://MY_BUCKET/MY_INPUT_FILE.jsonl
|
||||
```
|
||||
|
||||
### Step 2: Generate your presigned urls
|
||||
+135
-135
@@ -1,135 +1,135 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Test pause/resume with Data Parallel (DP) via HTTP API.
|
||||
|
||||
This example demonstrates coordinated pause/resume across multiple DP ranks.
|
||||
The pause synchronizes across all DP engines via all-reduce.
|
||||
|
||||
Prerequisites:
|
||||
Start a vLLM server with data parallelism:
|
||||
|
||||
$ VLLM_SERVER_DEV_MODE=1 vllm serve facebook/opt-125m \
|
||||
--enforce-eager \
|
||||
--data-parallel-size 4 \
|
||||
--tensor-parallel-size 1
|
||||
|
||||
Then run this script:
|
||||
|
||||
$ python data_parallel_pause_resume.py
|
||||
|
||||
The test verifies pause works by:
|
||||
1. Starting a streaming generation request
|
||||
2. Pausing the server mid-generation
|
||||
3. Sleeping for PAUSE_DURATION seconds
|
||||
4. Resuming the server
|
||||
5. Verifying there was a gap in token generation matching the pause duration
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import threading
|
||||
import time
|
||||
|
||||
import requests
|
||||
from openai import OpenAI
|
||||
|
||||
BASE_URL = "http://localhost:8000"
|
||||
MODEL_NAME = "facebook/opt-125m"
|
||||
PAUSE_DURATION = 3.0
|
||||
|
||||
|
||||
def pause_generation(base_url: str, mode: str = "keep") -> None:
|
||||
"""Pause generation via HTTP endpoint."""
|
||||
url = f"{base_url}/pause"
|
||||
response = requests.post(url, params={"mode": mode}, timeout=60)
|
||||
response.raise_for_status()
|
||||
print("Server paused")
|
||||
|
||||
|
||||
def resume_generation(base_url: str) -> None:
|
||||
"""Resume generation via HTTP endpoint."""
|
||||
url = f"{base_url}/resume"
|
||||
response = requests.post(url, timeout=60)
|
||||
response.raise_for_status()
|
||||
print("Server resumed")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--base-url", default=BASE_URL)
|
||||
parser.add_argument("--model", default=MODEL_NAME)
|
||||
args = parser.parse_args()
|
||||
|
||||
client = OpenAI(
|
||||
base_url=f"{args.base_url}/v1",
|
||||
api_key="EMPTY",
|
||||
)
|
||||
|
||||
prompt = "Write a long story about a dragon. Once upon a time"
|
||||
token_times: list[float] = []
|
||||
pause_token_idx = 0
|
||||
pause_triggered = threading.Event()
|
||||
|
||||
def generator_thread():
|
||||
"""Stream tokens and record timestamps."""
|
||||
stream = client.completions.create(
|
||||
model=args.model,
|
||||
prompt=prompt,
|
||||
max_tokens=50,
|
||||
stream=True,
|
||||
)
|
||||
for chunk in stream:
|
||||
if chunk.choices[0].text:
|
||||
token_times.append(time.monotonic())
|
||||
token_count = len(token_times)
|
||||
print(f"Token {token_count}: {chunk.choices[0].text!r}")
|
||||
|
||||
# Signal controller after some tokens
|
||||
if token_count >= 5 and not pause_triggered.is_set():
|
||||
pause_triggered.set()
|
||||
|
||||
def controller_thread():
|
||||
"""Pause and resume the server."""
|
||||
nonlocal pause_token_idx
|
||||
|
||||
# Wait for some tokens
|
||||
pause_triggered.wait()
|
||||
|
||||
print(f"\nPausing server (keep mode) at token {len(token_times)}...")
|
||||
pause_generation(args.base_url, mode="keep")
|
||||
pause_token_idx = len(token_times)
|
||||
print(f"Sleeping for {PAUSE_DURATION}s...")
|
||||
|
||||
time.sleep(PAUSE_DURATION)
|
||||
|
||||
print("Resuming server...")
|
||||
resume_generation(args.base_url)
|
||||
print("Resumed!\n")
|
||||
|
||||
# Run both threads
|
||||
gen_thread = threading.Thread(target=generator_thread)
|
||||
ctrl_thread = threading.Thread(target=controller_thread)
|
||||
|
||||
gen_thread.start()
|
||||
ctrl_thread.start()
|
||||
|
||||
gen_thread.join()
|
||||
ctrl_thread.join()
|
||||
|
||||
# Check gap at the pause point
|
||||
if pause_token_idx < len(token_times):
|
||||
pause_gap = token_times[pause_token_idx] - token_times[pause_token_idx - 1]
|
||||
print(
|
||||
f"\nGap after pause (token {pause_token_idx} -> "
|
||||
f"{pause_token_idx + 1}): {pause_gap:.3f}s"
|
||||
)
|
||||
if pause_gap >= PAUSE_DURATION * 0.9:
|
||||
print("Test passed! Pause synchronized across DP ranks.")
|
||||
else:
|
||||
print(f"Test failed! Expected ~{PAUSE_DURATION}s gap, got {pause_gap:.3f}s")
|
||||
else:
|
||||
print("Test failed! No tokens were generated after resuming.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Test pause/resume with Data Parallel (DP) via HTTP API.
|
||||
|
||||
This example demonstrates coordinated pause/resume across multiple DP ranks.
|
||||
The pause synchronizes across all DP engines via all-reduce.
|
||||
|
||||
Prerequisites:
|
||||
Start a vLLM server with data parallelism:
|
||||
|
||||
$ VLLM_SERVER_DEV_MODE=1 vllm serve facebook/opt-125m \
|
||||
--enforce-eager \
|
||||
--data-parallel-size 4 \
|
||||
--tensor-parallel-size 1
|
||||
|
||||
Then run this script:
|
||||
|
||||
$ python data_parallel_pause_resume.py
|
||||
|
||||
The test verifies pause works by:
|
||||
1. Starting a streaming generation request
|
||||
2. Pausing the server mid-generation
|
||||
3. Sleeping for PAUSE_DURATION seconds
|
||||
4. Resuming the server
|
||||
5. Verifying there was a gap in token generation matching the pause duration
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import threading
|
||||
import time
|
||||
|
||||
import requests
|
||||
from openai import OpenAI
|
||||
|
||||
BASE_URL = "http://localhost:8000"
|
||||
MODEL_NAME = "facebook/opt-125m"
|
||||
PAUSE_DURATION = 3.0
|
||||
|
||||
|
||||
def pause_generation(base_url: str, mode: str = "keep") -> None:
|
||||
"""Pause generation via HTTP endpoint."""
|
||||
url = f"{base_url}/pause"
|
||||
response = requests.post(url, params={"mode": mode}, timeout=60)
|
||||
response.raise_for_status()
|
||||
print("Server paused")
|
||||
|
||||
|
||||
def resume_generation(base_url: str) -> None:
|
||||
"""Resume generation via HTTP endpoint."""
|
||||
url = f"{base_url}/resume"
|
||||
response = requests.post(url, timeout=60)
|
||||
response.raise_for_status()
|
||||
print("Server resumed")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--base-url", default=BASE_URL)
|
||||
parser.add_argument("--model", default=MODEL_NAME)
|
||||
args = parser.parse_args()
|
||||
|
||||
client = OpenAI(
|
||||
base_url=f"{args.base_url}/v1",
|
||||
api_key="EMPTY",
|
||||
)
|
||||
|
||||
prompt = "Write a long story about a dragon. Once upon a time"
|
||||
token_times: list[float] = []
|
||||
pause_token_idx = 0
|
||||
pause_triggered = threading.Event()
|
||||
|
||||
def generator_thread():
|
||||
"""Stream tokens and record timestamps."""
|
||||
stream = client.completions.create(
|
||||
model=args.model,
|
||||
prompt=prompt,
|
||||
max_tokens=50,
|
||||
stream=True,
|
||||
)
|
||||
for chunk in stream:
|
||||
if chunk.choices[0].text:
|
||||
token_times.append(time.monotonic())
|
||||
token_count = len(token_times)
|
||||
print(f"Token {token_count}: {chunk.choices[0].text!r}")
|
||||
|
||||
# Signal controller after some tokens
|
||||
if token_count >= 5 and not pause_triggered.is_set():
|
||||
pause_triggered.set()
|
||||
|
||||
def controller_thread():
|
||||
"""Pause and resume the server."""
|
||||
nonlocal pause_token_idx
|
||||
|
||||
# Wait for some tokens
|
||||
pause_triggered.wait()
|
||||
|
||||
print(f"\nPausing server (keep mode) at token {len(token_times)}...")
|
||||
pause_generation(args.base_url, mode="keep")
|
||||
pause_token_idx = len(token_times)
|
||||
print(f"Sleeping for {PAUSE_DURATION}s...")
|
||||
|
||||
time.sleep(PAUSE_DURATION)
|
||||
|
||||
print("Resuming server...")
|
||||
resume_generation(args.base_url)
|
||||
print("Resumed!\n")
|
||||
|
||||
# Run both threads
|
||||
gen_thread = threading.Thread(target=generator_thread)
|
||||
ctrl_thread = threading.Thread(target=controller_thread)
|
||||
|
||||
gen_thread.start()
|
||||
ctrl_thread.start()
|
||||
|
||||
gen_thread.join()
|
||||
ctrl_thread.join()
|
||||
|
||||
# Check gap at the pause point
|
||||
if pause_token_idx < len(token_times):
|
||||
pause_gap = token_times[pause_token_idx] - token_times[pause_token_idx - 1]
|
||||
print(
|
||||
f"\nGap after pause (token {pause_token_idx} -> "
|
||||
f"{pause_token_idx + 1}): {pause_gap:.3f}s"
|
||||
)
|
||||
if pause_gap >= PAUSE_DURATION * 0.9:
|
||||
print("Test passed! Pause synchronized across DP ranks.")
|
||||
else:
|
||||
print(f"Test failed! Expected ~{PAUSE_DURATION}s gap, got {pause_gap:.3f}s")
|
||||
else:
|
||||
print("Test failed! No tokens were generated after resuming.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+1
-1
@@ -15,7 +15,7 @@ vllm serve meta-llama/Llama-3.2-1B-Instruct \
|
||||
--enable-prompt-embeds
|
||||
|
||||
Run the client:
|
||||
python examples/online_serving/prompt_embed_inference_with_openai_client.py
|
||||
python examples/features/prompt_embed/prompt_embed_inference_with_openai_client.py
|
||||
|
||||
Model: meta-llama/Llama-3.2-1B-Instruct
|
||||
Note: This model is gated on Hugging Face Hub.
|
||||
+1
-1
@@ -15,7 +15,7 @@ Requirements:
|
||||
- transformers
|
||||
|
||||
Run:
|
||||
python examples/offline_inference/prompt_embed_inference.py
|
||||
python examples/features/prompt_embed/prompt_embed_offline.py
|
||||
"""
|
||||
|
||||
import torch
|
||||
+3
-3
@@ -3,16 +3,16 @@
|
||||
"""
|
||||
Validates the loading of a model saved with the sharded_state format.
|
||||
This script demonstrates how to load a model that was previously saved
|
||||
using save_sharded_state.py and validates it by running inference.
|
||||
using save_sharded_state_offline.py and validates it by running inference.
|
||||
Example usage:
|
||||
(First need to save a sharded_state mode)
|
||||
|
||||
python save_sharded_state.py \
|
||||
python save_sharded_state_offline.py \
|
||||
--model /path/to/load \
|
||||
--tensor-parallel-size 8 \
|
||||
--output /path/to/save/sharded/model
|
||||
|
||||
python load_sharded_state.py \
|
||||
python load_sharded_state_offline.py \
|
||||
--model /path/to/saved/sharded/model \
|
||||
--load-format sharded_state \
|
||||
--tensor-parallel-size 8 \
|
||||
+1
-1
@@ -7,7 +7,7 @@ read its own shard rather than the entire checkpoint.
|
||||
|
||||
Example usage:
|
||||
|
||||
python save_sharded_state.py \
|
||||
python save_sharded_state_offline.py \
|
||||
--model /path/to/load \
|
||||
--tensor-parallel-size 8 \
|
||||
--output /path/to/save
|
||||
+5
-5
@@ -20,7 +20,7 @@ vllm serve deepseek-ai/DeepSeek-R1-Distill-Qwen-7B \
|
||||
If you want to run this script standalone with `uv`, you can use the following:
|
||||
|
||||
```bash
|
||||
uvx --from git+https://github.com/vllm-project/vllm#subdirectory=examples/online_serving/structured_outputs \
|
||||
uvx --from git+https://github.com/vllm-project/vllm#subdirectory=examples/features/structured_outputs \
|
||||
structured-outputs
|
||||
```
|
||||
|
||||
@@ -34,19 +34,19 @@ See [feature docs](https://docs.vllm.ai/en/latest/features/structured_outputs.ht
|
||||
Run all constraints, non-streaming:
|
||||
|
||||
```bash
|
||||
uv run structured_outputs.py
|
||||
uv run structured_outputs_offline.py
|
||||
```
|
||||
|
||||
Run all constraints, streaming:
|
||||
|
||||
```bash
|
||||
uv run structured_outputs.py --stream
|
||||
uv run structured_outputs_offline.py --stream
|
||||
```
|
||||
|
||||
Run certain constraints, for example `structural_tag` and `regex`, streaming:
|
||||
|
||||
```bash
|
||||
uv run structured_outputs.py \
|
||||
uv run structured_outputs_offline.py \
|
||||
--constraint structural_tag regex \
|
||||
--stream
|
||||
```
|
||||
@@ -54,5 +54,5 @@ uv run structured_outputs.py \
|
||||
Run all constraints, with reasoning models and streaming:
|
||||
|
||||
```bash
|
||||
uv run structured_outputs.py --reasoning --stream
|
||||
uv run structured_outputs_offline.py --reasoning --stream
|
||||
```
|
||||
+3
-3
@@ -7,15 +7,15 @@ no internal lb supported in external_launcher mode.
|
||||
|
||||
To run this example:
|
||||
```bash
|
||||
$ torchrun --nproc-per-node=2 examples/offline_inference/torchrun_dp_example.py
|
||||
$ torchrun --nproc-per-node=2 examples/features/torchrun/torchrun_dp_example_offline.py
|
||||
```
|
||||
|
||||
With custom parallelism settings:
|
||||
```bash
|
||||
$ torchrun --nproc-per-node=8 examples/offline_inference/torchrun_dp_example.py \
|
||||
$ torchrun --nproc-per-node=8 examples/features/torchrun/torchrun_dp_example_offline.py \
|
||||
--tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
|
||||
```
|
||||
"""
|
||||
""" # noqa: E501
|
||||
|
||||
import argparse
|
||||
|
||||
+1
-1
@@ -4,7 +4,7 @@
|
||||
experimental support for tensor-parallel inference with torchrun,
|
||||
see https://github.com/vllm-project/vllm/issues/11400 for
|
||||
the motivation and use case for this example.
|
||||
run the script with `torchrun --nproc-per-node=4 torchrun_example.py`,
|
||||
run the script with `torchrun --nproc-per-node=4 torchrun_example_offline.py`,
|
||||
the argument `4` should match the product of `tensor_parallel_size` and
|
||||
`pipeline_parallel_size` below. see `tests/distributed/test_torchrun_example.py`
|
||||
for the unit test.
|
||||
Executable → Regular
+1
-2
@@ -25,7 +25,6 @@ import os
|
||||
import pybase64 as base64
|
||||
import requests
|
||||
from openai import OpenAI
|
||||
from utils import get_first_model
|
||||
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
@@ -407,7 +406,7 @@ def parse_args():
|
||||
|
||||
def main(args) -> None:
|
||||
chat_type = args.chat_type
|
||||
model = get_first_model(client)
|
||||
model = client.models.list().data[0].id
|
||||
example_function_map[chat_type](model, args.max_completion_tokens)
|
||||
|
||||
|
||||
+6
-6
@@ -6,15 +6,15 @@ This folder provides several example scripts on how to inference Qwen2.5-Omni of
|
||||
|
||||
```bash
|
||||
# Audio + image + video
|
||||
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
|
||||
python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
|
||||
-q mixed_modalities
|
||||
|
||||
# Read vision and audio inputs from a single video file
|
||||
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
|
||||
python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
|
||||
-q use_audio_in_video
|
||||
|
||||
# Multiple audios
|
||||
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
|
||||
python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
|
||||
-q multi_audios
|
||||
```
|
||||
|
||||
@@ -24,16 +24,16 @@ You can also test Qwen2.5-Omni on a single modality:
|
||||
|
||||
```bash
|
||||
# Process audio inputs
|
||||
python examples/offline_inference/audio_language.py \
|
||||
python examples/generate/multimodal/audio_language_offline.py \
|
||||
--model-type qwen2_5_omni
|
||||
|
||||
# Process image inputs
|
||||
python examples/offline_inference/vision_language.py \
|
||||
python examples/generate/multimodal/vision_language_offline.py \
|
||||
--modality image \
|
||||
--model-type qwen2_5_omni
|
||||
|
||||
# Process video inputs
|
||||
python examples/offline_inference/vision_language.py \
|
||||
python examples/generate/multimodal/vision_language_offline.py \
|
||||
--modality video \
|
||||
--model-type qwen2_5_omni
|
||||
```
|
||||
Executable → Regular
Executable → Regular
+1
-1
@@ -1402,7 +1402,7 @@ def run_mantis(questions: list[str], modality: str) -> ModelRequestData:
|
||||
# MiniCPM-V
|
||||
def run_minicpmv_base(questions: list[str], modality: str, model_name):
|
||||
assert modality in ["image", "video", "image+video"]
|
||||
# If you want to use `MiniCPM-o-2_6` with audio inputs, check `audio_language.py` # noqa
|
||||
# If you want to use `MiniCPM-o-2_6` with audio inputs, check `audio_language_offline.py` # noqa
|
||||
|
||||
# 2.0
|
||||
# The official repo doesn't work yet, so we need to use a fork for now
|
||||
+1
-1
@@ -9,7 +9,7 @@ Validates that:
|
||||
3. Results are deterministic across runs (baseline vs reference).
|
||||
|
||||
Usage:
|
||||
python examples/offline_inference/routed_experts_e2e.py \
|
||||
python examples/rl/routed_experts_e2e.py \
|
||||
--model Qwen/Qwen3-30B-A3B \
|
||||
--tp 4 \
|
||||
--max-model-len 4096 \
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user