forked from Karylab-cklius/vllm
Compare commits
13
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5c846ccbda | ||
|
|
f2069b005b | ||
|
|
ccd49f6821 | ||
|
|
1c7bc18318 | ||
|
|
9a938df64e | ||
|
|
3da4a1b124 | ||
|
|
6871738777 | ||
|
|
aa4990a9a2 | ||
|
|
a4610da0c6 | ||
|
|
09cdcf34aa | ||
|
|
d2c671c29b | ||
|
|
b5a2adec4b | ||
|
|
78739e3bda |
@@ -34,7 +34,7 @@ case "${test_suite}" in
|
||||
pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py
|
||||
pytest -v -s v1/structured_output
|
||||
pytest -v -s v1/test_serial_utils.py
|
||||
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py
|
||||
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py
|
||||
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py
|
||||
;;
|
||||
server)
|
||||
|
||||
@@ -12,7 +12,6 @@ steps:
|
||||
- tests/basic_correctness/test_cpu_offload
|
||||
- tests/basic_correctness/test_mem.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s basic_correctness/test_mem.py
|
||||
- pytest -v -s basic_correctness/test_basic_correctness.py
|
||||
- pytest -v -s basic_correctness/test_cpu_offload.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s basic_correctness/test_mem.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s basic_correctness/test_basic_correctness.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s basic_correctness/test_cpu_offload.py
|
||||
|
||||
@@ -14,8 +14,7 @@ steps:
|
||||
- vllm/v1/cudagraph_dispatcher.py
|
||||
- tests/compile/correctness_e2e/test_sequence_parallel.py
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
|
||||
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
|
||||
|
||||
- label: Sequence Parallel Correctness Tests (2xH100)
|
||||
key: sequence-parallel-correctness-tests-2xh100
|
||||
@@ -25,8 +24,7 @@ steps:
|
||||
optional: true
|
||||
num_devices: 2
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
|
||||
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
|
||||
|
||||
- label: AsyncTP Correctness Tests (2xH100)
|
||||
key: asynctp-correctness-tests-2xh100
|
||||
@@ -36,8 +34,7 @@ steps:
|
||||
optional: true
|
||||
num_devices: 2
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
|
||||
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
|
||||
|
||||
- label: AsyncTP Correctness Tests (B200)
|
||||
key: asynctp-correctness-tests-b200
|
||||
@@ -47,8 +44,7 @@ steps:
|
||||
optional: true
|
||||
num_devices: 2
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
|
||||
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
|
||||
|
||||
- label: Distributed Compile Unit Tests (2xH100)
|
||||
key: distributed-compile-unit-tests-2xh100
|
||||
@@ -61,8 +57,7 @@ steps:
|
||||
- vllm/model_executor/layers
|
||||
- tests/compile/passes/distributed/
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- pytest -s -v tests/compile/passes/distributed
|
||||
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -s -v tests/compile/passes/distributed
|
||||
|
||||
- label: Fusion and Compile Unit Tests (2xB200)
|
||||
key: fusion-and-compile-unit-tests-2xb200
|
||||
|
||||
@@ -32,11 +32,10 @@ steps:
|
||||
- tests/entrypoints/openai/test_multi_api_servers.py
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
|
||||
- DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
|
||||
|
||||
- label: Distributed Compile + RPC Tests (2 GPUs)
|
||||
key: distributed-compile-rpc-tests-2-gpus
|
||||
@@ -56,10 +55,9 @@ steps:
|
||||
- tests/entrypoints/llm/test_collective_rpc.py
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- pytest -v -s entrypoints/llm/test_collective_rpc.py
|
||||
- pytest -v -s ./compile/fullgraph/test_basic_correctness.py
|
||||
- pytest -v -s ./compile/test_wrapper.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s entrypoints/llm/test_collective_rpc.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s ./compile/fullgraph/test_basic_correctness.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s ./compile/test_wrapper.py
|
||||
|
||||
- label: Distributed Torchrun + Shutdown Tests (2 GPUs)
|
||||
key: distributed-torchrun-shutdown-tests-2-gpus
|
||||
@@ -78,11 +76,10 @@ steps:
|
||||
- tests/v1/worker/test_worker_memory_snapshot.py
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
|
||||
- VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
|
||||
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
|
||||
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s v1/worker/test_worker_memory_snapshot.py
|
||||
|
||||
- label: Distributed Torchrun + Examples (4 GPUs)
|
||||
key: distributed-torchrun-examples-4-gpus
|
||||
@@ -97,24 +94,23 @@ steps:
|
||||
- tests/examples/features/data_parallel/data_parallel_offline.py
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
# test with torchrun tp=2 and external_dp=2
|
||||
- torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
|
||||
# test with torchrun tp=2 and pp=2
|
||||
- PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
|
||||
# test with torchrun tp=4 and dp=1
|
||||
- TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
# test with torchrun tp=2, pp=2 and dp=1
|
||||
- PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
# test with torchrun tp=1 and dp=4 with ep
|
||||
- DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
# test with torchrun tp=2 and dp=2 with ep
|
||||
- TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
# test with internal dp
|
||||
- python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager
|
||||
# rlhf examples
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py
|
||||
|
||||
- label: Distributed DP Tests (4 GPUs)
|
||||
key: distributed-dp-tests-4-gpus
|
||||
@@ -128,14 +124,13 @@ steps:
|
||||
- tests/distributed/test_utils
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
|
||||
- TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
|
||||
- TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
|
||||
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
|
||||
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
|
||||
- pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
|
||||
- pytest -v -s distributed/test_utils.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_utils.py
|
||||
|
||||
- label: Distributed Compile + Comm (4 GPUs)
|
||||
key: distributed-compile-comm-4-gpus
|
||||
@@ -151,13 +146,12 @@ steps:
|
||||
- tests/distributed/test_multiproc_executor.py
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- pytest -v -s compile/fullgraph/test_basic_correctness.py
|
||||
- pytest -v -s distributed/test_pynccl.py
|
||||
- pytest -v -s distributed/test_events.py
|
||||
- pytest -v -s distributed/test_symm_mem_allreduce.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s compile/fullgraph/test_basic_correctness.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_pynccl.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_events.py
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_symm_mem_allreduce.py
|
||||
# test multi-node TP with multiproc executor (simulated on single node)
|
||||
- pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
|
||||
|
||||
- label: Distributed Tests (8 GPUs)(H100)
|
||||
key: distributed-tests-8-gpus-h100
|
||||
@@ -176,9 +170,8 @@ steps:
|
||||
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
# test with torchrun tp=2 and dp=4 with ep
|
||||
- torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
|
||||
- NCCL_CUMEM_HOST_ENABLE=0 torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
|
||||
|
||||
- label: Distributed Tests (4 GPUs)(A100)
|
||||
key: distributed-tests-4-gpus-a100
|
||||
@@ -271,9 +264,7 @@ steps:
|
||||
- tests/distributed/test_pipeline_parallel.py
|
||||
- tests/basic_correctness/test_basic_correctness.py
|
||||
commands:
|
||||
- export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- pytest -v -s distributed/test_ray_v2_executor.py
|
||||
- pytest -v -s distributed/test_ray_v2_executor_e2e.py
|
||||
- pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
|
||||
- TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -k "ray"
|
||||
- VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_ray_v2_executor.py
|
||||
- VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_ray_v2_executor_e2e.py
|
||||
- VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
|
||||
- VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -k "ray"
|
||||
|
||||
@@ -22,10 +22,9 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/llm
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode
|
||||
- pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
|
||||
- pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -41,9 +40,8 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/serve
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -59,8 +57,7 @@ steps:
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -77,9 +74,8 @@ steps:
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/chat_completion
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -128,8 +124,7 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/speech_to_text
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/speech_to_text
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/speech_to_text
|
||||
|
||||
- label: Entrypoints Integration (Multimodal)
|
||||
device: h200_35gb
|
||||
@@ -140,8 +135,7 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/multimodal
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/multimodal
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/multimodal
|
||||
|
||||
- label: Entrypoints Integration (Pooling)
|
||||
key: entrypoints-integration-pooling
|
||||
@@ -151,8 +145,7 @@ steps:
|
||||
- vllm/
|
||||
- tests/entrypoints/pooling
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/pooling
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/pooling
|
||||
|
||||
- label: OpenAI API Correctness
|
||||
key: openai-api-correctness
|
||||
|
||||
@@ -50,8 +50,7 @@ steps:
|
||||
- csrc/
|
||||
- vllm/model_executor/layers/quantization
|
||||
commands:
|
||||
- export VLLM_USE_DEEP_GEMM=0 # We found Triton is faster than DeepGEMM for H100
|
||||
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4
|
||||
- VLLM_USE_DEEP_GEMM=0 pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4 # Triton is faster than DeepGEMM for H100
|
||||
|
||||
- label: LM Eval Small Models (B200)
|
||||
key: lm-eval-small-models-b200
|
||||
@@ -108,9 +107,7 @@ steps:
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_ROCM_ARCH=gfx942 pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt # Limit Quark compilation to save time
|
||||
|
||||
- label: MoE Refactor Integration Test (H100 - TEMPORARY)
|
||||
key: moe-refactor-integration-test-h100-temporary
|
||||
|
||||
@@ -36,14 +36,14 @@ steps:
|
||||
commands:
|
||||
# FIXIT: find out which code initialize cuda before running the test
|
||||
# before the fix, we need to use spawn to test it
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
#
|
||||
# Alot of these tests are on the edge of OOMing
|
||||
- export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
||||
#
|
||||
# There is some Tensor Parallelism related processing logic in LoRA that
|
||||
# requires multi-GPU testing for validation.
|
||||
- pytest -v -s -x lora/test_chatglm3_tp.py
|
||||
- pytest -v -s -x lora/test_llama_tp.py
|
||||
- pytest -v -s -x lora/test_qwen3_with_multi_loras.py
|
||||
- pytest -v -s -x lora/test_olmoe_tp.py
|
||||
- pytest -v -s -x lora/test_gptoss_tp.py
|
||||
- pytest -v -s -x lora/test_qwen35_densemodel_lora.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_chatglm3_tp.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_llama_tp.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_qwen3_with_multi_loras.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_olmoe_tp.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_gptoss_tp.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_qwen35_densemodel_lora.py
|
||||
@@ -18,9 +18,8 @@ steps:
|
||||
- vllm/v1/
|
||||
- tests/v1/spec_decode
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# TODO: create another `optional` test group for slow tests
|
||||
- pytest -v -s -m 'not slow_test' v1/spec_decode
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not slow_test' v1/spec_decode
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
@@ -50,12 +49,11 @@ steps:
|
||||
- tests/v1/test_request.py
|
||||
- tests/v1/test_outputs.py
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s v1/sample
|
||||
- pytest -v -s v1/logits_processors
|
||||
- pytest -v -s v1/test_oracle.py
|
||||
- pytest -v -s v1/test_request.py
|
||||
- pytest -v -s v1/test_outputs.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/sample
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/logits_processors
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/test_oracle.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/test_request.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/test_outputs.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -93,18 +91,17 @@ steps:
|
||||
- tests/entrypoints/openai/correctness/test_lmeval.py
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# split the test to avoid interference
|
||||
- pytest -v -s -m 'not cpu_test' v1/core
|
||||
- pytest -v -s v1/executor
|
||||
- pytest -v -s v1/kv_offload
|
||||
- pytest -v -s v1/simple_kv_offload
|
||||
- pytest -v -s v1/worker
|
||||
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'not cpu_test' v1/metrics
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not cpu_test' v1/core
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/executor
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/kv_offload
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/simple_kv_offload
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/worker
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not cpu_test' v1/metrics
|
||||
# Integration test for streaming correctness (requires special branch).
|
||||
- pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
|
||||
- pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -153,8 +150,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
|
||||
- tests/v1/kv_connector/extract_hidden_states_integration
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s v1/kv_connector/extract_hidden_states_integration
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/kv_connector/extract_hidden_states_integration
|
||||
|
||||
- label: Extract Hidden States Integration (2 GPUs)
|
||||
key: extract-hidden-states-integration-2-gpus
|
||||
@@ -167,8 +163,7 @@ steps:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
|
||||
- tests/v1/kv_connector/extract_hidden_states_integration
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration
|
||||
|
||||
- label: Regression
|
||||
key: regression
|
||||
@@ -360,10 +355,9 @@ steps:
|
||||
- vllm/model_executor/layers
|
||||
- tests/v1/determinism/
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pip install pytest-timeout pytest-forked
|
||||
- pytest -v -s v1/determinism/test_batch_invariance.py
|
||||
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_batch_invariance.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
|
||||
|
||||
- label: Batch Invariance (H100)
|
||||
key: batch-invariance-h100
|
||||
@@ -374,12 +368,11 @@ steps:
|
||||
- vllm/model_executor/layers
|
||||
- tests/v1/determinism/
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pip install pytest-timeout pytest-forked
|
||||
- pytest -v -s v1/determinism/test_batch_invariance.py
|
||||
- pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
|
||||
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
|
||||
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_batch_invariance.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
|
||||
|
||||
- label: Batch Invariance (B200)
|
||||
key: batch-invariance-b200
|
||||
@@ -390,14 +383,13 @@ steps:
|
||||
- vllm/model_executor/layers
|
||||
- tests/v1/determinism/
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pip install pytest-timeout pytest-forked
|
||||
- pytest -v -s v1/determinism/test_batch_invariance.py
|
||||
- pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
|
||||
- VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
|
||||
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
|
||||
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
|
||||
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_batch_invariance.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
|
||||
|
||||
- label: Acceptance Length Test (Large Models) # optional
|
||||
device: h200_35gb
|
||||
@@ -412,5 +404,4 @@ steps:
|
||||
- vllm/model_executor/models/mlp_speculator.py
|
||||
- tests/v1/spec_decode/test_acceptance_length.py
|
||||
commands:
|
||||
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
|
||||
- pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
|
||||
|
||||
@@ -13,13 +13,12 @@ steps:
|
||||
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
commands:
|
||||
- apt-get update && apt-get install -y curl libsodium23
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# Dump tracebacks of all threads if a test hangs, so a wedged GPU/CUDA
|
||||
# init surfaces a stack instead of silently stalling.
|
||||
- export PYTHONFAULTHANDLER=1
|
||||
# Per-test watchdog: a single hung test (e.g. stuck during engine/CUDA
|
||||
# init) fails fast with a traceback instead of running until the global
|
||||
# build timeout. The `thread` method also handles hangs inside C/CUDA
|
||||
# calls that the signal method cannot interrupt.
|
||||
- pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread
|
||||
- pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
|
||||
#
|
||||
# Env vars are inlined because CONTINUE_ON_FAILURE wraps each command
|
||||
# in a subshell, so a standalone `export` would be lost.
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTHONFAULTHANDLER=1 pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn PYTHONFAULTHANDLER=1 pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
|
||||
|
||||
@@ -16,15 +16,14 @@ steps:
|
||||
- tests/entrypoints/llm/test_struct_output_generate.py
|
||||
commands:
|
||||
- set -x
|
||||
- export VLLM_USE_V2_MODEL_RUNNER=1
|
||||
- pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
|
||||
# This requires eager until we sort out CG correctness issues.
|
||||
# TODO: remove ENFORCE_EAGER here after https://github.com/vllm-project/vllm/pull/32936 is merged.
|
||||
- ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
|
||||
- pytest -v -s v1/e2e/general/test_context_length.py
|
||||
- pytest -v -s v1/e2e/general/test_min_tokens.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/e2e/general/test_context_length.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/e2e/general/test_min_tokens.py
|
||||
# Temporary hack filter to exclude ngram spec decoding based tests.
|
||||
- pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
|
||||
|
||||
- label: Model Runner V2 Examples
|
||||
device: h200_35gb
|
||||
@@ -42,26 +41,25 @@ steps:
|
||||
- examples/features/tensorize_vllm_model.py
|
||||
commands:
|
||||
- set -x
|
||||
- export VLLM_USE_V2_MODEL_RUNNER=1
|
||||
- pip install tensorizer # for tensorizer test
|
||||
- python3 basic/offline_inference/chat.py # for basic
|
||||
- python3 basic/offline_inference/generate.py --model facebook/opt-125m
|
||||
#- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO
|
||||
#- python3 basic/offline_inference/embed.py # TODO
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/chat.py # for basic
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/generate.py --model facebook/opt-125m
|
||||
#- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO
|
||||
#- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/embed.py # TODO
|
||||
# for multi-modal models
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# for pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# for features demo
|
||||
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
|
||||
- python3 deployment/llm_engine_example.py
|
||||
- python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
|
||||
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 features/automatic_prefix_caching/prefix_caching_offline.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 deployment/llm_engine_example.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && VLLM_USE_V2_MODEL_RUNNER=1 python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
|
||||
# https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
|
||||
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
|
||||
|
||||
- label: Model Runner V2 Distributed (2 GPUs)
|
||||
key: model-runner-v2-distributed-2-gpus
|
||||
@@ -76,13 +74,11 @@ steps:
|
||||
- tests/v1/distributed/test_eagle_dp.py
|
||||
commands:
|
||||
- set -x
|
||||
- export VLLM_USE_V2_MODEL_RUNNER=1
|
||||
# The "and not True" here is a hacky way to exclude the prompt_embeds cases which aren't yet supported.
|
||||
- TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True"
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
|
||||
|
||||
- label: Model Runner V2 Pipeline Parallelism (4 GPUs)
|
||||
key: model-runner-v2-pipeline-parallelism-4-gpus
|
||||
@@ -97,10 +93,9 @@ steps:
|
||||
- tests/v1/distributed/test_pp_dp_v2.py
|
||||
commands:
|
||||
- set -x
|
||||
- export VLLM_USE_V2_MODEL_RUNNER=1
|
||||
- pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
|
||||
- pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
|
||||
- pytest -v -s v1/distributed/test_pp_dp_v2.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/distributed/test_pp_dp_v2.py
|
||||
|
||||
- label: Model Runner V2 Spec Decode
|
||||
device: h200_35gb
|
||||
@@ -115,8 +110,7 @@ steps:
|
||||
- tests/v1/e2e/spec_decode/test_spec_decode.py
|
||||
commands:
|
||||
- set -x
|
||||
- export VLLM_USE_V2_MODEL_RUNNER=1
|
||||
- pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
|
||||
- pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
|
||||
- pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
|
||||
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
|
||||
- VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
|
||||
|
||||
@@ -23,17 +23,15 @@ steps:
|
||||
# - tests/entrypoints/openai/test_uds.py
|
||||
- tests/v1/sample/test_logprobs_e2e.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
|
||||
- pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py
|
||||
# - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py
|
||||
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"
|
||||
|
||||
# - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
|
||||
- pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
|
||||
# - pytest -v -s entrypoints/openai/test_return_token_ids.py
|
||||
# - pytest -v -s entrypoints/openai/test_uds.py
|
||||
- pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"
|
||||
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
|
||||
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/test_return_token_ids.py
|
||||
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/test_uds.py
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"
|
||||
|
||||
- label: Rust Frontend Serve/Admin Coverage
|
||||
timeout_in_minutes: 60
|
||||
@@ -51,13 +49,11 @@ steps:
|
||||
- tests/entrypoints/serve/instrumentator/test_metrics.py
|
||||
# - tests/entrypoints/serve/dev/test_sleep.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
- pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
|
||||
- pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
|
||||
- pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
|
||||
# - pytest -v -s entrypoints/serve/dev/test_sleep.py
|
||||
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
|
||||
# - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/dev/test_sleep.py
|
||||
|
||||
- label: Rust Frontend Core Correctness
|
||||
timeout_in_minutes: 30
|
||||
@@ -69,9 +65,7 @@ steps:
|
||||
- tests/utils.py
|
||||
- tests/entrypoints/openai/correctness/test_lmeval.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
|
||||
- label: Rust Frontend Tool Use
|
||||
timeout_in_minutes: 60
|
||||
@@ -83,9 +77,7 @@ steps:
|
||||
- tests/utils.py
|
||||
- tests/tool_use/
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
|
||||
|
||||
- label: Rust Frontend Distributed
|
||||
timeout_in_minutes: 30
|
||||
@@ -103,9 +95,6 @@ steps:
|
||||
- tests/v1/distributed/test_hybrid_lb_dp.py
|
||||
- tests/v1/distributed/test_internal_lb_dp.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info"
|
||||
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info"
|
||||
- VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info"
|
||||
|
||||
@@ -153,8 +153,7 @@ steps:
|
||||
- vllm/model_executor/models/qwen3_dflash.py
|
||||
- tests/v1/spec_decode/test_speculators_correctness.py
|
||||
commands:
|
||||
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
|
||||
- pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
|
||||
|
||||
- label: Spec Decode MTP hybrid (B200)
|
||||
timeout_in_minutes: 30
|
||||
|
||||
@@ -135,6 +135,12 @@ Do not modify code in these areas without first reading and following the
|
||||
linked guide. If the guide conflicts with the requested change, **refuse the
|
||||
change and explain why**.
|
||||
|
||||
Security reviewers should start with [`SECURITY.md`](SECURITY.md),
|
||||
[`docs/usage/security.md`](docs/usage/security.md), and
|
||||
[`docs/contributing/vulnerability_management.md`](docs/contributing/vulnerability_management.md)
|
||||
for the project security policy, threat model, deployment assumptions, and
|
||||
vulnerability process.
|
||||
|
||||
- **Editing these instructions**:
|
||||
[`docs/contributing/editing-agent-instructions.md`](docs/contributing/editing-agent-instructions.md)
|
||||
— Rules for modifying AGENTS.md or any domain-specific guide it references.
|
||||
|
||||
@@ -4,6 +4,9 @@
|
||||
#ifdef CPU_CAPABILITY_AMXBF16
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_amx.hpp"
|
||||
#endif
|
||||
#if defined(__riscv_v)
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_rvv.hpp"
|
||||
#endif
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
|
||||
|
||||
#define VLLM_DISPATCH_CASE_16B_TYPES(...) \
|
||||
@@ -319,6 +322,8 @@ void cpu_gemm_wna16(
|
||||
return ISA::AMX;
|
||||
} else if (isa_hint == "vec") {
|
||||
return ISA::VEC;
|
||||
} else if (isa_hint == "rvv") {
|
||||
return ISA::RVV;
|
||||
} else {
|
||||
TORCH_CHECK(false, "unsupported isa hint: " + isa_hint);
|
||||
}
|
||||
@@ -397,6 +402,40 @@ void cpu_gemm_wna16(
|
||||
pack_factor);
|
||||
return;
|
||||
}
|
||||
} else if (isa == ISA::RVV) {
|
||||
using gemm_t = cpu_micro_gemm::MicroGemm<ISA::RVV, scalar_t>;
|
||||
if (has_zp) {
|
||||
using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, true, false>;
|
||||
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
|
||||
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
|
||||
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
|
||||
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
|
||||
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
|
||||
scales_group_stride, zeros_group_stride, group_num, group_size,
|
||||
pack_factor);
|
||||
return;
|
||||
}
|
||||
if (use_desc_act) {
|
||||
using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, true>;
|
||||
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
|
||||
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
|
||||
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
|
||||
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
|
||||
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
|
||||
scales_group_stride, zeros_group_stride, group_num, group_size,
|
||||
pack_factor);
|
||||
return;
|
||||
} else {
|
||||
using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, false>;
|
||||
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
|
||||
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
|
||||
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
|
||||
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
|
||||
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
|
||||
scales_group_stride, zeros_group_stride, group_num, group_size,
|
||||
pack_factor);
|
||||
return;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@@ -0,0 +1,228 @@
|
||||
#ifndef CPU_MICRO_GEMM_RVV_HPP
|
||||
#define CPU_MICRO_GEMM_RVV_HPP
|
||||
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
|
||||
|
||||
#if defined(__riscv_v)
|
||||
|
||||
namespace cpu_micro_gemm {
|
||||
namespace {
|
||||
|
||||
constexpr int32_t RVV_MGEMM_N8 = 8;
|
||||
constexpr int32_t RVV_MGEMM_B_GROUP_STRIDE = 16;
|
||||
|
||||
template <typename scalar_t>
|
||||
FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32(const scalar_t* ptr);
|
||||
|
||||
template <>
|
||||
FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32<float>(const float* ptr) {
|
||||
return RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, RVV_MGEMM_N8);
|
||||
}
|
||||
|
||||
template <>
|
||||
FORCE_INLINE fixed_fp32x8_t
|
||||
load_row8_b_as_f32<c10::Half>(const c10::Half* ptr) {
|
||||
#if defined(__riscv_zvfh)
|
||||
fixed_fp16x8_t vec = RVVI(__riscv_vle16_v_f16, LMUL_128)(
|
||||
reinterpret_cast<const _Float16*>(ptr), RVV_MGEMM_N8);
|
||||
return RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
|
||||
#else
|
||||
alignas(32) float values[RVV_MGEMM_N8];
|
||||
for (int32_t i = 0; i < RVV_MGEMM_N8; ++i) {
|
||||
values[i] = static_cast<float>(ptr[i]);
|
||||
}
|
||||
return RVVI(__riscv_vle32_v_f32, LMUL_256)(values, RVV_MGEMM_N8);
|
||||
#endif
|
||||
}
|
||||
|
||||
template <>
|
||||
FORCE_INLINE fixed_fp32x8_t
|
||||
load_row8_b_as_f32<c10::BFloat16>(const c10::BFloat16* ptr) {
|
||||
#if defined(__riscv_zvfbfmin)
|
||||
fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
|
||||
reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
|
||||
fixed_bf16x8_t vec =
|
||||
RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16, LMUL_128)(raw);
|
||||
return RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
|
||||
#else
|
||||
fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
|
||||
reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
|
||||
auto wide = RVVI(__riscv_vzext_vf2_u32, LMUL_256)(raw, RVV_MGEMM_N8);
|
||||
auto shifted = RVVI(__riscv_vsll_vx_u32, LMUL_256)(wide, 16, RVV_MGEMM_N8);
|
||||
return RVVI4(__riscv_vreinterpret_v_u32, LMUL_256, _f32, LMUL_256)(shifted);
|
||||
#endif
|
||||
}
|
||||
|
||||
// Mx8 RVV kernel. B points at one 8-channel half of a 16-channel packed group,
|
||||
// with rows separated by RVV_MGEMM_B_GROUP_STRIDE scalar elements.
|
||||
template <int32_t M, typename scalar_t>
|
||||
FORCE_INLINE void gemm_micro_rvv_fma_mx8_ku4(const scalar_t* __restrict__ a_ptr,
|
||||
const scalar_t* __restrict__ b_ptr,
|
||||
float* __restrict__ c_ptr,
|
||||
const int64_t lda,
|
||||
const int64_t ldc, const int32_t k,
|
||||
const bool accum_c) {
|
||||
static_assert(0 < M && M <= 8);
|
||||
|
||||
#define RVV_ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
|
||||
#define RVV_IF_M(i) if constexpr (M > (i))
|
||||
|
||||
#define RVV_DECL_A(i) const scalar_t* __restrict__ a##i = a_ptr + (i) * lda;
|
||||
RVV_ROWS_APPLY(RVV_DECL_A)
|
||||
#undef RVV_DECL_A
|
||||
|
||||
#define RVV_DECL_ACC(i) fixed_fp32x8_t acc##i;
|
||||
RVV_ROWS_APPLY(RVV_DECL_ACC)
|
||||
#undef RVV_DECL_ACC
|
||||
|
||||
#define RVV_INIT_ACC(i) \
|
||||
RVV_IF_M(i) { \
|
||||
if (accum_c) { \
|
||||
acc##i = RVVI(__riscv_vle32_v_f32, LMUL_256)(c_ptr + (i) * ldc, \
|
||||
RVV_MGEMM_N8); \
|
||||
} else { \
|
||||
acc##i = RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, RVV_MGEMM_N8); \
|
||||
} \
|
||||
}
|
||||
RVV_ROWS_APPLY(RVV_INIT_ACC)
|
||||
#undef RVV_INIT_ACC
|
||||
|
||||
int32_t k_idx = 0;
|
||||
for (; k_idx + 3 < k; k_idx += 4) {
|
||||
#define RVV_FMA_ROW(i, K_OFFSET) \
|
||||
RVV_IF_M(i) { \
|
||||
acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)( \
|
||||
acc##i, static_cast<float>(*(a##i + k_idx + (K_OFFSET))), b, \
|
||||
RVV_MGEMM_N8); \
|
||||
}
|
||||
|
||||
#define RVV_STEP_K(K_OFFSET) \
|
||||
{ \
|
||||
fixed_fp32x8_t b = load_row8_b_as_f32<scalar_t>( \
|
||||
b_ptr + (k_idx + (K_OFFSET)) * RVV_MGEMM_B_GROUP_STRIDE); \
|
||||
RVV_FMA_ROW(0, K_OFFSET) \
|
||||
RVV_FMA_ROW(1, K_OFFSET) \
|
||||
RVV_FMA_ROW(2, K_OFFSET) \
|
||||
RVV_FMA_ROW(3, K_OFFSET) \
|
||||
RVV_FMA_ROW(4, K_OFFSET) \
|
||||
RVV_FMA_ROW(5, K_OFFSET) \
|
||||
RVV_FMA_ROW(6, K_OFFSET) \
|
||||
RVV_FMA_ROW(7, K_OFFSET) \
|
||||
}
|
||||
|
||||
RVV_STEP_K(0)
|
||||
RVV_STEP_K(1)
|
||||
RVV_STEP_K(2)
|
||||
RVV_STEP_K(3)
|
||||
#undef RVV_STEP_K
|
||||
#undef RVV_FMA_ROW
|
||||
}
|
||||
|
||||
for (; k_idx < k; ++k_idx) {
|
||||
fixed_fp32x8_t b =
|
||||
load_row8_b_as_f32<scalar_t>(b_ptr + k_idx * RVV_MGEMM_B_GROUP_STRIDE);
|
||||
#define RVV_TAIL_ROW(i) \
|
||||
RVV_IF_M(i) { \
|
||||
acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)( \
|
||||
acc##i, static_cast<float>(*(a##i + k_idx)), b, RVV_MGEMM_N8); \
|
||||
}
|
||||
RVV_ROWS_APPLY(RVV_TAIL_ROW)
|
||||
#undef RVV_TAIL_ROW
|
||||
}
|
||||
|
||||
#define RVV_STORE_ROW(i) \
|
||||
RVV_IF_M(i) { \
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_256)(c_ptr + (i) * ldc, acc##i, \
|
||||
RVV_MGEMM_N8); \
|
||||
}
|
||||
RVV_ROWS_APPLY(RVV_STORE_ROW)
|
||||
#undef RVV_STORE_ROW
|
||||
|
||||
#undef RVV_ROWS_APPLY
|
||||
#undef RVV_IF_M
|
||||
}
|
||||
|
||||
template <int32_t M, typename scalar_t>
|
||||
FORCE_INLINE void gemm_micro_rvv_mx32_ku4(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
static_assert(0 < M && M <= 8);
|
||||
scalar_t* __restrict__ curr_b_0 = b_ptr;
|
||||
scalar_t* __restrict__ curr_b_1 = b_ptr + b_n_group_stride;
|
||||
|
||||
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0, c_ptr, lda, ldc, k, accum_c);
|
||||
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0 + RVV_MGEMM_N8,
|
||||
c_ptr + RVV_MGEMM_N8, lda, ldc, k, accum_c);
|
||||
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1, c_ptr + 16, lda, ldc, k,
|
||||
accum_c);
|
||||
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1 + RVV_MGEMM_N8, c_ptr + 24, lda,
|
||||
ldc, k, accum_c);
|
||||
}
|
||||
|
||||
class TileGemmRVV {
|
||||
public:
|
||||
template <typename scalar_t>
|
||||
FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
switch (m) {
|
||||
case 1:
|
||||
gemm_micro_rvv_mx32_ku4<1>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 2:
|
||||
gemm_micro_rvv_mx32_ku4<2>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 3:
|
||||
gemm_micro_rvv_mx32_ku4<3>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 4:
|
||||
gemm_micro_rvv_mx32_ku4<4>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 5:
|
||||
gemm_micro_rvv_mx32_ku4<5>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 6:
|
||||
gemm_micro_rvv_mx32_ku4<6>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 7:
|
||||
gemm_micro_rvv_mx32_ku4<7>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 8:
|
||||
gemm_micro_rvv_mx32_ku4<8>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
template <typename scalar_t>
|
||||
class MicroGemm<cpu_utils::ISA::RVV, scalar_t> {
|
||||
public:
|
||||
static constexpr int32_t MaxMSize = 8;
|
||||
static constexpr int32_t NSize = 32;
|
||||
|
||||
public:
|
||||
void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
TileGemmRVV::gemm<scalar_t>(CPU_MICRO_GEMM_PARAMS);
|
||||
}
|
||||
|
||||
static void pack_weight(const scalar_t* __restrict__ weight,
|
||||
scalar_t* __restrict__ packed_weight,
|
||||
const int32_t output_size, const int32_t input_size) {
|
||||
TORCH_CHECK_EQ(output_size % 16, 0);
|
||||
for (int32_t o_idx = 0; o_idx < output_size; ++o_idx) {
|
||||
const scalar_t* __restrict__ curr_weight = weight + o_idx * input_size;
|
||||
scalar_t* __restrict__ curr_packed_weight =
|
||||
packed_weight + (o_idx / 16) * (16 * input_size) + o_idx % 16;
|
||||
for (int32_t i_idx = 0; i_idx < input_size; ++i_idx) {
|
||||
*curr_packed_weight = *curr_weight;
|
||||
|
||||
curr_packed_weight += 16;
|
||||
++curr_weight;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace cpu_micro_gemm
|
||||
|
||||
#endif // defined(__riscv_v)
|
||||
|
||||
#endif // CPU_MICRO_GEMM_RVV_HPP
|
||||
+3
-1
@@ -8,13 +8,15 @@
|
||||
#include "cpu/cpu_types.hpp"
|
||||
|
||||
namespace cpu_utils {
|
||||
enum class ISA { AMX, VEC };
|
||||
enum class ISA { AMX, VEC, RVV };
|
||||
|
||||
inline ISA get_isa(const std::string& isa) {
|
||||
if (isa == "amx") {
|
||||
return ISA::AMX;
|
||||
} else if (isa == "vec") {
|
||||
return ISA::VEC;
|
||||
} else if (isa == "rvv") {
|
||||
return ISA::RVV;
|
||||
} else {
|
||||
TORCH_CHECK(false, "Invalid isa type: " + isa);
|
||||
}
|
||||
|
||||
@@ -136,8 +136,12 @@ typename T::Fmha::Arguments args_from_options(
|
||||
StrideQ stride_Q_pe = cute::make_tuple(
|
||||
static_cast<int64_t>(q_pe.stride(1)), _1{}, static_cast<int64_t>(q_pe.stride(0)));
|
||||
|
||||
// Read the token and page strides from the cache tensor instead of assuming
|
||||
// packed pages, so strided views (e.g. per-layer views into a cross-layer
|
||||
// block-major cache) are addressed correctly.
|
||||
StrideK stride_C = cute::make_tuple(
|
||||
static_cast<int64_t>(0 + D_latent + D_rope), _1{}, static_cast<int64_t>(page_size * (D_latent + D_rope)));
|
||||
static_cast<int64_t>(kv_c_and_k_pe_cache.stride(1)), _1{},
|
||||
static_cast<int64_t>(kv_c_and_k_pe_cache.stride(0)));
|
||||
StrideLSE stride_PT = cute::make_stride(_1{}, page_count_per_seq);
|
||||
StrideLSE stride_LSE = cute::make_tuple(_1{}, 0 + H);
|
||||
StrideO stride_O = cute::make_tuple(static_cast<int64_t>(0 + D_latent), _1{}, static_cast<int64_t>(0 + H * D_latent));
|
||||
|
||||
@@ -549,7 +549,7 @@ __global__ void indexer_k_quant_and_cache_kernel(
|
||||
const int head_dim, // dimension of each head
|
||||
const int quant_block_size, // quantization block size
|
||||
const int cache_block_size, // cache block size
|
||||
const int cache_stride, // stride for each token in kv_cache
|
||||
const int64_t cache_block_stride, // stride for each block in kv_cache
|
||||
|
||||
const bool use_ue8m0 // use ue8m0 scale format
|
||||
) {
|
||||
@@ -590,16 +590,15 @@ __global__ void indexer_k_quant_and_cache_kernel(
|
||||
scale = exp2f(ceilf(log2f(scale)));
|
||||
}
|
||||
|
||||
const int64_t dst_offset = block_idx * cache_block_size * cache_stride +
|
||||
block_offset * head_dim + head_dim_idx;
|
||||
const int64_t dst_offset =
|
||||
block_idx * cache_block_stride + block_offset * head_dim + head_dim_idx;
|
||||
for (int i = 0; i < VEC_SIZE; i++) {
|
||||
kv_cache[dst_offset + i] =
|
||||
fp8::scaled_convert<cache_t, scalar_t, kv_dt>(k_val_ptr[i], scale);
|
||||
}
|
||||
if (threadIdx.x == 0) {
|
||||
const int64_t dst_scale_idx =
|
||||
block_idx * cache_block_size * cache_stride +
|
||||
cache_block_size * head_dim +
|
||||
block_idx * cache_block_stride + cache_block_size * head_dim +
|
||||
(block_offset * head_dim + head_dim_idx) * 4 / quant_block_size;
|
||||
reinterpret_cast<float*>(kv_cache)[dst_scale_idx / 4] = scale;
|
||||
}
|
||||
@@ -1452,7 +1451,7 @@ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
reinterpret_cast<KV_T*>(k.data_ptr()), \
|
||||
reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
|
||||
slot_mapping.const_data_ptr<int64_t>(), head_dim, quant_block_size, \
|
||||
cache_block_size, cache_stride, use_ue8m0);
|
||||
cache_block_size, cache_block_stride, use_ue8m0);
|
||||
|
||||
void indexer_k_quant_and_cache(
|
||||
torch::stable::Tensor& k, // [num_tokens, head_dim]
|
||||
@@ -1463,7 +1462,7 @@ void indexer_k_quant_and_cache(
|
||||
int num_tokens = k.size(0);
|
||||
int head_dim = k.size(1);
|
||||
int cache_block_size = kv_cache.size(1);
|
||||
int cache_stride = kv_cache.size(2);
|
||||
int64_t cache_block_stride = kv_cache.stride(0);
|
||||
bool use_ue8m0 = scale_fmt == "ue8m0";
|
||||
|
||||
STD_TORCH_CHECK(k.device() == kv_cache.device(),
|
||||
|
||||
+10
-10
@@ -132,7 +132,7 @@ CMD ["/bin/bash"]
|
||||
FROM vllm-base AS ucx-nixl-build
|
||||
|
||||
ARG UCX_VERSION=v1.21.0-rc2
|
||||
ARG NIXL_VERSION=0.10.1
|
||||
ARG NIXL_VERSION=v1.2.0
|
||||
|
||||
# Build-time only: compiler, autotools, and verbs dev headers
|
||||
RUN apt-get update -y && apt-get install -y --no-install-recommends \
|
||||
@@ -149,25 +149,25 @@ RUN apt-get update -y && apt-get install -y --no-install-recommends \
|
||||
# patchelf (installed via uv) is used by the NIXL wheel build to rewrite
|
||||
# RPATH entries, making the wheel portable across stages.
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
git clone https://github.com/openucx/ucx /tmp/ucx_source && \
|
||||
cd /tmp/ucx_source && git checkout "${UCX_VERSION}" && \
|
||||
git clone --depth 1 --branch "${UCX_VERSION}" https://github.com/openucx/ucx /tmp/ucx_source && \
|
||||
cd /tmp/ucx_source && \
|
||||
bash autogen.sh && \
|
||||
./configure --prefix=/tmp/ucx_install --with-ze=yes --enable-examples --enable-mt && \
|
||||
make CFLAGS="-Wno-error=incompatible-pointer-types" -j8 && make install && \
|
||||
git clone https://github.com/ai-dynamo/nixl /tmp/nixl_source && \
|
||||
cd /tmp/nixl_source && git checkout "${NIXL_VERSION}" && \
|
||||
make CFLAGS="-Wno-error=incompatible-pointer-types" -j"$(nproc)" && make install && \
|
||||
git clone --depth 1 --branch "${NIXL_VERSION}" https://github.com/ai-dynamo/nixl /tmp/nixl_source && \
|
||||
cd /tmp/nixl_source && \
|
||||
uv pip install --upgrade meson pybind11 patchelf && \
|
||||
uv pip install -r requirements.txt && \
|
||||
PKG_CONFIG_PATH=/tmp/ucx_install/lib/pkgconfig \
|
||||
LD_LIBRARY_PATH=/tmp/ucx_install/lib \
|
||||
python -m pip wheel --no-deps . -w /tmp/nixl_wheels/ && \
|
||||
find /tmp/ucx_install -type f \( -name '*.a' -o -name '*.la' \) -delete && \
|
||||
rm -rf /tmp/ucx_install/include /tmp/ucx_install/share /tmp/ucx_install/etc /tmp/ucx_install/lib/cmake /tmp/ucx_install/bin && \
|
||||
rm -rf /tmp/ucx_source /tmp/nixl_source
|
||||
rm -rf /tmp/ucx_install/{include,share,etc,bin} /tmp/ucx_install/lib/cmake \
|
||||
/tmp/ucx_source /tmp/nixl_source
|
||||
|
||||
FROM vllm-base AS vllm-openai
|
||||
|
||||
ARG NIXL_VERSION=0.10.1
|
||||
ARG NIXL_VERSION=v1.2.0
|
||||
|
||||
# Copy compiled UCX runtime libraries and the pre-built NIXL wheel.
|
||||
# No compiler or autotools are installed in this stage.
|
||||
@@ -192,7 +192,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
librdmacm1t64 \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& uv pip install --no-deps /tmp/nixl_wheels/nixl*.whl \
|
||||
&& uv pip install nixl==${NIXL_VERSION} \
|
||||
&& uv pip install nixl==${NIXL_VERSION} && uv pip uninstall nixl-cu13 \
|
||||
&& rm -rf /tmp/nixl_wheels
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
|
||||
@@ -304,9 +304,15 @@ review process:
|
||||
resources. The reviewer will add `ready` label to the PR when the PR is
|
||||
ready to merge or a full CI run is needed.
|
||||
|
||||
### Escalating Stalled Contributions
|
||||
### Pull Request Limits and Escalation
|
||||
|
||||
If you have an important contribution that has not yet received maintainer attention, please email us at:
|
||||
vLLM uses GitHub's [pull request limit](https://github.blog/open-source/maintainers/how-pull-request-limits-are-cutting-down-the-noise/)
|
||||
for contributors without write access. The current cap is 6 open PRs. If this
|
||||
blocks well-intentioned critical work, contact a committer to request bypass
|
||||
list access.
|
||||
|
||||
If you need an expedited review for an important contribution, please email us
|
||||
at:
|
||||
|
||||
<pr-review-request@vllm.ai>
|
||||
|
||||
|
||||
@@ -212,3 +212,69 @@ def test_cutlass_mla_decode(
|
||||
print(
|
||||
f"{t:.3f} ms, {FLOPS / 10**9 / t:.0f} TFLOPS,", f"{bytes / 10**6 / t:.0f} GB/s"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.has_device_capability(100),
|
||||
reason=CUTLASS_MLA_UNSUPPORTED_REASON,
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_cutlass_mla_decode_cross_layer_view():
|
||||
"""The kernel must read the cache's page-dim stride instead of assuming
|
||||
pages are packed back-to-back. A per-layer view into a cross-layer
|
||||
(block-major) cache has stride(0) inflated by num_layers; outputs must
|
||||
match a contiguous cache holding the same data exactly."""
|
||||
device = torch.device("cuda:0")
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
torch.set_default_device(device)
|
||||
torch.manual_seed(42)
|
||||
|
||||
b, mean_sk, d, dv, block_size = 4, 512, 576, 512, 64
|
||||
num_layers, layer_idx = 3, 1
|
||||
scale = math.sqrt(d) ** (-1)
|
||||
|
||||
num_pages = b * (mean_sk // block_size)
|
||||
cache_seqlens = torch.full((b,), mean_sk, dtype=torch.int32)
|
||||
block_table = torch.arange(num_pages, dtype=torch.int32).view(
|
||||
b, mean_sk // block_size
|
||||
)
|
||||
|
||||
kv_contig = torch.randn(num_pages, block_size, d)
|
||||
# Neighbor layers hold random data so packed-pages addressing reads
|
||||
# garbage rather than zeros.
|
||||
kv_cross_layer = torch.randn(num_pages, num_layers, block_size, d)
|
||||
kv_view = kv_cross_layer[:, layer_idx]
|
||||
kv_view.copy_(kv_contig)
|
||||
assert kv_view.stride(0) == num_layers * block_size * d
|
||||
|
||||
q_nope = torch.randn(b, 128, dv)
|
||||
q_pe = torch.randn(b, 128, d - dv)
|
||||
sm_count = num_compute_units(device.index)
|
||||
workspace_size = ops.sm100_cutlass_mla_get_workspace_size(
|
||||
mean_sk, b, sm_count, num_kv_splits=1
|
||||
)
|
||||
workspace = torch.empty(workspace_size, dtype=torch.uint8)
|
||||
|
||||
def run(cache):
|
||||
out = torch.empty(b, 128, dv)
|
||||
lse = torch.empty(b, 128, dtype=torch.float32)
|
||||
ops.sm100_cutlass_mla_decode(
|
||||
out,
|
||||
lse,
|
||||
q_nope,
|
||||
q_pe,
|
||||
cache,
|
||||
cache_seqlens,
|
||||
block_table,
|
||||
workspace,
|
||||
scale,
|
||||
1,
|
||||
)
|
||||
return out, lse
|
||||
|
||||
out_contig, lse_contig = run(kv_contig)
|
||||
out_view, lse_view = run(kv_view)
|
||||
|
||||
# Same data and same compute order; only addressing differs.
|
||||
assert torch.equal(out_contig, out_view)
|
||||
assert torch.equal(lse_contig, lse_view)
|
||||
|
||||
@@ -0,0 +1,566 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Bit-exact kernel equivalence for MLA decode/write kernels on the
|
||||
cross-layer (block-major) KV cache layout.
|
||||
|
||||
The cross-layer layout carves each layer's per-block page out of a single
|
||||
unified slot, so the per-layer view has an inflated ``stride(0)`` (the full
|
||||
unified slot) and a non-zero storage offset. These tests confirm the MLA
|
||||
kernels behind the backends that opt in to the layout (FlashMLA dense,
|
||||
FlashInfer MLA dense, FlashMLA fp8 sparse, plus the ``concat_and_cache_mla``
|
||||
write) honor that strided view bit-identically to a contiguous per-layer
|
||||
cache, and that writes do not bleed into neighbouring layers' segments.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not torch.cuda.is_available(), reason="MLA cache kernels require CUDA"
|
||||
)
|
||||
|
||||
|
||||
def test_concat_and_cache_mla_into_unified_slot_view():
|
||||
"""concat_and_cache_mla must write correctly into a per-layer view whose
|
||||
block stride is the full unified slot (block-major), with zero bleed into
|
||||
the other layers' segments of the same slot."""
|
||||
from vllm import _custom_ops as ops
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
kv_lora_rank = 512
|
||||
pe = 64
|
||||
entry = kv_lora_rank + pe
|
||||
page = 64
|
||||
num_blocks = 32
|
||||
ntok = 200
|
||||
|
||||
kv_c = torch.randn(ntok, kv_lora_rank, device=dev, dtype=torch.bfloat16)
|
||||
k_pe = torch.randn(ntok, pe, device=dev, dtype=torch.bfloat16)
|
||||
slot = torch.randperm(num_blocks * page, device=dev, dtype=torch.int64)[:ntok]
|
||||
scale = torch.tensor(1.0, device=dev)
|
||||
|
||||
def write(cache):
|
||||
ops.concat_and_cache_mla(kv_c, k_pe, cache, slot, "auto", scale)
|
||||
|
||||
# Contiguous per-layer reference: (num_blocks, page, entry).
|
||||
ref = torch.zeros(num_blocks, page, entry, device=dev, dtype=torch.bfloat16)
|
||||
write(ref)
|
||||
|
||||
# Unified slot holding three layer pages per block. Carve the middle
|
||||
# layer's view (non-zero offset, block stride == full unified slot).
|
||||
layer_page_elems = page * entry
|
||||
n_layers = 3
|
||||
unified_slot_elems = n_layers * layer_page_elems
|
||||
big = torch.zeros(num_blocks, unified_slot_elems, device=dev, dtype=torch.bfloat16)
|
||||
flat = big.view(-1)
|
||||
offset = layer_page_elems # middle layer
|
||||
view = torch.as_strided(
|
||||
flat,
|
||||
size=(num_blocks, page, entry),
|
||||
stride=(unified_slot_elems, entry, 1),
|
||||
storage_offset=offset,
|
||||
)
|
||||
assert not view.is_contiguous()
|
||||
assert view.stride(0) == unified_slot_elems
|
||||
write(view)
|
||||
|
||||
# Bit-exact equivalence and zero bleed into the neighbour segments.
|
||||
max_diff = (ref.float() - view.float()).abs().max().item()
|
||||
assert max_diff == 0.0, f"max|Δ| = {max_diff}"
|
||||
|
||||
neighbour_lo = torch.as_strided(
|
||||
flat, (num_blocks, layer_page_elems), (unified_slot_elems, 1), 0
|
||||
)
|
||||
neighbour_hi = torch.as_strided(
|
||||
flat,
|
||||
(num_blocks, layer_page_elems),
|
||||
(unified_slot_elems, 1),
|
||||
2 * layer_page_elems,
|
||||
)
|
||||
assert neighbour_lo.abs().max().item() == 0.0
|
||||
assert neighbour_hi.abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_flashmla_dense_decode_unified_slot_view():
|
||||
"""FlashMLA dense decode (FLASHMLA backend, e.g. Kimi-K2-style dense MLA
|
||||
on Hopper) must read a unified-slot block-major view bit-identically to a
|
||||
contiguous per-layer cache."""
|
||||
import vllm.v1.attention.ops.flashmla as fm
|
||||
|
||||
ok, reason = fm.is_flashmla_dense_supported()
|
||||
if not ok:
|
||||
pytest.skip(reason)
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
dt = torch.bfloat16
|
||||
head_dim = 576
|
||||
hdv = 512
|
||||
h_q = 128
|
||||
page = 64
|
||||
num_blocks = 64
|
||||
bs = 4
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
|
||||
q = torch.randn(bs, 1, h_q, head_dim, device=dev, dtype=dt) * 0.1
|
||||
kv_data = torch.randn(num_blocks, page, 1, head_dim, device=dev, dtype=dt) * 0.1
|
||||
|
||||
# (A) contiguous per-layer reference.
|
||||
cache_contiguous = kv_data.clone().contiguous()
|
||||
|
||||
# (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
|
||||
unified = (
|
||||
torch.randn(num_blocks, n_layers, page, 1, head_dim, device=dev, dtype=dt) * 0.1
|
||||
)
|
||||
unified[:, layer].copy_(kv_data)
|
||||
cache_view = unified[:, layer]
|
||||
assert not cache_view.is_contiguous()
|
||||
assert cache_view.stride(0) == n_layers * page * 1 * head_dim
|
||||
|
||||
max_blk = num_blocks // bs
|
||||
block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
|
||||
bs, max_blk
|
||||
)
|
||||
cache_seqlens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
|
||||
|
||||
def run(kc):
|
||||
meta, num_splits = fm.get_mla_metadata()
|
||||
out, _ = fm.flash_mla_with_kvcache(
|
||||
q=q,
|
||||
k_cache=kc,
|
||||
block_table=block_table,
|
||||
cache_seqlens=cache_seqlens,
|
||||
head_dim_v=hdv,
|
||||
tile_scheduler_metadata=meta,
|
||||
num_splits=num_splits,
|
||||
softmax_scale=head_dim**-0.5,
|
||||
causal=True,
|
||||
)
|
||||
return out.clone().float()
|
||||
|
||||
out_ref = run(cache_contiguous)
|
||||
out_view = run(cache_view)
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert out_ref.abs().max().item() > 0.0
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_flashinfer_mla_dense_decode_unified_slot_view():
|
||||
"""FlashInfer MLA dense decode must read a unified-slot block-major view
|
||||
(inflated stride(0), non-zero storage offset) bit-identically to a
|
||||
contiguous per-layer cache."""
|
||||
try:
|
||||
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
|
||||
except ImportError:
|
||||
pytest.skip("flashinfer is not available")
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.is_device_capability_family(100):
|
||||
pytest.skip("FlashInfer trtllm-gen MLA requires sm100")
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
dt = torch.bfloat16
|
||||
kv_lora_rank = 512
|
||||
qk_rope_head_dim = 64
|
||||
qk_nope_head_dim = 128
|
||||
head_dim = kv_lora_rank + qk_rope_head_dim # 576
|
||||
num_qo_heads = 128
|
||||
page = 64
|
||||
num_blocks = 64
|
||||
bs = 4
|
||||
n_layers = 3 # >1 so the per-layer view's block stride is inflated.
|
||||
layer = 1
|
||||
|
||||
q = torch.randn(bs, 1, num_qo_heads, head_dim, device=dev, dtype=dt)
|
||||
kv_data = torch.randn(num_blocks, 1, page, head_dim, device=dev, dtype=dt)
|
||||
|
||||
# (A) contiguous per-layer reference.
|
||||
kv_contiguous = kv_data.clone().contiguous()
|
||||
|
||||
# (B) unified slot: block b of every layer packed together; view one layer
|
||||
# -> stride(0) is n_layers x larger and storage offset is non-zero.
|
||||
unified = torch.randn(num_blocks, n_layers, 1, page, head_dim, device=dev, dtype=dt)
|
||||
unified[:, layer].copy_(kv_data)
|
||||
kv_view = unified[:, layer]
|
||||
assert not kv_view.is_contiguous()
|
||||
assert kv_view.stride(0) == n_layers * 1 * page * head_dim
|
||||
|
||||
max_blk = num_blocks // bs
|
||||
block_tables = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
|
||||
bs, max_blk
|
||||
)
|
||||
seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
|
||||
ws = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=dev)
|
||||
scale = head_dim**-0.5
|
||||
|
||||
def run(kv):
|
||||
return trtllm_batch_decode_with_kv_cache_mla(
|
||||
query=q,
|
||||
kv_cache=kv,
|
||||
workspace_buffer=ws,
|
||||
qk_nope_head_dim=qk_nope_head_dim,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
block_tables=block_tables,
|
||||
seq_lens=seq_lens,
|
||||
max_seq_len=int(seq_lens.max().item()),
|
||||
bmm1_scale=scale,
|
||||
bmm2_scale=1.0,
|
||||
).clone()
|
||||
|
||||
out_ref = run(kv_contiguous).float()
|
||||
out_view = run(kv_view).float()
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_flashmla_fp8_sparse_decode_unified_slot_view():
|
||||
"""FlashMLA fp8 sparse decode (DeepSeek V3.2/V4 DSA path) must read a
|
||||
unified-slot block-major view bit-identically to a contiguous fp8_ds_mla
|
||||
cache, with finite nonzero output."""
|
||||
import vllm.v1.attention.ops.flashmla as fm
|
||||
|
||||
ok, reason = fm.is_flashmla_sparse_supported()
|
||||
if not ok:
|
||||
pytest.skip(reason)
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
entry = 656 # fp8_ds_mla bytes per token
|
||||
page = 64
|
||||
num_blocks = 32
|
||||
h_q = 128
|
||||
head_dim = 576
|
||||
hdv = 512
|
||||
batch = 2
|
||||
topk = 128
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
|
||||
q = torch.randn(batch, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1
|
||||
|
||||
# Structurally valid fp8 ds_mla payload: 512B fp8 + 16B f32 scales + 128B
|
||||
# bf16 rope (random bytes corrupt the scale region and yield NaNs).
|
||||
nope = (torch.randn(num_blocks, page, 1, 512, device=dev) * 0.1).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
scales = torch.ones(num_blocks, page, 1, 4, device=dev, dtype=torch.float32)
|
||||
rope = (torch.randn(num_blocks, page, 1, 64, device=dev) * 0.1).to(torch.bfloat16)
|
||||
payload = torch.cat(
|
||||
[
|
||||
nope.view(torch.uint8).view(num_blocks, page, 1, 512),
|
||||
scales.view(torch.uint8).view(num_blocks, page, 1, 16),
|
||||
rope.view(torch.uint8).view(num_blocks, page, 1, 128),
|
||||
],
|
||||
dim=-1,
|
||||
).contiguous()
|
||||
assert payload.shape[-1] == entry and payload.dtype == torch.uint8
|
||||
|
||||
# (A) contiguous reference.
|
||||
cache_contiguous = payload.clone().contiguous()
|
||||
|
||||
# (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
|
||||
unified = torch.randint(
|
||||
0, 256, (num_blocks, n_layers, page, 1, entry), device=dev, dtype=torch.uint8
|
||||
)
|
||||
unified[:, layer].copy_(payload)
|
||||
cache_view = unified[:, layer]
|
||||
assert not cache_view.is_contiguous()
|
||||
assert cache_view.stride(0) == n_layers * page * 1 * entry
|
||||
|
||||
# Sparse indices: each batch uses its own disjoint blocks.
|
||||
blocks_per_batch = num_blocks // batch
|
||||
idx = torch.full((batch, 1, topk), -1, device=dev, dtype=torch.int32)
|
||||
for b in range(batch):
|
||||
slots: list[int] = []
|
||||
for blk in range(b * blocks_per_batch, (b + 1) * blocks_per_batch):
|
||||
slots.extend(blk * page + off for off in range(page))
|
||||
slots_t = torch.tensor(slots[:topk], device=dev, dtype=torch.int32)
|
||||
idx[b, 0, : slots_t.numel()] = slots_t
|
||||
|
||||
def run(kc):
|
||||
meta, num_splits = fm.get_mla_metadata()
|
||||
out, _ = fm.flash_mla_with_kvcache(
|
||||
q=q,
|
||||
k_cache=kc,
|
||||
block_table=None,
|
||||
cache_seqlens=None,
|
||||
head_dim_v=hdv,
|
||||
tile_scheduler_metadata=meta,
|
||||
is_fp8_kvcache=True,
|
||||
indices=idx,
|
||||
softmax_scale=head_dim**-0.5,
|
||||
)
|
||||
return out.clone().float()
|
||||
|
||||
out_ref = run(cache_contiguous)
|
||||
out_view = run(cache_view)
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert out_ref.abs().max().item() > 0.0
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_indexer_k_quant_and_cache_into_unified_slot_view():
|
||||
"""indexer_k_quant_and_cache (DeepSeek V3.2/V4 DSA indexer K write) must
|
||||
write correctly into a per-layer view whose block stride is the full
|
||||
unified slot, with zero bleed into the other layers' segments."""
|
||||
from vllm import _custom_ops as ops
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
head_dim = 128
|
||||
quant_block_size = 128
|
||||
block_size = 64
|
||||
num_blocks = 16
|
||||
ntok = 100
|
||||
# Indexer cache layout per token: head_dim fp8 bytes followed by
|
||||
# head_dim * 4 / quant_block_size scale bytes.
|
||||
cache_stride = head_dim + head_dim * 4 // quant_block_size
|
||||
|
||||
k = torch.randn(ntok, head_dim, device=dev, dtype=torch.bfloat16)
|
||||
slot = torch.randperm(num_blocks * block_size, device=dev, dtype=torch.int64)[:ntok]
|
||||
|
||||
def write(cache):
|
||||
ops.indexer_k_quant_and_cache(k, cache, slot, quant_block_size, "ue8m0")
|
||||
|
||||
# Contiguous per-layer reference.
|
||||
ref = torch.zeros(
|
||||
num_blocks, block_size, cache_stride, device=dev, dtype=torch.uint8
|
||||
)
|
||||
write(ref)
|
||||
|
||||
# Unified slot holding three layer pages per block; carve the middle one.
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
unified = torch.zeros(
|
||||
num_blocks, n_layers, block_size, cache_stride, device=dev, dtype=torch.uint8
|
||||
)
|
||||
view = unified[:, layer]
|
||||
assert not view.is_contiguous()
|
||||
assert view.stride(0) == n_layers * block_size * cache_stride
|
||||
write(view)
|
||||
|
||||
assert torch.equal(ref, view.contiguous())
|
||||
# Zero bleed into the neighbour layers' segments.
|
||||
assert unified[:, 0].abs().max().item() == 0
|
||||
assert unified[:, 2].abs().max().item() == 0
|
||||
|
||||
|
||||
def test_flashattn_mla_dense_decode_unified_slot_view():
|
||||
"""FA3 decode (FLASH_ATTN_MLA backend) must read a unified-slot
|
||||
block-major view bit-identically to a contiguous per-layer cache."""
|
||||
try:
|
||||
from vllm.vllm_flash_attn import flash_attn_varlen_func
|
||||
except ImportError:
|
||||
pytest.skip("vllm_flash_attn is not available")
|
||||
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
|
||||
|
||||
if not flash_attn_supports_mla():
|
||||
pytest.skip("FA3 MLA requires a Hopper device")
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
dt = torch.bfloat16
|
||||
kv_lora_rank = 512
|
||||
rope_dim = 64
|
||||
entry = kv_lora_rank + rope_dim # 576
|
||||
h_q = 16
|
||||
page = 64
|
||||
num_blocks = 64
|
||||
bs = 4
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
|
||||
q_pe = torch.randn(bs, h_q, rope_dim, device=dev, dtype=dt) * 0.1
|
||||
q_nope = torch.randn(bs, h_q, kv_lora_rank, device=dev, dtype=dt) * 0.1
|
||||
kv_data = torch.randn(num_blocks, page, entry, device=dev, dtype=dt) * 0.1
|
||||
|
||||
# (A) contiguous per-layer reference.
|
||||
cache_contiguous = kv_data.clone().contiguous()
|
||||
|
||||
# (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
|
||||
unified = torch.randn(num_blocks, n_layers, page, entry, device=dev, dtype=dt) * 0.1
|
||||
unified[:, layer].copy_(kv_data)
|
||||
cache_view = unified[:, layer]
|
||||
assert not cache_view.is_contiguous()
|
||||
assert cache_view.stride(0) == n_layers * page * entry
|
||||
|
||||
max_blk = num_blocks // bs
|
||||
block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
|
||||
bs, max_blk
|
||||
)
|
||||
seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
|
||||
cu_seqlens_q = torch.arange(bs + 1, device=dev, dtype=torch.int32)
|
||||
|
||||
def run(cache):
|
||||
kv_c_cache = cache[..., :kv_lora_rank]
|
||||
k_pe_cache = cache[..., kv_lora_rank:]
|
||||
out = flash_attn_varlen_func(
|
||||
q=q_pe,
|
||||
k=k_pe_cache.unsqueeze(-2), # Add head dim of 1
|
||||
v=kv_c_cache.unsqueeze(-2), # Add head dim of 1
|
||||
q_v=q_nope,
|
||||
max_seqlen_q=1,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
max_seqlen_k=int(seq_lens.max().item()),
|
||||
seqused_k=seq_lens,
|
||||
block_table=block_table,
|
||||
softmax_scale=entry**-0.5,
|
||||
causal=True,
|
||||
fa_version=3,
|
||||
)
|
||||
return out.clone().float()
|
||||
|
||||
out_ref = run(cache_contiguous)
|
||||
out_view = run(cache_view)
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert out_ref.abs().max().item() > 0.0
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_flashmla_dense_fp8_decode_unified_slot_view():
|
||||
"""FlashMLA dense fp8 decode (FLASHMLA backend with quantized KV cache)
|
||||
must read a unified-slot block-major view bit-identically to a contiguous
|
||||
per-layer fp8 cache."""
|
||||
import vllm.v1.attention.ops.flashmla as fm
|
||||
|
||||
ok, reason = fm.is_flashmla_dense_supported()
|
||||
if not ok:
|
||||
pytest.skip(reason)
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
head_dim = 576
|
||||
hdv = 512
|
||||
h_q = 128
|
||||
page = 64
|
||||
num_blocks = 64
|
||||
bs = 4
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
|
||||
q = torch.randn(bs, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1
|
||||
kv_data = (torch.randn(num_blocks, page, head_dim, device=dev) * 0.1).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
|
||||
# (A) contiguous per-layer reference.
|
||||
cache_contiguous = kv_data.clone().contiguous()
|
||||
|
||||
# (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
|
||||
unified = (torch.randn(num_blocks, n_layers, page, head_dim, device=dev) * 0.1).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
unified[:, layer].copy_(kv_data)
|
||||
cache_view = unified[:, layer]
|
||||
assert not cache_view.is_contiguous()
|
||||
assert cache_view.stride(0) == n_layers * page * head_dim
|
||||
|
||||
max_blk = num_blocks // bs
|
||||
block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
|
||||
bs, max_blk
|
||||
)
|
||||
cache_seqlens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
|
||||
descale = torch.ones(1, device=dev, dtype=torch.float32)
|
||||
|
||||
def run(kc):
|
||||
tile_md, num_splits = fm.get_mla_metadata_dense_fp8(cache_seqlens, h_q, 1)
|
||||
out, _ = fm.flash_mla_with_kvcache_fp8(
|
||||
q=q,
|
||||
k_cache=kc.unsqueeze(-2), # Add head dim of 1
|
||||
block_table=block_table,
|
||||
cache_seqlens=cache_seqlens,
|
||||
head_dim_v=hdv,
|
||||
tile_scheduler_metadata=tile_md,
|
||||
num_splits=num_splits,
|
||||
softmax_scale=head_dim**-0.5,
|
||||
causal=True,
|
||||
descale_q=descale,
|
||||
descale_k=descale,
|
||||
)
|
||||
return out.clone().float()
|
||||
|
||||
out_ref = run(cache_contiguous)
|
||||
out_view = run(cache_view)
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert out_ref.abs().max().item() > 0.0
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
|
||||
|
||||
def test_flashinfer_mla_dense_fp8_decode_unified_slot_view():
|
||||
"""FlashInfer MLA dense decode with an fp8 KV cache must read a
|
||||
unified-slot block-major view bit-identically to a contiguous per-layer
|
||||
cache."""
|
||||
try:
|
||||
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
|
||||
except ImportError:
|
||||
pytest.skip("flashinfer is not available")
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.is_device_capability_family(100):
|
||||
pytest.skip("FlashInfer trtllm-gen MLA requires sm100")
|
||||
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
kv_lora_rank = 512
|
||||
qk_rope_head_dim = 64
|
||||
qk_nope_head_dim = 128
|
||||
head_dim = kv_lora_rank + qk_rope_head_dim # 576
|
||||
num_qo_heads = 128
|
||||
page = 64
|
||||
num_blocks = 64
|
||||
bs = 4
|
||||
n_layers = 3
|
||||
layer = 1
|
||||
|
||||
# With a quantized KV cache the decode query is quantized to fp8 as well
|
||||
# (trtllm-gen has no bf16-query x fp8-cache decode kernel).
|
||||
q = (torch.randn(bs, 1, num_qo_heads, head_dim, device=dev) * 0.1).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
kv_data = (torch.randn(num_blocks, 1, page, head_dim, device=dev) * 0.1).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
|
||||
# (A) contiguous per-layer reference.
|
||||
kv_contiguous = kv_data.clone().contiguous()
|
||||
|
||||
# (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
|
||||
unified = (
|
||||
torch.randn(num_blocks, n_layers, 1, page, head_dim, device=dev) * 0.1
|
||||
).to(torch.float8_e4m3fn)
|
||||
unified[:, layer].copy_(kv_data)
|
||||
kv_view = unified[:, layer]
|
||||
assert not kv_view.is_contiguous()
|
||||
assert kv_view.stride(0) == n_layers * 1 * page * head_dim
|
||||
|
||||
max_blk = num_blocks // bs
|
||||
block_tables = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
|
||||
bs, max_blk
|
||||
)
|
||||
seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
|
||||
ws = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=dev)
|
||||
scale = head_dim**-0.5
|
||||
|
||||
def run(kv):
|
||||
return trtllm_batch_decode_with_kv_cache_mla(
|
||||
query=q,
|
||||
kv_cache=kv,
|
||||
workspace_buffer=ws,
|
||||
qk_nope_head_dim=qk_nope_head_dim,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
block_tables=block_tables,
|
||||
seq_lens=seq_lens,
|
||||
max_seq_len=int(seq_lens.max().item()),
|
||||
bmm1_scale=scale,
|
||||
bmm2_scale=1.0,
|
||||
).clone()
|
||||
|
||||
out_ref = run(kv_contiguous).float()
|
||||
out_view = run(kv_view).float()
|
||||
assert torch.isfinite(out_ref).all()
|
||||
assert (out_ref - out_view).abs().max().item() == 0.0
|
||||
@@ -231,3 +231,95 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE)
|
||||
|
||||
# FP8 tolerances match test_mla_backends.py test_backend_correctness.
|
||||
torch.testing.assert_close(o_ref, o_fp8, atol=5e-1, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"H_Q,H_KV,D_QK,D_V,is_mla",
|
||||
[
|
||||
(16, 1, 576, 512, True), # MLA path (grouped kernel, v = trans(k))
|
||||
(32, 8, 128, 128, False), # GQA path (grouped kernel)
|
||||
(32, 32, 128, 128, False), # MHA path (normal kernel)
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("PAGE_SIZE", [16])
|
||||
def test_decode_attention_cross_layer_view(H_Q, H_KV, D_QK, D_V, is_mla, PAGE_SIZE):
|
||||
"""The kernel must honor the cache's page-dim stride, not assume pages are
|
||||
packed back-to-back. A per-layer view into a cross-layer (block-major)
|
||||
cache has stride(0) inflated by num_layers; outputs must match a
|
||||
contiguous cache holding the same data exactly."""
|
||||
B = 3
|
||||
seq_len = 1027
|
||||
CACHE_SIZE = 16384
|
||||
NUM_LAYERS = 3
|
||||
LAYER_IDX = 1
|
||||
dtype = torch.bfloat16
|
||||
sm_scale = 1.0 / (D_QK**0.5)
|
||||
num_kv_splits = 8
|
||||
num_pages = CACHE_SIZE // PAGE_SIZE
|
||||
|
||||
num_pages_per_batch = cdiv(seq_len, PAGE_SIZE)
|
||||
req_to_page = torch.randint(
|
||||
0, num_pages, (B, num_pages_per_batch), device=DEVICE_TYPE
|
||||
)
|
||||
|
||||
q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE)
|
||||
b_seq_len = torch.full((B,), seq_len, device=DEVICE_TYPE)
|
||||
|
||||
# Reference: contiguous paged cache.
|
||||
k_ref = torch.randn(
|
||||
num_pages, PAGE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE
|
||||
)
|
||||
if is_mla:
|
||||
v_ref = k_ref[..., :D_V]
|
||||
else:
|
||||
v_ref = torch.randn(
|
||||
num_pages, PAGE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE
|
||||
)
|
||||
|
||||
# Cross-layer cache: all layers' pages for a block are adjacent. The
|
||||
# per-layer view has the same shape as the contiguous cache but
|
||||
# stride(0) is NUM_LAYERS x larger. Neighbor layers hold random data so
|
||||
# any packed-pages addressing reads garbage rather than zeros.
|
||||
k_xl = torch.randn(
|
||||
num_pages, NUM_LAYERS, PAGE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE
|
||||
)
|
||||
k_view = k_xl[:, LAYER_IDX]
|
||||
k_view.copy_(k_ref)
|
||||
assert k_view.stride(0) == NUM_LAYERS * PAGE_SIZE * H_KV * D_QK
|
||||
if is_mla:
|
||||
v_view = k_view[..., :D_V]
|
||||
else:
|
||||
v_xl = torch.randn(
|
||||
num_pages, NUM_LAYERS, PAGE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE
|
||||
)
|
||||
v_view = v_xl[:, LAYER_IDX]
|
||||
v_view.copy_(v_ref)
|
||||
|
||||
def run(k_buffer, v_buffer):
|
||||
o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
|
||||
lse = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
|
||||
attn_logits = torch.empty(
|
||||
(B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE
|
||||
)
|
||||
decode_attention_fwd(
|
||||
q,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
o,
|
||||
lse,
|
||||
req_to_page,
|
||||
b_seq_len,
|
||||
attn_logits,
|
||||
num_kv_splits,
|
||||
sm_scale,
|
||||
PAGE_SIZE,
|
||||
is_mla=is_mla,
|
||||
)
|
||||
return o, lse
|
||||
|
||||
o_ref, lse_ref = run(k_ref, v_ref)
|
||||
o_xl, lse_xl = run(k_view, v_view)
|
||||
|
||||
# Same data and same compute order; only addressing differs.
|
||||
assert torch.equal(o_ref, o_xl)
|
||||
assert torch.equal(lse_ref, lse_xl)
|
||||
|
||||
@@ -124,6 +124,7 @@ def test_replace_submodules(default_vllm_config, dist_init, dummy_model):
|
||||
max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
|
||||
),
|
||||
torch.device(DEVICES[0]),
|
||||
default_vllm_config,
|
||||
)
|
||||
model = manager.model
|
||||
assert isinstance(model.get_submodule("dense1"), ColumnParallelLinearWithLoRA)
|
||||
@@ -152,6 +153,7 @@ def test_wrap_replicated_linear_subclasses(default_vllm_config, dist_init, dummy
|
||||
max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
|
||||
),
|
||||
torch.device(DEVICES[0]),
|
||||
default_vllm_config,
|
||||
)
|
||||
|
||||
assert isinstance(
|
||||
@@ -172,6 +174,7 @@ def test_wrap_gate_linear(default_vllm_config, dist_init, dummy_model):
|
||||
max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
|
||||
),
|
||||
torch.device(DEVICES[0]),
|
||||
default_vllm_config,
|
||||
)
|
||||
|
||||
assert isinstance(
|
||||
@@ -219,6 +222,7 @@ def test_dedup_shared_module_across_paths(default_vllm_config, dist_init, dummy_
|
||||
max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
|
||||
),
|
||||
torch.device(DEVICES[0]),
|
||||
default_vllm_config,
|
||||
)
|
||||
|
||||
canonical = manager.model.get_submodule("moe.gate")
|
||||
@@ -263,6 +267,7 @@ def test_lm_head_exempt_from_dedup(default_vllm_config, dist_init, dummy_model):
|
||||
max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
|
||||
),
|
||||
torch.device(DEVICES[0]),
|
||||
default_vllm_config,
|
||||
)
|
||||
|
||||
# lm_head's special handling still ran: logits_processor got wrapped
|
||||
@@ -293,6 +298,7 @@ def test_skip_unsupported_matched_modules(default_vllm_config, dist_init, dummy_
|
||||
max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
|
||||
),
|
||||
torch.device(DEVICES[0]),
|
||||
default_vllm_config,
|
||||
)
|
||||
|
||||
# Should not crash and should keep unsupported matched modules unchanged.
|
||||
@@ -325,6 +331,7 @@ def test_target_modules_fail_closed_on_unsupported_matched_modules(
|
||||
target_modules=["dense1"],
|
||||
),
|
||||
torch.device(DEVICES[0]),
|
||||
default_vllm_config,
|
||||
)
|
||||
|
||||
|
||||
@@ -374,6 +381,7 @@ def test_lora_model_manager(default_vllm_config, dist_init, dummy_model, device)
|
||||
max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE
|
||||
),
|
||||
device=device,
|
||||
vllm_config=default_vllm_config,
|
||||
)
|
||||
assert all(x is None for x in manager.lora_index_to_id)
|
||||
assert manager.add_adapter(model_lora1)
|
||||
@@ -442,6 +450,7 @@ def test_lora_lru_cache_model_manager(
|
||||
max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE
|
||||
),
|
||||
device=device,
|
||||
vllm_config=default_vllm_config,
|
||||
)
|
||||
assert all(x is None for x in manager.lora_index_to_id)
|
||||
assert manager.add_adapter(model_lora1)
|
||||
@@ -535,6 +544,7 @@ def test_lru_lora_model_manager(default_vllm_config, dist_init, dummy_model, dev
|
||||
max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE
|
||||
),
|
||||
device=device,
|
||||
vllm_config=default_vllm_config,
|
||||
)
|
||||
assert all(x is None for x in manager.lora_index_to_id)
|
||||
|
||||
@@ -642,9 +652,7 @@ def test_lru_lora_model_manager(default_vllm_config, dist_init, dummy_model, dev
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device", DEVICES)
|
||||
def test_lru_cache_worker_adapter_manager(
|
||||
default_vllm_config, dist_init, dummy_model, device, tmp_path
|
||||
):
|
||||
def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device, tmp_path):
|
||||
lora_config = LoRAConfig(
|
||||
max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE
|
||||
)
|
||||
@@ -670,7 +678,7 @@ def test_lru_cache_worker_adapter_manager(
|
||||
worker_adapter_manager.max_num_seqs = 4
|
||||
worker_adapter_manager.max_num_batched_tokens = 2
|
||||
|
||||
worker_adapter_manager.create_lora_manager(dummy_model)
|
||||
worker_adapter_manager.create_lora_manager(dummy_model, vllm_config)
|
||||
|
||||
mapping = LoRAMapping([], [])
|
||||
worker_adapter_manager.set_active_adapters(
|
||||
@@ -758,9 +766,7 @@ def test_lru_cache_worker_adapter_manager(
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device", DEVICES)
|
||||
def test_worker_adapter_manager(
|
||||
default_vllm_config, dist_init, dummy_model_gate_up, device, tmp_path
|
||||
):
|
||||
def test_worker_adapter_manager(dist_init, dummy_model_gate_up, device, tmp_path):
|
||||
# Should remove every LoRA not specified in the request.
|
||||
lora_config = LoRAConfig(
|
||||
max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE
|
||||
@@ -774,7 +780,7 @@ def test_worker_adapter_manager(
|
||||
|
||||
worker_adapter_manager = WorkerLoRAManager(vllm_config, device, EMBEDDING_MODULES)
|
||||
worker_adapter_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size
|
||||
worker_adapter_manager.create_lora_manager(dummy_model_gate_up)
|
||||
worker_adapter_manager.create_lora_manager(dummy_model_gate_up, vllm_config)
|
||||
|
||||
dummy_lora_files = f"{tmp_path}/lora_adapter"
|
||||
os.makedirs(dummy_lora_files, exist_ok=True)
|
||||
@@ -894,6 +900,7 @@ def test_packed_loras(default_vllm_config, dist_init, dummy_model_gate_up, devic
|
||||
max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE
|
||||
),
|
||||
device=device,
|
||||
vllm_config=default_vllm_config,
|
||||
)
|
||||
model = manager.model
|
||||
|
||||
@@ -944,6 +951,7 @@ def _test_target_modules(
|
||||
device: str,
|
||||
expected_lora: list[tuple[str, type]],
|
||||
expected_no_lora: list[tuple[str, type]],
|
||||
vllm_config,
|
||||
):
|
||||
"""Create a LoRAModelManager and assert which modules have LoRA applied."""
|
||||
LoRAModelManager(
|
||||
@@ -959,6 +967,7 @@ def _test_target_modules(
|
||||
target_modules=target_modules,
|
||||
),
|
||||
device=device,
|
||||
vllm_config=vllm_config,
|
||||
)
|
||||
for module_path, lora_cls in expected_lora:
|
||||
assert isinstance(model.get_submodule(module_path), lora_cls)
|
||||
@@ -981,6 +990,7 @@ def test_target_modules_config(default_vllm_config, dist_init, dummy_model, devi
|
||||
("dense2", RowParallelLinearWithLoRA),
|
||||
("layer1.dense2", RowParallelLinearWithLoRA),
|
||||
],
|
||||
vllm_config=default_vllm_config,
|
||||
)
|
||||
|
||||
|
||||
@@ -998,6 +1008,7 @@ def test_target_modules_multiple(default_vllm_config, dist_init, dummy_model, de
|
||||
("layer1.dense2", RowParallelLinearWithLoRA),
|
||||
],
|
||||
expected_no_lora=[],
|
||||
vllm_config=default_vllm_config,
|
||||
)
|
||||
|
||||
|
||||
@@ -1017,6 +1028,7 @@ def test_target_modules_none_uses_all(
|
||||
("layer1.dense2", RowParallelLinearWithLoRA),
|
||||
],
|
||||
expected_no_lora=[],
|
||||
vllm_config=default_vllm_config,
|
||||
)
|
||||
|
||||
|
||||
@@ -1036,4 +1048,5 @@ def test_target_modules_match_packed_runtime_modules(
|
||||
("layer1.dense1", ColumnParallelLinearWithLoRA),
|
||||
("layer1.dense2", RowParallelLinearWithLoRA),
|
||||
],
|
||||
vllm_config=default_vllm_config,
|
||||
)
|
||||
|
||||
@@ -43,8 +43,8 @@ MODELS = [
|
||||
pytest.param(
|
||||
"Intel/Qwen2-0.5B-Instruct-int4-sym-AutoRound",
|
||||
marks=pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="AWQ AutoRound model only supports CUDA backend for now.",
|
||||
not (current_platform.is_cuda() or current_platform.is_xpu()),
|
||||
reason="AWQ AutoRound model only supports CUDA/XPU backend for now.",
|
||||
),
|
||||
id="auto_round:auto_awq",
|
||||
),
|
||||
|
||||
@@ -74,6 +74,29 @@ def test_embed_dimensions(model_info: EmbedModelInfo):
|
||||
pooling_params.verify(model_config)
|
||||
|
||||
|
||||
@dataclass()
|
||||
class MockMatryoshkaModelConfig:
|
||||
pooler_config: PoolerConfig
|
||||
is_matryoshka: bool = True
|
||||
matryoshka_dimensions: list[int] | None = None
|
||||
served_model_name: str = "mock-matryoshka-model"
|
||||
embedding_size: int = 32
|
||||
|
||||
|
||||
def test_embed_dimensions_matryoshka_without_list_upper_bound():
|
||||
task = "embed"
|
||||
model_config = MockMatryoshkaModelConfig(
|
||||
pooler_config=PoolerConfig(seq_pooling_type="CLS"),
|
||||
matryoshka_dimensions=None,
|
||||
embedding_size=32,
|
||||
)
|
||||
|
||||
PoolingParams(task=task, dimensions=16).verify(model_config)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
PoolingParams(task=task, dimensions=64).verify(model_config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task", ["classify"])
|
||||
def test_classify(task):
|
||||
model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS"))
|
||||
|
||||
@@ -1,10 +1,13 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
def test_mla_backend_rejects_cross_layer_kv_cache():
|
||||
"""MLA backends return identity permutation (layers dim first)
|
||||
to signal cross-layer KV cache is unsupported."""
|
||||
|
||||
def test_mla_common_backend_rejects_cross_layer_kv_cache():
|
||||
"""MLACommonBackend defaults to the identity permutation (layers dim
|
||||
first) so MLA backends whose decode kernels are not verified to honor
|
||||
the cache's block-dim stride stay opted out of cross-layer KV cache."""
|
||||
from vllm.model_executor.layers.attention.mla_attention import (
|
||||
MLACommonBackend,
|
||||
)
|
||||
@@ -19,6 +22,35 @@ def test_mla_backend_rejects_cross_layer_kv_cache():
|
||||
) == (0, 1, 2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"backend_path",
|
||||
[
|
||||
"vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend",
|
||||
"vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend",
|
||||
"vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend",
|
||||
"vllm.v1.attention.backends.mla.flashmla.FlashMLABackend",
|
||||
"vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend",
|
||||
],
|
||||
)
|
||||
def test_verified_mla_backends_support_cross_layer_kv_cache(backend_path):
|
||||
"""Backends whose decode kernels honor the cache's block-dim stride opt
|
||||
in to the cross-layer layout with a non-identity permutation placing
|
||||
num_blocks first in physical layout."""
|
||||
module_path, name = backend_path.rsplit(".", 1)
|
||||
backend = getattr(
|
||||
pytest.importorskip(module_path, reason="backend deps unavailable"), name
|
||||
)
|
||||
|
||||
stride_order = backend.get_kv_cache_stride_order(include_num_layers_dimension=True)
|
||||
assert stride_order == (1, 0, 2, 3)
|
||||
assert stride_order[0] != 0 # num_blocks first => cross-layer supported
|
||||
assert backend.get_kv_cache_stride_order(include_num_layers_dimension=False) == (
|
||||
0,
|
||||
1,
|
||||
2,
|
||||
)
|
||||
|
||||
|
||||
def test_deepseek_v32_indexer_rejects_cross_layer_kv_cache():
|
||||
"""DeepseekV32Indexer returns identity permutation (layers dim first)
|
||||
to signal cross-layer KV cache is unsupported."""
|
||||
|
||||
@@ -0,0 +1,272 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Unit tests for OffloadingSpecFactory.
|
||||
|
||||
These tests verify:
|
||||
1. Pre-registration integrity — registered module paths can actually import
|
||||
and yield correct OffloadingSpec subclasses (CI sentinel against file moves).
|
||||
2. End-to-end factory → spec construction with real configs.
|
||||
3. Downstream collaboration — build_metric_definitions delegation.
|
||||
4. Error paths — unregistered specs, missing config, duplicate registration.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.config import KVTransferConfig
|
||||
from vllm.v1.kv_cache_interface import (
|
||||
FullAttentionSpec,
|
||||
KVCacheConfig,
|
||||
KVCacheGroupSpec,
|
||||
KVCacheTensor,
|
||||
)
|
||||
from vllm.v1.kv_offload.base import OffloadingSpec
|
||||
from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
|
||||
from vllm.v1.kv_offload.factory import OffloadingSpecFactory
|
||||
from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def restore_registry():
|
||||
"""Save and restore OffloadingSpecFactory._registry between tests."""
|
||||
original = dict(OffloadingSpecFactory._registry)
|
||||
yield
|
||||
OffloadingSpecFactory._registry = original
|
||||
|
||||
|
||||
def _make_vllm_config(
|
||||
spec_name: str | None = "CPUOffloadingSpec",
|
||||
cpu_bytes_to_use: int | None = None,
|
||||
store_threshold: int = 0,
|
||||
extra_config: dict | None = None,
|
||||
):
|
||||
"""Build a real VllmConfig with kv_transfer_config set for offloading."""
|
||||
from vllm.config import (
|
||||
CacheConfig,
|
||||
DeviceConfig,
|
||||
ModelConfig,
|
||||
SchedulerConfig,
|
||||
VllmConfig,
|
||||
)
|
||||
|
||||
model_config = ModelConfig(
|
||||
model="facebook/opt-125m",
|
||||
trust_remote_code=True,
|
||||
dtype="float16",
|
||||
seed=42,
|
||||
)
|
||||
scheduler_config = SchedulerConfig(
|
||||
max_num_seqs=16,
|
||||
max_num_batched_tokens=64,
|
||||
max_model_len=10000,
|
||||
enable_chunked_prefill=True,
|
||||
is_encoder_decoder=model_config.is_encoder_decoder,
|
||||
)
|
||||
cache_config = CacheConfig(
|
||||
block_size=16,
|
||||
gpu_memory_utilization=0.9,
|
||||
cache_dtype="auto",
|
||||
enable_prefix_caching=True,
|
||||
)
|
||||
|
||||
cfg = extra_config or {}
|
||||
if cpu_bytes_to_use is not None:
|
||||
cfg["cpu_bytes_to_use"] = cpu_bytes_to_use
|
||||
cfg["spec_name"] = spec_name
|
||||
if store_threshold > 0:
|
||||
cfg["store_threshold"] = store_threshold
|
||||
|
||||
kv_transfer_config = KVTransferConfig(
|
||||
kv_connector="OffloadingConnector",
|
||||
kv_role="kv_both",
|
||||
kv_connector_extra_config=cfg,
|
||||
)
|
||||
return VllmConfig(
|
||||
scheduler_config=scheduler_config,
|
||||
model_config=model_config,
|
||||
cache_config=cache_config,
|
||||
kv_transfer_config=kv_transfer_config,
|
||||
device_config=DeviceConfig("cpu"),
|
||||
)
|
||||
|
||||
|
||||
def _make_kv_cache_config():
|
||||
"""Build a minimal KVCacheConfig with one KV cache tensor."""
|
||||
num_blocks = 16
|
||||
num_kv_heads = 1
|
||||
head_size = 1
|
||||
dtype = torch.float32
|
||||
page_size = 2 * num_kv_heads * head_size * torch.finfo(dtype).bits // 8
|
||||
kv_tensor = KVCacheTensor(
|
||||
size=num_blocks * page_size, shared_by=["layer"], block_stride=0
|
||||
)
|
||||
return KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
kv_cache_tensors=[kv_tensor],
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec(
|
||||
["layer"],
|
||||
FullAttentionSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=num_kv_heads,
|
||||
head_size=head_size,
|
||||
dtype=dtype,
|
||||
),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pre-registration integrity (CI sentinel)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_pre_registered_specs_can_be_imported():
|
||||
"""If someone moves cpu/spec.py but forgets to update factory.py, CI fails."""
|
||||
for name in OffloadingSpecFactory._registry:
|
||||
cls = OffloadingSpecFactory._registry[name]()
|
||||
assert issubclass(cls, OffloadingSpec)
|
||||
|
||||
|
||||
def test_cpu_spec_registered():
|
||||
"""CPUOffloadingSpec is registered and importable."""
|
||||
cls = OffloadingSpecFactory._registry["CPUOffloadingSpec"]()
|
||||
assert cls is CPUOffloadingSpec
|
||||
|
||||
|
||||
def test_tiering_spec_registered():
|
||||
"""TieringOffloadingSpec is registered and importable."""
|
||||
cls = OffloadingSpecFactory._registry["TieringOffloadingSpec"]()
|
||||
assert cls is TieringOffloadingSpec
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Normal path — get_spec_cls
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_get_spec_cls_returns_registered_class():
|
||||
"""Registered spec_name returns correct class."""
|
||||
config = _make_vllm_config(spec_name="CPUOffloadingSpec")
|
||||
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
|
||||
assert spec_cls is CPUOffloadingSpec
|
||||
|
||||
|
||||
def test_get_spec_cls_default_to_cpu():
|
||||
"""Default spec_name (absent from config) resolves to CPUOffloadingSpec."""
|
||||
config = _make_vllm_config(spec_name=None)
|
||||
config.kv_transfer_config.kv_connector_extra_config.pop("spec_name", None)
|
||||
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
|
||||
assert spec_cls is CPUOffloadingSpec
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# End-to-end — create_spec
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_create_cpu_offloading_spec_end_to_end():
|
||||
"""Full factory → spec construction with real VllmConfig/KVCacheConfig.
|
||||
|
||||
Verifies:
|
||||
- cpu_bytes_to_use validation and num_blocks calculation
|
||||
- block_size % hash_block_size assertion
|
||||
- spec instance is CPUOffloadingSpec
|
||||
"""
|
||||
config = _make_vllm_config(cpu_bytes_to_use=65536)
|
||||
kv_cache_config = _make_kv_cache_config()
|
||||
spec = OffloadingSpecFactory.create_spec(config, kv_cache_config)
|
||||
assert isinstance(spec, CPUOffloadingSpec)
|
||||
assert spec.num_blocks > 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dynamic import via spec_module_path
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_dynamic_load_via_spec_module_path():
|
||||
"""External spec loaded via spec_module_path.
|
||||
|
||||
This is how external projects (e.g., llm-d-kv-cache SharedStorageOffloadingSpec)
|
||||
integrate with vLLM without being pre-registered in the factory.
|
||||
The fallback path: registry miss → spec_module_path → importlib.import_module.
|
||||
"""
|
||||
config = _make_vllm_config(spec_name="CPUOffloadingSpec")
|
||||
# Delete from registry to force the dynamic import path
|
||||
del OffloadingSpecFactory._registry["CPUOffloadingSpec"]
|
||||
# spec_name not in registry → falls through to spec_module_path
|
||||
config.kv_transfer_config.kv_connector_extra_config["spec_module_path"] = (
|
||||
"vllm.v1.kv_offload.cpu.spec"
|
||||
)
|
||||
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
|
||||
assert spec_cls is CPUOffloadingSpec
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Error paths
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_unregistered_spec_without_module_path_raises():
|
||||
"""spec_name not in registry + no spec_module_path → ValueError."""
|
||||
config = _make_vllm_config(spec_name="NonexistentSpec")
|
||||
with pytest.raises(ValueError, match="Unsupported spec type"):
|
||||
OffloadingSpecFactory.get_spec_cls(config)
|
||||
|
||||
# create_spec should also fail (calls get_spec_cls internally)
|
||||
kv_cache_config = _make_kv_cache_config()
|
||||
with pytest.raises(ValueError, match="Unsupported spec type"):
|
||||
OffloadingSpecFactory.create_spec(config, kv_cache_config)
|
||||
|
||||
|
||||
def test_cpu_spec_missing_cpu_bytes_to_use_raises():
|
||||
"""CPUOffloadingSpec requires cpu_bytes_to_use → Exception."""
|
||||
config = _make_vllm_config(cpu_bytes_to_use=None)
|
||||
config.kv_transfer_config.kv_connector_extra_config.pop("cpu_bytes_to_use", None)
|
||||
kv_cache_config = _make_kv_cache_config()
|
||||
with pytest.raises(Exception, match="cpu_bytes_to_use must be specified"):
|
||||
OffloadingSpecFactory.create_spec(config, kv_cache_config)
|
||||
|
||||
|
||||
def test_duplicate_registration_raises():
|
||||
"""register_spec with existing name → ValueError."""
|
||||
with pytest.raises(ValueError, match="is already registered"):
|
||||
OffloadingSpecFactory.register_spec(
|
||||
"CPUOffloadingSpec", "some.module", "SomeClass"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Downstream collaboration — build_metric_definitions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_build_metric_definitions_empty_below_threshold():
|
||||
"""store_threshold < 2 → only base metric (no stores_skipped)."""
|
||||
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
|
||||
|
||||
config = _make_vllm_config(store_threshold=1)
|
||||
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
|
||||
metrics = spec_cls.build_metric_definitions(
|
||||
config.kv_transfer_config.kv_connector_extra_config
|
||||
)
|
||||
assert CPUOffloadingMetrics.STORES_SKIPPED not in metrics
|
||||
|
||||
|
||||
def test_build_metric_definitions_returns_counter_at_threshold():
|
||||
"""store_threshold >= 2 → returns stores_skipped counter definition."""
|
||||
from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics
|
||||
|
||||
config = _make_vllm_config(store_threshold=2)
|
||||
spec_cls = OffloadingSpecFactory.get_spec_cls(config)
|
||||
metrics = spec_cls.build_metric_definitions(
|
||||
config.kv_transfer_config.kv_connector_extra_config
|
||||
)
|
||||
assert CPUOffloadingMetrics.STORES_SKIPPED in metrics
|
||||
@@ -0,0 +1,152 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Unit tests for SecondaryTierFactory.
|
||||
|
||||
These tests verify:
|
||||
1. Pre-registration integrity — registered tier module paths can import
|
||||
and yield correct SecondaryTierManager subclasses (CI sentinel).
|
||||
2. Multi-tier creation via factory with correct tier_type propagation.
|
||||
3. Error paths — missing tier_type, unknown tier_type, duplicate registration.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.v1.kv_offload.tiering.base import SecondaryTierManager
|
||||
from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager
|
||||
from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def restore_registry():
|
||||
"""Save and restore SecondaryTierFactory._registry between tests."""
|
||||
original = dict(SecondaryTierFactory._registry)
|
||||
yield
|
||||
SecondaryTierFactory._registry = original
|
||||
|
||||
|
||||
def _make_mock_args():
|
||||
"""Build common mock args for create_secondary_tier."""
|
||||
return MagicMock(), MagicMock() # primary_kv_view, offloading_spec
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pre-registration integrity (CI sentinel)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_pre_registered_tiers_can_be_imported():
|
||||
"""CI sentinel: example/fs/obj paths must import and yield SecondaryTierManager."""
|
||||
for tier_type in SecondaryTierFactory._registry:
|
||||
cls = SecondaryTierFactory._registry[tier_type]()
|
||||
assert issubclass(cls, SecondaryTierManager)
|
||||
|
||||
|
||||
def test_example_tier_registered():
|
||||
"""Example tier is registered."""
|
||||
cls = SecondaryTierFactory._registry["example"]()
|
||||
assert cls is ExampleSecondaryTierManager
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Normal path — create_secondary_tier
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_create_tier_from_registry():
|
||||
"""Registered tier_type creates instance with correct tier_type."""
|
||||
primary_kv_view, offloading_spec = _make_mock_args()
|
||||
tier_config = {"type": "example"}
|
||||
|
||||
tier = SecondaryTierFactory.create_secondary_tier(
|
||||
tier_config, primary_kv_view, offloading_spec
|
||||
)
|
||||
|
||||
assert isinstance(tier, SecondaryTierManager)
|
||||
assert tier.tier_type == "example"
|
||||
|
||||
|
||||
def test_create_multiple_tiers():
|
||||
"""Multiple tier configs can be created with correct tier_types."""
|
||||
primary_kv_view, offloading_spec = _make_mock_args()
|
||||
configs = [
|
||||
{"type": "example", "custom_param": 1},
|
||||
{"type": "example", "custom_param": 2},
|
||||
]
|
||||
|
||||
tiers = [
|
||||
SecondaryTierFactory.create_secondary_tier(
|
||||
cfg.copy(), primary_kv_view, offloading_spec
|
||||
)
|
||||
for cfg in configs
|
||||
]
|
||||
|
||||
assert len(tiers) == 2
|
||||
assert all(tier.tier_type == "example" for tier in tiers)
|
||||
assert all(isinstance(tier, ExampleSecondaryTierManager) for tier in tiers)
|
||||
|
||||
|
||||
def test_register_new_tier_type():
|
||||
"""Verify that new tier types can be registered and created.
|
||||
|
||||
This is how external projects add custom secondary tiers
|
||||
(e.g., llm-d FS backend was upstreamed as "fs" tier via this mechanism).
|
||||
"""
|
||||
# Register a new tier type (reuse example manager for simplicity)
|
||||
SecondaryTierFactory.register_tier(
|
||||
"custom_tier",
|
||||
"vllm.v1.kv_offload.tiering.example.manager",
|
||||
"ExampleSecondaryTierManager",
|
||||
)
|
||||
|
||||
primary_kv_view, offloading_spec = _make_mock_args()
|
||||
tier = SecondaryTierFactory.create_secondary_tier(
|
||||
{"type": "custom_tier", "custom_param": 99},
|
||||
primary_kv_view,
|
||||
offloading_spec,
|
||||
)
|
||||
|
||||
assert tier.tier_type == "custom_tier"
|
||||
assert isinstance(tier, ExampleSecondaryTierManager)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Error paths
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_missing_tier_type_raises():
|
||||
"""tier_config without 'type' → ValueError."""
|
||||
primary_kv_view, offloading_spec = _make_mock_args()
|
||||
tier_config: dict[str, str] = {}
|
||||
|
||||
with pytest.raises(ValueError, match="must include 'type'"):
|
||||
SecondaryTierFactory.create_secondary_tier(
|
||||
tier_config, primary_kv_view, offloading_spec
|
||||
)
|
||||
|
||||
|
||||
def test_unknown_tier_type_raises():
|
||||
"""Unrecognized tier_type → ValueError with supported types list."""
|
||||
primary_kv_view, offloading_spec = _make_mock_args()
|
||||
tier_config = {"type": "nonexistent_tier"}
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"Unknown secondary tier type.*Supported types:",
|
||||
):
|
||||
SecondaryTierFactory.create_secondary_tier(
|
||||
tier_config, primary_kv_view, offloading_spec
|
||||
)
|
||||
|
||||
|
||||
def test_duplicate_registration_raises():
|
||||
"""register_tier with existing type → ValueError."""
|
||||
with pytest.raises(ValueError, match="is already registered"):
|
||||
SecondaryTierFactory.register_tier("example", "some.module", "SomeClass")
|
||||
@@ -25,8 +25,6 @@ import regex as re
|
||||
# from "skip" to "silent", remove its directory from SEPARATE_GROUPS.
|
||||
SEPARATE_GROUPS = [
|
||||
"tests",
|
||||
# v0 related
|
||||
"vllm/lora",
|
||||
]
|
||||
|
||||
# TODO(woosuk): Include the code from Megatron and HuggingFace.
|
||||
|
||||
@@ -139,7 +139,7 @@ class RerankRequest(ScoringRequestMixin):
|
||||
# --8<-- [start:rerank-request-params]
|
||||
query: ScoreInput
|
||||
documents: ScoreInput | list[ScoreInput]
|
||||
top_n: int = Field(default_factory=lambda: 0)
|
||||
top_n: int = Field(default=0, ge=0)
|
||||
# --8<-- [end:rerank-request-params]
|
||||
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ from vllm.forward_context import (
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
LinearBase,
|
||||
QuantizeMethodBase,
|
||||
ReplicatedLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
@@ -182,6 +183,14 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
|
||||
lora_b, non_blocking=True
|
||||
)
|
||||
|
||||
def _get_quant_method(self) -> QuantizeMethodBase:
|
||||
quant_method = self.base_layer.quant_method
|
||||
if quant_method is None:
|
||||
raise RuntimeError(
|
||||
f"{type(self.base_layer).__name__} must define quant_method for LoRA."
|
||||
)
|
||||
return quant_method
|
||||
|
||||
def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
|
||||
# is_forward_context_available for tower modules
|
||||
if self._enable_aux_cuda_stream and is_forward_context_available():
|
||||
@@ -195,7 +204,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
|
||||
def _apply_sync(
|
||||
self, x: torch.Tensor, bias: torch.Tensor | None = None
|
||||
) -> torch.Tensor:
|
||||
output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
|
||||
output = self._get_quant_method().apply(self.base_layer, x, bias)
|
||||
return self._apply_lora_to_output(x, output)
|
||||
|
||||
def _apply_base_forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
@@ -242,7 +251,7 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):
|
||||
output_size = sum(self.output_slices)
|
||||
|
||||
def base_fn() -> torch.Tensor:
|
||||
return self.base_layer.quant_method.apply(self.base_layer, x, bias)
|
||||
return self._get_quant_method().apply(self.base_layer, x, bias)
|
||||
|
||||
def lora_fn() -> torch.Tensor:
|
||||
# Must be zeros, not empty: _lora_expand_kernel exits early (without
|
||||
|
||||
@@ -33,7 +33,7 @@ def _mcp_apply(x, bias, layer: "ColumnParallelLinearWithLoRA"):
|
||||
== len(layer.output_slices)
|
||||
)
|
||||
|
||||
output = layer.base_layer.quant_method.apply(layer.base_layer, x, bias)
|
||||
output = layer._get_quant_method().apply(layer.base_layer, x, bias)
|
||||
|
||||
x = x.view(-1, x.shape[-1])
|
||||
output, out_orig_shape = output.view(-1, output.shape[-1]), output.shape
|
||||
@@ -73,6 +73,8 @@ def _mcp_apply(x, bias, layer: "ColumnParallelLinearWithLoRA"):
|
||||
)
|
||||
|
||||
if not current_platform.can_update_inplace():
|
||||
if lora_output is None:
|
||||
raise RuntimeError("LoRA expand must return an output tensor.")
|
||||
output = lora_output
|
||||
|
||||
output = output.view(*out_orig_shape)
|
||||
@@ -327,12 +329,16 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
|
||||
|
||||
def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
|
||||
merged_cls = maybe_get_oot_by_class(MergedColumnParallelLinear)
|
||||
base_forward = getattr(type(self.base_layer), "forward", None)
|
||||
merged_forward = getattr(merged_cls, "forward", None)
|
||||
# Effectively unsharded subclasses can safely reuse their custom
|
||||
# forward() implementation before applying the LoRA delta.
|
||||
if (
|
||||
self.tp_size == 1
|
||||
and type(self.base_layer) is not merged_cls
|
||||
and type(self.base_layer).forward is not merged_cls.forward
|
||||
and base_forward is not None
|
||||
and merged_forward is not None
|
||||
and base_forward is not merged_forward
|
||||
):
|
||||
return self._apply_base_forward(x)
|
||||
return _mcp_apply(x, bias, self)
|
||||
@@ -482,6 +488,7 @@ class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA):
|
||||
lora_config: LoRAConfig,
|
||||
packed_modules_list: list,
|
||||
model_config: PretrainedConfig | None = None,
|
||||
decorate: bool = True,
|
||||
) -> bool:
|
||||
return (
|
||||
type(source_layer) is maybe_get_oot_by_class(QKVParallelLinear)
|
||||
@@ -523,6 +530,7 @@ class ColumnParallelLinearWithShardedLoRA(ColumnParallelLinearWithLoRA):
|
||||
lora_config: LoRAConfig,
|
||||
packed_modules_list: list,
|
||||
model_config: PretrainedConfig | None = None,
|
||||
decorate: bool = True,
|
||||
) -> bool:
|
||||
# specifying kwargs so they can be easily accessed in decorator
|
||||
return super().can_replace_layer(
|
||||
@@ -565,6 +573,7 @@ class MergedColumnParallelLinearWithShardedLoRA(MergedColumnParallelLinearWithLo
|
||||
lora_config: LoRAConfig,
|
||||
packed_modules_list: list,
|
||||
model_config: PretrainedConfig | None = None,
|
||||
decorate: bool = True,
|
||||
) -> bool:
|
||||
# specifying kwargs so they can be easily accessed in decorator
|
||||
return super().can_replace_layer(
|
||||
@@ -650,6 +659,7 @@ class MergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithLoRA):
|
||||
lora_config: LoRAConfig,
|
||||
packed_modules_list: list,
|
||||
model_config: PretrainedConfig | None = None,
|
||||
decorate: bool = True,
|
||||
) -> bool:
|
||||
# specifying kwargs so they can be easily accessed in decorator
|
||||
return super().can_replace_layer(
|
||||
@@ -678,6 +688,7 @@ class MergedColumnParallelLinearVariableSliceWithLoRA(
|
||||
lora_config: LoRAConfig,
|
||||
packed_modules_list: list,
|
||||
model_config: PretrainedConfig | None = None,
|
||||
decorate: bool = True,
|
||||
) -> bool:
|
||||
# Support MergedColumnParallelLinear with 3 or more slices
|
||||
# (2 slices are handled by MergedColumnParallelLinearWithLoRA)
|
||||
@@ -727,7 +738,7 @@ class MergedColumnParallelLinearVariableSliceWithLoRA(
|
||||
start_idx = 0
|
||||
for output_size in output_sizes:
|
||||
end_idx = start_idx + output_size
|
||||
lora_b_list.append(lora_b[start_idx:end_idx, :])
|
||||
lora_b_list.append(lora_b[start_idx:end_idx])
|
||||
start_idx = end_idx
|
||||
lora_b = lora_b_list
|
||||
|
||||
|
||||
@@ -12,10 +12,16 @@ from vllm.lora.layers.base import BaseLayerWithLoRA
|
||||
from vllm.model_executor.custom_op import maybe_get_oot_by_class
|
||||
from vllm.model_executor.layers.fused_moe import MoERunner
|
||||
from vllm.model_executor.layers.fused_moe.experts.lora_context import MoELoRAContext
|
||||
from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import (
|
||||
LoRAExpertsMixin,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
|
||||
FusedMoEModularMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
|
||||
from vllm.model_executor.layers.fused_moe.modular_kernel import (
|
||||
FusedMoEKernel,
|
||||
FusedMoEKernelModularImpl,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
|
||||
MoEPrepareAndFinalizeNoDPEPModular,
|
||||
)
|
||||
@@ -58,6 +64,13 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
|
||||
routed_experts._ensure_moe_quant_config_init()
|
||||
if getattr(routed_experts.quant_method, "supports_internal_mk", False):
|
||||
moe_kernel = routed_experts.quant_method.moe_kernel
|
||||
assert moe_kernel is not None, (
|
||||
"Fused MoE quant method must provide a moe_kernel."
|
||||
)
|
||||
# Don't let the kernel own shared experts so the runner can
|
||||
# overlap them with routed experts via a separate CUDA stream.
|
||||
assert isinstance(moe_kernel.impl, FusedMoEKernelModularImpl)
|
||||
moe_kernel.impl.shared_experts = None
|
||||
else:
|
||||
prepare_finalize = MoEPrepareAndFinalizeNoDPEPModular()
|
||||
moe_kernel = FusedMoEKernel(
|
||||
@@ -405,7 +418,11 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
|
||||
def set_mapping(self, punica_wrapper):
|
||||
super().set_mapping(punica_wrapper)
|
||||
lora_context = self._build_lora_context()
|
||||
self._moe_kernel.fused_experts.set_lora_context(lora_context)
|
||||
fused_experts = self._moe_kernel.fused_experts
|
||||
assert isinstance(fused_experts, LoRAExpertsMixin), (
|
||||
f"{type(fused_experts).__name__} does not support LoRA context setup."
|
||||
)
|
||||
fused_experts.set_lora_context(lora_context)
|
||||
prepare_finalize = self._moe_kernel.prepare_finalize
|
||||
if hasattr(prepare_finalize, "set_lora_context"):
|
||||
prepare_finalize.set_lora_context(lora_context)
|
||||
@@ -482,9 +499,13 @@ class FusedMoE3DWithLoRA(FusedMoEWithLoRA):
|
||||
) -> None:
|
||||
"""Initializes lora matrices."""
|
||||
|
||||
assert isinstance(model_config, PretrainedConfig)
|
||||
if model_config is None:
|
||||
raise ValueError("model_config must be provided for MoE LoRA.")
|
||||
architectures = model_config.architectures
|
||||
if not architectures:
|
||||
raise ValueError("model_config.architectures must be defined for MoE LoRA.")
|
||||
self._verify_ep_fs(lora_config)
|
||||
self._base_model = model_config.architectures[0]
|
||||
self._base_model = architectures[0]
|
||||
self.max_loras = lora_config.max_loras
|
||||
self.fully_sharded = lora_config.fully_sharded_loras
|
||||
|
||||
|
||||
@@ -116,7 +116,7 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):
|
||||
return lora_b
|
||||
|
||||
def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
|
||||
output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
|
||||
output = self._get_quant_method().apply(self.base_layer, x, bias)
|
||||
|
||||
x = x.view(-1, x.shape[-1])
|
||||
output, out_orig_shape = output.view(-1, output.shape[-1]), output.shape
|
||||
|
||||
@@ -111,19 +111,27 @@ def try_get_optimal_moe_lora_config(
|
||||
# base MoE weight's block-wise quantization, so block_shape is omitted
|
||||
# from the config lookup — the non-quantized branch in get_default_config
|
||||
# ignores it anyway.
|
||||
config = try_get_optimal_moe_config(w1_shape, w2_shape, top_k, dtype, M).copy()
|
||||
raw_config = try_get_optimal_moe_config(w1_shape, w2_shape, top_k, dtype, M)
|
||||
config: dict[str, int | None] = dict(raw_config)
|
||||
if op_type in [
|
||||
"fused_moe_lora_w13_shrink",
|
||||
"fused_moe_lora_w2_shrink",
|
||||
]:
|
||||
block_size_n = config.get("BLOCK_SIZE_N")
|
||||
config["BLOCK_SIZE_N"] = min(
|
||||
config.get("BLOCK_SIZE_N", 64), next_power_of_2(rank)
|
||||
block_size_n if block_size_n is not None else 64,
|
||||
next_power_of_2(rank),
|
||||
)
|
||||
elif op_type in [
|
||||
"fused_moe_lora_w13_expand",
|
||||
"fused_moe_lora_w2_expand",
|
||||
]:
|
||||
block_size_k = config.get("BLOCK_SIZE_K")
|
||||
config["BLOCK_SIZE_K"] = max(
|
||||
16, min(config.get("BLOCK_SIZE_K", 32), next_power_of_2(rank))
|
||||
16,
|
||||
min(
|
||||
block_size_k if block_size_k is not None else 32,
|
||||
next_power_of_2(rank),
|
||||
),
|
||||
)
|
||||
return config
|
||||
|
||||
@@ -245,9 +245,10 @@ class LoRAModel:
|
||||
from tensorizer import TensorDeserializer
|
||||
|
||||
tensorizer_config = TensorizerConfig(**tensorizer_config_dict)
|
||||
lora_tensor_path = os.path.join(
|
||||
tensorizer_config.tensorizer_dir, "adapter_model.tensors"
|
||||
)
|
||||
tensorizer_dir = tensorizer_config.tensorizer_dir
|
||||
if tensorizer_dir is None:
|
||||
raise ValueError("tensorizer_dir must be set in tensorizer config.")
|
||||
lora_tensor_path = os.path.join(tensorizer_dir, "adapter_model.tensors")
|
||||
tensorizer_args = tensorizer_config._construct_tensorizer_args()
|
||||
tensors = TensorDeserializer(
|
||||
lora_tensor_path,
|
||||
|
||||
+17
-10
@@ -34,6 +34,7 @@ from vllm.lora.utils import (
|
||||
from vllm.model_executor.layers.fused_moe import MoERunner
|
||||
from vllm.model_executor.models import (
|
||||
SupportsLoRA,
|
||||
SupportsMultiModal,
|
||||
is_pooling_model,
|
||||
supports_multimodal,
|
||||
)
|
||||
@@ -50,6 +51,12 @@ T = TypeVar("T")
|
||||
DEFAULT_LANGUAGE_WRAPPER_KEY = "language_model"
|
||||
|
||||
|
||||
class SupportsLoRAModel(nn.Module, SupportsLoRA): ...
|
||||
|
||||
|
||||
class SupportsLoRAMultiModalModel(SupportsLoRAModel, SupportsMultiModal): ...
|
||||
|
||||
|
||||
class AdapterLRUCache(LRUCache[int, T]):
|
||||
def __init__(self, capacity: int, deactivate_fn: Callable[[int], object]):
|
||||
super().__init__(capacity)
|
||||
@@ -66,13 +73,13 @@ class LoRAModelManager:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: SupportsLoRA,
|
||||
model: SupportsLoRAModel,
|
||||
max_num_seqs: int,
|
||||
max_num_batched_tokens: int,
|
||||
vocab_size: int,
|
||||
lora_config: LoRAConfig,
|
||||
device: torch.device,
|
||||
vllm_config: VllmConfig | None = None,
|
||||
vllm_config: VllmConfig,
|
||||
):
|
||||
"""Create a LoRAModelManager and adapter for a given model.
|
||||
|
||||
@@ -85,7 +92,7 @@ class LoRAModelManager:
|
||||
vocab_size: the vocab size of the model.
|
||||
lora_config: the LoRA configuration.
|
||||
"""
|
||||
self.model: SupportsLoRA = model
|
||||
self.model: SupportsLoRAModel = model
|
||||
self.supported_lora_modules = get_supported_lora_modules(self.model)
|
||||
assert self.supported_lora_modules, (
|
||||
f"No supported LoRA modules found in {self.model.__class__.__name__}."
|
||||
@@ -106,7 +113,6 @@ class LoRAModelManager:
|
||||
self.is_pooling_model = is_pooling_model(self.model)
|
||||
self.packed_modules: dict[str, list[str]] = {}
|
||||
self.modules: dict[str, BaseLayerWithLoRA] = {}
|
||||
# Dict instead of a set for compatibility with LRUCache.
|
||||
self._last_mapping: LoRAMapping | None = None
|
||||
is_moe = is_moe_model(self.model)
|
||||
self._is_moe = is_moe
|
||||
@@ -272,6 +278,7 @@ class LoRAModelManager:
|
||||
|
||||
@property
|
||||
def capacity(self) -> int:
|
||||
assert self.lora_config.max_cpu_loras is not None
|
||||
return self.lora_config.max_cpu_loras
|
||||
|
||||
@property
|
||||
@@ -1156,7 +1163,7 @@ class LoRAModelManager:
|
||||
|
||||
|
||||
class LoRALRUCache(AdapterLRUCache[LoRAModel]):
|
||||
def __init__(self, capacity: int, deactivate_lora_fn: Callable[[int], bool]):
|
||||
def __init__(self, capacity: int, deactivate_lora_fn: Callable[[int], object]):
|
||||
super().__init__(capacity, deactivate_lora_fn)
|
||||
|
||||
|
||||
@@ -1165,13 +1172,13 @@ class LRUCacheLoRAModelManager(LoRAModelManager):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: nn.Module,
|
||||
model: SupportsLoRAModel,
|
||||
max_num_seqs: int,
|
||||
max_num_batched_tokens: int,
|
||||
vocab_size: int,
|
||||
lora_config: LoRAConfig,
|
||||
device: torch.device,
|
||||
vllm_config: VllmConfig | None = None,
|
||||
vllm_config: VllmConfig,
|
||||
):
|
||||
super().__init__(
|
||||
model,
|
||||
@@ -1182,10 +1189,10 @@ class LRUCacheLoRAModelManager(LoRAModelManager):
|
||||
device,
|
||||
vllm_config,
|
||||
)
|
||||
self._registered_adapters: LoRALRUCache = LoRALRUCache(
|
||||
self._registered_adapters: LoRALRUCache = LoRALRUCache( # type: ignore[assignment]
|
||||
self.capacity, self.deactivate_adapter
|
||||
)
|
||||
self._active_adapters: LoRALRUCache = LoRALRUCache(
|
||||
self._active_adapters: LoRALRUCache = LoRALRUCache( # type: ignore[assignment]
|
||||
self.lora_slots, self._deactivate_adapter
|
||||
)
|
||||
|
||||
@@ -1248,7 +1255,7 @@ class LRUCacheLoRAModelManager(LoRAModelManager):
|
||||
|
||||
|
||||
def create_lora_manager(
|
||||
model: nn.Module,
|
||||
model: SupportsLoRAModel,
|
||||
max_num_seqs: int,
|
||||
max_num_batched_tokens: int,
|
||||
vocab_size: int,
|
||||
|
||||
@@ -91,9 +91,11 @@ class PEFTHelper:
|
||||
tensorizer_args = tensorizer_config._construct_tensorizer_args()
|
||||
from tensorizer.stream_io import open_stream
|
||||
|
||||
lora_config_path = os.path.join(
|
||||
tensorizer_config.tensorizer_dir, "adapter_config.json"
|
||||
)
|
||||
tensorizer_dir = tensorizer_config.tensorizer_dir
|
||||
if tensorizer_dir is None:
|
||||
raise ValueError("tensorizer_dir must be set in tensorizer config.")
|
||||
|
||||
lora_config_path = os.path.join(tensorizer_dir, "adapter_config.json")
|
||||
with open_stream(
|
||||
lora_config_path, mode="rb", **tensorizer_args.stream_kwargs
|
||||
) as f:
|
||||
|
||||
+16
-5
@@ -173,11 +173,18 @@ def parse_fine_tuned_lora_name(
|
||||
# mapping correctly.
|
||||
if name.startswith("base_model.model."):
|
||||
name = name.replace("base_model.model.", "")
|
||||
name = weights_mapper._map_name(name) if weights_mapper else name
|
||||
# recover the prefix `base_model.model.`
|
||||
name = "base_model.model." + name
|
||||
if weights_mapper:
|
||||
mapped_name = weights_mapper._map_name(name)
|
||||
if mapped_name is None:
|
||||
raise ValueError("Mapped LoRA weight name cannot be None.")
|
||||
# recover the prefix `base_model.model.`
|
||||
name = "base_model.model." + mapped_name
|
||||
else:
|
||||
name = weights_mapper._map_name(name) if weights_mapper else name
|
||||
if weights_mapper:
|
||||
mapped_name = weights_mapper._map_name(name)
|
||||
if mapped_name is None:
|
||||
raise ValueError("Mapped LoRA weight name cannot be None.")
|
||||
name = mapped_name
|
||||
|
||||
# In some situations, we may not start with `base_model.model.`.
|
||||
# If we don't (e.g., ibm-granite/granite-speech-3.3-8b),
|
||||
@@ -185,7 +192,11 @@ def parse_fine_tuned_lora_name(
|
||||
start_index = 2 if name.startswith("base_model.model.") else 0
|
||||
|
||||
parts = name.split(".")
|
||||
if parts[-1] == "weight" and (parts[-2] == "lora_A" or parts[-2] == "lora_B"):
|
||||
if (
|
||||
parts[-1] == "weight"
|
||||
and len(parts) >= 2
|
||||
and (parts[-2] == "lora_A" or parts[-2] == "lora_B")
|
||||
):
|
||||
new_name = ".".join(parts[start_index:-2])
|
||||
return new_name, parts[-2] == "lora_A"
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ from typing import Any, Literal
|
||||
import torch
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.config.lora import LoRAConfig
|
||||
from vllm.exceptions import LoRAAdapterNotFoundError
|
||||
from vllm.logger import init_logger
|
||||
from vllm.lora.lora_model import LoRAModel
|
||||
@@ -45,7 +46,10 @@ class WorkerLoRAManager:
|
||||
vllm_config.scheduler_config.max_num_batched_tokens
|
||||
)
|
||||
self.vocab_size = vllm_config.model_config.get_vocab_size()
|
||||
self.lora_config = vllm_config.lora_config
|
||||
lora_config = vllm_config.lora_config
|
||||
if lora_config is None:
|
||||
raise ValueError("LoRA config must be set for WorkerLoRAManager.")
|
||||
self.lora_config: LoRAConfig = lora_config
|
||||
|
||||
# Use get_text_config() in case of multimodal models
|
||||
text_config = vllm_config.model_config.hf_config.get_text_config()
|
||||
@@ -81,8 +85,10 @@ class WorkerLoRAManager:
|
||||
def create_lora_manager(
|
||||
self,
|
||||
model: torch.nn.Module,
|
||||
vllm_config: VllmConfig | None = None,
|
||||
vllm_config: VllmConfig,
|
||||
) -> Any:
|
||||
if vllm_config is None:
|
||||
raise ValueError("vllm_config must be provided to create a LoRA manager.")
|
||||
lora_manager = create_lora_manager(
|
||||
model,
|
||||
max_num_seqs=self.max_num_seqs,
|
||||
@@ -240,8 +246,10 @@ class LRUCacheWorkerLoRAManager(WorkerLoRAManager):
|
||||
def create_lora_manager(
|
||||
self,
|
||||
model: torch.nn.Module,
|
||||
vllm_config: VllmConfig | None = None,
|
||||
vllm_config: VllmConfig,
|
||||
) -> Any:
|
||||
if vllm_config is None:
|
||||
raise ValueError("vllm_config must be provided to create a LoRA manager.")
|
||||
lora_manager = create_lora_manager(
|
||||
model,
|
||||
lora_manager_cls=self._manager_cls,
|
||||
|
||||
@@ -9,7 +9,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
pack_quantized_values_into_int32,
|
||||
unpack_quantized_values_into_int32,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.platforms import CpuArchEnum, current_platform
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
|
||||
@@ -214,5 +214,7 @@ def _get_isa_hint(dtype: torch.dtype) -> str:
|
||||
supports_amx = torch.cpu._is_amx_tile_supported()
|
||||
if supports_amx and dtype in (torch.bfloat16,):
|
||||
return "amx"
|
||||
elif current_platform.get_cpu_architecture() == CpuArchEnum.RISCV:
|
||||
return "rvv"
|
||||
else:
|
||||
return "vec"
|
||||
|
||||
@@ -1212,9 +1212,9 @@ class MLACommonBackend(AttentionBackend):
|
||||
include_num_layers_dimension: bool = False,
|
||||
) -> tuple[int, ...]:
|
||||
if include_num_layers_dimension:
|
||||
# MLA kernels require contiguous per-layer KV cache views.
|
||||
# Identity permutation keeps num_layers first in physical
|
||||
# layout, signaling cross-layer allocation is unsupported.
|
||||
# Default to identity permutation to signal cross-layer allocation
|
||||
# is unsupported. Each MLA backend must opt in to support cross-layer
|
||||
# allocation by overriding this method.
|
||||
return (0, 1, 2, 3)
|
||||
return (0, 1, 2)
|
||||
|
||||
|
||||
@@ -1029,6 +1029,7 @@ class FusedMoEKernelModularImpl:
|
||||
):
|
||||
self.prepare_finalize = prepare_finalize
|
||||
self.fused_experts = fused_experts
|
||||
self.shared_experts: SharedExperts | None = None
|
||||
moe_parallel_config = fused_experts.moe_config.moe_parallel_config
|
||||
self.moe_parallel_config = moe_parallel_config
|
||||
self.is_dp_ep = (
|
||||
|
||||
@@ -185,11 +185,17 @@ class INCWNA16LinearScheme(INCLinearScheme):
|
||||
|
||||
|
||||
class INCXPULinearBase(INCLinearScheme):
|
||||
# AWQ packs nibbles within each int32 in the order [0, 2, 4, 6, 1, 3, 5, 7];
|
||||
# this permutation undoes that ordering so values can be repacked in
|
||||
# standard sequential (GPTQ) order.
|
||||
_REVERSE_AWQ_PACK_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
|
||||
|
||||
def __init__(self, layer_config: "INCLayerConfig") -> None:
|
||||
self.weight_bits = layer_config.bits
|
||||
self.group_size = layer_config.group_size
|
||||
self.sym = layer_config.sym
|
||||
self.pack_factor = 32 // self.weight_bits
|
||||
self.is_awq_packed = layer_config.is_awq
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
@@ -206,18 +212,34 @@ class INCXPULinearBase(INCLinearScheme):
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
scales_and_zp_size = input_size_per_partition // self.group_size
|
||||
|
||||
qweight = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition // self.pack_factor,
|
||||
output_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=0,
|
||||
packed_factor=self.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
if self.is_awq_packed:
|
||||
# AWQ: qweight [in, out // pack_factor] packed along output dim
|
||||
qweight = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition,
|
||||
output_size_per_partition // self.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=1,
|
||||
packed_factor=self.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
else:
|
||||
# GPTQ: qweight [in // pack_factor, out] packed along input dim
|
||||
qweight = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition // self.pack_factor,
|
||||
output_size_per_partition,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=0,
|
||||
packed_factor=self.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
scales = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
scales_and_zp_size,
|
||||
@@ -228,6 +250,8 @@ class INCXPULinearBase(INCLinearScheme):
|
||||
output_dim=1,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
# Both AWQ and GPTQ checkpoints store qzeros with this shape; for
|
||||
# symmetric quantization the values are ignored downstream.
|
||||
qzeros = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
scales_and_zp_size,
|
||||
@@ -255,6 +279,37 @@ class INCXPULinearBase(INCLinearScheme):
|
||||
)
|
||||
layer.register_parameter("g_idx", g_idx)
|
||||
|
||||
def _convert_awq_qweight_to_gptq(self, qw: torch.Tensor) -> torch.Tensor:
|
||||
"""Convert AWQ qweight [K, N // pf] to GPTQ qweight [K // pf, N].
|
||||
|
||||
AWQ packs along the output dim with a non-standard nibble order; GPTQ
|
||||
packs along the input dim with sequential nibble order. The conversion
|
||||
is lossless — it only reshuffles bits.
|
||||
"""
|
||||
size_bits = self.weight_bits
|
||||
pack_factor = self.pack_factor
|
||||
mask = (1 << size_bits) - 1
|
||||
device = qw.device
|
||||
reverse_order = torch.tensor(
|
||||
self._REVERSE_AWQ_PACK_ORDER, dtype=torch.long, device=device
|
||||
)
|
||||
shifts = torch.arange(0, 32, size_bits, dtype=torch.int32, device=device)
|
||||
|
||||
K, N_packed = qw.shape
|
||||
N = N_packed * pack_factor
|
||||
|
||||
# Unpack int32 → individual values, fix AWQ nibble ordering
|
||||
unpacked = (qw.unsqueeze(-1) >> shifts) & mask # (K, N_packed, pf)
|
||||
unpacked = unpacked[:, :, reverse_order]
|
||||
unpacked = unpacked.reshape(K, N) # (K, N)
|
||||
|
||||
# Repack along input dim (dim 0) in sequential nibble order
|
||||
unpacked = unpacked.reshape(K // pack_factor, pack_factor, N)
|
||||
new_qw = (unpacked.to(torch.int32) << shifts[None, :, None]).sum(
|
||||
dim=1, dtype=torch.int32
|
||||
)
|
||||
return new_qw.contiguous()
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
@@ -276,10 +331,24 @@ class INCXPULinearBase(INCLinearScheme):
|
||||
|
||||
|
||||
class INCXPULinearMethod(INCXPULinearBase):
|
||||
"""XPU linear method for INC w4a16 quantization (symmetric only).
|
||||
|
||||
Supports both GPTQ-packed (``auto_round:auto_gptq``) and AWQ-packed
|
||||
(``auto_round:auto_awq``) AutoRound checkpoints. AWQ-packed qweights are
|
||||
losslessly repacked into the GPTQ-style nibble layout during
|
||||
``process_weights_after_loading``, before the final oneDNN "NT" transpose
|
||||
that ``torch.ops._xpu_C.int4_gemm_w4a16`` expects.
|
||||
"""
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
device = layer.qweight.data.device
|
||||
|
||||
qweight_ct = layer.qweight.data.t().contiguous()
|
||||
qweight_data = layer.qweight.data
|
||||
if self.is_awq_packed:
|
||||
# Lossless repack: AWQ [K, N // pf] → GPTQ [K // pf, N]
|
||||
qweight_data = self._convert_awq_qweight_to_gptq(qweight_data)
|
||||
|
||||
qweight_ct = qweight_data.t().contiguous()
|
||||
layer.qweight = Parameter(qweight_ct.t(), requires_grad=False)
|
||||
layer.scales = Parameter(layer.scales.data, requires_grad=False)
|
||||
layer.qzeros = Parameter(
|
||||
@@ -370,7 +439,11 @@ class INCARKLinearMethod(INCXPULinearBase):
|
||||
ark_linear.to(layer.qweight.device)
|
||||
|
||||
with torch.no_grad():
|
||||
ark_linear.qweight.copy_(layer.qweight.detach())
|
||||
qweight_src = layer.qweight.detach()
|
||||
if self.is_awq_packed:
|
||||
# ARK consumes GPTQ-style packed nibbles; convert AWQ losslessly.
|
||||
qweight_src = self._convert_awq_qweight_to_gptq(qweight_src)
|
||||
ark_linear.qweight.copy_(qweight_src)
|
||||
if hasattr(layer, "qzeros") and layer.qzeros is not None:
|
||||
ark_linear.qzeros.copy_(layer.qzeros.detach())
|
||||
else:
|
||||
|
||||
@@ -42,6 +42,7 @@ from .interfaces_base import VllmModel
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.lora.model_manager import LoRAModelManager
|
||||
from vllm.model_executor.models.utils import WeightsMapper
|
||||
from vllm.multimodal.inputs import MultiModalFeatureSpec
|
||||
from vllm.multimodal.registry import _ProcessorFactories
|
||||
@@ -554,6 +555,7 @@ class SupportsLoRA(Protocol):
|
||||
packed_modules_mapping: dict[str, list[str]] = {}
|
||||
# Module prefixes to skip during LoRA loading (e.g., ["mtp."] for MTP layers)
|
||||
lora_skip_prefixes: ClassVar[list[str]] = []
|
||||
lora_manager: "LoRAModelManager | None"
|
||||
|
||||
|
||||
# We can't use runtime_checkable with ClassVar for issubclass checks
|
||||
|
||||
@@ -24,9 +24,7 @@ from vllm.distributed import (
|
||||
get_tensor_model_parallel_world_size,
|
||||
)
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
fused_moe_make_expert_params_mapping,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping
|
||||
from vllm.model_executor.layers.layernorm import RMSNorm
|
||||
from vllm.model_executor.layers.linear import ReplicatedLinear
|
||||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||||
|
||||
@@ -182,6 +182,11 @@ class PoolingParams(
|
||||
)
|
||||
elif self.dimensions < 1:
|
||||
raise ValueError("Dimensions must be greater than 0")
|
||||
elif self.dimensions > model_config.embedding_size:
|
||||
raise ValueError(
|
||||
"Dimensions must be less than or equal to the model's "
|
||||
f"embedding size ({model_config.embedding_size})"
|
||||
)
|
||||
|
||||
elif self.task in ["classify", "token_classify"]:
|
||||
if self.use_activation is None:
|
||||
|
||||
@@ -49,6 +49,14 @@ class CutlassMLABackend(MLACommonBackend):
|
||||
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
|
||||
return [128]
|
||||
|
||||
@staticmethod
|
||||
def get_kv_cache_stride_order(
|
||||
include_num_layers_dimension: bool = False,
|
||||
) -> tuple[int, ...]:
|
||||
if include_num_layers_dimension:
|
||||
return (1, 0, 2, 3)
|
||||
return (0, 1, 2)
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "CUTLASS_MLA"
|
||||
|
||||
@@ -52,6 +52,14 @@ class FlashAttnMLABackend(MLACommonBackend):
|
||||
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
|
||||
return [MultipleOf(16)]
|
||||
|
||||
@staticmethod
|
||||
def get_kv_cache_stride_order(
|
||||
include_num_layers_dimension: bool = False,
|
||||
) -> tuple[int, ...]:
|
||||
if include_num_layers_dimension:
|
||||
return (1, 0, 2, 3)
|
||||
return (0, 1, 2)
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "FLASH_ATTN_MLA"
|
||||
|
||||
@@ -49,6 +49,14 @@ class FlashInferMLABackend(MLACommonBackend):
|
||||
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
|
||||
return [32, 64]
|
||||
|
||||
@staticmethod
|
||||
def get_kv_cache_stride_order(
|
||||
include_num_layers_dimension: bool = False,
|
||||
) -> tuple[int, ...]:
|
||||
if include_num_layers_dimension:
|
||||
return (1, 0, 2, 3)
|
||||
return (0, 1, 2)
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "FLASHINFER_MLA"
|
||||
|
||||
@@ -58,6 +58,14 @@ class FlashMLABackend(MLACommonBackend):
|
||||
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
|
||||
return [64]
|
||||
|
||||
@staticmethod
|
||||
def get_kv_cache_stride_order(
|
||||
include_num_layers_dimension: bool = False,
|
||||
) -> tuple[int, ...]:
|
||||
if include_num_layers_dimension:
|
||||
return (1, 0, 2, 3)
|
||||
return (0, 1, 2)
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "FLASHMLA"
|
||||
|
||||
@@ -57,6 +57,14 @@ class TritonMLABackend(MLACommonBackend):
|
||||
return True
|
||||
return block_size % 16 == 0
|
||||
|
||||
@staticmethod
|
||||
def get_kv_cache_stride_order(
|
||||
include_num_layers_dimension: bool = False,
|
||||
) -> tuple[int, ...]:
|
||||
if include_num_layers_dimension:
|
||||
return (1, 0, 2, 3)
|
||||
return (0, 1, 2)
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "TRITON_MLA"
|
||||
|
||||
@@ -56,6 +56,15 @@ def tanh(x):
|
||||
return 2 * tl.sigmoid(2 * x) - 1
|
||||
|
||||
|
||||
def _page_stride(buf, page_size):
|
||||
# Stride between pages. 4D buffers have a page dim; 3D buffers pack pages
|
||||
# along the token dim, so split it out first. Read the real stride (a
|
||||
# cross-layer view has gaps), don't assume PAGE_SIZE * token stride.
|
||||
if buf.ndim == 3:
|
||||
buf = buf.unflatten(-3, (-1, page_size))
|
||||
return buf.stride(-4)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _fwd_kernel_stage1(
|
||||
Q,
|
||||
@@ -68,8 +77,10 @@ def _fwd_kernel_stage1(
|
||||
stride_req_to_tokens_b,
|
||||
stride_qbs,
|
||||
stride_qh,
|
||||
stride_buf_kpbs,
|
||||
stride_buf_kbs,
|
||||
stride_buf_kh,
|
||||
stride_buf_vpbs,
|
||||
stride_buf_vbs,
|
||||
stride_buf_vh,
|
||||
stride_mid_ob,
|
||||
@@ -123,9 +134,11 @@ def _fwd_kernel_stage1(
|
||||
mask=offs_n < split_kv_end,
|
||||
other=0,
|
||||
)
|
||||
kv_loc = kv_page_number * PAGE_SIZE + offs_n % PAGE_SIZE
|
||||
kv_in_page = offs_n % PAGE_SIZE
|
||||
offs_buf_k = (
|
||||
kv_loc[:, None] * stride_buf_kbs
|
||||
(kv_page_number * stride_buf_kpbs + kv_in_page * stride_buf_kbs)[
|
||||
:, None
|
||||
]
|
||||
+ cur_kv_head * stride_buf_kh
|
||||
+ offs_d[None, :]
|
||||
)
|
||||
@@ -145,7 +158,9 @@ def _fwd_kernel_stage1(
|
||||
qk = tl.where(offs_n < split_kv_end, qk, float("-inf"))
|
||||
|
||||
offs_buf_v = (
|
||||
kv_loc[:, None] * stride_buf_vbs
|
||||
(kv_page_number * stride_buf_vpbs + kv_in_page * stride_buf_vbs)[
|
||||
:, None
|
||||
]
|
||||
+ cur_kv_head * stride_buf_vh
|
||||
+ offs_dv[None, :]
|
||||
)
|
||||
@@ -235,8 +250,10 @@ def _decode_att_m_fwd(
|
||||
Req_to_tokens.stride(0),
|
||||
q.stride(0),
|
||||
q.stride(1),
|
||||
_page_stride(k_buffer, page_size),
|
||||
k_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
|
||||
k_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
|
||||
_page_stride(v_buffer, page_size),
|
||||
v_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
|
||||
v_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
|
||||
att_out.stride(0),
|
||||
@@ -270,8 +287,10 @@ def _fwd_grouped_kernel_stage1(
|
||||
stride_req_to_tokens_b,
|
||||
stride_qbs,
|
||||
stride_qh,
|
||||
stride_buf_kpbs,
|
||||
stride_buf_kbs,
|
||||
stride_buf_kh,
|
||||
stride_buf_vpbs,
|
||||
stride_buf_vbs,
|
||||
stride_buf_vh,
|
||||
stride_mid_ob,
|
||||
@@ -357,10 +376,12 @@ def _fwd_grouped_kernel_stage1(
|
||||
other=0,
|
||||
cache_modifier=".ca",
|
||||
)
|
||||
kv_loc = kv_page_number * PAGE_SIZE + offs_n % PAGE_SIZE
|
||||
kv_off_k = (
|
||||
kv_page_number * stride_buf_kpbs + (offs_n % PAGE_SIZE) * stride_buf_kbs
|
||||
)
|
||||
|
||||
# explicitly facilitate overlapping load/compute
|
||||
offs_buf_k = kv_loc[None, :] * stride_buf_kbs + base_offs_k
|
||||
offs_buf_k = kv_off_k[None, :] + base_offs_k
|
||||
k = tl.load(
|
||||
K_Buffer + offs_buf_k,
|
||||
mask=(offs_n[None, :] < split_kv_end) & (mask_d[:, None]),
|
||||
@@ -372,7 +393,7 @@ def _fwd_grouped_kernel_stage1(
|
||||
k = (k.to(tl.float32) * ks).to(q.dtype)
|
||||
qk = tl.dot(q, k.to(q.dtype))
|
||||
if BLOCK_DPE > 0:
|
||||
offs_buf_kpe = kv_loc[None, :] * stride_buf_kbs + base_offs_kpe
|
||||
offs_buf_kpe = kv_off_k[None, :] + base_offs_kpe
|
||||
kpe = tl.load(
|
||||
K_Buffer + offs_buf_kpe,
|
||||
mask=(offs_n[None, :] < split_kv_end) & (mask_dpe[:, None]),
|
||||
@@ -392,7 +413,11 @@ def _fwd_grouped_kernel_stage1(
|
||||
)
|
||||
|
||||
if not IS_MLA:
|
||||
offs_buf_v = kv_loc[:, None] * stride_buf_vbs + base_offs_v
|
||||
kv_off_v = (
|
||||
kv_page_number * stride_buf_vpbs
|
||||
+ (offs_n % PAGE_SIZE) * stride_buf_vbs
|
||||
)
|
||||
offs_buf_v = kv_off_v[:, None] + base_offs_v
|
||||
v = tl.load(
|
||||
V_Buffer + offs_buf_v,
|
||||
mask=(offs_n[:, None] < split_kv_end) & (mask_dv[None, :]),
|
||||
@@ -517,8 +542,10 @@ def _decode_grouped_att_m_fwd(
|
||||
Req_to_tokens.stride(0),
|
||||
q.stride(0),
|
||||
q.stride(1),
|
||||
_page_stride(k_buffer, page_size),
|
||||
k_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
|
||||
k_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
|
||||
_page_stride(v_buffer, page_size),
|
||||
v_buffer.stride(-3), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
|
||||
v_buffer.stride(-2), # Assume (..., PAGE_SIZE, NUM_HEADS, HEAD_DIM)
|
||||
att_out.stride(0),
|
||||
|
||||
@@ -17,6 +17,10 @@ def noop(*args: Any, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def fake_pin_memory(self: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
|
||||
return self
|
||||
|
||||
|
||||
class _EventPlaceholder:
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
self.record = noop
|
||||
@@ -41,6 +45,7 @@ torch.cuda.set_stream = noop
|
||||
torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder()
|
||||
torch.accelerator.synchronize = noop
|
||||
torch.accelerator.empty_cache = noop
|
||||
torch.Tensor.pin_memory = fake_pin_memory
|
||||
|
||||
# Patch vLLM torch utils
|
||||
import vllm.utils.torch_utils as torch_utils
|
||||
|
||||
Reference in New Issue
Block a user