forked from Karylab-cklius/vllm
Compare commits
72
Commits
v0.18.1rc0
...
sm103
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
257c0c5b50 | ||
|
|
e9855c5c19 | ||
|
|
7cd8824477 | ||
|
|
2e67fa756d | ||
|
|
e3c6c10cad | ||
|
|
16a664df24 | ||
|
|
7281199a8c | ||
|
|
b2dd75eb48 | ||
|
|
c59a132f96 | ||
|
|
de99d91ece | ||
|
|
83c9d525b6 | ||
|
|
8f4824b664 | ||
|
|
56777b5c89 | ||
|
|
2488a82f89 | ||
|
|
dc6908ac6a | ||
|
|
e85f8f0932 | ||
|
|
5bf3c42d4c | ||
|
|
38364a7e32 | ||
|
|
fafe76b4af | ||
|
|
ffb5b32b5f | ||
|
|
91fd695b75 | ||
|
|
1cbbcfe8a3 | ||
|
|
aceadb5ee1 | ||
|
|
ec2280611a | ||
|
|
7151ae6528 | ||
|
|
45bd5c8e75 | ||
|
|
10a1018c12 | ||
|
|
aec2dc6c0d | ||
|
|
7938d12119 | ||
|
|
debd6e768c | ||
|
|
9ace378a63 | ||
|
|
27d5ee3e6f | ||
|
|
35141a7eed | ||
|
|
e99fb98867 | ||
|
|
a16133a0f1 | ||
|
|
54ab804e87 | ||
|
|
02e6efe56d | ||
|
|
410d300893 | ||
|
|
d3fe857135 | ||
|
|
f85e479e66 | ||
|
|
1f0d210641 | ||
|
|
3bbe2e1e6e | ||
|
|
6e04e79326 | ||
|
|
e7767eccae | ||
|
|
43877a620b | ||
|
|
63f49b8bd4 | ||
|
|
a5e9d511de | ||
|
|
c058ff44d4 | ||
|
|
ce9b1d76cf | ||
|
|
e74c17e153 | ||
|
|
eaf4978621 | ||
|
|
77d24c4bfe | ||
|
|
b3e846017d | ||
|
|
cd1242d82a | ||
|
|
4383f1532e | ||
|
|
6eedec6e36 | ||
|
|
ffc8531524 | ||
|
|
6ecba840d7 | ||
|
|
3b06c55c78 | ||
|
|
b050700462 | ||
|
|
5dac719b2b | ||
|
|
c862481c02 | ||
|
|
c86b17cfe6 | ||
|
|
66f927f205 | ||
|
|
e78bc74268 | ||
|
|
6b2fa3a762 | ||
|
|
eeee5b262d | ||
|
|
5ad0446572 | ||
|
|
8cc700dd6a | ||
|
|
80b70884eb | ||
|
|
61e381dcf0 | ||
|
|
88f1b374f5 |
@@ -1,12 +0,0 @@
|
||||
# For vllm script, with -t option (tensor parallel size).
|
||||
# bash ./run-lm-eval-gsm-vllm-baseline.sh -m nm-testing/SparseLlama-3.1-8B-gsm8k-pruned.2of4-chnl_wts_per_tok_dyn_act_fp8-BitM -b "auto" -t 2
|
||||
model_name: "nm-testing/SparseLlama-3.1-8B-gsm8k-pruned.2of4-chnl_wts_per_tok_dyn_act_fp8-BitM"
|
||||
tasks:
|
||||
- name: "gsm8k"
|
||||
metrics:
|
||||
- name: "exact_match,strict-match"
|
||||
value: 0.6353
|
||||
- name: "exact_match,flexible-extract"
|
||||
value: 0.637
|
||||
limit: null
|
||||
num_fewshot: null
|
||||
@@ -12,7 +12,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-wheel-arm64-cuda-12-9
|
||||
agents:
|
||||
queue: arm64_cpu_queue_postmerge
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
# #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
|
||||
# https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
|
||||
@@ -27,7 +27,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-wheel-arm64-cuda-13-0
|
||||
agents:
|
||||
queue: arm64_cpu_queue_postmerge
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
# #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
|
||||
# https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
|
||||
@@ -42,7 +42,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-wheel-arm64-cpu
|
||||
agents:
|
||||
queue: arm64_cpu_queue_postmerge
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_BUILD_ACL=ON --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
|
||||
- "mkdir artifacts"
|
||||
@@ -55,7 +55,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-wheel-x86-cuda-12-9
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
@@ -68,7 +68,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-wheel-x86-cuda-13-0
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
@@ -81,7 +81,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-wheel-x86-cpu
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
|
||||
- "mkdir artifacts"
|
||||
@@ -97,7 +97,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-release-image-x86
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
@@ -110,7 +110,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64
|
||||
agents:
|
||||
queue: arm64_cpu_queue_postmerge
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
@@ -120,7 +120,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-release-image-x86-cuda-13-0
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
@@ -133,7 +133,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64-cuda-13-0
|
||||
agents:
|
||||
queue: arm64_cpu_queue_postmerge
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
# compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark
|
||||
@@ -149,7 +149,7 @@ steps:
|
||||
- block-cpu-release-image-build
|
||||
- input-release-version
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
|
||||
@@ -167,7 +167,7 @@ steps:
|
||||
- block-arm64-cpu-release-image-build
|
||||
- input-release-version
|
||||
agents:
|
||||
queue: arm64_cpu_queue_postmerge
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
|
||||
@@ -185,7 +185,7 @@ steps:
|
||||
- build-release-image-arm64
|
||||
id: create-multi-arch-manifest
|
||||
agents:
|
||||
queue: small_cpu_queue_postmerge
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend"
|
||||
@@ -196,7 +196,7 @@ steps:
|
||||
- create-multi-arch-manifest
|
||||
id: annotate-release-workflow
|
||||
agents:
|
||||
queue: small_cpu_queue_postmerge
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/annotate-release.sh"
|
||||
|
||||
@@ -206,7 +206,7 @@ steps:
|
||||
- build-release-image-arm64-cuda-13-0
|
||||
id: create-multi-arch-manifest-cuda-13-0
|
||||
agents:
|
||||
queue: small_cpu_queue_postmerge
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130 --amend"
|
||||
@@ -217,7 +217,7 @@ steps:
|
||||
- create-multi-arch-manifest
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
agents:
|
||||
queue: small_cpu_queue_postmerge
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/push-nightly-builds.sh"
|
||||
# Clean up old nightly builds (keep only last 14)
|
||||
@@ -235,7 +235,7 @@ steps:
|
||||
- create-multi-arch-manifest-cuda-13-0
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
agents:
|
||||
queue: small_cpu_queue_postmerge
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/push-nightly-builds.sh cu130"
|
||||
# Clean up old nightly builds (keep only last 14)
|
||||
@@ -262,7 +262,7 @@ steps:
|
||||
- block-upload-release-wheels
|
||||
id: upload-release-wheels
|
||||
agents:
|
||||
queue: small_cpu_queue_postmerge
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
|
||||
|
||||
@@ -323,7 +323,7 @@ steps:
|
||||
- step: input-rocm-config
|
||||
allow_failure: true # Allow failure so non-UI builds can proceed (input step is skipped)
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
# Set configuration and check cache
|
||||
- |
|
||||
@@ -465,7 +465,7 @@ steps:
|
||||
- step: build-rocm-base-wheels
|
||||
allow_failure: false
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
# Download artifacts and prepare Docker image
|
||||
@@ -575,7 +575,7 @@ steps:
|
||||
- step: build-rocm-vllm-wheel
|
||||
allow_failure: false
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
# Download all wheel artifacts and run upload
|
||||
@@ -624,7 +624,7 @@ steps:
|
||||
- step: input-release-version
|
||||
allow_failure: true
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/annotate-rocm-release.sh"
|
||||
env:
|
||||
@@ -641,7 +641,7 @@ steps:
|
||||
depends_on: block-generate-root-index-rocm-wheels
|
||||
id: generate-root-index-rocm-wheels
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
|
||||
env:
|
||||
@@ -655,7 +655,7 @@ steps:
|
||||
- step: build-rocm-base-wheels
|
||||
allow_failure: false
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- |
|
||||
|
||||
@@ -326,8 +326,7 @@ apply_rocm_test_overrides() {
|
||||
if [[ $cmds == *" kernels/moe"* ]]; then
|
||||
cmds="${cmds} \
|
||||
--ignore=kernels/moe/test_moe.py \
|
||||
--ignore=kernels/moe/test_cutlass_moe.py \
|
||||
--ignore=kernels/moe/test_triton_moe_ptpc_fp8.py"
|
||||
--ignore=kernels/moe/test_cutlass_moe.py"
|
||||
fi
|
||||
|
||||
# --- Entrypoint ignores ---
|
||||
|
||||
@@ -1435,7 +1435,7 @@ steps:
|
||||
- pytest -v -s entrypoints/offline_mode
|
||||
|
||||
|
||||
- label: Entrypoints Integration (API Server 1) # 1h 7m
|
||||
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
@@ -1448,10 +1448,43 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
|
||||
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 2) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/openai/speech_to_text/
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
|
||||
|
||||
|
||||
- label: Entrypoints Integration (API Server 2) #26.9m
|
||||
timeout_in_minutes: 45
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
@@ -2174,6 +2207,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2204,6 +2238,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2220,6 +2255,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2234,6 +2270,7 @@ steps:
|
||||
timeout_in_minutes: 106
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2249,6 +2286,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2262,6 +2300,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2935,7 +2974,7 @@ steps:
|
||||
# #
|
||||
#####################################################################################################################################
|
||||
|
||||
- label: Entrypoints Integration (API Server 1) # TBD
|
||||
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
@@ -2948,10 +2987,43 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
|
||||
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 2) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/openai/speech_to_text/
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses
|
||||
|
||||
|
||||
- label: Entrypoints Integration (API Server 2) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
@@ -3269,6 +3341,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -3284,6 +3357,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -3299,6 +3373,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -3315,6 +3390,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -3329,6 +3405,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -3344,6 +3421,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -3357,6 +3435,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -3370,6 +3449,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -257,6 +257,17 @@ steps:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
device: a100
|
||||
|
||||
@@ -25,8 +25,8 @@ steps:
|
||||
- pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
|
||||
- pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
|
||||
|
||||
- label: Entrypoints Integration (API Server 1)
|
||||
timeout_in_minutes: 130
|
||||
- label: Entrypoints Integration (API Server openai - Part 1)
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -34,7 +34,24 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 2)
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/openai/speech_to_text/
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
mirror:
|
||||
amd:
|
||||
@@ -42,6 +59,17 @@ steps:
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3)
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
|
||||
- label: Entrypoints Integration (API Server 2)
|
||||
timeout_in_minutes: 130
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
|
||||
@@ -90,6 +90,7 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
|
||||
|
||||
|
||||
- label: GPQA Eval (GPT-OSS) (H100)
|
||||
timeout_in_minutes: 120
|
||||
device: h100
|
||||
|
||||
@@ -8,7 +8,7 @@ steps:
|
||||
- vllm/lora
|
||||
- tests/lora
|
||||
commands:
|
||||
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemoel_lora.py
|
||||
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
|
||||
parallelism: 4
|
||||
|
||||
|
||||
@@ -31,4 +31,4 @@ steps:
|
||||
- pytest -v -s -x lora/test_llm_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_densemoel_lora.py
|
||||
- pytest -v -s -x lora/test_qwen35_densemodel_lora.py
|
||||
@@ -2,11 +2,54 @@ group: Miscellaneous
|
||||
depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: V1 Others
|
||||
timeout_in_minutes: 60
|
||||
- label: V1 Spec Decode
|
||||
timeout_in_minutes: 30
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/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
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: V1 Sample + Logits
|
||||
timeout_in_minutes: 30
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/v1/sample
|
||||
- tests/v1/logits_processors
|
||||
- tests/v1/test_oracle.py
|
||||
- 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
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: V1 Core + KV + Metrics
|
||||
timeout_in_minutes: 30
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/v1/core
|
||||
- tests/v1/executor
|
||||
- tests/v1/kv_offload
|
||||
- tests/v1/worker
|
||||
- tests/v1/kv_connector/unit
|
||||
- tests/v1/metrics
|
||||
- tests/entrypoints/openai/correctness/test_lmeval.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
@@ -14,16 +57,9 @@ steps:
|
||||
- pytest -v -s -m 'not cpu_test' v1/core
|
||||
- pytest -v -s v1/executor
|
||||
- pytest -v -s v1/kv_offload
|
||||
- pytest -v -s v1/sample
|
||||
- pytest -v -s v1/logits_processors
|
||||
- pytest -v -s v1/worker
|
||||
# TODO: create another `optional` test group for slow tests
|
||||
- pytest -v -s -m 'not slow_test' v1/spec_decode
|
||||
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'not cpu_test' v1/metrics
|
||||
- pytest -v -s v1/test_oracle.py
|
||||
- pytest -v -s v1/test_request.py
|
||||
- pytest -v -s v1/test_outputs.py
|
||||
# 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
|
||||
@@ -39,7 +75,7 @@ steps:
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/v1
|
||||
device: cpu
|
||||
device: cpu-small
|
||||
commands:
|
||||
# split the test to avoid interference
|
||||
- pytest -v -s -m 'cpu_test' v1/core
|
||||
@@ -141,7 +177,7 @@ steps:
|
||||
- tests/tool_parsers
|
||||
- tests/transformers_utils
|
||||
- tests/config
|
||||
device: cpu
|
||||
device: cpu-small
|
||||
commands:
|
||||
- python3 standalone_tests/lazy_imports.py
|
||||
- pytest -v -s test_inputs.py
|
||||
@@ -156,7 +192,7 @@ steps:
|
||||
- pytest -v -s config
|
||||
|
||||
- label: Batch Invariance (H100)
|
||||
timeout_in_minutes: 25
|
||||
timeout_in_minutes: 30
|
||||
device: h100
|
||||
source_file_dependencies:
|
||||
- vllm/v1/attention
|
||||
@@ -167,6 +203,8 @@ steps:
|
||||
- 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]
|
||||
|
||||
- label: Acceptance Length Test (Large Models) # optional
|
||||
timeout_in_minutes: 25
|
||||
|
||||
@@ -87,13 +87,12 @@ steps:
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
- tests/distributed/test_pipeline_parallel.py
|
||||
#- tests/distributed/test_pp_cudagraph.py
|
||||
- tests/distributed/test_pp_cudagraph.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"
|
||||
# TODO: Uncomment once https://github.com/vllm-project/vllm/pull/35162 is merged.
|
||||
#- pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
|
||||
- pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
|
||||
|
||||
- label: Model Runner V2 Spec Decode
|
||||
timeout_in_minutes: 30
|
||||
|
||||
@@ -51,7 +51,7 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/test_utils.py
|
||||
- tests/models/test_vision.py
|
||||
device: cpu
|
||||
device: cpu-small
|
||||
commands:
|
||||
- pytest -v -s models/test_utils.py models/test_vision.py
|
||||
|
||||
|
||||
@@ -70,7 +70,7 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
- tests/models/registry.py
|
||||
device: cpu
|
||||
device: cpu-medium
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
|
||||
|
||||
+1
-1
@@ -108,7 +108,7 @@ uv.lock
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
.python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
|
||||
+7
-26
@@ -343,7 +343,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
"csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu"
|
||||
"csrc/quantization/fp4/nvfp4_quant_entry.cu"
|
||||
"csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu"
|
||||
"csrc/sparse/cutlass/sparse_scaled_mm_entry.cu"
|
||||
"csrc/cutlass_extensions/common.cpp"
|
||||
"csrc/quantization/w8a8/fp8/per_token_group_quant.cu"
|
||||
"csrc/quantization/w8a8/int8/per_token_group_quant.cu")
|
||||
@@ -619,31 +618,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
#
|
||||
# 2:4 Sparse Kernels
|
||||
|
||||
# The 2:4 sparse kernels cutlass_scaled_sparse_mm and cutlass_compressor
|
||||
# require CUDA 12.2 or later (and only work on Hopper).
|
||||
cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;" "${CUDA_ARCHS}")
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.2 AND SCALED_MM_ARCHS)
|
||||
set(SRCS "csrc/sparse/cutlass/sparse_scaled_mm_c3x.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SRCS}"
|
||||
CUDA_ARCHS "${SCALED_MM_ARCHS}")
|
||||
list(APPEND VLLM_EXT_SRC "${SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SPARSE_SCALED_MM_C3X=1")
|
||||
message(STATUS "Building sparse_scaled_mm_c3x for archs: ${SCALED_MM_ARCHS}")
|
||||
else()
|
||||
if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.2 AND SCALED_MM_ARCHS)
|
||||
message(STATUS "Not building sparse_scaled_mm_c3x kernels as CUDA Compiler version is "
|
||||
"not >= 12.2, we recommend upgrading to CUDA 12.2 or later "
|
||||
"if you intend on running FP8 sparse quantized models on Hopper.")
|
||||
else()
|
||||
message(STATUS "Not building sparse_scaled_mm_c3x as no compatible archs found "
|
||||
"in CUDA target architectures")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# The nvfp4_scaled_mm_sm120 kernels for Geforce Blackwell SM120 require
|
||||
# CUDA 12.8 or later
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
@@ -690,6 +664,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
list(APPEND VLLM_EXT_SRC "${SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
|
||||
set(VLLM_NVFP4_SM100_ENABLED TRUE)
|
||||
message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building NVFP4 as no compatible archs were found.")
|
||||
@@ -985,6 +960,12 @@ define_extension_target(
|
||||
# Setting this variable sidesteps the issue by calling the driver directly.
|
||||
target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
|
||||
|
||||
# Propagate ENABLE_NVFP4_SM100 to all languages (including C++ files such as
|
||||
# torch_bindings.cpp) so that per-SM op registrations are compiled in.
|
||||
if(VLLM_NVFP4_SM100_ENABLED)
|
||||
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
|
||||
endif()
|
||||
|
||||
# add OR VLLM_GPU_LANG STREQUAL "HIP" here once
|
||||
# https://github.com/vllm-project/vllm/issues/35163 is resolved
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
|
||||
@@ -1,517 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import itertools
|
||||
import pickle as pkl
|
||||
import time
|
||||
from collections.abc import Callable, Iterable
|
||||
|
||||
import torch
|
||||
import torch.utils.benchmark as TBenchmark
|
||||
from torch.utils.benchmark import Measurement as TMeasurement
|
||||
from utils import make_rand_sparse_tensors
|
||||
from weight_shapes import WEIGHT_SHAPES
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())
|
||||
DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512]
|
||||
DEFAULT_TP_SIZES = [1]
|
||||
|
||||
|
||||
# bench
|
||||
def bench_fn(
|
||||
label: str, sub_label: str, description: str, fn: Callable, *args, **kwargs
|
||||
) -> TMeasurement:
|
||||
min_run_time = 1
|
||||
|
||||
globals = {
|
||||
"args": args,
|
||||
"kwargs": kwargs,
|
||||
"fn": fn,
|
||||
}
|
||||
return TBenchmark.Timer(
|
||||
stmt="fn(*args, **kwargs)",
|
||||
globals=globals,
|
||||
label=label,
|
||||
sub_label=sub_label,
|
||||
description=description,
|
||||
).blocked_autorange(min_run_time=min_run_time)
|
||||
|
||||
|
||||
def bench_int8(
|
||||
dtype: torch.dtype, m: int, k: int, n: int, label: str, sub_label: str
|
||||
) -> Iterable[TMeasurement]:
|
||||
assert dtype == torch.int8
|
||||
b_compressed, e, a, b = make_rand_sparse_tensors(torch.int8, m, n, k)
|
||||
scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32)
|
||||
scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32)
|
||||
bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
out = ops.cutlass_scaled_sparse_mm(
|
||||
a, b_compressed, e, scale_a, scale_b, torch.bfloat16
|
||||
)
|
||||
out_ref = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16)
|
||||
|
||||
if not torch.allclose(out, out_ref):
|
||||
print("Incorrect results")
|
||||
print(out)
|
||||
print(out_ref)
|
||||
else:
|
||||
print("Correct results")
|
||||
|
||||
timers = []
|
||||
# pytorch impl - bfloat16
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"pytorch_bf16_bf16_bf16_matmul-no-scales",
|
||||
torch.mm,
|
||||
a.to(dtype=torch.bfloat16),
|
||||
b.to(dtype=torch.bfloat16),
|
||||
)
|
||||
)
|
||||
|
||||
# pytorch impl - float16
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"pytorch_fp16_fp16_fp16_matmul-no-scales",
|
||||
torch.mm,
|
||||
a.to(dtype=torch.float16),
|
||||
b.to(dtype=torch.float16),
|
||||
)
|
||||
)
|
||||
|
||||
# cutlass impl
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"cutlass_i8_i8_bf16_scaled_mm",
|
||||
ops.cutlass_scaled_mm,
|
||||
a,
|
||||
b,
|
||||
scale_a,
|
||||
scale_b,
|
||||
torch.bfloat16,
|
||||
)
|
||||
)
|
||||
|
||||
# cutlass with bias
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"cutlass_i8_i8_bf16_scaled_mm_bias",
|
||||
ops.cutlass_scaled_mm,
|
||||
a,
|
||||
b,
|
||||
scale_a,
|
||||
scale_b,
|
||||
torch.bfloat16,
|
||||
bias,
|
||||
)
|
||||
)
|
||||
|
||||
# cutlass sparse impl
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"cutlass_i8_i8_bf16_scaled_sparse_mm",
|
||||
ops.cutlass_scaled_sparse_mm,
|
||||
a,
|
||||
b_compressed,
|
||||
e,
|
||||
scale_a,
|
||||
scale_b,
|
||||
torch.bfloat16,
|
||||
)
|
||||
)
|
||||
|
||||
# cutlass sparse with bias
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"cutlass_i8_i8_bf16_scaled_sparse_mm_bias",
|
||||
ops.cutlass_scaled_sparse_mm,
|
||||
a,
|
||||
b_compressed,
|
||||
e,
|
||||
scale_a,
|
||||
scale_b,
|
||||
torch.bfloat16,
|
||||
bias,
|
||||
)
|
||||
)
|
||||
|
||||
return timers
|
||||
|
||||
|
||||
def bench_fp8(
|
||||
dtype: torch.dtype, m: int, k: int, n: int, label: str, sub_label: str
|
||||
) -> Iterable[TMeasurement]:
|
||||
assert dtype == torch.float8_e4m3fn
|
||||
b_compressed, e, a, b = make_rand_sparse_tensors(torch.float8_e4m3fn, m, n, k)
|
||||
scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32)
|
||||
scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32)
|
||||
bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
out = ops.cutlass_scaled_sparse_mm(
|
||||
a, b_compressed, e, scale_a, scale_b, torch.bfloat16
|
||||
)
|
||||
out_ref = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16)
|
||||
|
||||
if not torch.allclose(out, out_ref):
|
||||
print("Incorrect results")
|
||||
print(out)
|
||||
print(out_ref)
|
||||
else:
|
||||
print("Correct results")
|
||||
|
||||
timers = []
|
||||
|
||||
# pytorch impl w. bf16
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"pytorch_bf16_bf16_bf16_matmul-no-scales",
|
||||
torch.mm,
|
||||
a.to(dtype=torch.bfloat16, device="cuda"),
|
||||
b.to(dtype=torch.bfloat16, device="cuda"),
|
||||
)
|
||||
)
|
||||
|
||||
# pytorch impl: bf16 output, without fp8 fast accum
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"pytorch_fp8_fp8_bf16_scaled_mm",
|
||||
torch._scaled_mm,
|
||||
a,
|
||||
b,
|
||||
scale_a=scale_a,
|
||||
scale_b=scale_b,
|
||||
out_dtype=torch.bfloat16,
|
||||
)
|
||||
)
|
||||
|
||||
# pytorch impl: bf16 output, with fp8 fast accum
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"pytorch_fp8_fp8_bf16_scaled_mm_fast_accum",
|
||||
torch._scaled_mm,
|
||||
a,
|
||||
b,
|
||||
scale_a=scale_a,
|
||||
scale_b=scale_b,
|
||||
out_dtype=torch.bfloat16,
|
||||
use_fast_accum=True,
|
||||
)
|
||||
)
|
||||
|
||||
# pytorch impl: fp16 output, without fp8 fast accum
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"pytorch_fp8_fp8_fp16_scaled_mm",
|
||||
torch._scaled_mm,
|
||||
a,
|
||||
b,
|
||||
scale_a=scale_a,
|
||||
scale_b=scale_b,
|
||||
out_dtype=torch.float16,
|
||||
)
|
||||
)
|
||||
|
||||
# pytorch impl: fp16 output, with fp8 fast accum
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"pytorch_fp8_fp8_fp16_scaled_mm_fast_accum",
|
||||
torch._scaled_mm,
|
||||
a,
|
||||
b,
|
||||
scale_a=scale_a,
|
||||
scale_b=scale_b,
|
||||
out_dtype=torch.float16,
|
||||
use_fast_accum=True,
|
||||
)
|
||||
)
|
||||
|
||||
# cutlass impl: bf16 output
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"cutlass_fp8_fp8_bf16_scaled_mm",
|
||||
ops.cutlass_scaled_mm,
|
||||
a,
|
||||
b,
|
||||
scale_a,
|
||||
scale_b,
|
||||
torch.bfloat16,
|
||||
)
|
||||
)
|
||||
|
||||
# cutlass impl: bf16 output
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"cutlass_fp8_fp8_bf16_scaled_sparse_mm",
|
||||
ops.cutlass_scaled_sparse_mm,
|
||||
a,
|
||||
b_compressed,
|
||||
e,
|
||||
scale_a,
|
||||
scale_b,
|
||||
torch.bfloat16,
|
||||
)
|
||||
)
|
||||
|
||||
# cutlass impl: fp16 output
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"cutlass_fp8_fp8_fp16_scaled_sparse_mm",
|
||||
ops.cutlass_scaled_sparse_mm,
|
||||
a,
|
||||
b_compressed,
|
||||
e,
|
||||
scale_a,
|
||||
scale_b,
|
||||
torch.float16,
|
||||
)
|
||||
)
|
||||
|
||||
# cutlass impl: bf16 output, with bias
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"cutlass_fp8_fp8_bf16_scaled_sparse_mm_bias",
|
||||
ops.cutlass_scaled_sparse_mm,
|
||||
a,
|
||||
b_compressed,
|
||||
e,
|
||||
scale_a,
|
||||
scale_b,
|
||||
torch.bfloat16,
|
||||
bias,
|
||||
)
|
||||
)
|
||||
|
||||
# cutlass impl: fp16 output, with bias
|
||||
timers.append(
|
||||
bench_fn(
|
||||
label,
|
||||
sub_label,
|
||||
"cutlass_fp8_fp8_fp16_scaled_sparse_mm_bias",
|
||||
ops.cutlass_scaled_sparse_mm,
|
||||
a,
|
||||
b_compressed,
|
||||
e,
|
||||
scale_a,
|
||||
scale_b,
|
||||
torch.float16,
|
||||
bias.to(dtype=torch.float16),
|
||||
)
|
||||
)
|
||||
|
||||
return timers
|
||||
|
||||
|
||||
def bench(
|
||||
dtype: torch.dtype, m: int, k: int, n: int, label: str, sub_label: str
|
||||
) -> Iterable[TMeasurement]:
|
||||
if dtype == torch.int8:
|
||||
return bench_int8(dtype, m, k, n, label, sub_label)
|
||||
if dtype == torch.float8_e4m3fn:
|
||||
return bench_fp8(dtype, m, k, n, label, sub_label)
|
||||
raise ValueError(
|
||||
f"Unsupported dtype {dtype}: should be one of torch.int8, torch.float8_e4m3fn."
|
||||
)
|
||||
|
||||
|
||||
# runner
|
||||
def print_timers(timers: Iterable[TMeasurement]):
|
||||
compare = TBenchmark.Compare(timers)
|
||||
compare.print()
|
||||
|
||||
|
||||
def run(
|
||||
dtype: torch.dtype, MKNs: Iterable[tuple[int, int, int]]
|
||||
) -> Iterable[TMeasurement]:
|
||||
results = []
|
||||
for m, k, n in MKNs:
|
||||
timers = bench(dtype, m, k, n, f"scaled-{dtype}-gemm", f"MKN=({m}x{k}x{n})")
|
||||
print_timers(timers)
|
||||
results.extend(timers)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# output makers
|
||||
def make_output(
|
||||
data: Iterable[TMeasurement],
|
||||
MKNs: Iterable[tuple[int, int, int]],
|
||||
base_description: str,
|
||||
timestamp=None,
|
||||
):
|
||||
print(f"== All Results {base_description} ====")
|
||||
print_timers(data)
|
||||
|
||||
# pickle all the results
|
||||
timestamp = int(time.time()) if timestamp is None else timestamp
|
||||
with open(f"{base_description}-{timestamp}.pkl", "wb") as f:
|
||||
pkl.dump(data, f)
|
||||
|
||||
|
||||
# argparse runners
|
||||
|
||||
|
||||
def run_square_bench(args):
|
||||
dim_sizes = list(range(args.dim_start, args.dim_end + 1, args.dim_increment))
|
||||
MKNs = list(zip(dim_sizes, dim_sizes, dim_sizes))
|
||||
data = run(args.dtype, MKNs)
|
||||
|
||||
make_output(data, MKNs, f"square_bench-{args.dtype}")
|
||||
|
||||
|
||||
def run_range_bench(args):
|
||||
dim_sizes = list(range(args.dim_start, args.dim_end, args.dim_increment))
|
||||
n = len(dim_sizes)
|
||||
Ms = [args.m_constant] * n if args.m_constant is not None else dim_sizes
|
||||
Ks = [args.k_constant] * n if args.k_constant is not None else dim_sizes
|
||||
Ns = [args.n_constant] * n if args.n_constant is not None else dim_sizes
|
||||
MKNs = list(zip(Ms, Ks, Ns))
|
||||
data = run(args.dtype, MKNs)
|
||||
|
||||
make_output(data, MKNs, f"range_bench-{args.dtype}")
|
||||
|
||||
|
||||
def run_model_bench(args):
|
||||
print("Benchmarking models:")
|
||||
for i, model in enumerate(args.models):
|
||||
print(f"[{i}] {model}")
|
||||
|
||||
def model_shapes(model_name: str, tp_size: int) -> list[tuple[int, int]]:
|
||||
KNs = []
|
||||
for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model_name]):
|
||||
KN[tp_split_dim] = KN[tp_split_dim] // tp_size
|
||||
KNs.append(KN)
|
||||
return KNs
|
||||
|
||||
model_bench_data = []
|
||||
models_tps = list(itertools.product(args.models, args.tp_sizes))
|
||||
for model, tp_size in models_tps:
|
||||
Ms = args.batch_sizes
|
||||
KNs = model_shapes(model, tp_size)
|
||||
MKNs = []
|
||||
for m in Ms:
|
||||
for k, n in KNs:
|
||||
MKNs.append((m, k, n))
|
||||
|
||||
data = run(args.dtype, MKNs)
|
||||
model_bench_data.append(data)
|
||||
|
||||
# Print all results
|
||||
for data, model_tp in zip(model_bench_data, models_tps):
|
||||
model, tp_size = model_tp
|
||||
print(f"== Results {args.dtype} {model}-TP{tp_size} ====")
|
||||
print_timers(data)
|
||||
|
||||
timestamp = int(time.time())
|
||||
|
||||
all_data = []
|
||||
for d in model_bench_data:
|
||||
all_data.extend(d)
|
||||
# pickle all data
|
||||
with open(f"model_bench-{args.dtype}-{timestamp}.pkl", "wb") as f:
|
||||
pkl.dump(all_data, f)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
def to_torch_dtype(dt):
|
||||
if dt == "int8":
|
||||
return torch.int8
|
||||
if dt == "fp8":
|
||||
return torch.float8_e4m3fn
|
||||
raise ValueError("unsupported dtype")
|
||||
|
||||
parser = FlexibleArgumentParser(
|
||||
description="""
|
||||
Benchmark Cutlass GEMM.
|
||||
|
||||
To run square GEMMs:
|
||||
python3 ./benchmarks/cutlass_benchmarks/sparse_benchmarks.py --dtype fp8 square_bench --dim-start 128 --dim-end 512 --dim-increment 64
|
||||
|
||||
To run constant N and K and sweep M:
|
||||
python3 ./benchmarks/cutlass_benchmarks/sparse_benchmarks.py --dtype fp8 range_bench --dim-start 128 --dim-end 512 --dim-increment 64 --n-constant 16384 --k-constant 16384
|
||||
|
||||
To run dimensions from a model:
|
||||
python3 ./benchmarks/cutlass_benchmarks/sparse_benchmarks.py --dtype fp8 model_bench --models meta-llama/Llama-2-7b-hf --batch-sizes 16 --tp-sizes 1
|
||||
|
||||
Output:
|
||||
- a .pkl file, that is a list of raw torch.benchmark.utils.Measurements for the pytorch and cutlass implementations for the various GEMMs.
|
||||
""", # noqa: E501
|
||||
formatter_class=argparse.RawTextHelpFormatter,
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
type=to_torch_dtype,
|
||||
required=True,
|
||||
help="Available options are ['int8', 'fp8']",
|
||||
)
|
||||
subparsers = parser.add_subparsers(dest="cmd")
|
||||
|
||||
square_parser = subparsers.add_parser("square_bench")
|
||||
square_parser.add_argument("--dim-start", type=int, required=True)
|
||||
square_parser.add_argument("--dim-end", type=int, required=True)
|
||||
square_parser.add_argument("--dim-increment", type=int, required=True)
|
||||
square_parser.set_defaults(func=run_square_bench)
|
||||
|
||||
range_parser = subparsers.add_parser("range_bench")
|
||||
range_parser.add_argument("--dim-start", type=int, required=True)
|
||||
range_parser.add_argument("--dim-end", type=int, required=True)
|
||||
range_parser.add_argument("--dim-increment", type=int, required=True)
|
||||
range_parser.add_argument("--m-constant", type=int, default=None)
|
||||
range_parser.add_argument("--n-constant", type=int, default=None)
|
||||
range_parser.add_argument("--k-constant", type=int, default=None)
|
||||
range_parser.set_defaults(func=run_range_bench)
|
||||
|
||||
model_parser = subparsers.add_parser("model_bench")
|
||||
model_parser.add_argument(
|
||||
"--models",
|
||||
nargs="+",
|
||||
type=str,
|
||||
default=DEFAULT_MODELS,
|
||||
choices=WEIGHT_SHAPES.keys(),
|
||||
)
|
||||
model_parser.add_argument(
|
||||
"--tp-sizes", nargs="+", type=int, default=DEFAULT_TP_SIZES
|
||||
)
|
||||
model_parser.add_argument(
|
||||
"--batch-sizes", nargs="+", type=int, default=DEFAULT_BATCH_SIZES
|
||||
)
|
||||
model_parser.set_defaults(func=run_model_bench)
|
||||
|
||||
args = parser.parse_args()
|
||||
args.func(args)
|
||||
@@ -5,8 +5,6 @@
|
||||
|
||||
import torch
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
|
||||
|
||||
def to_fp8(tensor: torch.Tensor) -> torch.Tensor:
|
||||
finfo = torch.finfo(torch.float8_e4m3fn)
|
||||
@@ -39,49 +37,3 @@ def make_rand_tensors(
|
||||
return to_fp8(a), to_fp8(b)
|
||||
|
||||
raise ValueError("unsupported dtype")
|
||||
|
||||
|
||||
def prune_to_2_4(tensor):
|
||||
# Reshape tensor to [N, 4] where N is number of groups of 4
|
||||
original_shape = tensor.shape
|
||||
reshaped = tensor.reshape(-1, 4)
|
||||
|
||||
# Get indices of top 2 absolute values in each group of 4
|
||||
_, indices = torch.topk(torch.abs(reshaped), k=2, dim=1)
|
||||
|
||||
# Create binary mask
|
||||
mask = torch.zeros_like(reshaped)
|
||||
mask.scatter_(dim=1, index=indices, src=torch.ones_like(indices, dtype=mask.dtype))
|
||||
|
||||
# Apply mask and reshape back
|
||||
pruned = reshaped * mask
|
||||
|
||||
# Turn all -0.0 to 0.0
|
||||
pruned[pruned == -0.0] = 0.0
|
||||
|
||||
return pruned.reshape(original_shape)
|
||||
|
||||
|
||||
def make_rand_sparse_tensors(
|
||||
dtype: torch.dtype, m: int, n: int, k: int
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
a = torch.randn((m, k), device="cuda") * 5
|
||||
b = torch.randn((n, k), device="cuda").t() * 5
|
||||
|
||||
b = prune_to_2_4(b.t()).t()
|
||||
|
||||
if dtype == torch.int8:
|
||||
a, b = to_int8(a), to_int8(b)
|
||||
elif dtype == torch.float8_e4m3fn:
|
||||
a, b = to_fp8(a), to_fp8(b)
|
||||
elif dtype == torch.float16:
|
||||
a, b = to_fp16(a), to_fp16(b)
|
||||
elif dtype == torch.bfloat16:
|
||||
a, b = to_bf16(a), to_bf16(b)
|
||||
else:
|
||||
raise ValueError("unsupported dtype")
|
||||
|
||||
b_compressed, e = ops.cutlass_sparse_compress(b.t())
|
||||
|
||||
# Compressed B, Metadata, Original A, B
|
||||
return b_compressed, e, a, b
|
||||
|
||||
@@ -0,0 +1,669 @@
|
||||
"""
|
||||
Benchmark: SM103 (B300) FP4 Ultra GEMM vs SM100 (B200) NVFP4 GEMM
|
||||
===================================================================
|
||||
|
||||
This benchmark compares the performance of the SM103-optimized FP4 Ultra
|
||||
GEMM kernel against the SM100 NVFP4 GEMM kernel, both running on B300
|
||||
hardware. It also benchmarks the effect of Programmatic Dependent Launch
|
||||
(PDL) on the quant->GEMM pipeline, where the GEMM consumer can begin
|
||||
before the quant producer finishes.
|
||||
|
||||
SM103 kernels use:
|
||||
- K=768 tile (vs K=256 on SM100)
|
||||
- FP4 Ultra MMA (UltraVs16) schedule
|
||||
- NoSmemWarpSpecialized epilogue
|
||||
- Sm103BlockScaledConfig scale factor layout
|
||||
|
||||
PDL kernels additionally set:
|
||||
- cudaLaunchAttributeProgrammaticStreamSerialization on quant (producer)
|
||||
- CUTLASS launch_with_pdl=true on GEMM (enables overlap with next kernel)
|
||||
|
||||
Usage:
|
||||
python benchmarks/kernels/benchmark_nvfp4_sm103.py [--mode gemm|quant|e2e|pdl|all]
|
||||
|
||||
Requirements:
|
||||
- B300 GPU (SM103 / compute capability 10.3)
|
||||
- CUDA >= 12.9
|
||||
- vLLM built with ENABLE_NVFP4_SM100=1 and SM103 support
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import vllm._C # noqa: F401 - registers ops into torch.ops._C
|
||||
|
||||
# ============================================================================
|
||||
# Helpers
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def round_up(x: int, y: int) -> int:
|
||||
return ((x + y - 1) // y) * y
|
||||
|
||||
|
||||
def get_sm_version() -> int:
|
||||
"""Return SM version as integer (e.g., 100, 103, 120)."""
|
||||
cap = torch.cuda.get_device_capability()
|
||||
return cap[0] * 10 + cap[1]
|
||||
|
||||
|
||||
def create_nvfp4_tensors(
|
||||
m: int, n: int, k: int, dtype: torch.dtype = torch.bfloat16
|
||||
) -> dict:
|
||||
"""
|
||||
Create synthetic NVFP4 GEMM input tensors (A, B, scales, alpha).
|
||||
|
||||
A: [m, k/2] uint8 (packed FP4)
|
||||
B: [n, k/2] uint8 (packed FP4, column-major)
|
||||
A_sf: [round_up(m,128), round_up(k/16,4)] float8_e4m3fn (SM100 swizzled)
|
||||
B_sf: [round_up(n,128), round_up(k/16,4)] float8_e4m3fn (SM100 swizzled)
|
||||
alpha: [1] float32
|
||||
D: [m, n] output
|
||||
"""
|
||||
# Packed FP4 data (random bytes -- content doesn't affect timing)
|
||||
A = torch.randint(0, 256, (m, k // 2), dtype=torch.uint8, device="cuda")
|
||||
B = torch.randint(0, 256, (n, k // 2), dtype=torch.uint8, device="cuda")
|
||||
|
||||
# Scale factors (SM100 swizzled layout)
|
||||
sf_m = round_up(m, 128)
|
||||
sf_n = round_up(n, 128)
|
||||
sf_k = round_up(k // 16, 4)
|
||||
|
||||
A_sf_sm100 = torch.randint(
|
||||
0, 256, (sf_m, sf_k), dtype=torch.uint8, device="cuda"
|
||||
).view(torch.float8_e4m3fn)
|
||||
B_sf_sm100 = torch.randint(
|
||||
0, 256, (sf_n, sf_k), dtype=torch.uint8, device="cuda"
|
||||
).view(torch.float8_e4m3fn)
|
||||
|
||||
# SM103 layout: convert from SM100 layout
|
||||
A_sf_sm103 = torch.empty_like(A_sf_sm100)
|
||||
B_sf_sm103 = torch.empty_like(B_sf_sm100)
|
||||
torch.ops._C.convert_sf_layout_sm100_to_sm103(A_sf_sm103, A_sf_sm100)
|
||||
torch.ops._C.convert_sf_layout_sm100_to_sm103(B_sf_sm103, B_sf_sm100)
|
||||
|
||||
# Global alpha
|
||||
alpha = torch.tensor([1.0], dtype=torch.float32, device="cuda")
|
||||
|
||||
# Output
|
||||
D = torch.empty(m, n, dtype=dtype, device="cuda")
|
||||
|
||||
return {
|
||||
"A": A,
|
||||
"B": B,
|
||||
"A_sf_sm100": A_sf_sm100,
|
||||
"B_sf_sm100": B_sf_sm100,
|
||||
"A_sf_sm103": A_sf_sm103,
|
||||
"B_sf_sm103": B_sf_sm103,
|
||||
"alpha": alpha,
|
||||
"D": D,
|
||||
}
|
||||
|
||||
|
||||
def create_quant_tensors(
|
||||
m: int, n: int, dtype: torch.dtype = torch.bfloat16
|
||||
) -> dict:
|
||||
"""Create inputs for activation quantization benchmark."""
|
||||
input_tensor = torch.randn(m, n, dtype=dtype, device="cuda")
|
||||
global_scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
return {"input": input_tensor, "global_scale": global_scale}
|
||||
|
||||
|
||||
def bench_fn(
|
||||
fn,
|
||||
warmup: int = 20,
|
||||
iters: int = 100,
|
||||
sync: bool = True,
|
||||
) -> float:
|
||||
"""Benchmark a function, returning median time in microseconds."""
|
||||
# Warmup
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
if sync:
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Timed iterations using CUDA events
|
||||
start_events = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
|
||||
end_events = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
|
||||
|
||||
for i in range(iters):
|
||||
start_events[i].record()
|
||||
fn()
|
||||
end_events[i].record()
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
times = [s.elapsed_time(e) * 1000 for s, e in zip(start_events, end_events)]
|
||||
times.sort()
|
||||
# Return median in microseconds
|
||||
return times[len(times) // 2]
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# GEMM Benchmark
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def benchmark_gemm(
|
||||
m_sizes: list[int],
|
||||
n: int = 7168,
|
||||
k: int = 7168,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
) -> list[dict]:
|
||||
"""
|
||||
Benchmark SM100 vs SM103 vs SM103+PDL NVFP4 GEMM kernels side by side.
|
||||
|
||||
PDL on the GEMM sets ProgrammaticStreamSerialization, allowing the NEXT
|
||||
kernel on the stream to overlap with the GEMM's tail. For isolated GEMM
|
||||
calls (no consumer kernel), the PDL overhead should be near-zero.
|
||||
"""
|
||||
vllm_ops = torch.ops._C
|
||||
|
||||
has_sm100a = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm100a")
|
||||
has_sm103a = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm103a")
|
||||
has_sm103a_pdl = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm103a_pdl")
|
||||
|
||||
if not has_sm100a and not has_sm103a:
|
||||
print("WARNING: Neither sm100a nor sm103a ops are available. "
|
||||
"Rebuild with ENABLE_NVFP4_SM100=1.")
|
||||
return []
|
||||
|
||||
results = []
|
||||
|
||||
for m in m_sizes:
|
||||
tensors = create_nvfp4_tensors(m, n, k, dtype)
|
||||
D = tensors["D"]
|
||||
A, B = tensors["A"], tensors["B"]
|
||||
A_sf_sm100, B_sf_sm100 = tensors["A_sf_sm100"], tensors["B_sf_sm100"]
|
||||
A_sf_sm103, B_sf_sm103 = tensors["A_sf_sm103"], tensors["B_sf_sm103"]
|
||||
alpha = tensors["alpha"]
|
||||
|
||||
flops = 2.0 * m * n * k
|
||||
|
||||
time_sm100: Optional[float] = None
|
||||
time_sm103: Optional[float] = None
|
||||
time_sm103_pdl: Optional[float] = None
|
||||
|
||||
if has_sm100a:
|
||||
def run_sm100():
|
||||
vllm_ops.cutlass_scaled_fp4_mm_sm100a(
|
||||
D, A, B, A_sf_sm100, B_sf_sm100, alpha
|
||||
)
|
||||
time_sm100 = bench_fn(run_sm100, warmup=20, iters=100)
|
||||
|
||||
if has_sm103a:
|
||||
def run_sm103():
|
||||
vllm_ops.cutlass_scaled_fp4_mm_sm103a(
|
||||
D, A, B, A_sf_sm103, B_sf_sm103, alpha
|
||||
)
|
||||
time_sm103 = bench_fn(run_sm103, warmup=20, iters=100)
|
||||
|
||||
if has_sm103a_pdl:
|
||||
def run_sm103_pdl():
|
||||
vllm_ops.cutlass_scaled_fp4_mm_sm103a_pdl(
|
||||
D, A, B, A_sf_sm103, B_sf_sm103, alpha
|
||||
)
|
||||
time_sm103_pdl = bench_fn(run_sm103_pdl, warmup=20, iters=100)
|
||||
|
||||
row: dict = {"M": m, "N": n, "K": k}
|
||||
|
||||
if time_sm100 is not None:
|
||||
row["sm100_us"] = time_sm100
|
||||
row["sm100_tflops"] = flops / (time_sm100 * 1e-6) / 1e12
|
||||
|
||||
if time_sm103 is not None:
|
||||
row["sm103_us"] = time_sm103
|
||||
row["sm103_tflops"] = flops / (time_sm103 * 1e-6) / 1e12
|
||||
|
||||
if time_sm103_pdl is not None:
|
||||
row["sm103pdl_us"] = time_sm103_pdl
|
||||
row["sm103pdl_tflops"] = flops / (time_sm103_pdl * 1e-6) / 1e12
|
||||
|
||||
if time_sm100 is not None and time_sm103 is not None:
|
||||
row["sm103_vs_100"] = time_sm100 / time_sm103
|
||||
|
||||
results.append(row)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Quantization Benchmark
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def benchmark_quant(
|
||||
m_sizes: list[int],
|
||||
n: int = 7168,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
) -> list[dict]:
|
||||
"""
|
||||
Benchmark SM100 vs SM103 activation quantization (BF16 -> NVFP4).
|
||||
"""
|
||||
vllm_ops = torch.ops._C
|
||||
results = []
|
||||
|
||||
has_sm103_quant = hasattr(vllm_ops, "scaled_fp4_quant_sm103")
|
||||
|
||||
for m in m_sizes:
|
||||
tensors = create_quant_tensors(m, n, dtype)
|
||||
input_t = tensors["input"]
|
||||
global_scale = tensors["global_scale"]
|
||||
|
||||
# SM100 quantization (swizzled layout)
|
||||
def run_sm100_quant():
|
||||
vllm_ops.scaled_fp4_quant(input_t, global_scale, True)
|
||||
|
||||
time_sm100 = bench_fn(run_sm100_quant, warmup=20, iters=100)
|
||||
|
||||
row: dict = {
|
||||
"M": m,
|
||||
"N": n,
|
||||
"sm100_us": time_sm100,
|
||||
"sm100_gb_s": (m * n * 2) / (time_sm100 * 1e-6) / 1e9,
|
||||
}
|
||||
|
||||
if has_sm103_quant:
|
||||
def run_sm103_quant():
|
||||
vllm_ops.scaled_fp4_quant_sm103(input_t, global_scale)
|
||||
|
||||
time_sm103 = bench_fn(run_sm103_quant, warmup=20, iters=100)
|
||||
row["sm103_us"] = time_sm103
|
||||
row["sm103_gb_s"] = (m * n * 2) / (time_sm103 * 1e-6) / 1e9
|
||||
|
||||
results.append(row)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# SF Layout Conversion Benchmark
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def benchmark_sf_conversion(
|
||||
m_sizes: list[int],
|
||||
k: int = 7168,
|
||||
) -> list[dict]:
|
||||
"""
|
||||
Benchmark the SM100 <-> SM103 scale factor layout conversion kernel.
|
||||
|
||||
This measures the overhead of converting scale factors between layouts,
|
||||
which happens once at model load time for weights.
|
||||
"""
|
||||
vllm_ops = torch.ops._C
|
||||
results = []
|
||||
|
||||
for m in m_sizes:
|
||||
sf_m = round_up(m, 128)
|
||||
sf_k = round_up(k // 16, 4)
|
||||
|
||||
# Create source SF tensor (SM100 layout)
|
||||
src = torch.randint(
|
||||
0, 256, (sf_m, sf_k), dtype=torch.uint8, device="cuda"
|
||||
).view(torch.float8_e4m3fn)
|
||||
|
||||
# Allocate destination (same shape)
|
||||
dst = torch.empty_like(src)
|
||||
|
||||
# Benchmark SM100 -> SM103 conversion
|
||||
def run_convert():
|
||||
vllm_ops.convert_sf_layout_sm100_to_sm103(dst, src)
|
||||
|
||||
time_us = bench_fn(run_convert, warmup=20, iters=200)
|
||||
|
||||
results.append({
|
||||
"M": m,
|
||||
"K": k,
|
||||
"sf_shape": f"{sf_m}x{sf_k}",
|
||||
"kernel": "SM100->SM103 SF convert",
|
||||
"time_us": time_us,
|
||||
"throughput_gb_s": (sf_m * sf_k) / (time_us * 1e-6) / 1e9,
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# End-to-End Benchmark (Quant + GEMM) with PDL comparison
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def benchmark_e2e(
|
||||
m_sizes: list[int],
|
||||
n: int = 7168,
|
||||
k: int = 7168,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
) -> list[dict]:
|
||||
"""
|
||||
Benchmark the full NVFP4 inference path: quantize activations + GEMM,
|
||||
comparing SM100, SM103, and SM103+PDL.
|
||||
|
||||
This measures what a real transformer linear layer does:
|
||||
1. Quantize BF16 activations to NVFP4 (with block scales)
|
||||
2. NVFP4 x NVFP4 GEMM
|
||||
|
||||
SM103+PDL enables ProgrammaticStreamSerialization on the quant kernel
|
||||
and launch_with_pdl on the GEMM, allowing the GEMM to begin executing
|
||||
while the quant kernel is still completing its last thread blocks.
|
||||
"""
|
||||
vllm_ops = torch.ops._C
|
||||
|
||||
has_sm100a = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm100a")
|
||||
has_sm103a = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm103a")
|
||||
has_sm103_quant = hasattr(vllm_ops, "scaled_fp4_quant_sm103")
|
||||
has_sm103_pdl_quant = hasattr(vllm_ops, "scaled_fp4_quant_sm103_pdl")
|
||||
has_sm103a_pdl = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm103a_pdl")
|
||||
|
||||
if not has_sm100a and not has_sm103a:
|
||||
print("WARNING: Neither sm100a nor sm103a ops are available. "
|
||||
"Rebuild with ENABLE_NVFP4_SM100=1.")
|
||||
return []
|
||||
|
||||
results = []
|
||||
|
||||
for m in m_sizes:
|
||||
# Create activation input
|
||||
activation = torch.randn(m, k, dtype=dtype, device="cuda")
|
||||
global_scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
|
||||
# Create weight (pre-quantized)
|
||||
B = torch.randint(0, 256, (n, k // 2), dtype=torch.uint8, device="cuda")
|
||||
sf_n = round_up(n, 128)
|
||||
sf_k = round_up(k // 16, 4)
|
||||
|
||||
# Weight SFs in SM100 layout (for SM100 kernel)
|
||||
B_sf_sm100 = torch.randint(
|
||||
0, 256, (sf_n, sf_k), dtype=torch.uint8, device="cuda"
|
||||
).view(torch.float8_e4m3fn)
|
||||
alpha = torch.tensor([1.0], dtype=torch.float32, device="cuda")
|
||||
|
||||
# Weight SFs in SM103 layout (pre-converted at load time)
|
||||
B_sf_sm103 = torch.empty_like(B_sf_sm100)
|
||||
vllm_ops.convert_sf_layout_sm100_to_sm103(B_sf_sm103, B_sf_sm100)
|
||||
|
||||
D = torch.empty(m, n, dtype=dtype, device="cuda")
|
||||
|
||||
flops = 2.0 * m * n * k
|
||||
row: dict = {"M": m, "N": n, "K": k}
|
||||
|
||||
# --- SM100 baseline: SM100 quant + SM100 GEMM ---
|
||||
if has_sm100a:
|
||||
def run_e2e_sm100():
|
||||
A_q, A_sf = vllm_ops.scaled_fp4_quant(
|
||||
activation, global_scale, True
|
||||
)
|
||||
A_sf = A_sf.view(torch.float8_e4m3fn)
|
||||
vllm_ops.cutlass_scaled_fp4_mm_sm100a(
|
||||
D, A_q, B, A_sf, B_sf_sm100, alpha
|
||||
)
|
||||
|
||||
time_sm100 = bench_fn(run_e2e_sm100, warmup=10, iters=50)
|
||||
row["sm100_us"] = time_sm100
|
||||
row["sm100_tflops"] = flops / (time_sm100 * 1e-6) / 1e12
|
||||
|
||||
# --- SM103 without PDL: SM103 quant + SM103 GEMM ---
|
||||
if has_sm103a and has_sm103_quant:
|
||||
def run_e2e_sm103():
|
||||
A_q, A_sf = vllm_ops.scaled_fp4_quant_sm103(
|
||||
activation, global_scale
|
||||
)
|
||||
A_sf = A_sf.view(torch.float8_e4m3fn)
|
||||
vllm_ops.cutlass_scaled_fp4_mm_sm103a(
|
||||
D, A_q, B, A_sf, B_sf_sm103, alpha
|
||||
)
|
||||
|
||||
time_sm103 = bench_fn(run_e2e_sm103, warmup=10, iters=50)
|
||||
row["sm103_us"] = time_sm103
|
||||
row["sm103_tflops"] = flops / (time_sm103 * 1e-6) / 1e12
|
||||
|
||||
# --- SM103 with PDL: PDL quant + PDL GEMM ---
|
||||
if has_sm103a_pdl and has_sm103_pdl_quant:
|
||||
def run_e2e_sm103_pdl():
|
||||
# PDL quant: ProgrammaticStreamSerialization allows GEMM to
|
||||
# begin before quant finishes.
|
||||
A_q, A_sf = vllm_ops.scaled_fp4_quant_sm103_pdl(
|
||||
activation, global_scale
|
||||
)
|
||||
A_sf = A_sf.view(torch.float8_e4m3fn)
|
||||
# PDL GEMM: ProgrammaticStreamSerialization allows the next
|
||||
# layer's kernel to begin before this GEMM finishes.
|
||||
vllm_ops.cutlass_scaled_fp4_mm_sm103a_pdl(
|
||||
D, A_q, B, A_sf, B_sf_sm103, alpha
|
||||
)
|
||||
|
||||
time_sm103_pdl = bench_fn(run_e2e_sm103_pdl, warmup=10, iters=50)
|
||||
row["sm103pdl_us"] = time_sm103_pdl
|
||||
row["sm103pdl_tflops"] = flops / (time_sm103_pdl * 1e-6) / 1e12
|
||||
|
||||
# Speedup columns
|
||||
if "sm100_us" in row and "sm103_us" in row:
|
||||
row["sm103_vs_100"] = row["sm100_us"] / row["sm103_us"]
|
||||
if "sm103_us" in row and "sm103pdl_us" in row:
|
||||
row["pdl_vs_nop"] = row["sm103_us"] / row["sm103pdl_us"]
|
||||
if "sm100_us" in row and "sm103pdl_us" in row:
|
||||
row["pdl_vs_100"] = row["sm100_us"] / row["sm103pdl_us"]
|
||||
|
||||
results.append(row)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# PDL Pipeline Benchmark (back-to-back quant+GEMM pairs)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def benchmark_pdl_pipeline(
|
||||
m_sizes: list[int],
|
||||
n: int = 7168,
|
||||
k: int = 7168,
|
||||
num_layers: int = 4,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
) -> list[dict]:
|
||||
"""
|
||||
Benchmark the PDL pipeline benefit for back-to-back layers.
|
||||
|
||||
In a real transformer, the same quant->GEMM pattern repeats for each
|
||||
linear layer. With PDL enabled on both quant and GEMM, each kernel
|
||||
launch overlaps with its predecessor's tail, creating a pipeline:
|
||||
|
||||
quant_1 -> GEMM_1 -> quant_2 -> GEMM_2 -> ...
|
||||
|
||||
This benchmark simulates `num_layers` consecutive quant+GEMM pairs
|
||||
to measure the cumulative pipeline benefit.
|
||||
"""
|
||||
vllm_ops = torch.ops._C
|
||||
|
||||
has_sm103_quant = hasattr(vllm_ops, "scaled_fp4_quant_sm103")
|
||||
has_sm103_pdl_quant = hasattr(vllm_ops, "scaled_fp4_quant_sm103_pdl")
|
||||
has_sm103a = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm103a")
|
||||
has_sm103a_pdl = hasattr(vllm_ops, "cutlass_scaled_fp4_mm_sm103a_pdl")
|
||||
|
||||
if not (has_sm103_quant and has_sm103a):
|
||||
print("WARNING: SM103 ops not available.")
|
||||
return []
|
||||
|
||||
results = []
|
||||
|
||||
for m in m_sizes:
|
||||
activation = torch.randn(m, k, dtype=dtype, device="cuda")
|
||||
global_scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
|
||||
B = torch.randint(0, 256, (n, k // 2), dtype=torch.uint8, device="cuda")
|
||||
sf_n = round_up(n, 128)
|
||||
sf_k = round_up(k // 16, 4)
|
||||
B_sf_sm100 = torch.randint(
|
||||
0, 256, (sf_n, sf_k), dtype=torch.uint8, device="cuda"
|
||||
).view(torch.float8_e4m3fn)
|
||||
B_sf_sm103 = torch.empty_like(B_sf_sm100)
|
||||
vllm_ops.convert_sf_layout_sm100_to_sm103(B_sf_sm103, B_sf_sm100)
|
||||
alpha = torch.tensor([1.0], dtype=torch.float32, device="cuda")
|
||||
D = torch.empty(m, n, dtype=dtype, device="cuda")
|
||||
|
||||
total_flops = 2.0 * m * n * k * num_layers
|
||||
|
||||
# SM103 without PDL: num_layers sequential quant+GEMM
|
||||
def run_pipeline_no_pdl():
|
||||
for _ in range(num_layers):
|
||||
A_q, A_sf = vllm_ops.scaled_fp4_quant_sm103(
|
||||
activation, global_scale
|
||||
)
|
||||
A_sf = A_sf.view(torch.float8_e4m3fn)
|
||||
vllm_ops.cutlass_scaled_fp4_mm_sm103a(
|
||||
D, A_q, B, A_sf, B_sf_sm103, alpha
|
||||
)
|
||||
|
||||
time_no_pdl = bench_fn(run_pipeline_no_pdl, warmup=5, iters=30)
|
||||
|
||||
row: dict = {
|
||||
"M": m, "layers": num_layers,
|
||||
"no_pdl_us": time_no_pdl,
|
||||
"no_pdl_tflops": total_flops / (time_no_pdl * 1e-6) / 1e12,
|
||||
}
|
||||
|
||||
# SM103 with PDL: num_layers pipelined quant+GEMM
|
||||
if has_sm103_pdl_quant and has_sm103a_pdl:
|
||||
def run_pipeline_pdl():
|
||||
for _ in range(num_layers):
|
||||
A_q, A_sf = vllm_ops.scaled_fp4_quant_sm103_pdl(
|
||||
activation, global_scale
|
||||
)
|
||||
A_sf = A_sf.view(torch.float8_e4m3fn)
|
||||
vllm_ops.cutlass_scaled_fp4_mm_sm103a_pdl(
|
||||
D, A_q, B, A_sf, B_sf_sm103, alpha
|
||||
)
|
||||
|
||||
time_pdl = bench_fn(run_pipeline_pdl, warmup=5, iters=30)
|
||||
row["pdl_us"] = time_pdl
|
||||
row["pdl_tflops"] = total_flops / (time_pdl * 1e-6) / 1e12
|
||||
row["pdl_speedup"] = time_no_pdl / time_pdl
|
||||
|
||||
results.append(row)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Main
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def print_results(results: list[dict], title: str):
|
||||
if not results:
|
||||
return
|
||||
|
||||
print(f"\n{'=' * 80}")
|
||||
print(f" {title}")
|
||||
print(f"{'=' * 80}")
|
||||
|
||||
# Determine columns from first result
|
||||
cols = list(results[0].keys())
|
||||
# Header
|
||||
header = " | ".join(f"{c:>15s}" for c in cols)
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
|
||||
for r in results:
|
||||
row = []
|
||||
for c in cols:
|
||||
v = r.get(c, "")
|
||||
if isinstance(v, float):
|
||||
row.append(f"{v:>15.2f}")
|
||||
elif isinstance(v, int):
|
||||
row.append(f"{v:>15d}")
|
||||
else:
|
||||
row.append(f"{v:>15s}")
|
||||
print(" | ".join(row))
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark NVFP4 SM103 vs SM100 kernels (with PDL)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=["gemm", "quant", "sf_convert", "e2e", "pdl", "all"],
|
||||
default="all",
|
||||
help="Which benchmark to run",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--n", type=int, default=7168,
|
||||
help="N dimension (default: 7168, DeepSeek)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--k", type=int, default=7168,
|
||||
help="K dimension (default: 7168, DeepSeek)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--layers", type=int, default=4,
|
||||
help="Number of back-to-back layers for PDL pipeline benchmark",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
sm = get_sm_version()
|
||||
print(f"GPU: {torch.cuda.get_device_name()}")
|
||||
print(f"SM version: {sm}")
|
||||
print(f"CUDA version: {torch.version.cuda}")
|
||||
|
||||
if sm < 100:
|
||||
print("ERROR: This benchmark requires SM100+ (Blackwell) GPU.")
|
||||
return
|
||||
|
||||
if sm == 103:
|
||||
print("NOTE: Running on SM103 (B300) -- all kernel variants will run.")
|
||||
else:
|
||||
print(f"NOTE: Running on SM{sm} -- SM100 kernel is native; "
|
||||
"SM103 kernel runs via forward compat (may be slower).")
|
||||
|
||||
# Problem sizes typical for LLM inference
|
||||
# Small M = decode, large M = prefill
|
||||
m_sizes = [1, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]
|
||||
|
||||
if args.mode in ("gemm", "all"):
|
||||
results = benchmark_gemm(m_sizes, n=args.n, k=args.k)
|
||||
print_results(
|
||||
results,
|
||||
f"NVFP4 GEMM: SM100 vs SM103 vs SM103+PDL (N={args.n}, K={args.k})",
|
||||
)
|
||||
|
||||
if args.mode in ("quant", "all"):
|
||||
results = benchmark_quant(m_sizes, n=args.k)
|
||||
print_results(results, f"NVFP4 Activation Quantization (N={args.k})")
|
||||
|
||||
if args.mode in ("sf_convert", "all"):
|
||||
sf_m_sizes = [1024, 2048, 4096, 7168, 8192, 14336, 16384]
|
||||
results = benchmark_sf_conversion(sf_m_sizes, k=args.k)
|
||||
print_results(results, "SF Layout Conversion SM100 <-> SM103")
|
||||
|
||||
if args.mode in ("e2e", "all"):
|
||||
results = benchmark_e2e(m_sizes, n=args.n, k=args.k)
|
||||
print_results(
|
||||
results,
|
||||
f"E2E NVFP4 (Quant+GEMM): SM100 vs SM103 vs SM103+PDL "
|
||||
f"(N={args.n}, K={args.k})",
|
||||
)
|
||||
print(
|
||||
"\nNOTE: sm103_vs_100 = SM100_time / SM103_time (>1 means SM103 faster)\n"
|
||||
" pdl_vs_nop = SM103_time / SM103+PDL_time (>1 means PDL faster)\n"
|
||||
" pdl_vs_100 = SM100_time / SM103+PDL_time (total speedup)"
|
||||
)
|
||||
|
||||
if args.mode in ("pdl", "all"):
|
||||
results = benchmark_pdl_pipeline(
|
||||
m_sizes, n=args.n, k=args.k, num_layers=args.layers
|
||||
)
|
||||
print_results(
|
||||
results,
|
||||
f"PDL Pipeline ({args.layers} layers): SM103 vs SM103+PDL "
|
||||
f"(N={args.n}, K={args.k})",
|
||||
)
|
||||
print(
|
||||
"\nNOTE: pdl_speedup = no_pdl_time / pdl_time\n"
|
||||
" PDL overlaps quant tail with GEMM head across layer boundaries.\n"
|
||||
" Benefit is most visible with multiple back-to-back layers."
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -232,6 +232,28 @@ void unmap_and_release(unsigned long long device, ssize_t size,
|
||||
}
|
||||
}
|
||||
|
||||
// ROCm workaround: hipMemRelease does not return physical VRAM to the
|
||||
// free pool while the virtual-address reservation is still held.
|
||||
// Cycling cuMemAddressFree → cuMemAddressReserve (at the same address)
|
||||
// forces the driver to actually release the physical pages while keeping
|
||||
// the same VA available for a later create_and_map.
|
||||
if (first_error == no_error) {
|
||||
first_error = cuMemAddressFree(d_mem, size);
|
||||
if (first_error == no_error) {
|
||||
CUdeviceptr d_mem_new = 0;
|
||||
first_error = cuMemAddressReserve(&d_mem_new, size, 0, d_mem, 0);
|
||||
if (first_error == no_error && d_mem_new != d_mem) {
|
||||
cuMemAddressFree(d_mem_new, size);
|
||||
snprintf(error_msg, sizeof(error_msg),
|
||||
"ROCm: VA re-reserve got %p instead of %p", (void*)d_mem_new,
|
||||
(void*)d_mem);
|
||||
error_code = CUresult(1);
|
||||
std::cerr << error_msg << std::endl;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (first_error != no_error) {
|
||||
CUDA_CHECK(first_error);
|
||||
}
|
||||
|
||||
+6
-10
@@ -237,6 +237,7 @@ void cutlass_scaled_fp4_mm(torch::Tensor& D, torch::Tensor const& A,
|
||||
torch::Tensor const& B_sf,
|
||||
torch::Tensor const& alpha);
|
||||
|
||||
|
||||
void cutlass_scaled_mm(torch::Tensor& out, torch::Tensor const& a,
|
||||
torch::Tensor const& b, torch::Tensor const& a_scales,
|
||||
torch::Tensor const& b_scales,
|
||||
@@ -285,16 +286,6 @@ void cutlass_scaled_mm_azp(torch::Tensor& out, torch::Tensor const& a,
|
||||
std::optional<torch::Tensor> const& azp,
|
||||
std::optional<torch::Tensor> const& bias);
|
||||
|
||||
bool cutlass_sparse_scaled_mm_supported(int64_t cuda_device_capability);
|
||||
|
||||
void cutlass_scaled_sparse_mm(torch::Tensor& out, torch::Tensor const& a,
|
||||
torch::Tensor const& b, torch::Tensor const& e,
|
||||
torch::Tensor const& a_scales,
|
||||
torch::Tensor const& b_scales,
|
||||
std::optional<torch::Tensor> const& bias);
|
||||
|
||||
std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a);
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
|
||||
torch::Tensor const& input, torch::Tensor const& input_scale,
|
||||
bool is_sf_swizzled_layout);
|
||||
@@ -316,6 +307,11 @@ void silu_and_mul_scaled_fp4_experts_quant(
|
||||
torch::Tensor const& input_offset_by_experts,
|
||||
torch::Tensor const& output_scale_offset_by_experts);
|
||||
|
||||
void convert_sf_layout_sm100_to_sm103(torch::Tensor& dst,
|
||||
torch::Tensor const& src);
|
||||
void convert_sf_layout_sm103_to_sm100(torch::Tensor& dst,
|
||||
torch::Tensor const& src);
|
||||
|
||||
void per_token_group_quant_fp8(const torch::Tensor& input,
|
||||
torch::Tensor& output_q, torch::Tensor& output_s,
|
||||
int64_t group_size, double eps, double fp8_min,
|
||||
|
||||
@@ -27,6 +27,18 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
||||
bool is_sf_swizzled_layout);
|
||||
#endif
|
||||
|
||||
#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
|
||||
void scaled_fp4_quant_sm103a(torch::Tensor const& output,
|
||||
torch::Tensor const& input,
|
||||
torch::Tensor const& output_sf,
|
||||
torch::Tensor const& input_sf);
|
||||
// PDL variant: launches quant with ProgrammaticStreamSerialization.
|
||||
void scaled_fp4_quant_sm103a_pdl(torch::Tensor const& output,
|
||||
torch::Tensor const& input,
|
||||
torch::Tensor const& output_sf,
|
||||
torch::Tensor const& input_sf);
|
||||
#endif
|
||||
|
||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
|
||||
void scaled_fp4_experts_quant_sm1xxa(
|
||||
@@ -132,3 +144,79 @@ void silu_and_mul_scaled_fp4_experts_quant(
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false, "No compiled silu_and_mul nvfp4 experts quantization kernel");
|
||||
}
|
||||
|
||||
// SM103-native quantization: writes SM103-layout scale factors directly,
|
||||
// eliminating the SM100->SM103 conversion step on the critical path.
|
||||
std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_sm103a_func(
|
||||
torch::Tensor const& input, torch::Tensor const& input_sf) {
|
||||
int64_t n = input.size(-1);
|
||||
int64_t m = input.numel() / n;
|
||||
auto device = input.device();
|
||||
|
||||
auto output = torch::empty(
|
||||
{m, n / 2}, torch::TensorOptions().device(device).dtype(torch::kUInt8));
|
||||
|
||||
auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
|
||||
auto output_sf = torch::empty(
|
||||
{sf_m, sf_n},
|
||||
torch::TensorOptions().device(device).dtype(torch::kInt32));
|
||||
|
||||
#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
|
||||
scaled_fp4_quant_sm103a(output, input, output_sf, input_sf);
|
||||
return {output, output_sf};
|
||||
#endif
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(false,
|
||||
"No compiled SM103 nvfp4 quantization kernel");
|
||||
}
|
||||
|
||||
void scaled_fp4_quant_sm103a_out(torch::Tensor const& input,
|
||||
torch::Tensor const& input_sf,
|
||||
torch::Tensor& output,
|
||||
torch::Tensor& output_sf) {
|
||||
#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
|
||||
scaled_fp4_quant_sm103a(output, input, output_sf, input_sf);
|
||||
return;
|
||||
#endif
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(false,
|
||||
"No compiled SM103 nvfp4 quantization kernel");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// PDL-enabled SM103 quantization entry points.
|
||||
//
|
||||
// These launch the quant kernel with ProgrammaticStreamSerialization,
|
||||
// allowing the subsequent GEMM to begin before quantization completes.
|
||||
// ============================================================================
|
||||
std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_sm103a_pdl_func(
|
||||
torch::Tensor const& input, torch::Tensor const& input_sf) {
|
||||
int64_t n = input.size(-1);
|
||||
int64_t m = input.numel() / n;
|
||||
auto device = input.device();
|
||||
|
||||
auto output = torch::empty(
|
||||
{m, n / 2}, torch::TensorOptions().device(device).dtype(torch::kUInt8));
|
||||
|
||||
auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
|
||||
auto output_sf = torch::empty(
|
||||
{sf_m, sf_n},
|
||||
torch::TensorOptions().device(device).dtype(torch::kInt32));
|
||||
|
||||
#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
|
||||
scaled_fp4_quant_sm103a_pdl(output, input, output_sf, input_sf);
|
||||
return {output, output_sf};
|
||||
#endif
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false, "No compiled SM103 PDL nvfp4 quantization kernel");
|
||||
}
|
||||
|
||||
void scaled_fp4_quant_sm103a_pdl_out(torch::Tensor const& input,
|
||||
torch::Tensor const& input_sf,
|
||||
torch::Tensor& output,
|
||||
torch::Tensor& output_sf) {
|
||||
#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
|
||||
scaled_fp4_quant_sm103a_pdl(output, input, output_sf, input_sf);
|
||||
return;
|
||||
#endif
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false, "No compiled SM103 PDL nvfp4 quantization kernel");
|
||||
}
|
||||
|
||||
@@ -171,8 +171,305 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// SM103 (B300) activation quantization kernel.
|
||||
//
|
||||
// Identical to the SM100 cvt_fp16_to_fp4 except it writes scale factors
|
||||
// in the SM103 swizzled layout (Sm103BlockScaledConfig).
|
||||
// ============================================================================
|
||||
template <class Type, bool UE8M0_SF = false>
|
||||
__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
cvt_fp16_to_fp4_sm103(int32_t numRows, int32_t numCols,
|
||||
int32_t num_padded_cols,
|
||||
Type const* __restrict__ in,
|
||||
float const* __restrict__ SFScale,
|
||||
uint32_t* __restrict__ out,
|
||||
uint32_t* __restrict__ SFout) {
|
||||
using PackedVec = vllm::PackedVec<Type, CVT_FP4_PACK16>;
|
||||
|
||||
static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
|
||||
(CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
|
||||
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
|
||||
"Vec size is not matched.");
|
||||
|
||||
int32_t const numKTiles = (numCols + 63) / 64;
|
||||
|
||||
int sf_m = round_up<int>(numRows, 128);
|
||||
int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
|
||||
int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
|
||||
|
||||
float const global_scale = (SFScale == nullptr) ? 1.0f : SFScale[0];
|
||||
|
||||
for (int rowIdx = blockIdx.x; rowIdx < sf_m; rowIdx += gridDim.x) {
|
||||
if (colIdx < num_padded_cols) {
|
||||
PackedVec in_vec;
|
||||
int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
|
||||
|
||||
bool valid = (rowIdx < numRows) && (elem_idx < numCols);
|
||||
if constexpr (CVT_FP4_PACK16) {
|
||||
ld256_cg_or_zero(reinterpret_cast<u32x8_t&>(in_vec),
|
||||
&reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
|
||||
valid);
|
||||
} else {
|
||||
ld128_cg_or_zero(reinterpret_cast<uint4&>(in_vec),
|
||||
&reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
|
||||
valid);
|
||||
}
|
||||
|
||||
// SM103: Use SM103-specific SF offset function
|
||||
auto sf_out =
|
||||
cvt_quant_to_fp4_get_sf_out_offset_sm103<uint32_t,
|
||||
CVT_FP4_NUM_THREADS_PER_SF>(
|
||||
rowIdx, colIdx, numKTiles, SFout);
|
||||
|
||||
auto out_val =
|
||||
cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
|
||||
in_vec, global_scale, sf_out);
|
||||
|
||||
if (valid) {
|
||||
if constexpr (CVT_FP4_PACK16) {
|
||||
int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;
|
||||
uint64_t packed64 =
|
||||
(uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
|
||||
reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
|
||||
} else {
|
||||
out[inOffset] = out_val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Scale factor layout conversion: SM100 <-> SM103
|
||||
//
|
||||
// Converts an already-swizzled SF tensor between SM100 and SM103 layouts.
|
||||
// Both layouts use the same 512-byte tile structure (128 M-rows x 4 K-cols)
|
||||
// but arrange bytes differently within each tile.
|
||||
//
|
||||
// SM100 offset: outerM(=mIdx%32)*16 + innerM(=(mIdx/32)%4)*4 + innerK
|
||||
// SM103 offset: m8(=(mIdx/16)%8)*16 + m4a(=(mIdx/4)%4)*128 + m4b(=mIdx%4)*4
|
||||
// + innerK
|
||||
// ============================================================================
|
||||
__global__ void convert_sf_sm100_to_sm103_kernel(
|
||||
const uint8_t* __restrict__ src,
|
||||
uint8_t* __restrict__ dst,
|
||||
int32_t numMTiles,
|
||||
int32_t numKTiles) {
|
||||
// Each thread converts one byte (one SF value).
|
||||
// Grid: numMTiles * numKTiles blocks, 512 threads per block.
|
||||
int32_t tile_idx = blockIdx.x;
|
||||
int32_t mTileIdx = tile_idx / numKTiles;
|
||||
int32_t kTileIdx = tile_idx % numKTiles;
|
||||
|
||||
// Each tile is 512 bytes: 128 M-positions x 4 K-positions.
|
||||
int32_t local_idx = threadIdx.x; // 0..511
|
||||
if (mTileIdx >= numMTiles) return;
|
||||
|
||||
int64_t tile_base = static_cast<int64_t>(tile_idx) << 9;
|
||||
|
||||
// Decode this thread's (mLocal, kLocal) from a simple linear index.
|
||||
int32_t mLocal = local_idx >> 2; // 0..127
|
||||
int32_t kLocal = local_idx & 3; // 0..3
|
||||
|
||||
// Compute SM100 source offset within tile.
|
||||
int32_t outerMIdx = mLocal & 31;
|
||||
int32_t innerMIdx = (mLocal >> 5) & 3;
|
||||
int32_t sm100_off = (outerMIdx << 4) | (innerMIdx << 2) | kLocal;
|
||||
|
||||
// Compute SM103 destination offset within tile.
|
||||
int32_t m4b = mLocal & 3;
|
||||
int32_t m4a = (mLocal >> 2) & 3;
|
||||
int32_t m8 = (mLocal >> 4) & 7;
|
||||
int32_t sm103_off = (m8 << 4) | (m4a << 7) | (m4b << 2) | kLocal;
|
||||
|
||||
dst[tile_base + sm103_off] = src[tile_base + sm100_off];
|
||||
}
|
||||
|
||||
__global__ void convert_sf_sm103_to_sm100_kernel(
|
||||
const uint8_t* __restrict__ src,
|
||||
uint8_t* __restrict__ dst,
|
||||
int32_t numMTiles,
|
||||
int32_t numKTiles) {
|
||||
int32_t tile_idx = blockIdx.x;
|
||||
int32_t mTileIdx = tile_idx / numKTiles;
|
||||
if (mTileIdx >= numMTiles) return;
|
||||
|
||||
int32_t local_idx = threadIdx.x;
|
||||
int64_t tile_base = static_cast<int64_t>(tile_idx) << 9;
|
||||
|
||||
int32_t mLocal = local_idx >> 2;
|
||||
int32_t kLocal = local_idx & 3;
|
||||
|
||||
// SM103 source offset
|
||||
int32_t m4b = mLocal & 3;
|
||||
int32_t m4a = (mLocal >> 2) & 3;
|
||||
int32_t m8 = (mLocal >> 4) & 7;
|
||||
int32_t sm103_off = (m8 << 4) | (m4a << 7) | (m4b << 2) | kLocal;
|
||||
|
||||
// SM100 destination offset
|
||||
int32_t outerMIdx = mLocal & 31;
|
||||
int32_t innerMIdx = (mLocal >> 5) & 3;
|
||||
int32_t sm100_off = (outerMIdx << 4) | (innerMIdx << 2) | kLocal;
|
||||
|
||||
dst[tile_base + sm100_off] = src[tile_base + sm103_off];
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
// ============================================================================
|
||||
// Host entry: SM103 activation quantization
|
||||
//
|
||||
// When use_pdl=true, the kernel is launched with
|
||||
// cudaLaunchAttributeProgrammaticStreamSerialization, allowing the next
|
||||
// kernel on the same stream (typically the GEMM consumer) to begin
|
||||
// executing before this quantization kernel fully completes. This
|
||||
// overlaps the tail of quantization with the head of the GEMM.
|
||||
// ============================================================================
|
||||
static void scaled_fp4_quant_sm103a_impl(torch::Tensor const& output,
|
||||
torch::Tensor const& input,
|
||||
torch::Tensor const& output_sf,
|
||||
torch::Tensor const& input_sf,
|
||||
bool use_pdl) {
|
||||
int32_t m = input.size(0);
|
||||
int32_t n = input.size(1);
|
||||
|
||||
TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16.");
|
||||
TORCH_CHECK(input.scalar_type() == at::ScalarType::Half ||
|
||||
input.scalar_type() == at::ScalarType::BFloat16,
|
||||
"Unsupported input data type for quantize_to_fp4.");
|
||||
|
||||
int multiProcessorCount =
|
||||
get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
|
||||
|
||||
auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr());
|
||||
auto sf_out = static_cast<int32_t*>(output_sf.data_ptr());
|
||||
auto output_ptr = static_cast<int64_t*>(output.data_ptr());
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
|
||||
auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
|
||||
|
||||
int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE);
|
||||
|
||||
dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
|
||||
int const numBlocksPerSM =
|
||||
vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
|
||||
|
||||
// SM103 always uses swizzled layout (the SM103 variant)
|
||||
int sf_n_int = int(vllm::round_up(sf_n_unpadded, 4) / 4);
|
||||
int32_t num_padded_cols =
|
||||
sf_n_int * 4 * CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
|
||||
|
||||
int grid_y = vllm::div_round_up(num_padded_cols, static_cast<int>(block.x));
|
||||
int grid_x =
|
||||
std::min(vllm::computeEffectiveRows(m),
|
||||
std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
|
||||
dim3 grid(grid_x, grid_y);
|
||||
|
||||
VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_sm103", [&] {
|
||||
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
|
||||
auto output_u32 = reinterpret_cast<uint32_t*>(output_ptr);
|
||||
auto sf_out_u32 = reinterpret_cast<uint32_t*>(sf_out);
|
||||
|
||||
if (use_pdl) {
|
||||
// PDL launch: set ProgrammaticStreamSerialization so the next kernel
|
||||
// (GEMM) can begin before this quant kernel fully completes.
|
||||
cudaLaunchConfig_t launch_config = {};
|
||||
launch_config.gridDim = grid;
|
||||
launch_config.blockDim = block;
|
||||
launch_config.dynamicSmemBytes = 0;
|
||||
launch_config.stream = stream;
|
||||
|
||||
cudaLaunchAttribute pdl_attr;
|
||||
pdl_attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
pdl_attr.val.programmaticStreamSerializationAllowed = 1;
|
||||
launch_config.numAttrs = 1;
|
||||
launch_config.attrs = &pdl_attr;
|
||||
|
||||
CUDA_CHECK(cudaLaunchKernelEx(
|
||||
&launch_config,
|
||||
vllm::cvt_fp16_to_fp4_sm103<cuda_type, false>,
|
||||
m, n, num_padded_cols, input_ptr, input_sf_ptr,
|
||||
output_u32, sf_out_u32));
|
||||
} else {
|
||||
vllm::cvt_fp16_to_fp4_sm103<cuda_type, false>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
m, n, num_padded_cols, input_ptr, input_sf_ptr,
|
||||
output_u32, sf_out_u32);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// Original entry point (no PDL).
|
||||
void scaled_fp4_quant_sm103a(torch::Tensor const& output,
|
||||
torch::Tensor const& input,
|
||||
torch::Tensor const& output_sf,
|
||||
torch::Tensor const& input_sf) {
|
||||
scaled_fp4_quant_sm103a_impl(output, input, output_sf, input_sf,
|
||||
/*use_pdl=*/false);
|
||||
}
|
||||
|
||||
// PDL-enabled entry point: launches quant kernel with
|
||||
// ProgrammaticStreamSerialization to overlap with a subsequent GEMM.
|
||||
void scaled_fp4_quant_sm103a_pdl(torch::Tensor const& output,
|
||||
torch::Tensor const& input,
|
||||
torch::Tensor const& output_sf,
|
||||
torch::Tensor const& input_sf) {
|
||||
scaled_fp4_quant_sm103a_impl(output, input, output_sf, input_sf,
|
||||
/*use_pdl=*/true);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Host entry: SF layout conversion SM100 <-> SM103
|
||||
// ============================================================================
|
||||
void convert_sf_layout_sm100_to_sm103(torch::Tensor& dst,
|
||||
torch::Tensor const& src) {
|
||||
TORCH_CHECK(src.is_contiguous(), "Source SF tensor must be contiguous");
|
||||
TORCH_CHECK(dst.is_contiguous(), "Destination SF tensor must be contiguous");
|
||||
TORCH_CHECK(src.numel() == dst.numel(),
|
||||
"Source and destination must have the same number of elements");
|
||||
|
||||
// SF tensors are stored as int32 with shape (rounded_m, rounded_k / 4)
|
||||
// Total bytes = rounded_m * (rounded_k / 4) * 4 = rounded_m * rounded_k
|
||||
int64_t total_bytes = src.numel() * src.element_size();
|
||||
int32_t numMTiles = src.size(0) / 128;
|
||||
int32_t numKTiles = total_bytes / (numMTiles * 512);
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(src));
|
||||
auto stream = at::cuda::getCurrentCUDAStream(src.get_device());
|
||||
|
||||
int32_t num_tiles = numMTiles * numKTiles;
|
||||
dim3 grid(num_tiles);
|
||||
dim3 block(512);
|
||||
|
||||
vllm::convert_sf_sm100_to_sm103_kernel<<<grid, block, 0, stream>>>(
|
||||
static_cast<const uint8_t*>(src.data_ptr()),
|
||||
static_cast<uint8_t*>(dst.data_ptr()),
|
||||
numMTiles, numKTiles);
|
||||
}
|
||||
|
||||
void convert_sf_layout_sm103_to_sm100(torch::Tensor& dst,
|
||||
torch::Tensor const& src) {
|
||||
TORCH_CHECK(src.is_contiguous() && dst.is_contiguous());
|
||||
TORCH_CHECK(src.numel() == dst.numel());
|
||||
|
||||
int64_t total_bytes = src.numel() * src.element_size();
|
||||
int32_t numMTiles = src.size(0) / 128;
|
||||
int32_t numKTiles = total_bytes / (numMTiles * 512);
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(src));
|
||||
auto stream = at::cuda::getCurrentCUDAStream(src.get_device());
|
||||
|
||||
int32_t num_tiles = numMTiles * numKTiles;
|
||||
vllm::convert_sf_sm103_to_sm100_kernel<<<dim3(num_tiles), dim3(512), 0, stream>>>(
|
||||
static_cast<const uint8_t*>(src.data_ptr()),
|
||||
static_cast<uint8_t*>(dst.data_ptr()),
|
||||
numMTiles, numKTiles);
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Original SM100 host entry
|
||||
// ============================================================================
|
||||
void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
||||
torch::Tensor const& input,
|
||||
torch::Tensor const& output_sf,
|
||||
|
||||
@@ -24,6 +24,20 @@ void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
|
||||
torch::Tensor const& A_sf,
|
||||
torch::Tensor const& B_sf,
|
||||
torch::Tensor const& alpha);
|
||||
// SM103 (B300) uses FP4 Ultra MMA -- separate entry point compiled from
|
||||
// the same source file, guarded by CUTLASS_ARCH_MMA_SM103_SUPPORTED.
|
||||
void cutlass_scaled_fp4_mm_sm103a(torch::Tensor& D, torch::Tensor const& A,
|
||||
torch::Tensor const& B,
|
||||
torch::Tensor const& A_sf,
|
||||
torch::Tensor const& B_sf,
|
||||
torch::Tensor const& alpha);
|
||||
// PDL variant: GEMM launched with ProgrammaticStreamSerialization.
|
||||
void cutlass_scaled_fp4_mm_sm103a_pdl(torch::Tensor& D,
|
||||
torch::Tensor const& A,
|
||||
torch::Tensor const& B,
|
||||
torch::Tensor const& A_sf,
|
||||
torch::Tensor const& B_sf,
|
||||
torch::Tensor const& alpha);
|
||||
#endif
|
||||
|
||||
#if defined ENABLE_NVFP4_SM120 && ENABLE_NVFP4_SM120
|
||||
@@ -43,6 +57,14 @@ void cutlass_scaled_fp4_mm(torch::Tensor& D, const torch::Tensor& A,
|
||||
const int32_t sm = get_sm_version_num();
|
||||
|
||||
#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
|
||||
// SM103 (B300): Use FP4 Ultra kernels with K=768 tiles for higher
|
||||
// throughput. Falls through to SM100 path if SM103 kernels weren’t compiled
|
||||
// (e.g., CUDA < 12.9).
|
||||
if (sm == 103) {
|
||||
cutlass_scaled_fp4_mm_sm103a(D, A, B, A_sf, B_sf, alpha);
|
||||
return;
|
||||
}
|
||||
|
||||
if (sm >= 100 && sm < 120) {
|
||||
cutlass_scaled_fp4_mm_sm100a(D, A, B, A_sf, B_sf, alpha);
|
||||
return;
|
||||
|
||||
@@ -36,6 +36,10 @@ using namespace cute;
|
||||
|
||||
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
|
||||
|
||||
// ============================================================================
|
||||
// SM100 (B200) Tile Configurations
|
||||
// ============================================================================
|
||||
|
||||
// Configuration for M in (256, inf)
|
||||
struct sm100_fp4_config_default {
|
||||
using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
|
||||
@@ -63,6 +67,51 @@ struct sm100_fp4_config_M16 {
|
||||
using PerSmTileShape_MNK = Shape<_128, _128, _256>;
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// SM103 (B300 / Blackwell Ultra) Tile Configurations
|
||||
//
|
||||
// Key differences from SM100:
|
||||
// - Tile K = 768 is MANDATORY (CUTLASS static_assert)
|
||||
// - Uses FP4 Ultra MMA instructions (UltraVs16) for higher throughput
|
||||
// - Uses NoSmem epilogue (saves shared memory for mainloop)
|
||||
// - 1SM for small M, 2SM for large M (cooperative SM pairs)
|
||||
// ============================================================================
|
||||
#if defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED)
|
||||
|
||||
// SM103 configuration for M in (256, inf) -- 2SM cooperative execution
|
||||
struct sm103_fp4_config_default {
|
||||
// 2SM schedule: two SMs cooperate on one tile for higher throughput
|
||||
using KernelSchedule = cutlass::gemm::
|
||||
KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103;
|
||||
using EpilogueSchedule = cutlass::epilogue::NoSmemWarpSpecialized2Sm;
|
||||
using TileShape = Shape<_256, _256, Int<768>>;
|
||||
using ClusterShape = Shape<_2, _2, _1>;
|
||||
using PerSmTileShape_MNK = Shape<_128, _256, Int<768>>;
|
||||
};
|
||||
|
||||
// SM103 configuration for M in (16, 256] -- 2SM with smaller N tile
|
||||
struct sm103_fp4_config_M256 {
|
||||
using KernelSchedule = cutlass::gemm::
|
||||
KernelTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103;
|
||||
using EpilogueSchedule = cutlass::epilogue::NoSmemWarpSpecialized2Sm;
|
||||
using TileShape = Shape<_256, _128, Int<768>>;
|
||||
using ClusterShape = Shape<_2, _1, _1>;
|
||||
using PerSmTileShape_MNK = Shape<_128, _128, Int<768>>;
|
||||
};
|
||||
|
||||
// SM103 configuration for M in [1, 16] -- 1SM (decode / small batch)
|
||||
struct sm103_fp4_config_M16 {
|
||||
// 1SM schedule: single SM per tile, lower latency for small problems
|
||||
using KernelSchedule = cutlass::gemm::
|
||||
KernelTmaWarpSpecialized1SmBlockScaledMxNvf4UltraVs16Sm103;
|
||||
using EpilogueSchedule = cutlass::epilogue::NoSmemWarpSpecialized1Sm;
|
||||
using TileShape = Shape<_128, _128, Int<768>>;
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
using PerSmTileShape_MNK = Shape<_128, _128, Int<768>>;
|
||||
};
|
||||
|
||||
#endif // CUTLASS_ARCH_MMA_SM103_SUPPORTED
|
||||
|
||||
template <typename Config, typename OutType>
|
||||
struct Fp4GemmSm100 {
|
||||
// A matrix configuration
|
||||
@@ -125,6 +174,99 @@ struct Fp4GemmSm100 {
|
||||
using LayoutD = decltype(cute::make_layout(make_shape(0, 0, 0), StrideD{}));
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// SM103 GEMM Definition (FP4 Ultra)
|
||||
//
|
||||
// SM103 differs from SM100 in several fundamental ways:
|
||||
// 1. Uses cutlass::arch::Sm103 (separate CollectiveBuilder specialization)
|
||||
// 2. Element types passed as cute::tuple<DataType, ScaleFactorType>
|
||||
// (SM100 uses nv_float4_t<float_e2m1_t> wrapper instead)
|
||||
// 3. Tile K = 768 (SM100 uses K = 256)
|
||||
// 4. Epilogue uses NoSmemWarpSpecialized (SM100 uses TmaWarpSpecialized)
|
||||
// 5. Scale factor memory layout uses Sm103BlockScaledConfig
|
||||
// (different swizzle pattern from SM100's Sm1xxBlockScaledConfig)
|
||||
//
|
||||
// IMPORTANT: Scale factor layout compatibility
|
||||
// SM103 and SM100 use DIFFERENT physical scale factor layouts in memory.
|
||||
// The activation quantization kernel (scaled_fp4_quant) and the weight
|
||||
// scale factors in NVFP4 checkpoints must produce/store data in the
|
||||
// SM103-expected layout when using these kernels. Passing SM100-format
|
||||
// scale factors to SM103 kernels will produce incorrect results.
|
||||
// See Sm103BlockScaledConfig::tile_atom_to_shape_SFA for the expected
|
||||
// layout.
|
||||
// ============================================================================
|
||||
#if defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED)
|
||||
|
||||
template <typename Config, typename OutType>
|
||||
struct Fp4GemmSm103 {
|
||||
// A matrix configuration -- bare float_e2m1_t (not nv_float4_t wrapper)
|
||||
using ElementA = cutlass::float_e2m1_t;
|
||||
using ElementSFA = cutlass::float_ue4m3_t;
|
||||
using LayoutATag = cutlass::layout::RowMajor;
|
||||
static constexpr int AlignmentA = 32;
|
||||
|
||||
// B matrix configuration
|
||||
using ElementB = cutlass::float_e2m1_t;
|
||||
using ElementSFB = cutlass::float_ue4m3_t;
|
||||
using LayoutBTag = cutlass::layout::ColumnMajor;
|
||||
static constexpr int AlignmentB = 32;
|
||||
|
||||
// C/D matrix configuration
|
||||
using ElementD = OutType;
|
||||
using ElementC = OutType;
|
||||
using LayoutCTag = cutlass::layout::RowMajor;
|
||||
using LayoutDTag = cutlass::layout::RowMajor;
|
||||
static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
|
||||
static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
|
||||
|
||||
// Kernel functional config
|
||||
using ElementAccumulator = float;
|
||||
using ArchTag = cutlass::arch::Sm103;
|
||||
using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
|
||||
|
||||
// Use config's tile shapes (K=768 mandatory for SM103)
|
||||
using MmaTileShape = typename Config::TileShape;
|
||||
using ClusterShape = typename Config::ClusterShape;
|
||||
using PerSmTileShape_MNK = typename Config::PerSmTileShape_MNK;
|
||||
|
||||
// Epilogue: SM103 uses NoSmem variant with OpClassTensorOp
|
||||
// Note: epilogue builder uses Sm100 arch tag (shared epilogue HW)
|
||||
using CollectiveEpilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm100, cutlass::arch::OpClassTensorOp,
|
||||
PerSmTileShape_MNK, ClusterShape,
|
||||
cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
|
||||
ElementAccumulator, ElementC, LayoutCTag, AlignmentC, ElementD,
|
||||
LayoutDTag, AlignmentD,
|
||||
typename Config::EpilogueSchedule>::CollectiveOp;
|
||||
|
||||
// Mainloop: SM103 passes element+SF types as tuples to the builder
|
||||
using CollectiveMainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
ArchTag, OperatorClass, cute::tuple<ElementA, ElementSFA>, LayoutATag,
|
||||
AlignmentA, cute::tuple<ElementB, ElementSFB>, LayoutBTag, AlignmentB,
|
||||
ElementAccumulator, MmaTileShape, ClusterShape,
|
||||
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
|
||||
sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
typename Config::KernelSchedule>::CollectiveOp;
|
||||
|
||||
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
|
||||
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
using StrideA = typename Gemm::GemmKernel::StrideA;
|
||||
using LayoutA = decltype(cute::make_layout(make_shape(0, 0, 0), StrideA{}));
|
||||
using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA;
|
||||
using StrideB = typename Gemm::GemmKernel::StrideB;
|
||||
using LayoutB = decltype(cute::make_layout(make_shape(0, 0, 0), StrideB{}));
|
||||
using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB;
|
||||
using StrideC = typename Gemm::GemmKernel::StrideC;
|
||||
using LayoutC = decltype(cute::make_layout(make_shape(0, 0, 0), StrideC{}));
|
||||
using StrideD = typename Gemm::GemmKernel::StrideD;
|
||||
using LayoutD = decltype(cute::make_layout(make_shape(0, 0, 0), StrideD{}));
|
||||
};
|
||||
|
||||
#endif // CUTLASS_ARCH_MMA_SM103_SUPPORTED
|
||||
|
||||
template <typename Config>
|
||||
typename Config::Gemm::Arguments args_from_options(
|
||||
at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
|
||||
@@ -177,7 +319,7 @@ template <typename Config>
|
||||
void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
|
||||
at::Tensor const& A_sf, at::Tensor const& B_sf,
|
||||
at::Tensor const& alpha, int64_t m, int64_t n, int64_t k,
|
||||
cudaStream_t stream) {
|
||||
cudaStream_t stream, bool launch_with_pdl = false) {
|
||||
typename Config::Gemm gemm;
|
||||
|
||||
auto arguments =
|
||||
@@ -192,7 +334,12 @@ void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
|
||||
|
||||
CUTLASS_CHECK(gemm.initialize(arguments, workspace.data_ptr(), stream));
|
||||
|
||||
CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream));
|
||||
// When launch_with_pdl=true, CUTLASS sets
|
||||
// cudaLaunchAttributeProgrammaticStreamSerialization on the GEMM kernel,
|
||||
// allowing the next kernel on the stream to begin before this GEMM
|
||||
// fully completes.
|
||||
CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream,
|
||||
/*cuda_adapter=*/nullptr, launch_with_pdl));
|
||||
}
|
||||
|
||||
// Dispatch function to select appropriate config based on M
|
||||
@@ -220,6 +367,39 @@ void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// SM103 Dispatch
|
||||
// ============================================================================
|
||||
#if defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED)
|
||||
|
||||
template <typename OutType>
|
||||
void cutlass_fp4_gemm_sm103_dispatch(torch::Tensor& D, torch::Tensor const& A,
|
||||
torch::Tensor const& B,
|
||||
torch::Tensor const& A_sf,
|
||||
torch::Tensor const& B_sf,
|
||||
torch::Tensor const& alpha, int64_t m,
|
||||
int64_t n, int64_t k,
|
||||
cudaStream_t stream,
|
||||
bool launch_with_pdl = false) {
|
||||
uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m));
|
||||
|
||||
if (mp2 <= 16) {
|
||||
// m in [1, 16] -- 1SM, low-latency decode
|
||||
runGemm<Fp4GemmSm103<sm103_fp4_config_M16, OutType>>(
|
||||
D, A, B, A_sf, B_sf, alpha, m, n, k, stream, launch_with_pdl);
|
||||
} else if (mp2 <= 256) {
|
||||
// m in (16, 256] -- 2SM, small tile
|
||||
runGemm<Fp4GemmSm103<sm103_fp4_config_M256, OutType>>(
|
||||
D, A, B, A_sf, B_sf, alpha, m, n, k, stream, launch_with_pdl);
|
||||
} else {
|
||||
// m in (256, inf) -- 2SM, large tile
|
||||
runGemm<Fp4GemmSm103<sm103_fp4_config_default, OutType>>(
|
||||
D, A, B, A_sf, B_sf, alpha, m, n, k, stream, launch_with_pdl);
|
||||
}
|
||||
}
|
||||
|
||||
#endif // CUTLASS_ARCH_MMA_SM103_SUPPORTED
|
||||
|
||||
#else
|
||||
template <typename OutType>
|
||||
void cutlass_fp4_gemm_dispatch(torch::Tensor& D, torch::Tensor const& A,
|
||||
@@ -315,3 +495,107 @@ void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
|
||||
")");
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// SM103 Entry Point (B300 / Blackwell Ultra)
|
||||
//
|
||||
// Uses FP4 Ultra MMA instructions with K=768 tiles for higher throughput.
|
||||
// Scale factors must be in Sm103BlockScaledConfig layout (different from SM100).
|
||||
//
|
||||
// When launch_with_pdl=true, the CUTLASS GEMM is launched with
|
||||
// ProgrammaticStreamSerialization, allowing the next kernel on the stream
|
||||
// to begin before this GEMM completes. Combined with a PDL-enabled
|
||||
// quantization producer, this creates a pipelined quant->GEMM overlap.
|
||||
// ============================================================================
|
||||
#if defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED)
|
||||
|
||||
static void cutlass_scaled_fp4_mm_sm103a_impl(
|
||||
torch::Tensor& D, torch::Tensor const& A, torch::Tensor const& B,
|
||||
torch::Tensor const& A_sf, torch::Tensor const& B_sf,
|
||||
torch::Tensor const& alpha, bool launch_with_pdl) {
|
||||
CHECK_INPUT(A, FLOAT4_E2M1X2, "a");
|
||||
CHECK_INPUT(B, FLOAT4_E2M1X2, "b");
|
||||
|
||||
CHECK_INPUT(A_sf, SF_DTYPE, "scale_a");
|
||||
CHECK_INPUT(B_sf, SF_DTYPE, "scale_b");
|
||||
|
||||
CHECK_INPUT(alpha, at::ScalarType::Float, "alpha");
|
||||
|
||||
TORCH_CHECK(A.dim() == 2, "a must be a matrix");
|
||||
TORCH_CHECK(B.dim() == 2, "b must be a matrix");
|
||||
TORCH_CHECK(A.sizes()[1] == B.sizes()[1],
|
||||
"a and b shapes cannot be multiplied (", A.sizes()[0], "x",
|
||||
A.sizes()[1], " and ", B.sizes()[0], "x", B.sizes()[1], ")");
|
||||
|
||||
auto const m = A.sizes()[0];
|
||||
auto const n = B.sizes()[0];
|
||||
auto const k = A.sizes()[1] * 2;
|
||||
|
||||
constexpr int alignment = 32;
|
||||
TORCH_CHECK(k % alignment == 0, "Expected k to be divisible by ", alignment,
|
||||
", but got a shape: (", A.sizes()[0], "x", A.sizes()[1],
|
||||
"), k: ", k, ".");
|
||||
TORCH_CHECK(n % alignment == 0, "Expected n to be divisible by ", alignment,
|
||||
", but got b shape: (", B.sizes()[0], "x", B.sizes()[1], ").");
|
||||
|
||||
// SM103 scale factor shape validation.
|
||||
// Physical dimensions are the same as SM100 (padded to 128 x ceil(k/16,4)),
|
||||
// but the internal swizzle pattern (Sm103BlockScaledConfig) differs.
|
||||
auto round_up = [](int x, int y) { return (x + y - 1) / y * y; };
|
||||
int rounded_m = round_up(m, 128);
|
||||
int rounded_n = round_up(n, 128);
|
||||
int rounded_k = round_up(k / 16, 4);
|
||||
|
||||
TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix");
|
||||
TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix");
|
||||
TORCH_CHECK(A_sf.sizes()[1] == B_sf.sizes()[1],
|
||||
"scale_a and scale_b shapes cannot be multiplied (",
|
||||
A_sf.sizes()[0], "x", A_sf.sizes()[1], " and ", B_sf.sizes()[0],
|
||||
"x", B_sf.sizes()[1], ")");
|
||||
TORCH_CHECK(A_sf.sizes()[0] == rounded_m && A_sf.sizes()[1] == rounded_k,
|
||||
"scale_a must be padded and swizzled to a shape (", rounded_m,
|
||||
"x", rounded_k, "), but got a shape (", A_sf.sizes()[0], "x",
|
||||
A_sf.sizes()[1], ")");
|
||||
TORCH_CHECK(B_sf.sizes()[0] == rounded_n && B_sf.sizes()[1] == rounded_k,
|
||||
"scale_b must be padded and swizzled to a shape (", rounded_n,
|
||||
"x", rounded_k, "), but got a shape (", B_sf.sizes()[0], "x",
|
||||
B_sf.sizes()[1], ")");
|
||||
|
||||
auto out_dtype = D.dtype();
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(A));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(A.get_device());
|
||||
|
||||
if (out_dtype == at::ScalarType::Half) {
|
||||
cutlass_fp4_gemm_sm103_dispatch<cutlass::half_t>(
|
||||
D, A, B, A_sf, B_sf, alpha, m, n, k, stream, launch_with_pdl);
|
||||
} else if (out_dtype == at::ScalarType::BFloat16) {
|
||||
cutlass_fp4_gemm_sm103_dispatch<cutlass::bfloat16_t>(
|
||||
D, A, B, A_sf, B_sf, alpha, m, n, k, stream, launch_with_pdl);
|
||||
} else {
|
||||
TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm (", out_dtype,
|
||||
")");
|
||||
}
|
||||
}
|
||||
|
||||
// Original entry point (no PDL).
|
||||
void cutlass_scaled_fp4_mm_sm103a(torch::Tensor& D, torch::Tensor const& A,
|
||||
torch::Tensor const& B,
|
||||
torch::Tensor const& A_sf,
|
||||
torch::Tensor const& B_sf,
|
||||
torch::Tensor const& alpha) {
|
||||
cutlass_scaled_fp4_mm_sm103a_impl(D, A, B, A_sf, B_sf, alpha,
|
||||
/*launch_with_pdl=*/false);
|
||||
}
|
||||
|
||||
// PDL-enabled entry point: GEMM launched with ProgrammaticStreamSerialization
|
||||
// so the next kernel on the stream can overlap with this GEMM's tail.
|
||||
void cutlass_scaled_fp4_mm_sm103a_pdl(torch::Tensor& D, torch::Tensor const& A,
|
||||
torch::Tensor const& B,
|
||||
torch::Tensor const& A_sf,
|
||||
torch::Tensor const& B_sf,
|
||||
torch::Tensor const& alpha) {
|
||||
cutlass_scaled_fp4_mm_sm103a_impl(D, A, B, A_sf, B_sf, alpha,
|
||||
/*launch_with_pdl=*/true);
|
||||
}
|
||||
|
||||
#endif // CUTLASS_ARCH_MMA_SM103_SUPPORTED
|
||||
|
||||
@@ -199,6 +199,55 @@ __device__ __forceinline__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(
|
||||
return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// SM103 (Blackwell Ultra / B300) swizzled SF offset.
|
||||
//
|
||||
// SM103 uses Sm103BlockScaledConfig with a 3-level M decomposition:
|
||||
// M -> (m8, m4a, m4b) where mIdx = m4b + m4a*4 + m8*16
|
||||
// K -> (sfv16_broadcast, k4)
|
||||
//
|
||||
// Atom layout:
|
||||
// Shape: <Shape<_8, _4, _4>, Shape<SFVecSize=16, _4>>
|
||||
// Stride: <Stride<_16, _128, _4>, Stride<_0, _1>>
|
||||
//
|
||||
// Physical offset = m8*16 + m4a*128 + m4b*4 + k4
|
||||
// Each 128-row x 4-col tile occupies 512 bytes (same as SM100).
|
||||
// ============================================================================
|
||||
template <class SFType, int CVT_FP4_NUM_THREADS_PER_SF>
|
||||
__device__ __forceinline__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset_sm103(
|
||||
int rowIdx, int colIdx, int32_t numKTiles, SFType* SFout) {
|
||||
static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 ||
|
||||
CVT_FP4_NUM_THREADS_PER_SF == 2);
|
||||
|
||||
if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF != 0) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF;
|
||||
int32_t mIdx = rowIdx;
|
||||
|
||||
// SM103 tile decomposition (128 rows per M-tile, 4 K-positions per K-tile).
|
||||
int32_t mTileIdx = mIdx >> 7; // mIdx / 128
|
||||
int32_t mLocal = mIdx & 127; // mIdx % 128
|
||||
|
||||
// SM103 3-level M decomposition: mLocal = m4b + m4a*4 + m8*16
|
||||
int32_t m4b = mLocal & 3; // mLocal % 4
|
||||
int32_t m4a = (mLocal >> 2) & 3; // (mLocal / 4) % 4
|
||||
int32_t m8 = (mLocal >> 4) & 7; // (mLocal / 16) % 8
|
||||
|
||||
int32_t kTileIdx = kIdx >> 2; // kIdx / 4
|
||||
int32_t innerKIdx = kIdx & 3; // kIdx % 4
|
||||
|
||||
// Physical offset within the 512-byte tile:
|
||||
// m8 * 16 + m4a * 128 + m4b * 4 + innerKIdx
|
||||
// Tile base: (mTileIdx * numKTiles + kTileIdx) * 512
|
||||
int64_t SFOffset = (static_cast<int64_t>(mTileIdx) * numKTiles + kTileIdx)
|
||||
<< 9 |
|
||||
(m8 << 4) | (m4a << 7) | (m4b << 2) | innerKIdx;
|
||||
|
||||
return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
|
||||
}
|
||||
|
||||
template <class SFType>
|
||||
__device__ __forceinline__ uint8_t* sf_out_rowmajor_u8(int row, int pack,
|
||||
int packs_per_row_sf,
|
||||
|
||||
@@ -1,90 +0,0 @@
|
||||
#pragma once
|
||||
|
||||
// clang-format will break include orders
|
||||
// clang-format off
|
||||
#include <cudaTypedefs.h>
|
||||
|
||||
#if defined CUDA_VERSION && CUDA_VERSION >= 12020
|
||||
#include "sparse_scaled_mm_c3x.cuh"
|
||||
|
||||
#include "cutlass/numeric_conversion.h"
|
||||
#include "cutlass/transform/device/transform_universal_adapter.hpp"
|
||||
#include "cutlass/transform/kernel/sparse_gemm_compressor.hpp"
|
||||
#include "cutlass/epilogue/collective/default_epilogue.hpp"
|
||||
|
||||
// clang-format on
|
||||
|
||||
using namespace cute;
|
||||
using namespace vllm;
|
||||
|
||||
using CompressorResult = std::tuple<torch::Tensor, torch::Tensor>;
|
||||
/// Make A structured sparse by replacing elements with 0 and compress it
|
||||
template <typename Gemm>
|
||||
CompressorResult cutlass_sparse_compress(torch::Tensor const& a) {
|
||||
// Checks for conformality
|
||||
TORCH_CHECK(a.dtype() == torch::kInt8 || a.dtype() == torch::kFloat8_e4m3fn ||
|
||||
a.dtype() == torch::kFloat16 || a.dtype() == torch::kBFloat16);
|
||||
TORCH_CHECK(a.dim() == 2)
|
||||
// Check for strides and alignment
|
||||
TORCH_CHECK(a.stride(0) % 4 == 0) // Required for semi-structured sparsity
|
||||
TORCH_CHECK(a.stride(1) == 1)
|
||||
|
||||
using GemmKernel = typename Gemm::KernelType;
|
||||
using ElementA = typename Gemm::ElementAB;
|
||||
using ElementE = typename GemmKernel::CollectiveMainloop::ElementE;
|
||||
|
||||
int m = a.size(0);
|
||||
int k = a.size(1);
|
||||
using ProblemShape = typename GemmKernel::ProblemShape;
|
||||
ProblemShape prob_shape{m, 1, k, 1};
|
||||
|
||||
int64_t lda = a.stride(0);
|
||||
using StrideA = Stride<int64_t, Int<1>, int64_t>;
|
||||
StrideA a_stride{lda, Int<1>{}, 0};
|
||||
|
||||
using CompressorUtility = typename Gemm::CompressorUtility;
|
||||
CompressorUtility compressor_utility(prob_shape, a_stride);
|
||||
|
||||
// Allocate buffers for the metadata E and the compressed matrix A
|
||||
int ME = compressor_utility.get_metadata_m_physical();
|
||||
int KE = compressor_utility.get_metadata_k_physical();
|
||||
int MC = compressor_utility.get_tensorA_m_physical();
|
||||
int KC = compressor_utility.get_tensorA_k_physical();
|
||||
|
||||
auto const a_meta_options =
|
||||
torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
|
||||
auto const a_nzs_options =
|
||||
torch::TensorOptions().dtype(a.dtype()).device(a.device());
|
||||
|
||||
auto a_meta = torch::zeros({ME, KE}, a_meta_options);
|
||||
auto a_nzs = torch::zeros({MC, KC}, a_nzs_options);
|
||||
|
||||
auto a_ptr = static_cast<ElementA*>(a.data_ptr());
|
||||
auto a_nzs_ptr = static_cast<ElementA*>(a_nzs.data_ptr());
|
||||
auto a_meta_ptr = static_cast<ElementE*>(a_meta.data_ptr());
|
||||
|
||||
cutlass::KernelHardwareInfo hw_info;
|
||||
hw_info.device_id = a.device().index();
|
||||
hw_info.sm_count =
|
||||
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
|
||||
hw_info.device_id);
|
||||
|
||||
using Compressor = typename Gemm::Compressor;
|
||||
typename Compressor::Arguments arguments{
|
||||
prob_shape, {a_ptr, a_stride, a_nzs_ptr, a_meta_ptr}, {hw_info}};
|
||||
|
||||
Compressor compressor_op;
|
||||
size_t workspace_size = Compressor::get_workspace_size(arguments);
|
||||
auto const workspace_options =
|
||||
torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
|
||||
auto workspace = torch::empty(workspace_size, workspace_options);
|
||||
|
||||
CUTLASS_CHECK(compressor_op.can_implement(arguments));
|
||||
CUTLASS_CHECK(compressor_op.initialize(arguments, workspace.data_ptr()));
|
||||
CUTLASS_CHECK(compressor_op.run());
|
||||
CUDA_CHECK(cudaDeviceSynchronize());
|
||||
|
||||
return {a_meta, a_nzs};
|
||||
}
|
||||
|
||||
#endif
|
||||
@@ -1,307 +0,0 @@
|
||||
// clang-format will break include orders
|
||||
// clang-format off
|
||||
#include <cudaTypedefs.h>
|
||||
|
||||
#if defined CUDA_VERSION && CUDA_VERSION >= 12020
|
||||
#include "sparse_scaled_mm_c3x.cuh"
|
||||
// clang-format on
|
||||
|
||||
using namespace cute;
|
||||
using namespace vllm;
|
||||
|
||||
struct GemmCallerTraits {
|
||||
using return_type = void;
|
||||
|
||||
template <typename GemmConfig, typename... Args>
|
||||
static return_type invoke(Args&&... args) {
|
||||
return cutlass_sparse_gemm_caller<GemmConfig>(std::forward<Args>(args)...);
|
||||
}
|
||||
};
|
||||
|
||||
struct GemmCompressorTraits {
|
||||
using return_type = CompressorResult;
|
||||
|
||||
template <typename GemmConfig, typename... Args>
|
||||
static return_type invoke(Args&&... args) {
|
||||
return cutlass_sparse_compress<GemmConfig>(std::forward<Args>(args)...);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue,
|
||||
typename DispatchFunc, typename... Args>
|
||||
typename DispatchFunc::return_type cutlass_gemm_sm90_fp8_dispatch(
|
||||
uint32_t m, uint32_t n, Args&&... args) {
|
||||
static_assert(std::is_same_v<InType, cutlass::float_e4m3_t>);
|
||||
|
||||
using Cutlass3xGemmDefault =
|
||||
typename sm90_config_default<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
using Cutlass3xGemmM64 =
|
||||
typename sm90_fp8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
using Cutlass3xGemmM128 =
|
||||
typename sm90_fp8_config_M128<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
using Cutlass3xGemmM256 =
|
||||
typename sm90_fp8_config_M256<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
using Cutlass3xGemmM512 =
|
||||
typename sm90_fp8_config_M512<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
|
||||
using Cutlass3xGemm1 =
|
||||
typename sm90_fp8_config_1<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
using Cutlass3xGemm2 =
|
||||
typename sm90_fp8_config_2<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
using Cutlass3xGemm3 =
|
||||
typename sm90_fp8_config_3<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
using Cutlass3xGemm4 =
|
||||
typename sm90_fp8_config_4<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
using Cutlass3xGemm5 =
|
||||
typename sm90_fp8_config_5<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
using Cutlass3xGemm6 =
|
||||
typename sm90_fp8_config_6<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
using Cutlass3xGemm7 =
|
||||
typename sm90_fp8_config_7<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
using Cutlass3xGemm8 =
|
||||
typename sm90_fp8_config_8<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
|
||||
uint32_t const mp2 =
|
||||
std::max(static_cast<uint32_t>(64), next_pow_2(m)); // next power of 2
|
||||
|
||||
if (mp2 <= 64) {
|
||||
if (n == 28672) {
|
||||
return DispatchFunc::template invoke<Cutlass3xGemm2>(
|
||||
std::forward<Args>(args)...);
|
||||
} else if (n == 4096 || n == 6144) {
|
||||
return DispatchFunc::template invoke<Cutlass3xGemm1>(
|
||||
std::forward<Args>(args)...);
|
||||
}
|
||||
} else if (mp2 <= 128) {
|
||||
if (n == 4096) {
|
||||
return DispatchFunc::template invoke<Cutlass3xGemm3>(
|
||||
std::forward<Args>(args)...);
|
||||
} else if (n == 28672) {
|
||||
return DispatchFunc::template invoke<Cutlass3xGemm5>(
|
||||
std::forward<Args>(args)...);
|
||||
} else if (n == 6144) {
|
||||
return DispatchFunc::template invoke<Cutlass3xGemm4>(
|
||||
std::forward<Args>(args)...);
|
||||
}
|
||||
} else if (mp2 <= 256) {
|
||||
if (n == 4096) {
|
||||
return DispatchFunc::template invoke<Cutlass3xGemm6>(
|
||||
std::forward<Args>(args)...);
|
||||
} else if (n == 28672) {
|
||||
return DispatchFunc::template invoke<Cutlass3xGemm8>(
|
||||
std::forward<Args>(args)...);
|
||||
} else if (n == 6144) {
|
||||
return DispatchFunc::template invoke<Cutlass3xGemm7>(
|
||||
std::forward<Args>(args)...);
|
||||
}
|
||||
} else {
|
||||
if (n == 6144 || n == 28672) {
|
||||
return DispatchFunc::template invoke<Cutlass3xGemm8>(
|
||||
std::forward<Args>(args)...);
|
||||
} else if (n == 4096) {
|
||||
return DispatchFunc::template invoke<Cutlass3xGemm7>(
|
||||
std::forward<Args>(args)...);
|
||||
}
|
||||
}
|
||||
|
||||
// Otherwise the default heuristic
|
||||
if (mp2 <= 64) {
|
||||
// n in [1, 64]
|
||||
return DispatchFunc::template invoke<Cutlass3xGemmM64>(
|
||||
std::forward<Args>(args)...);
|
||||
} else if (mp2 <= 128) {
|
||||
// n in (64, 128]
|
||||
return DispatchFunc::template invoke<Cutlass3xGemmM128>(
|
||||
std::forward<Args>(args)...);
|
||||
} else if (mp2 <= 256) {
|
||||
// n in (128, 256]
|
||||
return DispatchFunc::template invoke<Cutlass3xGemmM256>(
|
||||
std::forward<Args>(args)...);
|
||||
} else {
|
||||
// n in (256, inf)
|
||||
return DispatchFunc::template invoke<Cutlass3xGemmM512>(
|
||||
std::forward<Args>(args)...);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue,
|
||||
typename DispatchFunc, typename... Args>
|
||||
typename DispatchFunc::return_type cutlass_gemm_sm90_16bit_dispatch(
|
||||
uint32_t m, uint32_t n, Args&&... args) {
|
||||
using Cutlass3xGemmDefault =
|
||||
typename sm90_config_default<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
|
||||
return DispatchFunc::template invoke<Cutlass3xGemmDefault>(
|
||||
std::forward<Args>(args)...);
|
||||
}
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue,
|
||||
typename DispatchFunc, typename... Args>
|
||||
typename DispatchFunc::return_type cutlass_gemm_sm90_int8_dispatch(
|
||||
uint32_t m, uint32_t n, Args&&... args) {
|
||||
static_assert(std::is_same_v<InType, int8_t>);
|
||||
|
||||
using Cutlass3xGemmDefault =
|
||||
typename sm90_config_default<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
using Cutlass3xGemmM128 =
|
||||
typename sm90_int8_config_M128<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
using Cutlass3xGemmM64 =
|
||||
typename sm90_int8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
using Cutlass3xGemmM32NBig =
|
||||
typename sm90_int8_config_M32_NBig<InType, OutType,
|
||||
Epilogue>::Cutlass3xGemm;
|
||||
using Cutlass3xGemmM32NSmall =
|
||||
typename sm90_int8_config_M32_NSmall<InType, OutType,
|
||||
Epilogue>::Cutlass3xGemm;
|
||||
|
||||
bool const is_small_n = n < 8192;
|
||||
uint32_t const mp2 =
|
||||
std::max(static_cast<uint32_t>(32), next_pow_2(m)); // next power of 2
|
||||
|
||||
if (mp2 <= 32) {
|
||||
// m in [1, 32]
|
||||
if (is_small_n) {
|
||||
return DispatchFunc::template invoke<Cutlass3xGemmM32NSmall>(
|
||||
std::forward<Args>(args)...);
|
||||
} else {
|
||||
return DispatchFunc::template invoke<Cutlass3xGemmM32NBig>(
|
||||
std::forward<Args>(args)...);
|
||||
}
|
||||
} else if (mp2 <= 64) {
|
||||
// m in (32, 64]
|
||||
return DispatchFunc::template invoke<Cutlass3xGemmM64>(
|
||||
std::forward<Args>(args)...);
|
||||
} else if (mp2 <= 128) {
|
||||
// m in (64, 128]
|
||||
return DispatchFunc::template invoke<Cutlass3xGemmM128>(
|
||||
std::forward<Args>(args)...);
|
||||
} else {
|
||||
// m in (128, inf)
|
||||
return DispatchFunc::template invoke<Cutlass3xGemmDefault>(
|
||||
std::forward<Args>(args)...);
|
||||
}
|
||||
}
|
||||
|
||||
// Dispatch to GEMM implementations based on element types
|
||||
template <template <typename, typename, typename> typename Epilogue,
|
||||
typename... EpilogueArgs>
|
||||
void cutlass_scaled_sparse_mm_sm90_epilogue(torch::Tensor& out,
|
||||
torch::Tensor const& a,
|
||||
torch::Tensor const& bt_nzs,
|
||||
torch::Tensor const& bt_meta,
|
||||
EpilogueArgs&&... epilogue_args) {
|
||||
uint32_t const m = out.size(0);
|
||||
uint32_t const n = out.size(1);
|
||||
|
||||
// TODO: add dispatch functions to all of these
|
||||
TORCH_CHECK(bt_meta.dtype() == torch::kUInt8);
|
||||
if (a.dtype() == torch::kInt8) {
|
||||
TORCH_CHECK(bt_nzs.dtype() == torch::kInt8);
|
||||
|
||||
if (out.dtype() == torch::kBFloat16) {
|
||||
return cutlass_gemm_sm90_int8_dispatch<int8_t, cutlass::bfloat16_t,
|
||||
Epilogue, GemmCallerTraits>(
|
||||
m, n, out, a, bt_nzs, bt_meta,
|
||||
std::forward<EpilogueArgs>(epilogue_args)...);
|
||||
} else {
|
||||
TORCH_CHECK(out.dtype() == torch::kFloat16);
|
||||
return cutlass_gemm_sm90_int8_dispatch<int8_t, cutlass::half_t, Epilogue,
|
||||
GemmCallerTraits>(
|
||||
m, n, out, a, bt_nzs, bt_meta,
|
||||
std::forward<EpilogueArgs>(epilogue_args)...);
|
||||
}
|
||||
} else if (a.dtype() == torch::kFloat8_e4m3fn) {
|
||||
TORCH_CHECK(bt_nzs.dtype() == torch::kFloat8_e4m3fn);
|
||||
|
||||
if (out.dtype() == torch::kBFloat16) {
|
||||
return cutlass_gemm_sm90_fp8_dispatch<cutlass::float_e4m3_t,
|
||||
cutlass::bfloat16_t, Epilogue,
|
||||
GemmCallerTraits>(
|
||||
m, n, out, a, bt_nzs, bt_meta,
|
||||
std::forward<EpilogueArgs>(epilogue_args)...);
|
||||
} else {
|
||||
TORCH_CHECK(out.dtype() == torch::kFloat16);
|
||||
return cutlass_gemm_sm90_fp8_dispatch<
|
||||
cutlass::float_e4m3_t, cutlass::half_t, Epilogue, GemmCallerTraits>(
|
||||
m, n, out, a, bt_nzs, bt_meta,
|
||||
std::forward<EpilogueArgs>(epilogue_args)...);
|
||||
}
|
||||
} else if (a.dtype() == torch::kFloat16) {
|
||||
TORCH_CHECK(bt_nzs.dtype() == torch::kFloat16);
|
||||
TORCH_CHECK(out.dtype() == torch::kFloat16);
|
||||
|
||||
return cutlass_gemm_sm90_16bit_dispatch<cutlass::half_t, cutlass::half_t,
|
||||
Epilogue, GemmCallerTraits>(
|
||||
m, n, out, a, bt_nzs, bt_meta,
|
||||
std::forward<EpilogueArgs>(epilogue_args)...);
|
||||
} else { // a.dtype() == torch::kBFloat16
|
||||
TORCH_CHECK(a.dtype() == torch::kBFloat16);
|
||||
TORCH_CHECK(bt_nzs.dtype() == torch::kBFloat16);
|
||||
TORCH_CHECK(out.dtype() == torch::kBFloat16);
|
||||
|
||||
return cutlass_gemm_sm90_16bit_dispatch<
|
||||
cutlass::bfloat16_t, cutlass::bfloat16_t, Epilogue, GemmCallerTraits>(
|
||||
m, n, out, a, bt_nzs, bt_meta,
|
||||
std::forward<EpilogueArgs>(epilogue_args)...);
|
||||
}
|
||||
}
|
||||
|
||||
void cutlass_scaled_sparse_mm_sm90(torch::Tensor& out, torch::Tensor const& a,
|
||||
torch::Tensor const& bt_nzs,
|
||||
torch::Tensor const& bt_meta,
|
||||
torch::Tensor const& a_scales,
|
||||
torch::Tensor const& b_scales,
|
||||
std::optional<torch::Tensor> const& bias) {
|
||||
TORCH_CHECK(bt_meta.dtype() == torch::kUInt8);
|
||||
TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
|
||||
TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
|
||||
|
||||
if (bias) {
|
||||
TORCH_CHECK(bias->dtype() == out.dtype(),
|
||||
"CUTLASS scaled_mm bias dtype must match output dtype ",
|
||||
out.dtype());
|
||||
return cutlass_scaled_sparse_mm_sm90_epilogue<
|
||||
c3x::ScaledEpilogueColumnBias>(out, a, bt_nzs, bt_meta, b_scales,
|
||||
a_scales, *bias);
|
||||
} else {
|
||||
return cutlass_scaled_sparse_mm_sm90_epilogue<c3x::ScaledEpilogue>(
|
||||
out, a, bt_nzs, bt_meta, b_scales, a_scales);
|
||||
}
|
||||
}
|
||||
|
||||
CompressorResult cutlass_sparse_compress_sm90(torch::Tensor const& a) {
|
||||
// These m and n variables are fordispatching to different GEMM algorithms.
|
||||
uint32_t const m = 1; // Set M to 1 for compression
|
||||
uint32_t const n = a.size(1);
|
||||
|
||||
// Note: For correctness, the compressed format must be invariant in:
|
||||
// - M, the flattened number of tokens
|
||||
// - Whether output dtype is fp16 or bf16
|
||||
// - CUTLASS epilogues
|
||||
|
||||
if (a.dtype() == torch::kInt8) {
|
||||
return cutlass_gemm_sm90_int8_dispatch<int8_t, cutlass::bfloat16_t,
|
||||
c3x::TrivialEpilogue,
|
||||
GemmCompressorTraits>(m, n, a);
|
||||
} else if (a.dtype() == torch::kFloat8_e4m3fn) {
|
||||
return cutlass_gemm_sm90_fp8_dispatch<
|
||||
cutlass::float_e4m3_t, cutlass::bfloat16_t, c3x::TrivialEpilogue,
|
||||
GemmCompressorTraits>(m, n, a);
|
||||
} else if (a.dtype() == torch::kFloat16) {
|
||||
return cutlass_gemm_sm90_16bit_dispatch<
|
||||
cutlass::bfloat16_t, cutlass::bfloat16_t, c3x::TrivialEpilogue,
|
||||
GemmCompressorTraits>(m, n, a);
|
||||
} else {
|
||||
TORCH_CHECK(a.dtype() == torch::kBFloat16,
|
||||
"cutlass_sparse_compress only supports int8, fp8_e4m3, fp16, "
|
||||
"and bf16 datatypes");
|
||||
return cutlass_gemm_sm90_16bit_dispatch<cutlass::half_t, cutlass::half_t,
|
||||
c3x::TrivialEpilogue,
|
||||
GemmCompressorTraits>(m, n, a);
|
||||
}
|
||||
}
|
||||
|
||||
#endif
|
||||
@@ -1,570 +0,0 @@
|
||||
#pragma once
|
||||
|
||||
// clang-format will break include orders
|
||||
// clang-format off
|
||||
#include <cudaTypedefs.h>
|
||||
|
||||
#include <torch/all.h>
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
|
||||
#include "cuda_utils.h"
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
|
||||
#include "cutlass/gemm/device/gemm_universal_adapter.h"
|
||||
#include "cutlass/epilogue/collective/collective_builder.hpp"
|
||||
#include "cutlass/gemm/collective/collective_builder.hpp"
|
||||
|
||||
#include "cutlass/transform/device/transform_universal_adapter.hpp"
|
||||
#include "cutlass/transform/kernel/sparse_gemm_compressor.hpp"
|
||||
|
||||
#include "core/math.hpp"
|
||||
#include "cutlass_extensions/cute_utils.cuh"
|
||||
#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
#include "cutlass_extensions/torch_utils.hpp"
|
||||
// clang-format on
|
||||
|
||||
using namespace cute;
|
||||
|
||||
/*
|
||||
This file defines 2:4 sparse GEMM operations using the CUTLASS 3.x API,
|
||||
for NVIDIA GPUs with sm90a (Hopper) or later.
|
||||
*/
|
||||
|
||||
namespace {
|
||||
|
||||
// A wrapper for the GEMM kernel that is used to guard against compilation on
|
||||
// architectures that will never use the kernel. The purpose of this is to
|
||||
// reduce the size of the compiled binary.
|
||||
// __CUDA_ARCH__ is not defined in host code, so this lets us smuggle the ifdef
|
||||
// into code that will be executed on the device where it is defined.
|
||||
template <typename Kernel>
|
||||
struct enable_sm90_or_later : Kernel {
|
||||
template <typename... Args>
|
||||
CUTLASS_DEVICE void operator()(Args&&... args) {
|
||||
#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 900
|
||||
Kernel::operator()(std::forward<Args>(args)...);
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
using GemmUniversalMode = cutlass::gemm::GemmUniversalMode;
|
||||
|
||||
/*
|
||||
* cutlass_sparse_3x_gemm defines a 2:4 sparse GEMM kernel via CUTLASS
|
||||
* for SM90 Hopper systems.
|
||||
*/
|
||||
template <typename ElementAB_, typename ElementD_,
|
||||
template <typename, typename, typename> typename Epilogue_,
|
||||
typename TileShape, typename ClusterShape, typename KernelSchedule,
|
||||
typename EpilogueSchedule>
|
||||
struct cutlass_sparse_3x_gemm {
|
||||
using ElementAB = ElementAB_;
|
||||
using ElementD = ElementD_;
|
||||
using ElementAcc =
|
||||
typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
|
||||
float>::type;
|
||||
|
||||
using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
|
||||
|
||||
using ElementC = void;
|
||||
using LayoutC = cutlass::layout::RowMajor;
|
||||
using LayoutC_Transpose =
|
||||
typename cutlass::layout::LayoutTranspose<LayoutC>::type;
|
||||
|
||||
using EVTCompute = typename Epilogue::EVTCompute;
|
||||
|
||||
// These are the minimum alignments needed for the kernels to compile
|
||||
static constexpr int AlignmentAB =
|
||||
128 / cutlass::sizeof_bits<ElementAB>::value;
|
||||
static constexpr int AlignmentCD =
|
||||
128 / cutlass::sizeof_bits<ElementD>::value;
|
||||
|
||||
using CollectiveEpilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, TileShape,
|
||||
ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto,
|
||||
ElementAcc, float, ElementC, LayoutC_Transpose, AlignmentCD, ElementD,
|
||||
LayoutC_Transpose, AlignmentCD, EpilogueSchedule,
|
||||
EVTCompute>::CollectiveOp;
|
||||
|
||||
static constexpr size_t CEStorageSize =
|
||||
sizeof(typename CollectiveEpilogue::SharedStorage);
|
||||
using Stages = typename cutlass::gemm::collective::StageCountAutoCarveout<
|
||||
static_cast<int>(CEStorageSize)>;
|
||||
|
||||
// clang-format off
|
||||
using CollectiveMainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
cutlass::arch::Sm90, cutlass::arch::OpClassSparseTensorOp,
|
||||
ElementAB, cutlass::layout::RowMajor, AlignmentAB,
|
||||
ElementAB, cutlass::layout::ColumnMajor, AlignmentAB,
|
||||
ElementAcc, TileShape, ClusterShape,
|
||||
Stages,
|
||||
KernelSchedule>::CollectiveOp;
|
||||
// clang-format on
|
||||
|
||||
using KernelType = enable_sm90_or_later<cutlass::gemm::kernel::GemmUniversal<
|
||||
cute::Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
|
||||
cutlass::gemm::PersistentScheduler>>;
|
||||
|
||||
struct GemmKernel : public KernelType {};
|
||||
|
||||
// Sparse compressor definitions
|
||||
using SparseConfig = typename GemmKernel::CollectiveMainloop::SparseConfig;
|
||||
using LayoutTagA = cutlass::layout::RowMajor;
|
||||
using CompressorUtility =
|
||||
cutlass::transform::kernel::StructuredSparseCompressorUtility<
|
||||
typename GemmKernel::ProblemShape, ElementAB, LayoutTagA,
|
||||
SparseConfig>;
|
||||
using CompressorKernel =
|
||||
cutlass::transform::kernel::StructuredSparseCompressor<
|
||||
typename GemmKernel::ProblemShape, ElementAB, LayoutTagA,
|
||||
SparseConfig, cutlass::arch::Sm90>;
|
||||
using Compressor =
|
||||
cutlass::transform::device::TransformUniversalAdapter<CompressorKernel>;
|
||||
};
|
||||
|
||||
/*
|
||||
* This class defines kernel to compress a 2:4 sparse matrix.
|
||||
* The particular format is defined by the Gemm template parameter,
|
||||
* which is a cutlass_sparse_3x_gemm.
|
||||
*/
|
||||
using CompressorResult = std::tuple<torch::Tensor, torch::Tensor>;
|
||||
/// Make A structured sparse by replacing elements with 0 and compress it
|
||||
template <typename Gemm>
|
||||
CompressorResult cutlass_sparse_compress(torch::Tensor const& a) {
|
||||
// Checks for conformality
|
||||
TORCH_CHECK(a.dtype() == torch::kInt8 || a.dtype() == torch::kFloat8_e4m3fn ||
|
||||
a.dtype() == torch::kFloat16 || a.dtype() == torch::kBFloat16);
|
||||
TORCH_CHECK(a.dim() == 2)
|
||||
// Check for strides and alignment
|
||||
TORCH_CHECK(a.stride(0) % 4 == 0) // Required for semi-structured sparsity
|
||||
TORCH_CHECK(a.stride(1) == 1)
|
||||
|
||||
using GemmKernel = typename Gemm::KernelType;
|
||||
using ElementA = typename Gemm::ElementAB;
|
||||
using ElementE = typename GemmKernel::CollectiveMainloop::ElementE;
|
||||
|
||||
int m = a.size(0);
|
||||
int k = a.size(1);
|
||||
using ProblemShape = typename GemmKernel::ProblemShape;
|
||||
ProblemShape prob_shape{m, 1, k, 1};
|
||||
|
||||
int64_t lda = a.stride(0);
|
||||
using StrideA = Stride<int64_t, Int<1>, int64_t>;
|
||||
StrideA a_stride{lda, Int<1>{}, 0};
|
||||
|
||||
using CompressorUtility = typename Gemm::CompressorUtility;
|
||||
CompressorUtility compressor_utility(prob_shape, a_stride);
|
||||
|
||||
// Allocate buffers for the metadata E and the compressed matrix A
|
||||
int ME = compressor_utility.get_metadata_m_physical();
|
||||
int KE = compressor_utility.get_metadata_k_physical();
|
||||
int MC = compressor_utility.get_tensorA_m_physical();
|
||||
int KC = compressor_utility.get_tensorA_k_physical();
|
||||
|
||||
auto const a_meta_options =
|
||||
torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
|
||||
auto const a_nzs_options =
|
||||
torch::TensorOptions().dtype(a.dtype()).device(a.device());
|
||||
|
||||
auto a_meta = torch::zeros({ME, KE}, a_meta_options);
|
||||
auto a_nzs = torch::zeros({MC, KC}, a_nzs_options);
|
||||
|
||||
auto a_ptr = static_cast<ElementA*>(a.data_ptr());
|
||||
auto a_nzs_ptr = static_cast<ElementA*>(a_nzs.data_ptr());
|
||||
auto a_meta_ptr = static_cast<ElementE*>(a_meta.data_ptr());
|
||||
|
||||
cutlass::KernelHardwareInfo hw_info;
|
||||
hw_info.device_id = a.device().index();
|
||||
hw_info.sm_count =
|
||||
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
|
||||
hw_info.device_id);
|
||||
|
||||
using Compressor = typename Gemm::Compressor;
|
||||
typename Compressor::Arguments arguments{
|
||||
prob_shape, {a_ptr, a_stride, a_nzs_ptr, a_meta_ptr}, {hw_info}};
|
||||
|
||||
Compressor compressor_op;
|
||||
size_t workspace_size = Compressor::get_workspace_size(arguments);
|
||||
auto const workspace_options =
|
||||
torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
|
||||
auto workspace = torch::empty(workspace_size, workspace_options);
|
||||
|
||||
CUTLASS_CHECK(compressor_op.can_implement(arguments));
|
||||
CUTLASS_CHECK(compressor_op.initialize(arguments, workspace.data_ptr()));
|
||||
CUTLASS_CHECK(compressor_op.run());
|
||||
CUDA_CHECK(cudaDeviceSynchronize());
|
||||
|
||||
return {a_meta, a_nzs};
|
||||
}
|
||||
|
||||
template <typename Gemm, typename... EpilogueArgs>
|
||||
void cutlass_sparse_gemm_caller(torch::Tensor& out, torch::Tensor const& a,
|
||||
torch::Tensor const& bt_nzs,
|
||||
torch::Tensor const& bt_meta,
|
||||
EpilogueArgs&&... epilogue_params) {
|
||||
using ElementAB = typename Gemm::ElementAB;
|
||||
using ElementD = typename Gemm::ElementD;
|
||||
|
||||
// Interface stride expected from the argument a (will get transposed)
|
||||
// We compute C^T = B^T * A^T, but we assume B is transposed before
|
||||
// compression and hence the bt_* naming
|
||||
using LayoutB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutA;
|
||||
using LayoutE = typename Gemm::GemmKernel::CollectiveMainloop::LayoutE;
|
||||
|
||||
// M, N, K after transposition
|
||||
int32_t m = out.size(1);
|
||||
int32_t n = out.size(0);
|
||||
int32_t k = a.size(1);
|
||||
|
||||
int64_t lda = a.stride(0);
|
||||
int64_t ldc = out.stride(0);
|
||||
|
||||
using StrideA = Stride<int64_t, Int<1>, int64_t>;
|
||||
using StrideC = Stride<Int<1>, int64_t, int64_t>;
|
||||
|
||||
StrideA a_stride{lda, Int<1>{}, Int<0>{}};
|
||||
StrideC c_stride{Int<1>{}, ldc, Int<0>{}};
|
||||
|
||||
using GemmKernel = typename Gemm::GemmKernel;
|
||||
typename GemmKernel::ProblemShape prob_shape{m, n, k, 1};
|
||||
|
||||
using ElementE = typename GemmKernel::CollectiveMainloop::ElementE;
|
||||
using SparseConfig = typename GemmKernel::CollectiveMainloop::SparseConfig;
|
||||
|
||||
LayoutB b_layout = SparseConfig::fill_layoutA(prob_shape);
|
||||
LayoutE e_layout = SparseConfig::fill_layoutE(prob_shape);
|
||||
|
||||
auto a_ptr = static_cast<ElementAB*>(a.data_ptr());
|
||||
auto b_ptr = static_cast<ElementAB*>(bt_nzs.data_ptr());
|
||||
auto e_ptr = static_cast<ElementE*>(bt_meta.data_ptr());
|
||||
typename GemmKernel::MainloopArguments mainloop_args{
|
||||
b_ptr, b_layout, a_ptr, a_stride, e_ptr, e_layout};
|
||||
|
||||
auto c_ptr = static_cast<ElementD*>(out.data_ptr());
|
||||
typename GemmKernel::EpilogueArguments epilogue_args{
|
||||
Gemm::Epilogue::prepare_args(
|
||||
std::forward<EpilogueArgs>(epilogue_params)...),
|
||||
c_ptr, c_stride, c_ptr, c_stride};
|
||||
|
||||
typename GemmKernel::Arguments args{cutlass::gemm::GemmUniversalMode::kGemm,
|
||||
prob_shape, mainloop_args, epilogue_args};
|
||||
|
||||
// Launch the CUTLASS GEMM kernel.
|
||||
using GemmOp = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
|
||||
GemmOp gemm_op;
|
||||
CUTLASS_CHECK(gemm_op.can_implement(args));
|
||||
|
||||
size_t workspace_size = gemm_op.get_workspace_size(args);
|
||||
auto const workspace_options =
|
||||
torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
|
||||
auto workspace = torch::empty(workspace_size, workspace_options);
|
||||
|
||||
auto stream = at::cuda::getCurrentCUDAStream(a.get_device());
|
||||
|
||||
cutlass::Status status = gemm_op.run(args, workspace.data_ptr(), stream);
|
||||
CUTLASS_CHECK(status);
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////
|
||||
// Gemm Configs are defined below
|
||||
//////////////////////////////////////////////////
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_config_default {};
|
||||
|
||||
template <typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_config_default<half_t, OutType, Epilogue> {
|
||||
using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecialized;
|
||||
using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
|
||||
using TileShape = Shape<_128, _128, _128>;
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<half_t, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_config_default<cutlass::bfloat16_t, OutType, Epilogue> {
|
||||
using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecialized;
|
||||
using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
|
||||
using TileShape = Shape<_128, _128, _128>;
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<cutlass::bfloat16_t, OutType, Epilogue, TileShape,
|
||||
ClusterShape, KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
//////////////////////// Cherry-Picking Kernels ////////////////////////
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_fp8_config_1 {
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
|
||||
using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
|
||||
using TileShape = Shape<_64, _64, _256>;
|
||||
using ClusterShape = Shape<_8, _1, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_fp8_config_2 {
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
using KernelSchedule =
|
||||
cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
|
||||
using EpilogueSchedule =
|
||||
typename cutlass::epilogue::TmaWarpSpecializedCooperative;
|
||||
using TileShape = Shape<_128, _64, _256>;
|
||||
using ClusterShape = Shape<_8, _1, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_fp8_config_3 {
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
|
||||
using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
|
||||
using TileShape = Shape<_64, _64, _256>;
|
||||
using ClusterShape = Shape<_1, _2, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_fp8_config_4 {
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
|
||||
using EpilogueSchedule =
|
||||
typename cutlass::epilogue::TmaWarpSpecializedCooperative;
|
||||
using TileShape = Shape<_64, _128, _256>;
|
||||
using ClusterShape = Shape<_8, _1, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_fp8_config_5 {
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
using KernelSchedule =
|
||||
cutlass::gemm::KernelTmaWarpSpecializedPingpongFP8FastAccum;
|
||||
using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
|
||||
using TileShape = Shape<_128, _128, _256>;
|
||||
using ClusterShape = Shape<_8, _1, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_fp8_config_6 {
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
|
||||
using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
|
||||
using TileShape = Shape<_64, _128, _256>;
|
||||
using ClusterShape = Shape<_1, _2, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_fp8_config_7 {
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
using KernelSchedule =
|
||||
cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
|
||||
using EpilogueSchedule =
|
||||
typename cutlass::epilogue::TmaWarpSpecializedCooperative;
|
||||
using TileShape = Shape<_128, _128, _256>;
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_fp8_config_8 {
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
using KernelSchedule =
|
||||
cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
|
||||
using EpilogueSchedule =
|
||||
typename cutlass::epilogue::TmaWarpSpecializedCooperative;
|
||||
using TileShape = Shape<_128, _256, _128>;
|
||||
using ClusterShape = Shape<_8, _1, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_config_default<cutlass::float_e4m3_t, OutType, Epilogue> {
|
||||
// M in (128, inf)
|
||||
using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
|
||||
using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
|
||||
using TileShape = Shape<_128, _128, _128>;
|
||||
using ClusterShape = Shape<_1, _2, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<cutlass::float_e4m3_t, OutType, Epilogue,
|
||||
TileShape, ClusterShape, KernelSchedule,
|
||||
EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_fp8_config_M64 {
|
||||
// M in [1, 64]
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
|
||||
using EpilogueSchedule =
|
||||
typename cutlass::epilogue::TmaWarpSpecializedCooperative;
|
||||
using TileShape = Shape<_64, _64, _256>;
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_fp8_config_M128 {
|
||||
// M in (64, 128]
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
using KernelSchedule =
|
||||
cutlass::gemm::KernelTmaWarpSpecializedPingpongFP8FastAccum;
|
||||
using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
|
||||
using TileShape = Shape<_64, _128, _256>;
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_fp8_config_M256 {
|
||||
// M in (128, 256]
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
using KernelSchedule =
|
||||
cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
|
||||
using EpilogueSchedule =
|
||||
typename cutlass::epilogue::TmaWarpSpecializedCooperative;
|
||||
using TileShape = Shape<_128, _128, _256>;
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_fp8_config_M512 {
|
||||
// M in (256, ]
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
using KernelSchedule =
|
||||
cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
|
||||
using EpilogueSchedule =
|
||||
typename cutlass::epilogue::TmaWarpSpecializedCooperative;
|
||||
using TileShape = Shape<_128, _128, _256>;
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_config_default<int8_t, OutType, Epilogue> {
|
||||
// For M > 128 and any N
|
||||
using KernelSchedule =
|
||||
typename cutlass::gemm::KernelTmaWarpSpecializedPingpong;
|
||||
using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
|
||||
using TileShape = Shape<_128, _128, _128>;
|
||||
using ClusterShape = Shape<_2, _1, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<int8_t, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_int8_config_M128 {
|
||||
// For M in (64, 128] and any N
|
||||
static_assert(std::is_same<InType, int8_t>());
|
||||
using KernelSchedule =
|
||||
typename cutlass::gemm::KernelTmaWarpSpecializedPingpong;
|
||||
using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
|
||||
using TileShape = Shape<_64, _128, _128>;
|
||||
using ClusterShape = Shape<_2, _1, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_int8_config_M64 {
|
||||
// For M in (32, 64] and any N
|
||||
static_assert(std::is_same<InType, int8_t>());
|
||||
using KernelSchedule = typename cutlass::gemm::KernelTmaWarpSpecialized;
|
||||
using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
|
||||
using TileShape = Shape<_64, _64, _256>;
|
||||
using ClusterShape = Shape<_1, _1, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_int8_config_M32_NBig {
|
||||
// For M in [1, 32] and N >= 8192
|
||||
static_assert(std::is_same<InType, int8_t>());
|
||||
using KernelSchedule = typename cutlass::gemm::KernelTmaWarpSpecialized;
|
||||
using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
|
||||
using TileShape = Shape<_64, _128, _256>;
|
||||
using ClusterShape = Shape<_1, _4, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue>
|
||||
struct sm90_int8_config_M32_NSmall {
|
||||
// For M in [1, 32] and N < 8192
|
||||
static_assert(std::is_same<InType, int8_t>());
|
||||
using KernelSchedule = typename cutlass::gemm::KernelTmaWarpSpecialized;
|
||||
using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
|
||||
using TileShape = Shape<_64, _64, _256>;
|
||||
using ClusterShape = Shape<_1, _8, _1>;
|
||||
using Cutlass3xGemm =
|
||||
cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
|
||||
KernelSchedule, EpilogueSchedule>;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
@@ -1,104 +0,0 @@
|
||||
#include <cudaTypedefs.h>
|
||||
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#include "cutlass_extensions/common.hpp"
|
||||
|
||||
bool cutlass_sparse_scaled_mm_supported(int64_t cuda_device_capability) {
|
||||
// sparse CUTLASS kernels need exactly hopper and are not forward compatible
|
||||
// CUDA 12.2 and SM90 (Hopper)
|
||||
|
||||
#if defined CUDA_VERSION
|
||||
return CUDA_VERSION >= 12020 && cuda_device_capability == 90;
|
||||
#endif
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
#if defined ENABLE_SPARSE_SCALED_MM_C3X && ENABLE_SPARSE_SCALED_MM_C3X
|
||||
void cutlass_scaled_sparse_mm_sm90(torch::Tensor& c, torch::Tensor const& a,
|
||||
torch::Tensor const& b,
|
||||
torch::Tensor const& e,
|
||||
torch::Tensor const& a_scales,
|
||||
torch::Tensor const& b_scales,
|
||||
std::optional<torch::Tensor> const& bias);
|
||||
|
||||
using CompressorResult = std::tuple<torch::Tensor, torch::Tensor>;
|
||||
CompressorResult cutlass_sparse_compress_sm90(torch::Tensor const& a);
|
||||
#endif
|
||||
|
||||
void cutlass_scaled_sparse_mm(torch::Tensor& c, torch::Tensor const& a,
|
||||
torch::Tensor const& bt_nzs,
|
||||
torch::Tensor const& bt_meta,
|
||||
torch::Tensor const& a_scales,
|
||||
torch::Tensor const& b_scales,
|
||||
std::optional<torch::Tensor> const& bias) {
|
||||
// Checks for conformality
|
||||
TORCH_CHECK(a.dim() == 2 && bt_nzs.dim() == 2 && c.dim() == 2);
|
||||
TORCH_CHECK(c.size(1) == bt_nzs.size(0) && bt_nzs.size(1) * 2 == a.size(1) &&
|
||||
a.size(0) == c.size(0));
|
||||
TORCH_CHECK(a_scales.numel() == 1 || a_scales.numel() == a.size(0));
|
||||
TORCH_CHECK(b_scales.numel() == 1 || b_scales.numel() == bt_nzs.size(0));
|
||||
|
||||
// Check for strides and alignment
|
||||
TORCH_CHECK(a.stride(1) == 1 && bt_nzs.stride(1) == 1 &&
|
||||
c.stride(1) == 1); // Row-major
|
||||
TORCH_CHECK(c.stride(0) % 16 == 0); // 16 Byte Alignment
|
||||
TORCH_CHECK(bt_nzs.stride(0) % 16 == 0); // 16 Byte Alignment
|
||||
TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
|
||||
|
||||
if (bias) {
|
||||
TORCH_CHECK(bias->numel() == bt_nzs.size(0) && bias->is_contiguous() &&
|
||||
bias->dim() == 1);
|
||||
}
|
||||
|
||||
at::cuda::OptionalCUDAGuard const device_guard(device_of(a));
|
||||
int32_t version_num = get_sm_version_num();
|
||||
|
||||
// Guard against compilation issues for sm90 kernels
|
||||
#if defined ENABLE_SPARSE_SCALED_MM_C3X && ENABLE_SPARSE_SCALED_MM_C3X
|
||||
// We build for 9.0a which is not forward compatible, so restrict this to
|
||||
// Hopper only
|
||||
if (version_num == 90) {
|
||||
cutlass_scaled_sparse_mm_sm90(c, a, bt_nzs, bt_meta, a_scales, b_scales,
|
||||
bias);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false,
|
||||
"No compiled cutlass_scaled_sparse_mm for a compute capability less than "
|
||||
"CUDA device capability: ",
|
||||
version_num);
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a) {
|
||||
// Check for strides and alignment
|
||||
TORCH_CHECK(a.stride(1) == 1); // Row-major
|
||||
TORCH_CHECK(a.stride(0) % 8 == 0); // 8 Byte Alignment for Compression
|
||||
|
||||
at::cuda::OptionalCUDAGuard const device_guard(device_of(a));
|
||||
int32_t version_num = get_sm_version_num();
|
||||
|
||||
// Guard against compilation issues for sm90 kernels
|
||||
#if defined ENABLE_SPARSE_SCALED_MM_C3X && ENABLE_SPARSE_SCALED_MM_C3X
|
||||
// We build for 9.0a which is not forward compatible, so restrict this to
|
||||
// Hopper only
|
||||
if (version_num == 90) {
|
||||
std::vector<torch::Tensor> result_tensors;
|
||||
|
||||
auto [a_meta, a_nzs] = cutlass_sparse_compress_sm90(a);
|
||||
result_tensors.push_back(std::move(a_nzs));
|
||||
result_tensors.push_back(std::move(a_meta));
|
||||
return result_tensors;
|
||||
}
|
||||
#endif
|
||||
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false,
|
||||
"No compiled cutlass_sparse_compress for a compute capability equal to "
|
||||
"CUDA device capability: ",
|
||||
version_num);
|
||||
}
|
||||
+104
-20
@@ -6,6 +6,41 @@
|
||||
#include <torch/library.h>
|
||||
#include <torch/version.h>
|
||||
|
||||
// Forward declarations for per-SM NVFP4 GEMM and quantization entry points.
|
||||
// Defined in nvfp4_scaled_mm_kernels.cu / nvfp4_quant_entry.cu and only
|
||||
// compiled when ENABLE_NVFP4_SM100 is set.
|
||||
#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
|
||||
void cutlass_scaled_fp4_mm_sm100a(torch::Tensor& D, torch::Tensor const& A,
|
||||
torch::Tensor const& B,
|
||||
torch::Tensor const& A_sf,
|
||||
torch::Tensor const& B_sf,
|
||||
torch::Tensor const& alpha);
|
||||
void cutlass_scaled_fp4_mm_sm103a(torch::Tensor& D, torch::Tensor const& A,
|
||||
torch::Tensor const& B,
|
||||
torch::Tensor const& A_sf,
|
||||
torch::Tensor const& B_sf,
|
||||
torch::Tensor const& alpha);
|
||||
std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_sm103a_func(
|
||||
torch::Tensor const& input, torch::Tensor const& input_sf);
|
||||
void scaled_fp4_quant_sm103a_out(torch::Tensor const& input,
|
||||
torch::Tensor const& input_sf,
|
||||
torch::Tensor& output,
|
||||
torch::Tensor& output_sf);
|
||||
// PDL-enabled variants (ProgrammaticStreamSerialization).
|
||||
void cutlass_scaled_fp4_mm_sm103a_pdl(torch::Tensor& D,
|
||||
torch::Tensor const& A,
|
||||
torch::Tensor const& B,
|
||||
torch::Tensor const& A_sf,
|
||||
torch::Tensor const& B_sf,
|
||||
torch::Tensor const& alpha);
|
||||
std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_sm103a_pdl_func(
|
||||
torch::Tensor const& input, torch::Tensor const& input_sf);
|
||||
void scaled_fp4_quant_sm103a_pdl_out(torch::Tensor const& input,
|
||||
torch::Tensor const& input_sf,
|
||||
torch::Tensor& output,
|
||||
torch::Tensor& output_sf);
|
||||
#endif
|
||||
|
||||
// Note on op signatures:
|
||||
// The X_meta signatures are for the meta functions corresponding to op X.
|
||||
// They must be kept in sync with the signature for X. Generally, only
|
||||
@@ -416,6 +451,65 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
" Tensor alpha) -> ()");
|
||||
ops.impl("cutlass_scaled_fp4_mm", torch::kCUDA, &cutlass_scaled_fp4_mm);
|
||||
|
||||
#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100
|
||||
// SM100-specific entry point (B200 / Blackwell, SM100 SF layout)
|
||||
ops.def(
|
||||
"cutlass_scaled_fp4_mm_sm100a(Tensor! out, Tensor a, Tensor b,"
|
||||
" Tensor block_scale_a, Tensor block_scale_b,"
|
||||
" Tensor alpha) -> ()");
|
||||
ops.impl("cutlass_scaled_fp4_mm_sm100a", torch::kCUDA,
|
||||
&cutlass_scaled_fp4_mm_sm100a);
|
||||
|
||||
// SM103-specific entry point (B300 / Blackwell Ultra, SM103 SF layout)
|
||||
ops.def(
|
||||
"cutlass_scaled_fp4_mm_sm103a(Tensor! out, Tensor a, Tensor b,"
|
||||
" Tensor block_scale_a, Tensor block_scale_b,"
|
||||
" Tensor alpha) -> ()");
|
||||
ops.impl("cutlass_scaled_fp4_mm_sm103a", torch::kCUDA,
|
||||
&cutlass_scaled_fp4_mm_sm103a);
|
||||
|
||||
// SM103-native quantization: produces SM103-layout scale factors directly.
|
||||
ops.def(
|
||||
"scaled_fp4_quant_sm103(Tensor input,"
|
||||
" Tensor input_scale) -> (Tensor, Tensor)");
|
||||
ops.impl("scaled_fp4_quant_sm103", torch::kCUDA, &scaled_fp4_quant_sm103a_func);
|
||||
|
||||
ops.def(
|
||||
"scaled_fp4_quant_sm103.out(Tensor input,"
|
||||
" Tensor input_scale,"
|
||||
" *, Tensor(a!) output, Tensor(b!) output_scale)"
|
||||
" -> ()");
|
||||
ops.impl("scaled_fp4_quant_sm103.out", torch::kCUDA,
|
||||
&scaled_fp4_quant_sm103a_out);
|
||||
|
||||
// PDL-enabled SM103 GEMM: launched with ProgrammaticStreamSerialization
|
||||
// so the next kernel on the stream can overlap with this GEMM's tail.
|
||||
ops.def(
|
||||
"cutlass_scaled_fp4_mm_sm103a_pdl(Tensor! out, Tensor a, Tensor b,"
|
||||
" Tensor block_scale_a,"
|
||||
" Tensor block_scale_b,"
|
||||
" Tensor alpha) -> ()");
|
||||
ops.impl("cutlass_scaled_fp4_mm_sm103a_pdl", torch::kCUDA,
|
||||
&cutlass_scaled_fp4_mm_sm103a_pdl);
|
||||
|
||||
// PDL-enabled SM103 quantization: launched with
|
||||
// ProgrammaticStreamSerialization so the subsequent GEMM can begin
|
||||
// before this quant kernel completes.
|
||||
ops.def(
|
||||
"scaled_fp4_quant_sm103_pdl(Tensor input,"
|
||||
" Tensor input_scale) -> (Tensor, Tensor)");
|
||||
ops.impl("scaled_fp4_quant_sm103_pdl", torch::kCUDA,
|
||||
&scaled_fp4_quant_sm103a_pdl_func);
|
||||
|
||||
ops.def(
|
||||
"scaled_fp4_quant_sm103_pdl.out(Tensor input,"
|
||||
" Tensor input_scale,"
|
||||
" *, Tensor(a!) output,"
|
||||
" Tensor(b!) output_scale) -> ()");
|
||||
ops.impl("scaled_fp4_quant_sm103_pdl.out", torch::kCUDA,
|
||||
&scaled_fp4_quant_sm103a_pdl_out);
|
||||
#endif
|
||||
|
||||
// cutlass nvfp4 block scaled group GEMM
|
||||
ops.def(
|
||||
"cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b,"
|
||||
@@ -523,26 +617,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.impl("cutlass_scaled_mm_supports_block_fp8",
|
||||
&cutlass_scaled_mm_supports_block_fp8);
|
||||
|
||||
// Check if cutlass sparse scaled_mm is supported for CUDA devices of the
|
||||
// given capability
|
||||
ops.def(
|
||||
"cutlass_sparse_scaled_mm_supported(int cuda_device_capability) -> bool");
|
||||
ops.impl("cutlass_sparse_scaled_mm_supported",
|
||||
&cutlass_sparse_scaled_mm_supported);
|
||||
|
||||
// CUTLASS sparse GEMM, supporting symmetric per-tensor or per-row/column
|
||||
// quantization, as well as bias
|
||||
ops.def(
|
||||
"cutlass_scaled_sparse_mm(Tensor! out, Tensor a,"
|
||||
" Tensor bt_nzs,"
|
||||
" Tensor bt_meta, Tensor a_scales,"
|
||||
" Tensor b_scales, Tensor? bias) -> ()");
|
||||
ops.impl("cutlass_scaled_sparse_mm", torch::kCUDA, &cutlass_scaled_sparse_mm);
|
||||
|
||||
// CUTLASS sparse matrix compressor
|
||||
ops.def("cutlass_sparse_compress(Tensor a) -> Tensor[]");
|
||||
ops.impl("cutlass_sparse_compress", &cutlass_sparse_compress);
|
||||
|
||||
// SM100 CUTLASS MLA decode
|
||||
ops.def(
|
||||
"sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope,"
|
||||
@@ -593,6 +667,16 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.impl("silu_and_mul_scaled_fp4_experts_quant", torch::kCUDA,
|
||||
&silu_and_mul_scaled_fp4_experts_quant);
|
||||
|
||||
// SM100 <-> SM103 scale factor layout conversion (B300 / Blackwell Ultra)
|
||||
ops.def(
|
||||
"convert_sf_layout_sm100_to_sm103(Tensor(a!) dst, Tensor src) -> ()");
|
||||
ops.impl("convert_sf_layout_sm100_to_sm103", torch::kCUDA,
|
||||
&convert_sf_layout_sm100_to_sm103);
|
||||
ops.def(
|
||||
"convert_sf_layout_sm103_to_sm100(Tensor(a!) dst, Tensor src) -> ()");
|
||||
ops.impl("convert_sf_layout_sm103_to_sm100", torch::kCUDA,
|
||||
&convert_sf_layout_sm103_to_sm100);
|
||||
|
||||
// Check if cutlass_scaled_mm_fp4 is supported for CUDA devices
|
||||
// of the given capability
|
||||
ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool");
|
||||
|
||||
@@ -244,6 +244,7 @@ statistics relating to that iteration:
|
||||
prefill in this iteration. However, we calculate this interval
|
||||
relative to when the request was first received by the frontend
|
||||
(`arrival_time`) in order to account for input processing time.
|
||||
Currently `arrival_time` starts when tokenization begins.
|
||||
|
||||
For any requests that were completed in a given iteration, we also
|
||||
record:
|
||||
@@ -587,7 +588,7 @@ see:
|
||||
- [Benchmarking LLM Workloads for Performance Evaluation and Autoscaling in Kubernetes](https://docs.google.com/document/d/1k4Q4X14hW4vftElIuYGDu5KDe2LtV1XammoG-Xi3bbQ)
|
||||
- [Inference Perf](https://github.com/kubernetes-sigs/wg-serving/tree/main/proposals/013-inference-perf)
|
||||
- <https://github.com/vllm-project/vllm/issues/5041> and <https://github.com/vllm-project/vllm/pull/12726>.
|
||||
|
||||
|
||||
This is a non-trivial topic. Consider this comment from Rob:
|
||||
|
||||
> I think this metric should focus on trying to estimate what the max
|
||||
|
||||
@@ -33,10 +33,10 @@ th {
|
||||
| Backend | Output act. format | Quant. types | Quant. format | Async | Apply Weight On Input | Subclass |
|
||||
| ------- | ------------------ | ------------ | ------------- | ----- | --------------------- | --------- |
|
||||
| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE] |
|
||||
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |
|
||||
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_nvlink_two_sided_prepare_finalize.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_nvlink_one_sided_prepare_finalize.FlashInferNVLinkOneSidedPrepareAndFinalize] |
|
||||
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
|
||||
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
|
||||
|
||||
!!! info "Table key"
|
||||
1. All types: mxfp4, nvfp4, int4, int8, fp8
|
||||
|
||||
@@ -137,6 +137,3 @@ llm = LLM("facebook/opt-125m", quantization="fp8")
|
||||
result = llm.generate("Hello, my name is")
|
||||
print(result[0].outputs[0].text)
|
||||
```
|
||||
|
||||
!!! warning
|
||||
Currently, we load the model at original precision before quantizing down to 8-bits, so you need enough memory to load the whole model.
|
||||
|
||||
@@ -535,7 +535,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
|
||||
| ------------ | ------ | ------ | ----------------- | -------------------- | ------------------------- |
|
||||
| `AriaForConditionalGeneration` | Aria | T + I<sup>+</sup> | `rhymes-ai/Aria` | | |
|
||||
| `AudioFlamingo3ForConditionalGeneration` | AudioFlamingo3 | T + A | `nvidia/audio-flamingo-3-hf`, `nvidia/music-flamingo-2601-hf` | ✅︎ | ✅︎ |
|
||||
| `AudioFlamingo3ForConditionalGeneration` | AudioFlamingo3 | T + A | `nvidia/audio-flamingo-3-hf`, `nvidia/music-flamingo-hf` | ✅︎ | ✅︎ |
|
||||
| `AyaVisionForConditionalGeneration` | Aya Vision | T + I<sup>+</sup> | `CohereLabs/aya-vision-8b`, `CohereLabs/aya-vision-32b`, etc. | | ✅︎ |
|
||||
| `BagelForConditionalGeneration` | BAGEL | T + I<sup>+</sup> | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ |
|
||||
| `BeeForConditionalGeneration` | Bee-8B | T + I<sup>E+</sup> | `Open-Bee/Bee-8B-RL`, `Open-Bee/Bee-8B-SFT` | | ✅︎ |
|
||||
@@ -586,6 +586,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
|
||||
| `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
|
||||
| `Molmo2ForConditionalGeneration` | Molmo2 | T + I<sup>+</sup> / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B` | ✅︎ | ✅︎ |
|
||||
| `MusicFlamingoForConditionalGeneration` | MusicFlamingo | T + A | `nvidia/music-flamingo-2601-hf`, `nvidia/music-flamingo-think-2601-hf` | ✅︎ | ✅︎ |
|
||||
| `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
|
||||
| `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + I<sup>E+</sup> | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ |
|
||||
| `OpenPanguVLForConditionalGeneration` | openpangu-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `FreedomIntelligence/openPangu-VL-7B` | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -104,12 +104,22 @@ def run_musicflamingo(question: str, audio_count: int) -> ModelRequestData:
|
||||
enforce_eager=True,
|
||||
)
|
||||
|
||||
# MusicFlamingo uses <sound> token for audio
|
||||
# MusicFlamingo prompt placeholders use <sound>; vLLM's MusicFlamingo
|
||||
# multimodal processor expands each one into <|sound_bos|> + audio tokens +
|
||||
# <|sound_eos|> based on extracted audio feature lengths.
|
||||
audio_placeholder = "<sound>" * audio_count
|
||||
system_prompt = (
|
||||
"You are Music Flamingo, a multimodal assistant for language and music. "
|
||||
"On each turn you receive an audio clip which contains music and optional "
|
||||
"text, you will receive at least one or both; use your world knowledge and "
|
||||
"reasoning to help the user with any task. Interpret the entirety of the "
|
||||
"content any input music--regardlenss of whether the user calls it audio, "
|
||||
"music, or sound."
|
||||
)
|
||||
|
||||
prompt = (
|
||||
"<|im_start|>system\n"
|
||||
"You are a helpful assistant.<|im_end|>\n"
|
||||
f"{system_prompt}<|im_end|>\n"
|
||||
"<|im_start|>user\n"
|
||||
f"{audio_placeholder}{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
|
||||
@@ -47,6 +47,8 @@ pystemmer==3.0.0
|
||||
# Multi-modal processing
|
||||
av==16.1.0
|
||||
# required for audio_in_video tests
|
||||
resampy==0.4.3
|
||||
# audio processing, required for audio_in_video tests
|
||||
blobfile==3.0.0
|
||||
# Multi-Modal Models Test
|
||||
decord==0.6.0
|
||||
|
||||
+208
@@ -0,0 +1,208 @@
|
||||
# SM103 NVFP4 Programmatic Dependent Launch (PDL) Summary
|
||||
|
||||
## Overview
|
||||
|
||||
This document summarizes the addition of **Programmatic Dependent Launch (PDL)** to the SM103 (B300 Blackwell Ultra) NVFP4 quantization and CUTLASS GEMM kernels in vLLM. PDL is a CUDA 12+ feature that allows consecutive kernels on the same stream to overlap execution, reducing the gap between a producer kernel's tail and a consumer kernel's head.
|
||||
|
||||
## What is PDL?
|
||||
|
||||
In standard CUDA stream semantics, Kernel B cannot begin until Kernel A fully completes. PDL relaxes this constraint:
|
||||
|
||||
```
|
||||
Without PDL:
|
||||
[==== quant kernel ====] [==== GEMM kernel ====]
|
||||
^ idle gap
|
||||
|
||||
With PDL:
|
||||
[==== quant kernel ====]
|
||||
[==== GEMM kernel ====]
|
||||
^ overlap region
|
||||
```
|
||||
|
||||
The CUDA attribute `cudaLaunchAttributeProgrammaticStreamSerialization` is set on the **producer** kernel, telling the driver that the next kernel on the stream may begin before the producer fully completes. This is safe when the consumer's early thread blocks operate on data that the producer has already finished writing (which is the typical case for tile-based execution).
|
||||
|
||||
## Changes Made
|
||||
|
||||
### 1. Quant kernel PDL (`csrc/quantization/fp4/nvfp4_quant_kernels.cu`)
|
||||
|
||||
- Refactored `scaled_fp4_quant_sm103a` into `scaled_fp4_quant_sm103a_impl` with a `use_pdl` parameter
|
||||
- When `use_pdl=true`, the kernel is launched via `cudaLaunchKernelEx` with `ProgrammaticStreamSerialization = 1`
|
||||
- Added `scaled_fp4_quant_sm103a_pdl()` entry point
|
||||
|
||||
### 2. CUTLASS GEMM PDL (`csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu`)
|
||||
|
||||
- Added `launch_with_pdl` parameter to `runGemm()` template, forwarded to CUTLASS's `GemmUniversalAdapter::run(..., launch_with_pdl)`
|
||||
- CUTLASS internally sets `ProgrammaticStreamSerialization` on the GEMM launch via `ClusterLauncher`
|
||||
- Added `cutlass_scaled_fp4_mm_sm103a_pdl()` entry point
|
||||
|
||||
### 3. Op registration (`csrc/torch_bindings.cpp`, `csrc/quantization/fp4/nvfp4_*_entry.cu`)
|
||||
|
||||
New torch ops registered:
|
||||
- `torch.ops._C.scaled_fp4_quant_sm103_pdl` -- PDL-enabled SM103 quant
|
||||
- `torch.ops._C.scaled_fp4_quant_sm103_pdl.out` -- out-variant
|
||||
- `torch.ops._C.cutlass_scaled_fp4_mm_sm103a_pdl` -- PDL-enabled SM103 GEMM
|
||||
|
||||
### 4. Benchmark (`benchmarks/kernels/benchmark_nvfp4_sm103.py`)
|
||||
|
||||
Updated benchmark with new modes:
|
||||
- `--mode gemm`: SM100 vs SM103 vs SM103+PDL GEMM-only
|
||||
- `--mode e2e`: End-to-end quant+GEMM with PDL comparison
|
||||
- `--mode pdl`: Multi-layer pipeline benchmark (back-to-back quant+GEMM pairs)
|
||||
|
||||
## PDL Pipeline Analysis
|
||||
|
||||
### Single kernel pair (quant + GEMM)
|
||||
|
||||
For a single quant->GEMM pair, PDL enables:
|
||||
1. **Quant kernel** with `ProgrammaticStreamSerialization`: GEMM can start before quant finishes
|
||||
2. **GEMM kernel** with `ProgrammaticStreamSerialization` (via CUTLASS): the next layer's kernel can start before GEMM finishes
|
||||
|
||||
### Multi-layer pipeline
|
||||
|
||||
In a real transformer, the pattern repeats:
|
||||
```
|
||||
Layer 1: quant_1 -> GEMM_1
|
||||
Layer 2: quant_2 -> GEMM_2
|
||||
...
|
||||
```
|
||||
|
||||
With PDL on both kernels, each transition overlaps:
|
||||
```
|
||||
[quant_1]--->[GEMM_1]--->[quant_2]--->[GEMM_2]---> (without PDL)
|
||||
|
||||
[quant_1]--[GEMM_1]--[quant_2]--[GEMM_2]-- (with PDL)
|
||||
^^ ^^ ^^
|
||||
overlap at each transition
|
||||
```
|
||||
|
||||
### Measured performance impact (B300 SXM6, CUDA 12.9)
|
||||
|
||||
| Scenario | PDL benefit | Explanation |
|
||||
|----------|-------------|-------------|
|
||||
| GEMM-only (isolated) | ~0-3% | Marginal; no meaningful consumer overlap for a single kernel |
|
||||
| Single quant+GEMM (decode, M=1-16) | 0-3% | Quant is tiny, GEMM dominates wall time |
|
||||
| Single quant+GEMM (prefill, M=64-512) | 2-3% | Moderate overlap window |
|
||||
| Single quant+GEMM (prefill, M=1024+) | 2-5% | Larger quant = more overlap-able tail |
|
||||
| 4-layer pipeline (decode, M=1-16) | 4-6% | Cumulative overlap across 8 kernel transitions |
|
||||
| 4-layer pipeline (M=1024) | **12%** | Best case: sustained overlap on compute-heavy layers |
|
||||
| 4-layer pipeline (prefill, M=4096) | 4% | GEMM dominates; quant tail is proportionally smaller |
|
||||
|
||||
The PDL benefit is proportional to the **ratio of overlap-able tail time to total kernel time**. The sweet spot is M=256-1024 where the quant kernel is large enough to provide meaningful overlap but doesn't yet dominate the pipeline.
|
||||
|
||||
## SM100 vs SM103 vs SM103+PDL Comparison
|
||||
|
||||
### Kernel architecture differences
|
||||
|
||||
| Aspect | SM100 (B200) | SM103 (B300) | SM103+PDL |
|
||||
|--------|-------------|-------------|-----------|
|
||||
| MMA instructions | FP4 BlockScaled | FP4 Ultra (UltraVs16) | Same as SM103 |
|
||||
| Tile K size | 256 | 768 (3x larger) | Same as SM103 |
|
||||
| Cooperative SMs | 1-2 per tile | 2 per tile (default) | Same as SM103 |
|
||||
| SF layout | Sm1xxBlockScaledConfig | Sm103BlockScaledConfig | Same as SM103 |
|
||||
| Kernel overlap | None (stream-serialized) | None | Quant tail overlaps GEMM head |
|
||||
| Epilogue | TmaWarpSpecialized | NoSmemWarpSpecialized | Same as SM103 |
|
||||
|
||||
### Measured performance on B300 SXM6, N=K=7168
|
||||
|
||||
**SM103 vs SM100 (GEMM-only):** SM103 Ultra MMA provides 3-7% higher throughput for small-to-medium M (decode/small batch). At large M (2048+), both achieve similar throughput as the problem becomes compute-bound on both paths. The SM103 advantage comes from:
|
||||
- 3x larger K-tile (768 vs 256): fewer mainloop iterations
|
||||
- UltraVs16 instructions: higher throughput per clock
|
||||
- NoSmem epilogue: more shared memory for mainloop double-buffering
|
||||
|
||||
**PDL benefit (E2E):** PDL provides a consistent 2-5% speedup on the end-to-end quant+GEMM path for most M sizes, with a peak of 5% at M=1024 where the quant and GEMM are well-balanced.
|
||||
|
||||
**PDL pipeline benefit (4 layers):** The multi-layer pipeline shows 4-12% speedup, with the best result at M=1024 (12% speedup) where cumulative overlap across 8 kernel transitions provides maximum benefit.
|
||||
|
||||
## Files Modified
|
||||
|
||||
| File | Change |
|
||||
|------|--------|
|
||||
| `csrc/quantization/fp4/nvfp4_quant_kernels.cu` | PDL launch via `cudaLaunchKernelEx` for SM103 quant |
|
||||
| `csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu` | `launch_with_pdl` parameter forwarded to CUTLASS |
|
||||
| `csrc/quantization/fp4/nvfp4_quant_entry.cu` | PDL entry points and forward declarations |
|
||||
| `csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu` | PDL GEMM forward declaration |
|
||||
| `csrc/torch_bindings.cpp` | Op registration for `_pdl` variants |
|
||||
| `benchmarks/kernels/benchmark_nvfp4_sm103.py` | PDL benchmark modes (gemm, e2e, pdl pipeline) |
|
||||
|
||||
## How to Run
|
||||
|
||||
```bash
|
||||
# Build vLLM with SM103 support
|
||||
python setup.py build_ext --inplace
|
||||
|
||||
# Run all benchmarks
|
||||
python benchmarks/kernels/benchmark_nvfp4_sm103.py --mode all
|
||||
|
||||
# Run only the PDL pipeline benchmark with 8 layers
|
||||
python benchmarks/kernels/benchmark_nvfp4_sm103.py --mode pdl --layers 8
|
||||
|
||||
# Run end-to-end comparison
|
||||
python benchmarks/kernels/benchmark_nvfp4_sm103.py --mode e2e
|
||||
```
|
||||
|
||||
## Benchmark Results
|
||||
|
||||
**Hardware:** NVIDIA B300 SXM6 AC (SM103), CUDA 12.9
|
||||
**Problem:** N=7168, K=7168 (DeepSeek-style dimensions), BF16 output
|
||||
|
||||
### GEMM-Only: SM100 vs SM103 vs SM103+PDL
|
||||
|
||||
| M | SM100 (us) | SM100 TFLOPS | SM103 (us) | SM103 TFLOPS | SM103+PDL (us) | SM103+PDL TFLOPS | SM103 vs SM100 |
|
||||
|---|-----------|-------------|-----------|-------------|---------------|-----------------|----------------|
|
||||
| 1 | 26.43 | 3.89 | 25.09 | 4.10 | 24.35 | 4.22 | 1.05x |
|
||||
| 16 | 27.30 | 60.23 | 24.83 | 66.21 | 25.31 | 64.96 | 1.10x |
|
||||
| 128 | 28.26 | 465.51 | 26.43 | 497.63 | 26.34 | 499.44 | 1.07x |
|
||||
| 512 | 28.19 | 1866.25 | 27.36 | 1923.00 | 26.53 | 1983.31 | 1.03x |
|
||||
| 1024 | 30.72 | 3425.35 | 34.72 | 3030.72 | 34.78 | 3025.15 | 0.88x |
|
||||
| 4096 | 93.12 | 4520.05 | 90.02 | 4675.91 | 89.89 | 4682.57 | 1.03x |
|
||||
|
||||
**Observations:**
|
||||
- SM103 is faster than SM100 for M <= 512 (up to 10% at M=16)
|
||||
- At M=1024, SM100 is faster (tile configuration tradeoff)
|
||||
- PDL on GEMM-only has marginal effect (expected: no consumer kernel to overlap)
|
||||
|
||||
### End-to-End: Quant + GEMM
|
||||
|
||||
| M | SM100 (us) | SM103 (us) | SM103+PDL (us) | SM103 vs SM100 | PDL vs no-PDL | PDL vs SM100 |
|
||||
|---|-----------|-----------|---------------|----------------|---------------|--------------|
|
||||
| 1 | 36.54 | 35.74 | 35.90 | 1.02x | 1.00x | 1.02x |
|
||||
| 8 | 36.58 | 34.53 | 33.73 | 1.06x | **1.02x** | **1.08x** |
|
||||
| 64 | 38.30 | 36.64 | 35.84 | 1.05x | **1.02x** | **1.07x** |
|
||||
| 256 | 38.24 | 36.74 | 36.03 | 1.04x | **1.02x** | **1.06x** |
|
||||
| 1024 | 38.66 | 40.90 | 38.78 | 0.95x | **1.05x** | 1.00x |
|
||||
| 2048 | 65.63 | 65.41 | 63.04 | 1.00x | **1.04x** | **1.04x** |
|
||||
| 4096 | 112.38 | 110.78 | 108.38 | 1.01x | **1.02x** | **1.04x** |
|
||||
|
||||
**Observations:**
|
||||
- PDL consistently improves E2E by 2-5% over non-PDL SM103
|
||||
- Best PDL improvement at M=1024: 5% (40.90 us -> 38.78 us)
|
||||
- Total SM103+PDL vs SM100 speedup: up to 8% at M=8
|
||||
|
||||
### PDL Pipeline: 4 Back-to-Back Layers
|
||||
|
||||
| M | No PDL (us) | No PDL TFLOPS | PDL (us) | PDL TFLOPS | PDL Speedup |
|
||||
|---|------------|--------------|---------|-----------|-------------|
|
||||
| 1 | 111.65 | 3.68 | 107.10 | 3.84 | 1.04x |
|
||||
| 16 | 112.42 | 58.50 | 107.07 | 61.42 | **1.05x** |
|
||||
| 64 | 120.51 | 218.29 | 115.65 | 227.47 | 1.04x |
|
||||
| 256 | 120.83 | 870.85 | 115.33 | 912.41 | **1.05x** |
|
||||
| 1024 | 151.52 | 2777.90 | 134.94 | 3119.12 | **1.12x** |
|
||||
| 2048 | 250.27 | 3363.59 | 236.35 | 3561.69 | **1.06x** |
|
||||
| 4096 | 432.38 | 3893.82 | 416.32 | 4044.07 | 1.04x |
|
||||
|
||||
**Key finding:** At M=1024, PDL provides **12% speedup** over non-PDL SM103 in the 4-layer pipeline benchmark. This is where the quant and GEMM kernels are well-balanced in execution time, maximizing the overlap benefit across 8 kernel transitions (4 quant + 4 GEMM).
|
||||
|
||||
### Activation Quantization: SM100 vs SM103
|
||||
|
||||
| M | SM100 (us) | SM100 GB/s | SM103 (us) | SM103 GB/s |
|
||||
|---|-----------|-----------|-----------|-----------|
|
||||
| 1 | 17.06 | 0.84 | 17.09 | 0.84 |
|
||||
| 256 | 17.25 | 212.78 | 17.18 | 213.57 |
|
||||
| 1024 | 17.25 | 851.12 | 16.90 | 868.85 |
|
||||
| 4096 | 19.17 | 3063.45 | 18.24 | 3219.31 |
|
||||
|
||||
**Observations:** SM103 quant is 2-5% faster than SM100 quant, primarily from the different SF swizzle pattern being more cache-friendly on SM103.
|
||||
|
||||
## Conclusion
|
||||
|
||||
PDL is a low-cost optimization that provides consistent 2-5% E2E improvement for single quant+GEMM pairs, scaling to **12% in multi-layer pipelines** at the M=1024 sweet spot. The implementation adds no correctness risk (PDL is a scheduling hint) and no overhead when the GPU decides not to overlap. It should be enabled by default for SM103 production workloads.
|
||||
@@ -127,7 +127,7 @@ class AttentionQuantPatternModel(torch.nn.Module):
|
||||
raw_tensor = raw_tensor.view(kv_cache_shape)
|
||||
kv_cache = raw_tensor.permute(*inv_order)
|
||||
|
||||
self.attn.kv_cache = [kv_cache]
|
||||
self.attn.kv_cache = kv_cache
|
||||
|
||||
# Build attn metadata
|
||||
self.attn_metadata = self.builder.build(
|
||||
|
||||
@@ -148,7 +148,7 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
|
||||
raw_tensor = raw_tensor.view(kv_cache_shape)
|
||||
kv_cache = raw_tensor.permute(*inv_order)
|
||||
|
||||
self.attn.kv_cache = [kv_cache]
|
||||
self.attn.kv_cache = kv_cache
|
||||
|
||||
# Build attn metadata
|
||||
attn_metadata = self.builder.build(
|
||||
@@ -295,7 +295,7 @@ def test_rope_kvcache_fusion(
|
||||
}
|
||||
q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused)
|
||||
attn_layer = forward_context.no_compile_layers[model.layer_name]
|
||||
kv_cache_unfused = attn_layer.kv_cache[0]
|
||||
kv_cache_unfused = attn_layer.kv_cache
|
||||
del dummy
|
||||
|
||||
torch._dynamo.mark_dynamic(qkv, 0)
|
||||
@@ -309,7 +309,7 @@ def test_rope_kvcache_fusion(
|
||||
}
|
||||
q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos)
|
||||
attn_layer = forward_context.no_compile_layers[model.layer_name]
|
||||
kv_cache_fused = attn_layer.kv_cache[0]
|
||||
kv_cache_fused = attn_layer.kv_cache
|
||||
del dummy
|
||||
|
||||
assert fusion_pass.matched_count == 1
|
||||
|
||||
@@ -15,6 +15,7 @@ from torch._dynamo.utils import counters
|
||||
import vllm.envs as envs
|
||||
from vllm.compilation.counter import compilation_counter
|
||||
from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode
|
||||
from vllm.utils.torch_utils import is_torch_equal_or_newer
|
||||
|
||||
from ..utils import fork_new_process_for_each_test
|
||||
|
||||
@@ -71,7 +72,10 @@ def test_moe_startup(monkeypatch, vllm_runner, fresh_vllm_cache, mega_aot_artifa
|
||||
num_compiled_artifacts_saved=0,
|
||||
):
|
||||
_run_vllm(vllm_runner)
|
||||
if envs.VLLM_USE_MEGA_AOT_ARTIFACT:
|
||||
mega_aot_active = envs.VLLM_USE_MEGA_AOT_ARTIFACT and is_torch_equal_or_newer(
|
||||
"2.10.0"
|
||||
)
|
||||
if mega_aot_active:
|
||||
# MEGA_AOT_ARTIFACT is enabled, so we expect no aot_autograd running on
|
||||
# subgraphs.
|
||||
assert counters["aot_autograd"]["total"] == 0
|
||||
|
||||
@@ -182,7 +182,7 @@ async def test_streaming_response(foscolo, client_and_model, server):
|
||||
# being very close semantically.
|
||||
assert (
|
||||
sum([x == y for x, y in zip(res_stream, res_no_stream.text.split())])
|
||||
>= len(res_stream) * 0.9
|
||||
>= len(res_stream) * 0.87
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -1,12 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def default_tokenizer() -> TokenizerLike:
|
||||
return AutoTokenizer.from_pretrained("gpt2")
|
||||
@@ -1,18 +1,9 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import json
|
||||
import random
|
||||
from typing import Any
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
from vllm.entrypoints.openai.engine.protocol import (
|
||||
DeltaMessage,
|
||||
)
|
||||
from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser
|
||||
|
||||
from ....utils import RemoteOpenAIServer
|
||||
|
||||
@@ -38,137 +29,6 @@ def server():
|
||||
yield server
|
||||
|
||||
|
||||
def create_complex_input(create_string_args: bool):
|
||||
coord_arg: dict | str = {
|
||||
"coordinates": [[23.54, 43.1], [-12.2, 54.3], [4, 5]],
|
||||
"coordinate_type": "latlong",
|
||||
}
|
||||
if create_string_args:
|
||||
# test granite behavior
|
||||
coord_arg = json.dumps(coord_arg)
|
||||
return [
|
||||
{"name": "find_bbox", "arguments": coord_arg},
|
||||
{
|
||||
"name": "get_stock_price",
|
||||
"arguments": {
|
||||
"symbol": "AAPL",
|
||||
"start_date": "2021-01-01",
|
||||
"end_date": "2021-12-31",
|
||||
},
|
||||
},
|
||||
{"name": "find_bbox", "arguments": coord_arg},
|
||||
]
|
||||
|
||||
|
||||
def random_chunks(s: str, min_len: int, max_len: int):
|
||||
chunks = []
|
||||
i = 0
|
||||
n = len(s)
|
||||
|
||||
while i < n:
|
||||
size = random.randint(min_len, max_len)
|
||||
chunks.append(s[i : i + size])
|
||||
i += size
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def tokenizer():
|
||||
return AutoTokenizer.from_pretrained(MODEL)
|
||||
|
||||
|
||||
# create a variety of input chunk sizes
|
||||
@pytest.mark.parametrize(
|
||||
"min_chunk, max_chunk",
|
||||
[
|
||||
(1, 1),
|
||||
(1, 2),
|
||||
(5, 7),
|
||||
(6, 20),
|
||||
],
|
||||
)
|
||||
def test_tool_call_parser_complex(min_chunk: int, max_chunk: int, tokenizer):
|
||||
input_dicts = create_complex_input(True)
|
||||
|
||||
formatted_tcs = [
|
||||
"<tool_call> " + json.dumps(call) + " </tool_call>" for call in input_dicts
|
||||
]
|
||||
|
||||
text_messages = [
|
||||
"Here goes the bbox call: \n",
|
||||
" Now the stock price call: \n ",
|
||||
" Now another bbox call: \n ",
|
||||
" See? I'm a helpful assistant.",
|
||||
]
|
||||
|
||||
test_input = (
|
||||
text_messages[0]
|
||||
+ formatted_tcs[0]
|
||||
+ text_messages[1]
|
||||
+ formatted_tcs[1]
|
||||
+ text_messages[2]
|
||||
+ formatted_tcs[2]
|
||||
+ text_messages[3]
|
||||
)
|
||||
|
||||
any_chat_request = ChatCompletionRequest(
|
||||
seed=42,
|
||||
model=MODEL,
|
||||
messages=[],
|
||||
)
|
||||
|
||||
parser = Granite4ToolParser(tokenizer=tokenizer)
|
||||
|
||||
delta_messages = list[DeltaMessage]()
|
||||
for text in random_chunks(test_input, min_chunk, max_chunk):
|
||||
delta = parser.extract_tool_calls_streaming(
|
||||
previous_text="",
|
||||
current_text="",
|
||||
delta_text=text,
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[],
|
||||
delta_token_ids=[],
|
||||
request=any_chat_request,
|
||||
)
|
||||
if delta is not None:
|
||||
delta_messages.append(delta)
|
||||
|
||||
content = ""
|
||||
tool_calls = list[dict[str, Any]]()
|
||||
|
||||
current_name = "__start__"
|
||||
current_args = ""
|
||||
|
||||
for msg in delta_messages:
|
||||
if msg.content:
|
||||
content += msg.content
|
||||
for tool_call in msg.tool_calls:
|
||||
if delta_func := tool_call.function:
|
||||
if delta_func.name is not None:
|
||||
if current_name == "__start__":
|
||||
current_name = delta_func.name
|
||||
|
||||
if delta_func.name != current_name:
|
||||
tool_calls.append(
|
||||
{
|
||||
"name": current_name,
|
||||
"arguments": json.loads(current_args),
|
||||
}
|
||||
)
|
||||
current_name = delta_func.name
|
||||
current_args = ""
|
||||
|
||||
if delta_func.arguments:
|
||||
current_args += delta_func.arguments
|
||||
|
||||
if current_name != "__start__":
|
||||
tool_calls.append({"name": current_name, "arguments": json.loads(current_args)})
|
||||
|
||||
assert content == "".join(text_messages)
|
||||
assert tool_calls == create_complex_input(False)
|
||||
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
|
||||
@@ -9,8 +9,6 @@ import pytest_asyncio
|
||||
from huggingface_hub import snapshot_download
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.tool_parsers.abstract_tool_parser import ToolParser
|
||||
from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser
|
||||
from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
|
||||
@@ -325,202 +323,3 @@ async def test_streaming_product_tool_call(
|
||||
print("\n[Streaming Product Test Passed]")
|
||||
print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
|
||||
print(f"Reconstructed Arguments: {arguments}")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def qwen_tokenizer() -> TokenizerLike:
|
||||
from vllm.tokenizers import get_tokenizer
|
||||
|
||||
return get_tokenizer("Qwen/Qwen3-32B")
|
||||
|
||||
|
||||
@pytest.fixture(params=CONFIGS.keys())
|
||||
def hermes_parser(request, qwen_tokenizer: TokenizerLike) -> ToolParser:
|
||||
config = CONFIGS[request.param]
|
||||
return config["tool_parser"](qwen_tokenizer)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def any_chat_request() -> ChatCompletionRequest:
|
||||
return ChatCompletionRequest(
|
||||
seed=42,
|
||||
model="Qwen/Qwen3-32B",
|
||||
messages=[],
|
||||
)
|
||||
|
||||
|
||||
def test_hermes_parser_streaming_just_forward_text(
|
||||
qwen_tokenizer: TokenizerLike,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """This is some prior text that has nothing to do with tool calling."""
|
||||
tokens = qwen_tokenizer.encode(text)
|
||||
previous_text = ""
|
||||
delta_messages = []
|
||||
for token in tokens:
|
||||
delta_text = qwen_tokenizer.decode([token])
|
||||
current_text = previous_text + delta_text
|
||||
delta = hermes_parser.extract_tool_calls_streaming(
|
||||
previous_text=previous_text,
|
||||
current_text=current_text,
|
||||
delta_text=delta_text,
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[],
|
||||
delta_token_ids=[],
|
||||
request=any_chat_request,
|
||||
)
|
||||
previous_text = current_text
|
||||
delta_messages.append(delta)
|
||||
|
||||
for delta in delta_messages:
|
||||
assert delta is not None
|
||||
assert not delta.tool_calls
|
||||
|
||||
print(delta_messages)
|
||||
assert "".join([delta.content for delta in delta_messages]) == text
|
||||
|
||||
|
||||
def test_hermes_parser_streaming_failure_case_bug_19056(
|
||||
qwen_tokenizer: TokenizerLike,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """<tool_call>
|
||||
{"name": "final_answer", "arguments": {"trigger": true}}
|
||||
</tool_call>"""
|
||||
tokens = qwen_tokenizer.encode(text)
|
||||
previous_text = ""
|
||||
delta_messages = []
|
||||
for token in tokens:
|
||||
text = qwen_tokenizer.decode([token])
|
||||
current_text = previous_text + text
|
||||
delta = hermes_parser.extract_tool_calls_streaming(
|
||||
previous_text=previous_text,
|
||||
current_text=current_text,
|
||||
delta_text=text,
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[],
|
||||
delta_token_ids=[],
|
||||
request=any_chat_request,
|
||||
)
|
||||
previous_text = current_text
|
||||
if delta is not None:
|
||||
delta_messages.append(delta)
|
||||
|
||||
assert delta_messages[0].tool_calls[0].function.name == "final_answer"
|
||||
tool_call_args = "".join(
|
||||
delta.tool_calls[0].function.arguments or "" for delta in delta_messages
|
||||
)
|
||||
assert tool_call_args == '{"trigger": true}'
|
||||
|
||||
|
||||
def test_hermes_parser_streaming(
|
||||
qwen_tokenizer: TokenizerLike,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = '<tool_call>\
|
||||
{"name": "get_current_temperature",\
|
||||
"arguments": {"location":\
|
||||
"San Francisco, California, United States", "unit": "celsius"}}\
|
||||
</tool_call>'
|
||||
|
||||
tokens = qwen_tokenizer.encode(text)
|
||||
previous_text = ""
|
||||
delta_messages = []
|
||||
for token in tokens:
|
||||
text = qwen_tokenizer.decode([token])
|
||||
current_text = previous_text + text
|
||||
delta = hermes_parser.extract_tool_calls_streaming(
|
||||
previous_text=previous_text,
|
||||
current_text=current_text,
|
||||
delta_text=text,
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[],
|
||||
delta_token_ids=[],
|
||||
request=any_chat_request,
|
||||
)
|
||||
previous_text = current_text
|
||||
if delta is not None:
|
||||
delta_messages.append(delta)
|
||||
print(delta_messages)
|
||||
assert delta_messages[0].tool_calls[0].function.name == "get_current_temperature"
|
||||
# load to normalize whitespace
|
||||
tool_call_args = json.loads(
|
||||
"".join(
|
||||
delta.tool_calls[0].function.arguments or "" for delta in delta_messages
|
||||
)
|
||||
)
|
||||
assert tool_call_args == {
|
||||
"location": "San Francisco, California, United States",
|
||||
"unit": "celsius",
|
||||
}
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_no_tool_call(
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """This is not a tool call."""
|
||||
tool_call = hermes_parser.extract_tool_calls(
|
||||
model_output=text,
|
||||
request=any_chat_request,
|
||||
)
|
||||
|
||||
assert tool_call is not None
|
||||
assert not tool_call.tools_called
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_tool_call_between_tags(
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """<tool_call>
|
||||
{"name": "final_answer", "arguments": {"trigger": true}}
|
||||
</tool_call>"""
|
||||
tool_call = hermes_parser.extract_tool_calls(
|
||||
model_output=text,
|
||||
request=any_chat_request,
|
||||
)
|
||||
|
||||
assert tool_call is not None
|
||||
assert tool_call.tools_called
|
||||
assert tool_call.tool_calls[0].function.name == "final_answer"
|
||||
assert tool_call.tool_calls[0].function.arguments == '{"trigger": true}'
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_tool_call_until_eos(
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
if isinstance(hermes_parser, Granite4ToolParser):
|
||||
pytest.skip(reason="The Granite4 tool parser enforces a complete response")
|
||||
|
||||
text = """<tool_call>
|
||||
{"name": "final_answer", "arguments": {"trigger": true}}"""
|
||||
tool_call = hermes_parser.extract_tool_calls(
|
||||
model_output=text,
|
||||
request=any_chat_request,
|
||||
)
|
||||
|
||||
assert tool_call is not None
|
||||
assert tool_call.tools_called
|
||||
assert tool_call.tool_calls[0].function.name == "final_answer"
|
||||
assert tool_call.tool_calls[0].function.arguments == '{"trigger": true}'
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_tool_call_invalid_json(
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
# Missing closing brace to trigger exception
|
||||
text = """<tool_call>
|
||||
{"name": "final_answer", "arguments": {"trigger": true}"""
|
||||
tool_call = hermes_parser.extract_tool_calls(
|
||||
model_output=text,
|
||||
request=any_chat_request,
|
||||
)
|
||||
|
||||
assert tool_call is not None
|
||||
assert not tool_call.tools_called
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
model_name: "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16"
|
||||
accuracy_threshold: 0.93
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
startup_max_wait_seconds: 1200
|
||||
server_args: >-
|
||||
--enforce-eager
|
||||
--max-model-len 4096
|
||||
--tensor-parallel-size 8
|
||||
--enable-expert-parallel
|
||||
--speculative-config '{"method":"mtp","num_speculative_tokens":5}'
|
||||
@@ -0,0 +1,11 @@
|
||||
model_name: "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8"
|
||||
accuracy_threshold: 0.93
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
startup_max_wait_seconds: 1200
|
||||
server_args: >-
|
||||
--enforce-eager
|
||||
--max-model-len 4096
|
||||
--tensor-parallel-size 8
|
||||
--enable-expert-parallel
|
||||
--speculative-config '{"method":"mtp","num_speculative_tokens":5}'
|
||||
@@ -0,0 +1,11 @@
|
||||
model_name: "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4"
|
||||
accuracy_threshold: 0.93
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
startup_max_wait_seconds: 1200
|
||||
server_args: >-
|
||||
--enforce-eager
|
||||
--max-model-len 4096
|
||||
--tensor-parallel-size 2
|
||||
--enable-expert-parallel
|
||||
--speculative-config '{"method":"mtp","num_speculative_tokens":5}'
|
||||
@@ -5,3 +5,4 @@ DeepSeek-V2-Lite-Instruct-FP8.yaml
|
||||
Qwen3-30B-A3B-NVFP4.yaml
|
||||
Qwen3-Next-80B-A3B-NVFP4-EP2.yaml
|
||||
Qwen3-Next-FP8-EP2.yaml
|
||||
Nemotron-3-Super-120B-A12B-NVFP4.yaml
|
||||
|
||||
@@ -2,3 +2,5 @@ DeepSeek-R1-TP.yaml
|
||||
DeepSeek-R1-DP.yaml
|
||||
DeepSeek-V3.2-TP.yaml
|
||||
DeepSeek-V3.2-DP.yaml
|
||||
Nemotron-3-Super-120B-A12B-BF16.yaml
|
||||
Nemotron-3-Super-120B-A12B-FP8.yaml
|
||||
|
||||
@@ -55,37 +55,37 @@ def _clear_supports_cache():
|
||||
# supports_trtllm_attention
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=True)
|
||||
def test_supports_batch_invariant_disables(_mock):
|
||||
@patch("vllm.envs.VLLM_BATCH_INVARIANT", True)
|
||||
def test_supports_batch_invariant_disables():
|
||||
assert supports_trtllm_attention() is False
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=False)
|
||||
@patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
|
||||
@patch(
|
||||
"vllm.utils.flashinfer.current_platform.is_device_capability_family",
|
||||
return_value=True,
|
||||
)
|
||||
@patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=True)
|
||||
def test_supports_sm100_with_artifactory(_art, _cap, _bi):
|
||||
def test_supports_sm100_with_artifactory(_art, _cap):
|
||||
assert supports_trtllm_attention() is True
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=False)
|
||||
@patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
|
||||
@patch(
|
||||
"vllm.utils.flashinfer.current_platform.is_device_capability_family",
|
||||
return_value=False,
|
||||
)
|
||||
def test_supports_non_sm100_platform(_cap, _bi):
|
||||
def test_supports_non_sm100_platform(_cap):
|
||||
assert supports_trtllm_attention() is False
|
||||
|
||||
|
||||
@patch("vllm.utils.flashinfer.vllm_is_batch_invariant", return_value=False)
|
||||
@patch("vllm.envs.VLLM_BATCH_INVARIANT", False)
|
||||
@patch(
|
||||
"vllm.utils.flashinfer.current_platform.is_device_capability_family",
|
||||
return_value=True,
|
||||
)
|
||||
@patch("vllm.utils.flashinfer.has_nvidia_artifactory", return_value=False)
|
||||
def test_supports_sm100_without_artifactory(_art, _cap, _bi):
|
||||
def test_supports_sm100_without_artifactory(_art, _cap):
|
||||
assert supports_trtllm_attention() is False
|
||||
|
||||
|
||||
|
||||
@@ -199,10 +199,10 @@ register_experts(
|
||||
|
||||
# Disable on blackwell for now
|
||||
if has_deep_ep() and not current_platform.has_device_capability(100):
|
||||
from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import (
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import (
|
||||
DeepEPHTPrepareAndFinalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import (
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import (
|
||||
DeepEPLLPrepareAndFinalize,
|
||||
)
|
||||
|
||||
@@ -240,7 +240,7 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability
|
||||
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
|
||||
FlashInferExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.flashinfer_nvlink_two_sided_prepare_finalize import ( # noqa: E501
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided import ( # noqa: E501
|
||||
FlashInferNVLinkTwoSidedPrepareAndFinalize,
|
||||
)
|
||||
|
||||
@@ -271,7 +271,7 @@ if (
|
||||
and has_flashinfer_cutlass_fused_moe()
|
||||
and current_platform.has_device_capability(100)
|
||||
):
|
||||
from vllm.model_executor.layers.fused_moe.flashinfer_nvlink_one_sided_prepare_finalize import ( # noqa: E501
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided import ( # noqa: E501
|
||||
FlashInferNVLinkOneSidedPrepareAndFinalize,
|
||||
)
|
||||
|
||||
|
||||
@@ -19,10 +19,10 @@ from vllm.utils.import_utils import has_deep_ep
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
|
||||
if has_deep_ep():
|
||||
from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import (
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import (
|
||||
DeepEPHTPrepareAndFinalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import (
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import (
|
||||
DeepEPLLPrepareAndFinalize,
|
||||
)
|
||||
|
||||
|
||||
@@ -17,7 +17,7 @@ from flashinfer import fp4_quantize
|
||||
from torch.nn import functional as F
|
||||
|
||||
from vllm.model_executor.layers.activation import SiluAndMul
|
||||
from vllm.model_executor.layers.fused_moe.flashinfer_cutedsl_moe import (
|
||||
from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import (
|
||||
flashinfer_cutedsl_moe_masked,
|
||||
)
|
||||
from vllm.utils.flashinfer import (
|
||||
|
||||
@@ -37,10 +37,10 @@ from .parallel_utils import ProcessGroupInfo, parallel_launch
|
||||
from .utils import make_dummy_moe_config, make_test_weights
|
||||
|
||||
if has_deep_ep():
|
||||
from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import (
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import (
|
||||
DeepEPHTPrepareAndFinalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import (
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import (
|
||||
DeepEPLLPrepareAndFinalize,
|
||||
)
|
||||
|
||||
|
||||
@@ -32,10 +32,10 @@ from ...utils import multi_gpu_test
|
||||
from .parallel_utils import ProcessGroupInfo, parallel_launch
|
||||
|
||||
if has_deep_ep():
|
||||
from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import (
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import (
|
||||
DeepEPHTPrepareAndFinalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import (
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import (
|
||||
DeepEPLLPrepareAndFinalize,
|
||||
)
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@ Run `pytest tests/kernels/moe/test_grouped_topk.py`.
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm.model_executor.layers.batch_invariant as batch_invariant
|
||||
import vllm.envs as envs
|
||||
from vllm.config import (
|
||||
CompilationConfig,
|
||||
VllmConfig,
|
||||
@@ -69,7 +69,7 @@ def test_grouped_topk(
|
||||
|
||||
with set_current_vllm_config(vllm_config), monkeypatch.context() as m:
|
||||
m.setenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "0")
|
||||
m.setattr(batch_invariant, "VLLM_BATCH_INVARIANT", True)
|
||||
m.setattr(envs, "VLLM_BATCH_INVARIANT", True)
|
||||
grouped_topk = GroupedTopk(
|
||||
topk=topk,
|
||||
renormalize=renormalize,
|
||||
|
||||
@@ -1,238 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for sparse cutlass kernels
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_cutlass_2of4_sparse.py`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import baseline_scaled_mm, to_fp8, to_int8
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
|
||||
sparse_cutlass_supported,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
CUDA_DEVICES = [
|
||||
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
|
||||
]
|
||||
|
||||
capability = current_platform.get_device_capability()
|
||||
capability = capability[0] * 10 + capability[1]
|
||||
|
||||
|
||||
def to_bf16(tensor: torch.Tensor) -> torch.Tensor:
|
||||
return tensor.to(dtype=torch.bfloat16)
|
||||
|
||||
|
||||
def to_fp16(tensor: torch.Tensor) -> torch.Tensor:
|
||||
return tensor.to(dtype=torch.float16)
|
||||
|
||||
|
||||
def prune_to_2_4(tensor):
|
||||
# Reshape tensor to [N, 4] where N is number of groups of 4
|
||||
original_shape = tensor.shape
|
||||
reshaped = tensor.reshape(-1, 4)
|
||||
|
||||
# Get indices of top 2 absolute values in each group of 4
|
||||
_, indices = torch.topk(torch.abs(reshaped), k=2, dim=1)
|
||||
|
||||
# Create binary mask
|
||||
mask = torch.zeros_like(reshaped)
|
||||
mask.scatter_(dim=1, index=indices, src=torch.ones_like(indices, dtype=mask.dtype))
|
||||
|
||||
# Apply mask and reshape back
|
||||
pruned = reshaped * mask
|
||||
|
||||
# Turn all -0.0 to 0.0
|
||||
pruned[pruned == -0.0] = 0.0
|
||||
|
||||
return pruned.reshape(original_shape)
|
||||
|
||||
|
||||
# This function checks that applying an identity matrix multiplication
|
||||
# to the compressed weights yields the original uncompressed weights.
|
||||
def check_compress_decompress_invariance(
|
||||
dtype: torch.dtype,
|
||||
b: torch.Tensor,
|
||||
b_compressed: torch.Tensor,
|
||||
b_metadata: torch.Tensor,
|
||||
):
|
||||
# For float16 and bfloat16, cutlass_scaled_sparse_mm's output must be the
|
||||
# same dtype as its inputs. This line addresses that constraint while
|
||||
# arbitrarily using bfloat16 for the int8/fp8 cases.
|
||||
out_dtype = torch.float16 if dtype is torch.float16 else torch.bfloat16
|
||||
|
||||
eye = torch.eye(b.shape[0], device="cuda", dtype=dtype)
|
||||
eye_scale = torch.ones(1, device="cuda", dtype=torch.float32)
|
||||
b_decomp = ops.cutlass_scaled_sparse_mm(
|
||||
eye, b_compressed, b_metadata, eye_scale, eye_scale, out_dtype=out_dtype
|
||||
)
|
||||
|
||||
torch.testing.assert_close(b.to(dtype=out_dtype), b_decomp)
|
||||
|
||||
|
||||
def make_rand_sparse_tensors(
|
||||
dtype: torch.dtype, m: int, n: int, k: int
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
a = torch.randn((m, k), device="cuda")
|
||||
b = torch.randn((n, k), device="cuda").t()
|
||||
|
||||
if dtype == torch.int8:
|
||||
# ensure A and B aren't all zeros after rounding
|
||||
a = a * 5.0
|
||||
b = b * 5.0
|
||||
|
||||
b = prune_to_2_4(b.t()).t()
|
||||
|
||||
if dtype == torch.int8:
|
||||
a, b = to_int8(a), to_int8(b)
|
||||
elif dtype == torch.float8_e4m3fn:
|
||||
a, b = to_fp8(a), to_fp8(b)
|
||||
elif dtype == torch.float16:
|
||||
a, b = to_fp16(a), to_fp16(b)
|
||||
elif dtype == torch.bfloat16:
|
||||
a, b = to_bf16(a), to_bf16(b)
|
||||
else:
|
||||
raise ValueError("unsupported dtype")
|
||||
|
||||
b_compressed, e = ops.cutlass_sparse_compress(b.t())
|
||||
check_compress_decompress_invariance(dtype, b, b_compressed, e)
|
||||
|
||||
# Compressed B, Metadata, Original A, B
|
||||
return b_compressed, e, a, b
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not sparse_cutlass_supported(),
|
||||
reason="Sparse CUTLASS is not supported on this GPU type.",
|
||||
)
|
||||
# Test working with a subset of A and B for sparse matmul
|
||||
def test_cutlass_sparse_subset():
|
||||
big_m = 1024
|
||||
m, n, k = 512, 512, 512
|
||||
|
||||
# Create tensors
|
||||
b_comp, e, whole_a, b = make_rand_sparse_tensors(torch.float8_e4m3fn, big_m, n, k)
|
||||
a = whole_a[0:m, 0:k]
|
||||
scale_a = torch.randn((1, 1), device="cuda", dtype=torch.float32) / 10
|
||||
scale_b = torch.randn((1, 1), device="cuda", dtype=torch.float32) / 10
|
||||
|
||||
out = ops.cutlass_scaled_sparse_mm(
|
||||
a, b_comp, e, scale_a, scale_b, out_dtype=torch.bfloat16
|
||||
)
|
||||
baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype=torch.bfloat16)
|
||||
|
||||
torch.testing.assert_close(out, baseline, rtol=1e-1, atol=1e0)
|
||||
|
||||
|
||||
MNK_FACTORS = [
|
||||
(1, 256, 128),
|
||||
(1, 16384, 1024),
|
||||
(1, 24576, 512),
|
||||
(16, 256, 512),
|
||||
(16, 16384, 128),
|
||||
(16, 24576, 4096),
|
||||
(32, 8192, 4096),
|
||||
(32, 16384, 4096),
|
||||
(33, 1024, 1024),
|
||||
(33, 8192, 128),
|
||||
(64, 2048, 512),
|
||||
(64, 16384, 1024),
|
||||
(100, 8192, 512),
|
||||
(128, 32768, 4096),
|
||||
(256, 4096, 4096),
|
||||
(512, 256, 1024),
|
||||
(512, 8192, 4096),
|
||||
(512, 16384, 128),
|
||||
(512, 24576, 128),
|
||||
]
|
||||
|
||||
|
||||
# Test working with a subset of A and B for sparse matmul
|
||||
@pytest.mark.skipif(
|
||||
not sparse_cutlass_supported(),
|
||||
reason="Sparse CUTLASS is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("m, n, k", MNK_FACTORS)
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
def test_cutlass_sparse_gemm(
|
||||
m: int, k: int, n: int, dtype: type[torch.dtype], use_bias: bool
|
||||
):
|
||||
# Create tensors
|
||||
b_comp, e, a, b = make_rand_sparse_tensors(dtype, m, n, k)
|
||||
scale_a = torch.ones((1, 1), device="cuda", dtype=torch.float32)
|
||||
scale_b = torch.ones((1, 1), device="cuda", dtype=torch.float32)
|
||||
|
||||
bias = torch.rand((n,), device="cuda", dtype=dtype) if use_bias else None
|
||||
|
||||
out = ops.cutlass_scaled_sparse_mm(
|
||||
a, b_comp, e, scale_a, scale_b, out_dtype=dtype, bias=bias
|
||||
)
|
||||
|
||||
baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype=dtype, bias=bias)
|
||||
|
||||
torch.testing.assert_close(out, baseline, rtol=1e-2, atol=3e-1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not sparse_cutlass_supported(),
|
||||
reason="Sparse CUTLASS is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("m, k, n", MNK_FACTORS)
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.has_device_capability(89),
|
||||
reason="FP8 is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
def test_cutlass_sparse_fp8_gemm(m: int, n: int, k: int, use_bias: bool):
|
||||
# Create tensors
|
||||
b_comp, e, a, b = make_rand_sparse_tensors(torch.float8_e4m3fn, m, n, k)
|
||||
scale_a = torch.randn((1, 1), device="cuda", dtype=torch.float32)
|
||||
scale_b = torch.randn((1, 1), device="cuda", dtype=torch.float32)
|
||||
out_dtype = torch.bfloat16
|
||||
|
||||
bias = torch.rand((n,), device="cuda", dtype=out_dtype) * 10 if use_bias else None
|
||||
|
||||
out = ops.cutlass_scaled_sparse_mm(
|
||||
a, b_comp, e, scale_a, scale_b, out_dtype=out_dtype, bias=bias
|
||||
)
|
||||
|
||||
baseline = baseline_scaled_mm(
|
||||
a, b, scale_a, scale_b, out_dtype=out_dtype, bias=bias
|
||||
)
|
||||
|
||||
torch.testing.assert_close(out, baseline, rtol=1e-2, atol=3e-1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not sparse_cutlass_supported(),
|
||||
reason="Sparse CUTLASS is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("m,k,n", MNK_FACTORS)
|
||||
@pytest.mark.parametrize("per_act_token", [True, False])
|
||||
@pytest.mark.parametrize("per_out_ch", [True, False])
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
def test_cutlass_sparse_int8_gemm(
|
||||
m: int, n: int, k: int, per_act_token: bool, per_out_ch: bool, use_bias: bool
|
||||
):
|
||||
# Create tensors
|
||||
b_comp, e, a, b = make_rand_sparse_tensors(torch.int8, m, n, k)
|
||||
scale_a = torch.randn((1, 1), device="cuda", dtype=torch.float32)
|
||||
scale_b = torch.randn((1, 1), device="cuda", dtype=torch.float32)
|
||||
out_dtype = torch.bfloat16
|
||||
|
||||
bias = torch.rand((n,), device="cuda", dtype=out_dtype) * 10 if use_bias else None
|
||||
|
||||
out = ops.cutlass_scaled_sparse_mm(
|
||||
a, b_comp, e, scale_a, scale_b, out_dtype=out_dtype, bias=bias
|
||||
)
|
||||
|
||||
baseline = baseline_scaled_mm(
|
||||
a, b, scale_a, scale_b, out_dtype=out_dtype, bias=bias
|
||||
)
|
||||
|
||||
torch.testing.assert_close(out, baseline, rtol=1e0, atol=2e0)
|
||||
@@ -247,3 +247,68 @@ def test_quantize_to_fp4_padded_no_sf_swizzled(pad_shape: tuple[int, int]) -> No
|
||||
out_ans = cast_from_fp4(out, m, n)
|
||||
torch.testing.assert_close(out_ans, out_ref)
|
||||
torch.testing.assert_close(scale_ans, scale_ref)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# SM103-native quantization correctness
|
||||
# ============================================================================
|
||||
|
||||
_SM103_QUANT_SHAPES = SHAPES + PAD_SHAPES
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not hasattr(torch.ops._C, "scaled_fp4_quant_sm103"),
|
||||
reason="scaled_fp4_quant_sm103 op not available "
|
||||
"(rebuild without VLLM_USE_PRECOMPILED=1)",
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("shape", _SM103_QUANT_SHAPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@torch.inference_mode()
|
||||
def test_scaled_fp4_quant_sm103_matches_sm100(
|
||||
dtype: torch.dtype,
|
||||
shape: tuple[int, int],
|
||||
seed: int,
|
||||
) -> None:
|
||||
"""
|
||||
Verify scaled_fp4_quant_sm103 (SM103-native layout) against the SM100 path.
|
||||
|
||||
Two invariants:
|
||||
1. Packed FP4 data is identical — both kernels quantize to the same
|
||||
e2m1 values; only the SF memory layout differs.
|
||||
2. SM103 native SFs are byte-identical to SM100 SFs run through
|
||||
convert_sf_layout_sm100_to_sm103 — confirms the in-kernel swizzle
|
||||
matches the standalone conversion kernel.
|
||||
"""
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device("cuda:0")
|
||||
m, n = shape
|
||||
|
||||
x = torch.randn((m, n), dtype=dtype)
|
||||
tensor_amax = torch.abs(x).max().to(torch.float32)
|
||||
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
|
||||
|
||||
# SM100 reference: ops wrapper returns fp4 (uint8) and sf (float8_e4m3fn)
|
||||
fp4_sm100, sf_sm100 = ops.scaled_fp4_quant(
|
||||
x, global_scale, is_sf_swizzled_layout=True
|
||||
)
|
||||
|
||||
# SM103 native: raw C++ op returns sf as int32; view as float8 to match
|
||||
fp4_sm103, sf_sm103_i32 = torch.ops._C.scaled_fp4_quant_sm103(x, global_scale)
|
||||
sf_sm103 = sf_sm103_i32.view(torch.float8_e4m3fn)
|
||||
|
||||
# 1. FP4 quantized data must be identical
|
||||
assert torch.equal(fp4_sm103, fp4_sm100), (
|
||||
f"FP4 data mismatch between SM100 and SM103 quant kernels "
|
||||
f"(shape={shape}, dtype={dtype})"
|
||||
)
|
||||
|
||||
# 2. SM103 native SFs must match SM100 SFs converted to SM103 layout
|
||||
sf_sm100_converted = torch.empty_like(sf_sm100)
|
||||
torch.ops._C.convert_sf_layout_sm100_to_sm103(sf_sm100_converted, sf_sm100)
|
||||
|
||||
assert torch.equal(sf_sm103.view(torch.uint8), sf_sm100_converted.view(torch.uint8)), (
|
||||
f"SM103 native SF layout doesn't match "
|
||||
f"convert_sf_layout_sm100_to_sm103(SM100 SFs) "
|
||||
f"(shape={shape}, dtype={dtype})"
|
||||
)
|
||||
|
||||
@@ -295,10 +295,15 @@ def whisper_lora_files():
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def qwen35_dense_model_lora_files():
|
||||
def qwen35_text_lora_files():
|
||||
return snapshot_download(repo_id="jeeejeee/qwen35-4b-text-only-sql-lora")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def qwen35_vl_lora_files():
|
||||
return snapshot_download(repo_id="jeeejeee/qwen35-4b-all-linear-pokemon-lora")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def reset_default_device():
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,361 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
import vllm
|
||||
import vllm.config
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.lora.request import LoRARequest
|
||||
|
||||
from ..utils import create_new_process_for_each_test, multi_gpu_test
|
||||
|
||||
MODEL_PATH = "Qwen/Qwen3.5-4B"
|
||||
TEXT_LORA_ID = 1
|
||||
VL_LORA_ID = 2
|
||||
|
||||
# text-only task
|
||||
TEXT_PROMPT_TEMPLATE = """Write a SQL query for the given database.\nSchema:\nTables:\n - stadium(Stadium_ID, Location, Name, Capacity, Highest, Lowest, Average)\n - singer(Singer_ID, Name, Country, Song_Name, Song_release_year, Age, Is_male)\n - concert(concert_ID, concert_Name, Theme, Stadium_ID, Year)\n - singer_in_concert(concert_ID, Singer_ID)\n\nQuestion:\n{query}""" # noqa: E501
|
||||
|
||||
TEXT_EXPECTED_LORA_OUTPUT = [
|
||||
"SELECT count(*) FROM singer",
|
||||
"SELECT avg(age) , min(age) , max(age) FROM singer WHERE country = 'France'",
|
||||
"SELECT name FROM stadium WHERE stadium_id NOT IN (SELECT stadium_id FROM concert)",
|
||||
]
|
||||
|
||||
|
||||
# visual caption
|
||||
VL_QUESTION = "What is in the image?"
|
||||
VL_TEST_IMAGES = [
|
||||
ImageAsset("stop_sign"),
|
||||
ImageAsset("cherry_blossom"),
|
||||
]
|
||||
VL_EXPECTED_LORA_OUTPUT = [
|
||||
'A red STOP sign stands prominently in the foreground, with a traditional Chinese gate adorned with red lanterns and the Chinese characters "中華門" in the background, signaling the entrance to a Chinatown. A black car passes by on the street, and stone lion statues guard the entrance to the culturally rich area.', # noqa: E501
|
||||
"A vibrant blue sky serves as a backdrop for the iconic Tokyo Skytree, partially obscured by the delicate pink blossoms of cherry trees in full bloom.", # noqa: E501
|
||||
]
|
||||
|
||||
TOKENIZER = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
|
||||
|
||||
|
||||
def _assert_exact_outputs(
|
||||
generated_texts: list[str], expected_outputs: list[str]
|
||||
) -> None:
|
||||
assert generated_texts == expected_outputs
|
||||
|
||||
|
||||
def _assert_prefix_outputs(
|
||||
generated_texts: list[str],
|
||||
expected_outputs: list[str],
|
||||
) -> None:
|
||||
assert len(generated_texts) == len(expected_outputs)
|
||||
for generated_text, expected_text in zip(generated_texts, expected_outputs):
|
||||
assert expected_text.startswith(generated_text), (
|
||||
f"Generated {generated_text!r} is not a prefix of expected "
|
||||
f"{expected_text!r}"
|
||||
)
|
||||
|
||||
|
||||
def _run_text_lora_sample(
|
||||
llm: vllm.LLM,
|
||||
lora_path: str,
|
||||
lora_id: int,
|
||||
) -> list[str]:
|
||||
prompts = [
|
||||
TEXT_PROMPT_TEMPLATE.format(query="How many singers do we have?"),
|
||||
TEXT_PROMPT_TEMPLATE.format(
|
||||
query=(
|
||||
"What is the average, minimum, and maximum "
|
||||
"age of all singers from France?"
|
||||
)
|
||||
),
|
||||
TEXT_PROMPT_TEMPLATE.format(
|
||||
query="What are the names of the stadiums without any concerts?"
|
||||
),
|
||||
]
|
||||
input_templates = []
|
||||
for prompt_text in prompts:
|
||||
messages = [{"role": "user", "content": prompt_text}]
|
||||
prompt = TOKENIZER.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=False, # disable thinking
|
||||
)
|
||||
input_templates.append(prompt)
|
||||
|
||||
outputs = llm.generate(
|
||||
input_templates,
|
||||
vllm.SamplingParams(temperature=0, max_tokens=512),
|
||||
lora_request=LoRARequest(str(lora_id), lora_id, lora_path),
|
||||
)
|
||||
|
||||
generated_texts: list[str] = []
|
||||
for output in outputs:
|
||||
generated_text = output.outputs[0].text.strip()
|
||||
generated_texts.append(generated_text)
|
||||
print(f"Prompt: {output.prompt!r}, Generated text: {generated_text!r}")
|
||||
return generated_texts
|
||||
|
||||
|
||||
def _run_vl_lora_sample(
|
||||
llm: vllm.LLM,
|
||||
lora_path: str | None = None,
|
||||
lora_id: int = 0,
|
||||
) -> list[str]:
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image"},
|
||||
{"type": "text", "text": VL_QUESTION},
|
||||
],
|
||||
}
|
||||
]
|
||||
prompt = TOKENIZER.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
prompts = [
|
||||
{
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": asset.pil_image},
|
||||
}
|
||||
for asset in VL_TEST_IMAGES
|
||||
]
|
||||
outputs = llm.generate(
|
||||
prompts,
|
||||
vllm.SamplingParams(temperature=0, max_tokens=128),
|
||||
lora_request=(
|
||||
LoRARequest(str(lora_id), lora_id, lora_path)
|
||||
if lora_path is not None
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
generated_texts: list[str] = []
|
||||
for output in outputs:
|
||||
generated_text = output.outputs[0].text.strip()
|
||||
generated_texts.append(generated_text)
|
||||
print(f"Prompt: {output.prompt!r}, Generated text: {generated_text!r}")
|
||||
return generated_texts
|
||||
|
||||
|
||||
def _build_text_prompts() -> list[str]:
|
||||
prompts = [
|
||||
TEXT_PROMPT_TEMPLATE.format(query="How many singers do we have?"),
|
||||
TEXT_PROMPT_TEMPLATE.format(
|
||||
query=(
|
||||
"What is the average, minimum, and maximum "
|
||||
"age of all singers from France?"
|
||||
)
|
||||
),
|
||||
TEXT_PROMPT_TEMPLATE.format(
|
||||
query="What are the names of the stadiums without any concerts?"
|
||||
),
|
||||
]
|
||||
input_templates = []
|
||||
for prompt_text in prompts:
|
||||
messages = [{"role": "user", "content": prompt_text}]
|
||||
prompt = TOKENIZER.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
input_templates.append(prompt)
|
||||
return input_templates
|
||||
|
||||
|
||||
def _build_vl_prompts() -> list[dict]:
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image"},
|
||||
{"type": "text", "text": VL_QUESTION},
|
||||
],
|
||||
}
|
||||
]
|
||||
prompt = TOKENIZER.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
return [
|
||||
{
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": asset.pil_image},
|
||||
}
|
||||
for asset in VL_TEST_IMAGES
|
||||
]
|
||||
|
||||
|
||||
def _run_mixed_lora_sample(
|
||||
llm: vllm.LLM,
|
||||
text_lora_path: str,
|
||||
vl_lora_path: str,
|
||||
text_lora_id: int,
|
||||
vl_lora_id: int,
|
||||
) -> list[str]:
|
||||
text_prompts = _build_text_prompts()[:2]
|
||||
vl_prompts = _build_vl_prompts()
|
||||
prompts = [
|
||||
text_prompts[0],
|
||||
vl_prompts[0],
|
||||
text_prompts[1],
|
||||
vl_prompts[1],
|
||||
]
|
||||
lora_requests = [
|
||||
LoRARequest("qwen35-text", text_lora_id, text_lora_path),
|
||||
LoRARequest("qwen35-vl", vl_lora_id, vl_lora_path),
|
||||
LoRARequest("qwen35-text", text_lora_id, text_lora_path),
|
||||
LoRARequest("qwen35-vl", vl_lora_id, vl_lora_path),
|
||||
]
|
||||
outputs = llm.generate(
|
||||
prompts,
|
||||
vllm.SamplingParams(temperature=0, max_tokens=256),
|
||||
lora_request=lora_requests,
|
||||
)
|
||||
|
||||
generated_texts: list[str] = []
|
||||
for output in outputs:
|
||||
generated_text = output.outputs[0].text.strip()
|
||||
generated_texts.append(generated_text)
|
||||
print(f"Prompt: {output.prompt!r}, Generated text: {generated_text!r}")
|
||||
return generated_texts
|
||||
|
||||
|
||||
def _run_mixed_lora_and_base_sample(
|
||||
llm: vllm.LLM,
|
||||
text_lora_path: str,
|
||||
vl_lora_path: str,
|
||||
text_lora_id: int,
|
||||
vl_lora_id: int,
|
||||
) -> list[str]:
|
||||
text_prompt = _build_text_prompts()[0]
|
||||
vl_prompt = _build_vl_prompts()[0]
|
||||
prompts = [
|
||||
text_prompt,
|
||||
vl_prompt,
|
||||
text_prompt,
|
||||
vl_prompt,
|
||||
]
|
||||
lora_requests = [
|
||||
LoRARequest("qwen35-text", text_lora_id, text_lora_path),
|
||||
LoRARequest("qwen35-vl", vl_lora_id, vl_lora_path),
|
||||
None,
|
||||
None,
|
||||
]
|
||||
outputs = llm.generate(
|
||||
prompts,
|
||||
vllm.SamplingParams(temperature=0, max_tokens=256),
|
||||
lora_request=lora_requests,
|
||||
)
|
||||
|
||||
generated_texts: list[str] = []
|
||||
for output in outputs:
|
||||
generated_text = output.outputs[0].text.strip()
|
||||
generated_texts.append(generated_text)
|
||||
print(f"Prompt: {output.prompt!r}, Generated text: {generated_text!r}")
|
||||
return generated_texts
|
||||
|
||||
|
||||
def _assert_qwen35_text_vl_and_mixed_lora(
|
||||
llm: vllm.LLM,
|
||||
qwen35_text_lora_files: str,
|
||||
qwen35_vl_lora_files: str,
|
||||
) -> None:
|
||||
generated_texts = _run_text_lora_sample(
|
||||
llm,
|
||||
qwen35_text_lora_files,
|
||||
TEXT_LORA_ID,
|
||||
)
|
||||
|
||||
_assert_exact_outputs(generated_texts, TEXT_EXPECTED_LORA_OUTPUT)
|
||||
|
||||
generated_texts = _run_vl_lora_sample(
|
||||
llm,
|
||||
qwen35_vl_lora_files,
|
||||
VL_LORA_ID,
|
||||
)
|
||||
_assert_prefix_outputs(generated_texts, VL_EXPECTED_LORA_OUTPUT)
|
||||
|
||||
generated_texts = _run_mixed_lora_sample(
|
||||
llm,
|
||||
qwen35_text_lora_files,
|
||||
qwen35_vl_lora_files,
|
||||
text_lora_id=TEXT_LORA_ID,
|
||||
vl_lora_id=VL_LORA_ID,
|
||||
)
|
||||
assert generated_texts[0] == TEXT_EXPECTED_LORA_OUTPUT[0]
|
||||
assert generated_texts[2] == TEXT_EXPECTED_LORA_OUTPUT[1]
|
||||
_assert_prefix_outputs([generated_texts[1]], [VL_EXPECTED_LORA_OUTPUT[0]])
|
||||
_assert_prefix_outputs([generated_texts[3]], [VL_EXPECTED_LORA_OUTPUT[1]])
|
||||
|
||||
generated_texts = _run_mixed_lora_and_base_sample(
|
||||
llm,
|
||||
qwen35_text_lora_files,
|
||||
qwen35_vl_lora_files,
|
||||
text_lora_id=TEXT_LORA_ID,
|
||||
vl_lora_id=VL_LORA_ID,
|
||||
)
|
||||
assert generated_texts[0] == TEXT_EXPECTED_LORA_OUTPUT[0]
|
||||
_assert_prefix_outputs([generated_texts[1]], [VL_EXPECTED_LORA_OUTPUT[0]])
|
||||
assert generated_texts[2] != TEXT_EXPECTED_LORA_OUTPUT[0]
|
||||
assert not VL_EXPECTED_LORA_OUTPUT[0].startswith(generated_texts[3]), (
|
||||
"Non-LoRA vision output unexpectedly matches the LoRA expectation."
|
||||
)
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_qwen35_text_lora(qwen35_text_lora_files, qwen35_vl_lora_files):
|
||||
llm = vllm.LLM(
|
||||
model=MODEL_PATH,
|
||||
max_model_len=4096,
|
||||
enable_lora=True,
|
||||
max_loras=2,
|
||||
max_num_seqs=4,
|
||||
max_lora_rank=8,
|
||||
enforce_eager=True,
|
||||
trust_remote_code=True,
|
||||
enable_tower_connector_lora=True,
|
||||
mm_processor_cache_gb=0,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
)
|
||||
|
||||
_assert_qwen35_text_vl_and_mixed_lora(
|
||||
llm,
|
||||
qwen35_text_lora_files,
|
||||
qwen35_vl_lora_files,
|
||||
)
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=4)
|
||||
def test_qwen35_text_lora_tp4(qwen35_text_lora_files, qwen35_vl_lora_files):
|
||||
llm = vllm.LLM(
|
||||
model=MODEL_PATH,
|
||||
max_model_len=4096,
|
||||
enable_lora=True,
|
||||
max_loras=2,
|
||||
max_lora_rank=8,
|
||||
max_num_seqs=4,
|
||||
enforce_eager=True,
|
||||
tensor_parallel_size=4,
|
||||
trust_remote_code=True,
|
||||
enable_tower_connector_lora=True,
|
||||
mm_processor_cache_gb=0,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
compilation_config=vllm.config.CompilationConfig(
|
||||
cudagraph_specialize_lora=False,
|
||||
),
|
||||
)
|
||||
|
||||
_assert_qwen35_text_vl_and_mixed_lora(
|
||||
llm,
|
||||
qwen35_text_lora_files,
|
||||
qwen35_vl_lora_files,
|
||||
)
|
||||
@@ -1,132 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
import vllm
|
||||
import vllm.config
|
||||
from vllm.lora.request import LoRARequest
|
||||
|
||||
from ..utils import create_new_process_for_each_test, multi_gpu_test
|
||||
|
||||
MODEL_PATH = "Qwen/Qwen3.5-4B"
|
||||
|
||||
PROMPT_TEMPLATE = """Write a SQL query for the given database.\nSchema:\nTables:\n - stadium(Stadium_ID, Location, Name, Capacity, Highest, Lowest, Average)\n - singer(Singer_ID, Name, Country, Song_Name, Song_release_year, Age, Is_male)\n - concert(concert_ID, concert_Name, Theme, Stadium_ID, Year)\n - singer_in_concert(concert_ID, Singer_ID)\n\nQuestion:\n{query}""" # noqa: E501
|
||||
|
||||
EXPECTED_LORA_OUTPUT = [
|
||||
"SELECT count(*) FROM singer",
|
||||
"SELECT avg(age) , min(age) , max(age) FROM singer WHERE country = 'France'",
|
||||
"SELECT name FROM stadium WHERE stadium_id NOT IN (SELECT stadium_id FROM concert)",
|
||||
]
|
||||
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
|
||||
|
||||
|
||||
def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
|
||||
prompts = [
|
||||
PROMPT_TEMPLATE.format(query="How many singers do we have?"),
|
||||
PROMPT_TEMPLATE.format(
|
||||
query=(
|
||||
"What is the average, minimum, and maximum "
|
||||
"age of all singers from France?"
|
||||
)
|
||||
),
|
||||
PROMPT_TEMPLATE.format(
|
||||
query=("What are the names of the stadiums without any concerts?")
|
||||
),
|
||||
]
|
||||
input_templates = []
|
||||
for prmpt in prompts:
|
||||
messages = [{"role": "user", "content": prmpt}]
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=False, # disable thinking
|
||||
)
|
||||
input_templates.append(prompt)
|
||||
sampling_params = vllm.SamplingParams(temperature=0, max_tokens=512)
|
||||
outputs = llm.generate(
|
||||
input_templates,
|
||||
sampling_params,
|
||||
lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
|
||||
)
|
||||
|
||||
generated_texts: list[str] = []
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text.strip()
|
||||
generated_texts.append(generated_text)
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
return generated_texts
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_qwen35_dense_model_lora(qwen35_dense_model_lora_files):
|
||||
llm = vllm.LLM(
|
||||
MODEL_PATH,
|
||||
max_model_len=512,
|
||||
enable_lora=True,
|
||||
max_loras=2,
|
||||
max_num_seqs=16,
|
||||
max_lora_rank=8,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
|
||||
output1 = do_sample(llm, qwen35_dense_model_lora_files, lora_id=1)
|
||||
for i in range(len(EXPECTED_LORA_OUTPUT)):
|
||||
assert output1[i] == EXPECTED_LORA_OUTPUT[i]
|
||||
output2 = do_sample(llm, qwen35_dense_model_lora_files, lora_id=2)
|
||||
for i in range(len(EXPECTED_LORA_OUTPUT)):
|
||||
assert output2[i] == EXPECTED_LORA_OUTPUT[i]
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=4)
|
||||
def test_qwen35_dense_model_lora_tp4(qwen35_dense_model_lora_files):
|
||||
llm = vllm.LLM(
|
||||
MODEL_PATH,
|
||||
max_model_len=1024,
|
||||
enable_lora=True,
|
||||
max_loras=2,
|
||||
max_lora_rank=8,
|
||||
max_num_seqs=16,
|
||||
tensor_parallel_size=4,
|
||||
trust_remote_code=True,
|
||||
fully_sharded_loras=False,
|
||||
compilation_config=vllm.config.CompilationConfig( # Avoid OOM
|
||||
cudagraph_specialize_lora=False,
|
||||
),
|
||||
)
|
||||
|
||||
output1 = do_sample(llm, qwen35_dense_model_lora_files, lora_id=1)
|
||||
print(output1)
|
||||
for i in range(len(EXPECTED_LORA_OUTPUT)):
|
||||
assert output1[i] == EXPECTED_LORA_OUTPUT[i]
|
||||
output2 = do_sample(llm, qwen35_dense_model_lora_files, lora_id=2)
|
||||
for i in range(len(EXPECTED_LORA_OUTPUT)):
|
||||
assert output2[i] == EXPECTED_LORA_OUTPUT[i]
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=4)
|
||||
def test_qwen35_dense_model_lora_tp4_fully_sharded_loras(qwen35_dense_model_lora_files):
|
||||
llm = vllm.LLM(
|
||||
MODEL_PATH,
|
||||
max_model_len=512,
|
||||
enable_lora=True,
|
||||
max_loras=2,
|
||||
max_lora_rank=8,
|
||||
tensor_parallel_size=4,
|
||||
trust_remote_code=True,
|
||||
fully_sharded_loras=True,
|
||||
gpu_memory_utilization=0.8,
|
||||
compilation_config=vllm.config.CompilationConfig( # Avoid OOM
|
||||
cudagraph_specialize_lora=False,
|
||||
),
|
||||
)
|
||||
output1 = do_sample(llm, qwen35_dense_model_lora_files, lora_id=1)
|
||||
for i in range(len(EXPECTED_LORA_OUTPUT)):
|
||||
assert output1[i] == EXPECTED_LORA_OUTPUT[i]
|
||||
output2 = do_sample(llm, qwen35_dense_model_lora_files, lora_id=2)
|
||||
for i in range(len(EXPECTED_LORA_OUTPUT)):
|
||||
assert output2[i] == EXPECTED_LORA_OUTPUT[i]
|
||||
@@ -26,7 +26,7 @@ def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch):
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
|
||||
monkeypatch.setattr(utils, "get_cu_count", lambda: 120)
|
||||
monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
|
||||
|
||||
wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
|
||||
monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
|
||||
@@ -50,7 +50,7 @@ def test_rocm_unquantized_gemm_gfx1x_n_gt_4_falls_back(monkeypatch):
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
|
||||
monkeypatch.setattr(utils, "get_cu_count", lambda: 120)
|
||||
monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
|
||||
|
||||
wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
|
||||
monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
|
||||
@@ -74,7 +74,7 @@ def test_rocm_unquantized_gemm_gfx950_wvsplitkrc_path(monkeypatch):
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: False)
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
|
||||
monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True)
|
||||
monkeypatch.setattr(utils, "get_cu_count", lambda: 120)
|
||||
monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
|
||||
|
||||
wvsplitkrc_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
|
||||
monkeypatch.setattr(utils.ops, "wvSplitKrc", wvsplitkrc_mock)
|
||||
|
||||
@@ -1 +1 @@
|
||||
{"transcriptions": ["The content of the input audio is 'you can ask why over and over and over again forever even if one day we explain every physical interaction and scientific law and hope and dream and regret with a single elegant equation'."], "token_ids": [[785, 2213, 315, 279, 1946, 7699, 374, 364, 9330, 646, 2548, 3170, 916, 323, 916, 323, 916, 1549, 15683, 1496, 421, 825, 1899, 582, 10339, 1449, 6961, 16230, 323, 12344, 2329, 323, 3900, 323, 7904, 323, 22231, 448, 264, 3175, 25777, 23606, 4427, 151645]]}
|
||||
{"transcriptions": ["There is no clear relationship between the barking and the music, as they seem to be independent of each other."], "token_ids": [[3862, 374, 902, 2797, 5025, 1948, 279, 293, 33452, 323, 279, 4627, 11, 438, 807, 2803, 311, 387, 9489, 315, 1817, 1008, 13, 151645]]}
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
{"transcriptions": ["This track is an energetic Eurodance / Dance‑Pop anthem that blends the bright, melodic sensibilities of mainstream pop with the driving, club‑ready pulse of classic Eurodance. The duration of the piece is ", "**Verse 1**\nMidnight cravings in bloom, lights flicker in the room, pepperoni dreams arise, pizza party on your skies\n\n**Verse 2**\nCheese melts on the crust, in flavor we trust, boxes stacked to the"], "token_ids": [[1986, 3754, 374, 458, 44855, 19461, 98875, 378, 107, 14, 378, 107, 35, 681, 55964, 11598, 55564, 429, 57843, 279, 9906, 11, 10581, 52760, 6097, 13450, 315, 20729, 2420, 448, 279, 9842, 11, 6335, 55964, 2307, 27235, 315, 11416, 19461, 98875, 13, 220, 576, 8090, 315, 279, 6573, 374, 220], [334, 68043, 220, 16, 1019, 33648, 9287, 88828, 304, 51454, 11, 12711, 28347, 261, 304, 279, 3054, 11, 24353, 20783, 18707, 30789, 11, 22502, 4614, 389, 697, 49293, 271, 334, 68043, 220, 17, 1019, 26843, 2367, 98091, 389, 279, 39612, 11, 304, 17172, 582, 6950, 11, 14697, 41315, 311, 279]]}
|
||||
@@ -0,0 +1 @@
|
||||
{"transcriptions": ["This track is an energetic Eurodance / Dance‑Pop anthem that blends the bright, melodic sensibilities of mainstream pop with the driving, club‑ready pulse of classic Eurodance. The duration of the piece is "], "token_ids": [[1986, 3754, 374, 458, 44855, 19461, 98875, 378, 107, 14, 378, 107, 35, 681, 55964, 11598, 55564, 429, 57843, 279, 9906, 11, 10581, 52760, 6097, 13450, 315, 20729, 2420, 448, 279, 9842, 11, 6335, 55964, 2307, 27235, 315, 11416, 19461, 98875, 13, 220, 576, 8090, 315, 279, 6573, 374, 220]]}
|
||||
@@ -26,6 +26,54 @@ from tests.models.registry import HF_EXAMPLE_MODELS
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
MODEL_NAME = "nvidia/audio-flamingo-3-hf"
|
||||
SINGLE_CONVERSATION = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What is surprising about the relationship between "
|
||||
"the barking and the music?",
|
||||
},
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {
|
||||
"url": "https://huggingface.co/datasets/nvidia/AudioSkills/"
|
||||
"resolve/main/assets/"
|
||||
"dogs_barking_in_sync_with_the_music.wav",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
BATCHED_CONVERSATIONS = [
|
||||
SINGLE_CONVERSATION,
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Why is the philosopher's name mentioned in the "
|
||||
"lyrics? (A) To express a sense of nostalgia "
|
||||
"(B) To indicate that language cannot express clearly, "
|
||||
"satirizing the inversion of black and white in the world "
|
||||
"(C) To add depth and complexity to the lyrics "
|
||||
"(D) To showcase the wisdom and influence of the "
|
||||
"philosopher",
|
||||
},
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {
|
||||
"url": "https://huggingface.co/datasets/nvidia/"
|
||||
"AudioSkills/resolve/main/assets/"
|
||||
"Ch6Ae9DT6Ko_00-04-03_00-04-31.wav",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
]
|
||||
|
||||
|
||||
def get_fixture_path(filename):
|
||||
@@ -34,21 +82,29 @@ def get_fixture_path(filename):
|
||||
)
|
||||
|
||||
|
||||
def assert_output_matches(output, expected_text, expected_token_ids):
|
||||
generated = output.outputs[0]
|
||||
assert generated.text.strip() == expected_text
|
||||
actual_token_ids = list(generated.token_ids)
|
||||
assert (
|
||||
actual_token_ids == expected_token_ids
|
||||
or actual_token_ids == expected_token_ids[:-1]
|
||||
or actual_token_ids[:-1] == expected_token_ids
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def llm():
|
||||
# Check if the model is supported by the current transformers version
|
||||
model_info = HF_EXAMPLE_MODELS.get_hf_info("AudioFlamingo3ForConditionalGeneration")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
try:
|
||||
llm = LLM(
|
||||
return LLM(
|
||||
model=MODEL_NAME,
|
||||
trust_remote_code=True,
|
||||
dtype="bfloat16",
|
||||
enforce_eager=True,
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
)
|
||||
return llm
|
||||
except Exception as e:
|
||||
pytest.skip(f"Failed to load model {MODEL_NAME}: {e}")
|
||||
|
||||
@@ -61,29 +117,17 @@ def test_single_generation(llm):
|
||||
with open(fixture_path) as f:
|
||||
expected = json.load(f)
|
||||
|
||||
audio_url = "https://huggingface.co/datasets/nvidia/AudioSkills/resolve/main/assets/Why_do_we_ask_questions_converted.wav"
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "audio_url", "audio_url": {"url": audio_url}},
|
||||
{"type": "text", "text": "Transcribe the input speech."},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=128)
|
||||
|
||||
outputs = llm.chat(
|
||||
messages=messages,
|
||||
messages=SINGLE_CONVERSATION,
|
||||
sampling_params=sampling_params,
|
||||
)
|
||||
generated_text = outputs[0].outputs[0].text.strip()
|
||||
|
||||
expected_text = expected["transcriptions"][0]
|
||||
|
||||
assert expected_text in generated_text or generated_text in expected_text
|
||||
assert_output_matches(
|
||||
outputs[0],
|
||||
expected["transcriptions"][0],
|
||||
expected["token_ids"][0],
|
||||
)
|
||||
|
||||
|
||||
def test_batched_generation(llm):
|
||||
@@ -94,49 +138,34 @@ def test_batched_generation(llm):
|
||||
with open(fixture_path) as f:
|
||||
expected = json.load(f)
|
||||
|
||||
items = [
|
||||
{
|
||||
"audio_url": "https://huggingface.co/datasets/nvidia/AudioSkills/resolve/main/assets/dogs_barking_in_sync_with_the_music.wav",
|
||||
"question": "What is surprising about the relationship "
|
||||
"between the barking and the music?",
|
||||
"expected_idx": 0,
|
||||
},
|
||||
{
|
||||
"audio_url": "https://huggingface.co/datasets/nvidia/AudioSkills/resolve/main/assets/Ch6Ae9DT6Ko_00-04-03_00-04-31.wav",
|
||||
"question": (
|
||||
"Why is the philosopher's name mentioned in the lyrics? "
|
||||
"(A) To express a sense of nostalgia "
|
||||
"(B) To indicate that language cannot express clearly, "
|
||||
"satirizing the inversion of black and white in the world "
|
||||
"(C) To add depth and complexity to the lyrics "
|
||||
"(D) To showcase the wisdom and influence of the philosopher"
|
||||
),
|
||||
"expected_idx": 1,
|
||||
},
|
||||
]
|
||||
|
||||
conversations = []
|
||||
for item in items:
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "audio_url", "audio_url": {"url": item["audio_url"]}},
|
||||
{"type": "text", "text": item["question"]},
|
||||
],
|
||||
}
|
||||
]
|
||||
conversations.append(messages)
|
||||
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=128)
|
||||
|
||||
outputs = llm.chat(
|
||||
messages=conversations,
|
||||
messages=BATCHED_CONVERSATIONS,
|
||||
sampling_params=sampling_params,
|
||||
)
|
||||
|
||||
for i, output in enumerate(outputs):
|
||||
generated_text = output.outputs[0].text.strip()
|
||||
expected_text = expected["transcriptions"][i]
|
||||
assert_output_matches(
|
||||
output,
|
||||
expected["transcriptions"][i],
|
||||
expected["token_ids"][i],
|
||||
)
|
||||
|
||||
assert expected_text in generated_text or generated_text in expected_text
|
||||
|
||||
def test_single_and_batched_generation_match(llm):
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=128)
|
||||
|
||||
single_output = llm.chat(
|
||||
messages=SINGLE_CONVERSATION,
|
||||
sampling_params=sampling_params,
|
||||
)[0]
|
||||
batched_output = llm.chat(
|
||||
messages=BATCHED_CONVERSATIONS,
|
||||
sampling_params=sampling_params,
|
||||
)[0]
|
||||
|
||||
assert single_output.outputs[0].text == batched_output.outputs[0].text
|
||||
assert list(single_output.outputs[0].token_ids) == list(
|
||||
batched_output.outputs[0].token_ids
|
||||
)
|
||||
|
||||
@@ -545,8 +545,12 @@ VLM_TEST_SETTINGS = {
|
||||
auto_cls=AutoModelForImageTextToText,
|
||||
),
|
||||
"isaac": VLMTestInfo(
|
||||
# NOTE: PerceptronAI/Isaac-0.1 removed because the upstream HF
|
||||
# repo has a stale model.safetensors.index.json that references
|
||||
# shard files which no longer exist (consolidated into a single
|
||||
# model.safetensors on 2026-03-20). Re-add once upstream fixes
|
||||
# the index file.
|
||||
models=[
|
||||
"PerceptronAI/Isaac-0.1",
|
||||
"PerceptronAI/Isaac-0.2-2B-Preview",
|
||||
],
|
||||
test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.models.registry import HF_EXAMPLE_MODELS
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
MODEL_NAME = "nvidia/music-flamingo-2601-hf"
|
||||
SINGLE_CONVERSATION = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Describe this track in full detail - tell me the "
|
||||
"genre, tempo, and key, then dive into the instruments, "
|
||||
"production style, and overall mood it creates.",
|
||||
},
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {
|
||||
"url": "https://huggingface.co/datasets/nvidia/AudioSkills/"
|
||||
"resolve/main/assets/song_1.mp3",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
BATCHED_CONVERSATIONS = [
|
||||
SINGLE_CONVERSATION,
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Generate a structured lyric sheet from the input music.",
|
||||
},
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {
|
||||
"url": "https://huggingface.co/datasets/nvidia/"
|
||||
"AudioSkills/resolve/main/assets/song_2.mp3",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
]
|
||||
|
||||
|
||||
def get_fixture_path(filename):
|
||||
return os.path.join(
|
||||
os.path.dirname(__file__), "../../fixtures/musicflamingo", filename
|
||||
)
|
||||
|
||||
|
||||
def assert_output_matches(output, expected_text, expected_token_ids):
|
||||
generated = output.outputs[0]
|
||||
assert generated.text == expected_text
|
||||
actual_token_ids = list(generated.token_ids)
|
||||
assert (
|
||||
actual_token_ids == expected_token_ids
|
||||
or actual_token_ids == expected_token_ids[:-1]
|
||||
or actual_token_ids[:-1] == expected_token_ids
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def llm():
|
||||
model_info = HF_EXAMPLE_MODELS.get_hf_info("MusicFlamingoForConditionalGeneration")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
try:
|
||||
return LLM(
|
||||
model=MODEL_NAME,
|
||||
dtype="bfloat16",
|
||||
enforce_eager=True,
|
||||
max_model_len=8192,
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
)
|
||||
except Exception as e:
|
||||
pytest.skip(f"Failed to load model {MODEL_NAME}: {e}")
|
||||
|
||||
|
||||
def test_single_generation(llm):
|
||||
fixture_path = get_fixture_path("expected_results_single.json")
|
||||
if not os.path.exists(fixture_path):
|
||||
pytest.skip(f"Fixture not found: {fixture_path}")
|
||||
|
||||
with open(fixture_path) as f:
|
||||
expected = json.load(f)
|
||||
|
||||
outputs = llm.chat(
|
||||
messages=SINGLE_CONVERSATION,
|
||||
sampling_params=SamplingParams(temperature=0.0, max_tokens=50),
|
||||
)
|
||||
|
||||
assert_output_matches(
|
||||
outputs[0],
|
||||
expected["transcriptions"][0],
|
||||
expected["token_ids"][0],
|
||||
)
|
||||
|
||||
|
||||
def test_batched_generation(llm):
|
||||
fixture_path = get_fixture_path("expected_results_batched.json")
|
||||
if not os.path.exists(fixture_path):
|
||||
pytest.skip(f"Fixture not found: {fixture_path}")
|
||||
|
||||
with open(fixture_path) as f:
|
||||
expected = json.load(f)
|
||||
|
||||
outputs = llm.chat(
|
||||
messages=BATCHED_CONVERSATIONS,
|
||||
sampling_params=SamplingParams(temperature=0.0, max_tokens=50),
|
||||
)
|
||||
|
||||
for i, output in enumerate(outputs):
|
||||
assert_output_matches(
|
||||
output,
|
||||
expected["transcriptions"][i],
|
||||
expected["token_ids"][i],
|
||||
)
|
||||
|
||||
|
||||
def test_single_and_batched_generation_match(llm):
|
||||
sampling_params = SamplingParams(temperature=0.0, max_tokens=50)
|
||||
|
||||
single_output = llm.chat(
|
||||
messages=SINGLE_CONVERSATION,
|
||||
sampling_params=sampling_params,
|
||||
)[0]
|
||||
batched_output = llm.chat(
|
||||
messages=BATCHED_CONVERSATIONS,
|
||||
sampling_params=sampling_params,
|
||||
)[0]
|
||||
|
||||
assert single_output.outputs[0].text == batched_output.outputs[0].text
|
||||
assert list(single_output.outputs[0].token_ids) == list(
|
||||
batched_output.outputs[0].token_ids
|
||||
)
|
||||
@@ -40,6 +40,7 @@ class MockAudioFlamingo3Processor:
|
||||
def __init__(self):
|
||||
self.audio_token = "<sound>"
|
||||
self.audio_token_id = 12345
|
||||
self.max_audio_len = 60
|
||||
self.feature_extractor = MockFeatureExtractor()
|
||||
|
||||
def __call__(self, text=None, audios=None, **kwargs):
|
||||
@@ -65,7 +66,6 @@ def mock_ctx():
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def check_transformers_version():
|
||||
# Check if the model is supported by the current transformers version
|
||||
model_info = HF_EXAMPLE_MODELS.get_hf_info("AudioFlamingo3ForConditionalGeneration")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
@@ -84,7 +84,7 @@ def test_audio_chunk_counting(mock_ctx):
|
||||
|
||||
sr = 16000
|
||||
audio_1 = np.zeros(30 * sr)
|
||||
audio_2 = np.zeros(45 * sr)
|
||||
audio_2 = np.zeros(75 * sr)
|
||||
|
||||
mm_data = {"audio": [audio_1, audio_2]}
|
||||
prompt = "<|user|>Listen.<|end|>"
|
||||
@@ -121,5 +121,107 @@ def test_dummy_data_generation(mock_ctx):
|
||||
assert "audio" in dummy_data
|
||||
assert len(dummy_data["audio"]) == 2
|
||||
|
||||
expected_len = 600 * 16000
|
||||
expected_len = 60 * 16000
|
||||
assert len(dummy_data["audio"][0]) == expected_len
|
||||
|
||||
|
||||
def test_audio_token_count_matches_hf_processor_math():
|
||||
from vllm.model_executor.models.audioflamingo3 import (
|
||||
_count_audio_tokens_from_mask,
|
||||
)
|
||||
|
||||
feature_attention_mask = torch.zeros((3, 3000), dtype=torch.long)
|
||||
feature_attention_mask[0, :2999] = 1
|
||||
feature_attention_mask[1, :2999] = 1
|
||||
feature_attention_mask[2, :1500] = 1
|
||||
chunk_counts = torch.tensor([2, 1], dtype=torch.long)
|
||||
|
||||
assert (
|
||||
_count_audio_tokens_from_mask(feature_attention_mask, chunk_counts, 0) == 1499
|
||||
)
|
||||
assert _count_audio_tokens_from_mask(feature_attention_mask, chunk_counts, 1) == 375
|
||||
|
||||
|
||||
def test_audio_feature_pipeline_matches_hf_small_config():
|
||||
from transformers.models.audioflamingo3 import (
|
||||
modeling_audioflamingo3 as hf_audioflamingo3_modeling,
|
||||
)
|
||||
from transformers.models.audioflamingo3.configuration_audioflamingo3 import (
|
||||
AudioFlamingo3Config,
|
||||
)
|
||||
|
||||
from vllm.model_executor.models.audioflamingo3 import (
|
||||
AudioFlamingo3Encoder,
|
||||
AudioFlamingo3MultiModalProjector,
|
||||
_build_audio_encoder_attention_mask,
|
||||
_flatten_valid_audio_embeddings,
|
||||
)
|
||||
|
||||
text_config = {
|
||||
"model_type": "qwen2",
|
||||
"intermediate_size": 64,
|
||||
"initializer_range": 0.02,
|
||||
"hidden_size": 32,
|
||||
"max_position_embeddings": 1024,
|
||||
"num_hidden_layers": 2,
|
||||
"num_attention_heads": 4,
|
||||
"num_key_value_heads": 2,
|
||||
"vocab_size": 128,
|
||||
"pad_token_id": 1,
|
||||
"use_mrope": False,
|
||||
}
|
||||
audio_config = {
|
||||
"hidden_size": 16,
|
||||
"num_attention_heads": 4,
|
||||
"intermediate_size": 32,
|
||||
"num_hidden_layers": 2,
|
||||
"num_mel_bins": 80,
|
||||
"max_source_positions": 1500,
|
||||
"dropout": 0.0,
|
||||
"attention_dropout": 0.0,
|
||||
"activation_dropout": 0.0,
|
||||
"encoder_layerdrop": 0.0,
|
||||
}
|
||||
|
||||
torch.manual_seed(0)
|
||||
config = AudioFlamingo3Config(
|
||||
text_config=text_config,
|
||||
audio_config=audio_config,
|
||||
audio_token_id=0,
|
||||
)
|
||||
hf_model = hf_audioflamingo3_modeling.AudioFlamingo3ForConditionalGeneration(
|
||||
config
|
||||
).eval()
|
||||
|
||||
vllm_encoder = AudioFlamingo3Encoder(config.audio_config).eval()
|
||||
vllm_encoder.load_state_dict(hf_model.audio_tower.state_dict())
|
||||
|
||||
vllm_projector = AudioFlamingo3MultiModalProjector(config).eval()
|
||||
vllm_projector.load_state_dict(hf_model.multi_modal_projector.state_dict())
|
||||
|
||||
input_features = torch.randn(3, 80, 3000)
|
||||
feature_attention_mask = torch.zeros(3, 3000, dtype=torch.bool)
|
||||
feature_attention_mask[0, :3000] = True
|
||||
feature_attention_mask[1, :2500] = True
|
||||
feature_attention_mask[2, :1500] = True
|
||||
|
||||
hf_output = hf_model.get_audio_features(
|
||||
input_features,
|
||||
feature_attention_mask,
|
||||
return_dict=True,
|
||||
).pooler_output
|
||||
vllm_attention_mask = _build_audio_encoder_attention_mask(
|
||||
feature_attention_mask,
|
||||
dtype=vllm_encoder.conv1.weight.dtype,
|
||||
device=vllm_encoder.conv1.weight.device,
|
||||
)
|
||||
vllm_hidden_states = vllm_encoder(
|
||||
input_features,
|
||||
attention_mask=vllm_attention_mask,
|
||||
)
|
||||
vllm_output, _ = _flatten_valid_audio_embeddings(
|
||||
vllm_projector(vllm_hidden_states),
|
||||
feature_attention_mask,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(vllm_output, hf_output)
|
||||
|
||||
@@ -0,0 +1,222 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Copyright 2026 The vLLM team.
|
||||
# Copyright 2026 NVIDIA CORPORATION and the HuggingFace Inc. team. All rights
|
||||
# reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from tests.models.registry import HF_EXAMPLE_MODELS
|
||||
|
||||
|
||||
class MockMusicFlamingoConfig(PretrainedConfig):
|
||||
model_type = "musicflamingo"
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.audio_config = PretrainedConfig()
|
||||
self.text_config = PretrainedConfig()
|
||||
|
||||
|
||||
class MockMusicFlamingoProcessor:
|
||||
def __init__(self):
|
||||
self.audio_token = "<sound>"
|
||||
self.audio_token_id = 12345
|
||||
self.audio_bos_token = "<|sound_bos|>"
|
||||
self.audio_bos_token_id = 12346
|
||||
self.audio_eos_token = "<|sound_eos|>"
|
||||
self.audio_eos_token_id = 12347
|
||||
self.max_audio_len = 1200
|
||||
self.feature_extractor = MockFeatureExtractor()
|
||||
|
||||
|
||||
class MockFeatureExtractor:
|
||||
def __init__(self):
|
||||
self.sampling_rate = 16000
|
||||
self.chunk_length = 30
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_ctx():
|
||||
config = MockMusicFlamingoConfig()
|
||||
|
||||
ctx = MagicMock()
|
||||
ctx.get_hf_config.return_value = config
|
||||
ctx.get_hf_processor.return_value = MockMusicFlamingoProcessor()
|
||||
ctx.model_config.hf_config = config
|
||||
return ctx
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def check_transformers_version():
|
||||
model_info = HF_EXAMPLE_MODELS.get_hf_info("MusicFlamingoForConditionalGeneration")
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
|
||||
def test_musicflamingo_chunk_counting_uses_rote_timestamps(mock_ctx, monkeypatch):
|
||||
from vllm.model_executor.models.musicflamingo import (
|
||||
MusicFlamingoDummyInputsBuilder,
|
||||
MusicFlamingoMultiModalProcessor,
|
||||
MusicFlamingoProcessingInfo,
|
||||
)
|
||||
|
||||
info = MusicFlamingoProcessingInfo(mock_ctx)
|
||||
processor = MusicFlamingoMultiModalProcessor(
|
||||
info, MusicFlamingoDummyInputsBuilder(info)
|
||||
)
|
||||
|
||||
sr = 16000
|
||||
audio_1 = np.zeros(30 * sr)
|
||||
audio_2 = np.zeros(45 * sr)
|
||||
|
||||
mm_data = {"audio": [audio_1, audio_2]}
|
||||
prompt = "<|user|>Listen.<|end|>"
|
||||
|
||||
from vllm.multimodal.processing import BaseMultiModalProcessor
|
||||
|
||||
def mock_base_call(self, prompt, mm_data, mm_kwargs, tok_kwargs):
|
||||
del self, prompt, mm_data, mm_kwargs, tok_kwargs
|
||||
return {
|
||||
"input_ids": [1, 2, 3],
|
||||
"input_features": torch.randn(3, 80, 3000),
|
||||
"rote_timestamps": torch.randn(3, 750),
|
||||
}
|
||||
|
||||
monkeypatch.setattr(BaseMultiModalProcessor, "_call_hf_processor", mock_base_call)
|
||||
|
||||
processed = processor._call_hf_processor(prompt, mm_data, {}, {})
|
||||
|
||||
chunk_counts = processed["chunk_counts"]
|
||||
|
||||
assert chunk_counts.tolist() == [1, 2]
|
||||
assert "rote_timestamps" in processed
|
||||
|
||||
|
||||
def test_musicflamingo_dummy_text_uses_plain_audio_tokens(mock_ctx):
|
||||
from vllm.model_executor.models.musicflamingo import (
|
||||
MusicFlamingoDummyInputsBuilder,
|
||||
MusicFlamingoProcessingInfo,
|
||||
)
|
||||
|
||||
info = MusicFlamingoProcessingInfo(mock_ctx)
|
||||
builder = MusicFlamingoDummyInputsBuilder(info)
|
||||
|
||||
assert builder.get_dummy_text({"audio": 2}) == "<sound><sound>"
|
||||
|
||||
|
||||
def test_musicflamingo_audio_feature_pipeline_matches_hf_small_config():
|
||||
from transformers.models.musicflamingo import (
|
||||
modeling_musicflamingo as hf_musicflamingo_modeling,
|
||||
)
|
||||
from transformers.models.musicflamingo.configuration_musicflamingo import (
|
||||
MusicFlamingoConfig,
|
||||
)
|
||||
|
||||
from vllm.model_executor.models.audioflamingo3 import (
|
||||
_build_audio_encoder_attention_mask,
|
||||
_flatten_valid_audio_embeddings,
|
||||
)
|
||||
from vllm.model_executor.models.musicflamingo import (
|
||||
MusicFlamingoEncoder,
|
||||
MusicFlamingoMultiModalProjector,
|
||||
MusicFlamingoRotaryEmbedding,
|
||||
apply_rotary_time_emb,
|
||||
)
|
||||
|
||||
text_config = {
|
||||
"model_type": "qwen2",
|
||||
"intermediate_size": 64,
|
||||
"initializer_range": 0.02,
|
||||
"hidden_size": 32,
|
||||
"max_position_embeddings": 1024,
|
||||
"num_hidden_layers": 2,
|
||||
"num_attention_heads": 4,
|
||||
"num_key_value_heads": 2,
|
||||
"vocab_size": 128,
|
||||
"pad_token_id": 1,
|
||||
"use_mrope": False,
|
||||
}
|
||||
audio_config = {
|
||||
"hidden_size": 16,
|
||||
"num_attention_heads": 4,
|
||||
"intermediate_size": 32,
|
||||
"num_hidden_layers": 2,
|
||||
"num_mel_bins": 80,
|
||||
"max_source_positions": 1500,
|
||||
"dropout": 0.0,
|
||||
"attention_dropout": 0.0,
|
||||
"activation_dropout": 0.0,
|
||||
"encoder_layerdrop": 0.0,
|
||||
}
|
||||
|
||||
torch.manual_seed(0)
|
||||
config = MusicFlamingoConfig(
|
||||
text_config=text_config,
|
||||
audio_config=audio_config,
|
||||
audio_token_id=0,
|
||||
head_dim=8,
|
||||
rope_parameters={"rope_type": "default", "rope_theta": 2048},
|
||||
)
|
||||
hf_model = hf_musicflamingo_modeling.MusicFlamingoForConditionalGeneration(
|
||||
config
|
||||
).eval()
|
||||
|
||||
vllm_encoder = MusicFlamingoEncoder(config.audio_config).eval()
|
||||
vllm_encoder.load_state_dict(hf_model.audio_tower.state_dict())
|
||||
|
||||
vllm_projector = MusicFlamingoMultiModalProjector(config).eval()
|
||||
vllm_projector.load_state_dict(hf_model.multi_modal_projector.state_dict())
|
||||
|
||||
vllm_rope = MusicFlamingoRotaryEmbedding(config).eval()
|
||||
vllm_rope.load_state_dict(hf_model.pos_emb.state_dict(), strict=False)
|
||||
|
||||
input_features = torch.randn(3, 80, 3000)
|
||||
feature_attention_mask = torch.zeros(3, 3000, dtype=torch.bool)
|
||||
feature_attention_mask[0, :3000] = True
|
||||
feature_attention_mask[1, :2500] = True
|
||||
feature_attention_mask[2, :1500] = True
|
||||
rote_timestamps = (
|
||||
torch.arange(750, dtype=torch.float32).unsqueeze(0).repeat(3, 1) * 0.04
|
||||
)
|
||||
|
||||
hf_output = hf_model.get_audio_features(
|
||||
input_features,
|
||||
feature_attention_mask,
|
||||
rote_timestamps=rote_timestamps,
|
||||
return_dict=True,
|
||||
).pooler_output
|
||||
vllm_attention_mask = _build_audio_encoder_attention_mask(
|
||||
feature_attention_mask,
|
||||
dtype=vllm_encoder.conv1.weight.dtype,
|
||||
device=vllm_encoder.conv1.weight.device,
|
||||
)
|
||||
vllm_hidden_states = vllm_encoder(
|
||||
input_features,
|
||||
attention_mask=vllm_attention_mask,
|
||||
)
|
||||
cos, sin = vllm_rope(rote_timestamps, seq_len=vllm_hidden_states.shape[-2])
|
||||
vllm_hidden_states = apply_rotary_time_emb(vllm_hidden_states, cos, sin)
|
||||
vllm_output, _ = _flatten_valid_audio_embeddings(
|
||||
vllm_projector(vllm_hidden_states),
|
||||
feature_attention_mask,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(vllm_output, hf_output)
|
||||
@@ -752,7 +752,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
"nvidia/audio-flamingo-3-hf", min_transformers_version="5.0.0"
|
||||
),
|
||||
"MusicFlamingoForConditionalGeneration": _HfExamplesInfo(
|
||||
"nvidia/music-flamingo-2601-hf", min_transformers_version="5.0.0.dev"
|
||||
"nvidia/music-flamingo-2601-hf", min_transformers_version="5.3.0"
|
||||
),
|
||||
"AyaVisionForConditionalGeneration": _HfExamplesInfo("CohereLabs/aya-vision-8b"),
|
||||
"BagelForConditionalGeneration": _HfExamplesInfo("ByteDance-Seed/BAGEL-7B-MoT"),
|
||||
|
||||
+6
-19
@@ -65,21 +65,10 @@ class BgeM3SparseEmbeddingsProcessor(
|
||||
f"Unsupported task {raw_embed_request}, "
|
||||
f"Supported tasks are {EMBED_TASKS}"
|
||||
)
|
||||
has_dense_embed = True
|
||||
if raw_embed_request.embed_task == "dense":
|
||||
params.task = "embed"
|
||||
params.skip_reading_prefix_cache = False
|
||||
elif raw_embed_request.embed_task == "sparse":
|
||||
params.task = "token_classify"
|
||||
has_dense_embed = False
|
||||
else:
|
||||
params.task = "embed&token_classify"
|
||||
params.task = "embed&token_classify"
|
||||
params.use_activation = raw_embed_request.use_activation
|
||||
if params.use_activation is None:
|
||||
params.use_activation = True
|
||||
if not has_dense_embed:
|
||||
params.dimensions = None
|
||||
return params
|
||||
|
||||
params.dimensions = raw_embed_request.dimensions
|
||||
|
||||
@@ -170,13 +159,11 @@ class BgeM3SparseEmbeddingsProcessor(
|
||||
raw_request = self._get_sparse_embedding_request(request_id)
|
||||
has_dense_embed = raw_request.embed_task in ["dense", "dense&sparse"]
|
||||
has_sparse_embed = raw_request.embed_task in ["sparse", "dense&sparse"]
|
||||
embed_dimensions = 0
|
||||
if has_dense_embed:
|
||||
embed_dimensions = (
|
||||
self.embed_dimensions
|
||||
if raw_request.dimensions is None
|
||||
else raw_request.dimensions
|
||||
)
|
||||
embed_dimensions = (
|
||||
self.embed_dimensions
|
||||
if raw_request.dimensions is None
|
||||
else raw_request.dimensions
|
||||
)
|
||||
for idx in range(len(model_output)):
|
||||
mo = model_output[idx]
|
||||
sparse_embedding_dict: dict[int, float] = {}
|
||||
|
||||
@@ -12,7 +12,6 @@ from compressed_tensors.quantization import QuantizationType
|
||||
from tests.models.utils import check_logprobs_close
|
||||
from vllm.model_executor.layers.fused_moe import UnquantizedFusedMoEMethod
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
|
||||
CompressedTensors24,
|
||||
CompressedTensorsLinearMethod,
|
||||
CompressedTensorsW4A4Fp4,
|
||||
CompressedTensorsW4A8Fp8,
|
||||
@@ -27,9 +26,6 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8
|
||||
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
|
||||
cutlass_fp4_supported,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
|
||||
sparse_cutlass_supported,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
|
||||
|
||||
@@ -362,283 +358,6 @@ def test_compressed_tensors_kv_cache_fp8_per_attn_head(vllm_runner):
|
||||
assert output
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not sparse_cutlass_supported(),
|
||||
reason="Sparse FP8 is not yet supported on this GPU type.",
|
||||
)
|
||||
def _test_2of4_quant_models(qkv_proj, weight_strategy, input_strategy, format="dense"):
|
||||
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
|
||||
assert isinstance(qkv_proj.scheme, CompressedTensors24)
|
||||
|
||||
assert qkv_proj.scheme.weight_quant.strategy == weight_strategy
|
||||
assert qkv_proj.scheme.input_quant.strategy == input_strategy
|
||||
assert qkv_proj.scheme.quantized
|
||||
assert qkv_proj.quant_method.quantization_config.sparsity_scheme_map
|
||||
sparsity_map = qkv_proj.quant_method.quantization_config.sparsity_scheme_map # noqa: E501
|
||||
assert sparsity_map.get("Linear").format == format
|
||||
assert sparsity_map.get("Linear").sparsity_structure == "2:4"
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda() or not current_platform.has_device_capability(90),
|
||||
reason="Sparse FP8 is not yet supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"args_2of4",
|
||||
[
|
||||
(
|
||||
"nm-testing/Meta-Llama-3-8B-Instruct-FP8-Dynamic-2of4-testing",
|
||||
"channel",
|
||||
"token",
|
||||
),
|
||||
(
|
||||
"nm-testing/Meta-Llama-3-8B-Instruct-FP8-Static-Per-Tensor-testing",
|
||||
"channel",
|
||||
"tensor",
|
||||
),
|
||||
(
|
||||
"nm-testing/Meta-Llama-3-8B-Instruct-FP8-Static-testing",
|
||||
"tensor",
|
||||
"tensor",
|
||||
),
|
||||
(
|
||||
"nm-testing/Meta-Llama-3-8B-Instruct-FP8-Dynamic-IA-Per-Tensor-Weight-testing",
|
||||
"tensor",
|
||||
"token",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_compressed_tensors_2of4_quant_fp8(vllm_runner, args_2of4):
|
||||
model, weight_strategy, input_strategy = args_2of4
|
||||
with vllm_runner(model, enforce_eager=True) as llm:
|
||||
|
||||
def check_model(model):
|
||||
layer = model.model.layers[0]
|
||||
|
||||
qkv_proj = layer.self_attn.qkv_proj
|
||||
assert qkv_proj.scheme.weights_dtype == torch.float8_e4m3fn
|
||||
_test_2of4_quant_models(qkv_proj, weight_strategy, input_strategy)
|
||||
|
||||
llm.apply_model(check_model)
|
||||
|
||||
output = llm.generate_greedy("Hello my name is", max_tokens=4)
|
||||
print(output)
|
||||
assert output
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda() or not current_platform.has_device_capability(90),
|
||||
reason="Sparse FP8 is not yet supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"args_2of4",
|
||||
[
|
||||
(
|
||||
"nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_per_tok_dyn_act_fp8-BitM",
|
||||
"channel",
|
||||
"token",
|
||||
),
|
||||
(
|
||||
"nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_tensor_act_fp8-BitM",
|
||||
"channel",
|
||||
"tensor",
|
||||
),
|
||||
(
|
||||
"nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_per_tok_dyn_act_fp8-BitM",
|
||||
"tensor",
|
||||
"token",
|
||||
),
|
||||
(
|
||||
"nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_tensor_act_fp8-BitM",
|
||||
"tensor",
|
||||
"tensor",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_compressed_tensors_2of4_quant_fp8_compressed(vllm_runner, args_2of4):
|
||||
model, weight_strategy, input_strategy = args_2of4
|
||||
with vllm_runner(model, enforce_eager=True) as llm:
|
||||
|
||||
def check_model(model):
|
||||
layer = model.model.layers[0]
|
||||
|
||||
qkv_proj = layer.self_attn.qkv_proj
|
||||
assert qkv_proj.scheme.weights_dtype == torch.float8_e4m3fn
|
||||
_test_2of4_quant_models(
|
||||
qkv_proj,
|
||||
weight_strategy,
|
||||
input_strategy,
|
||||
format="sparse-24-bitmask",
|
||||
)
|
||||
|
||||
llm.apply_model(check_model)
|
||||
|
||||
output = llm.generate_greedy("Hello my name is", max_tokens=4)
|
||||
print(output)
|
||||
assert output
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not sparse_cutlass_supported(),
|
||||
reason="cutlass is not yet supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"args_2of4",
|
||||
[
|
||||
(
|
||||
"nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_per_tok_dyn_act_int8-BitM",
|
||||
"channel",
|
||||
"token",
|
||||
),
|
||||
(
|
||||
"nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_tensor_act_int8-BitM",
|
||||
"channel",
|
||||
"tensor",
|
||||
),
|
||||
(
|
||||
"nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_per_tok_dyn_act_int8-BitM",
|
||||
"tensor",
|
||||
"token",
|
||||
),
|
||||
(
|
||||
"nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_tensor_act_int8-BitM",
|
||||
"tensor",
|
||||
"tensor",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_compressed_tensors_2of4_quant_int8_compressed(vllm_runner, args_2of4):
|
||||
model, weight_strategy, input_strategy = args_2of4
|
||||
with vllm_runner(model, enforce_eager=True) as llm:
|
||||
|
||||
def check_model(model):
|
||||
layer = model.model.layers[0]
|
||||
|
||||
qkv_proj = layer.self_attn.qkv_proj
|
||||
assert qkv_proj.scheme.weights_dtype == torch.int8
|
||||
_test_2of4_quant_models(
|
||||
qkv_proj,
|
||||
weight_strategy,
|
||||
input_strategy,
|
||||
format="sparse-24-bitmask",
|
||||
)
|
||||
|
||||
llm.apply_model(check_model)
|
||||
|
||||
output = llm.generate_greedy("Hello my name is", max_tokens=4)
|
||||
print(output)
|
||||
assert output
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not sparse_cutlass_supported(),
|
||||
reason="Sparse FP8 is not yet supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"args_2of4",
|
||||
[
|
||||
(
|
||||
"nm-testing/TinyLlama-1.1B-Chat-v1.0-INT8-Dynamic-IA-Per-Channel-Weight-testing",
|
||||
"channel",
|
||||
"token",
|
||||
),
|
||||
(
|
||||
"nm-testing/TinyLlama-1.1B-Chat-v1.0-INT8-Static-testing",
|
||||
"tensor",
|
||||
"tensor",
|
||||
),
|
||||
(
|
||||
"nm-testing/TinyLlama-1.1B-Chat-v1.0-INT8-Dynamic-IA-Per-Tensor-Weight-testing",
|
||||
"tensor",
|
||||
"token",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_compressed_tensors_2of4_quant_int8(vllm_runner, args_2of4):
|
||||
model, weight_strategy, input_strategy = args_2of4
|
||||
with vllm_runner(model, enforce_eager=True) as llm:
|
||||
|
||||
def check_model(model):
|
||||
layer = model.model.layers[0]
|
||||
|
||||
qkv_proj = layer.self_attn.qkv_proj
|
||||
assert qkv_proj.scheme.weights_dtype == torch.int8
|
||||
_test_2of4_quant_models(qkv_proj, weight_strategy, input_strategy)
|
||||
|
||||
llm.apply_model(check_model)
|
||||
|
||||
output = llm.generate_greedy("Hello my name is", max_tokens=4)
|
||||
print(output)
|
||||
assert output
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not sparse_cutlass_supported(),
|
||||
reason="2of4 Sparse is not yet supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"args_2of4",
|
||||
[("nm-testing/TinyLlama-1.1B-Chat-v1.0-2of4-Sparse-Dense-Compressor")],
|
||||
)
|
||||
def test_compressed_tensors_2of4_sparse(vllm_runner, args_2of4):
|
||||
model = args_2of4
|
||||
with vllm_runner(model, enforce_eager=True) as llm:
|
||||
|
||||
def check_model(model):
|
||||
layer = model.model.layers[0]
|
||||
|
||||
qkv_proj = layer.self_attn.qkv_proj
|
||||
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
|
||||
assert isinstance(qkv_proj.scheme, CompressedTensors24)
|
||||
|
||||
assert qkv_proj.scheme.weight_quant is None
|
||||
assert qkv_proj.scheme.input_quant is None
|
||||
assert not qkv_proj.scheme.quantized
|
||||
assert qkv_proj.quant_method.quantization_config.sparsity_scheme_map
|
||||
sparsity_map = qkv_proj.quant_method.quantization_config.sparsity_scheme_map # noqa: E501
|
||||
assert sparsity_map.get("Linear").format == "dense"
|
||||
assert sparsity_map.get("Linear").sparsity_structure == "2:4"
|
||||
|
||||
llm.apply_model(check_model)
|
||||
|
||||
output = llm.generate_greedy("Hello my name is", max_tokens=4)
|
||||
print(output)
|
||||
assert output
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not sparse_cutlass_supported(),
|
||||
reason="Cutlass is not yet supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"args_2of4", [("nm-testing/llama2.c-stories42M-pruned2.4-compressed")]
|
||||
)
|
||||
def test_compressed_tensors_2of4_sparse_compressed(vllm_runner, args_2of4):
|
||||
model = args_2of4
|
||||
with vllm_runner(model, enforce_eager=True) as llm:
|
||||
|
||||
def check_model(model):
|
||||
layer = model.model.layers[0]
|
||||
|
||||
qkv_proj = layer.self_attn.qkv_proj
|
||||
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
|
||||
assert isinstance(qkv_proj.scheme, CompressedTensors24)
|
||||
|
||||
assert qkv_proj.scheme.weight_quant is None
|
||||
assert qkv_proj.scheme.input_quant is None
|
||||
assert not qkv_proj.scheme.quantized
|
||||
assert qkv_proj.quant_method.quantization_config.sparsity_scheme_map
|
||||
sparsity_map = qkv_proj.quant_method.quantization_config.sparsity_scheme_map # noqa: E501
|
||||
assert sparsity_map.get("Linear").format == "sparse-24-bitmask"
|
||||
assert sparsity_map.get("Linear").sparsity_structure == "2:4"
|
||||
|
||||
llm.apply_model(check_model)
|
||||
|
||||
output = llm.generate_greedy("Hello my name is", max_tokens=4)
|
||||
print(output)
|
||||
assert output
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
|
||||
)
|
||||
|
||||
@@ -1,57 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests whether PTPC w8a8 FP8 computation is enabled correctly.
|
||||
|
||||
Run `pytest tests/quantization/test_ptpc_fp8.py --forked`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.quantization.utils import is_quant_method_supported
|
||||
from vllm.model_executor.layers.quantization.fp8 import Fp8KVCacheMethod
|
||||
from vllm.model_executor.layers.quantization.ptpc_fp8 import PTPCFp8LinearMethod
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
@pytest.fixture(scope="function", autouse=True)
|
||||
def enable_pickle(monkeypatch):
|
||||
"""`LLM.apply_model` requires pickling a function."""
|
||||
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_quant_method_supported("ptpc_fp8"),
|
||||
reason="PTPC FP8 is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.skipif(not current_platform.is_rocm(), reason="This test is for ROCm GPU.")
|
||||
@pytest.mark.parametrize("dtype", ["bfloat16"])
|
||||
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
|
||||
def test_ptpc_fp8_rocm(vllm_runner, dtype: str, kv_cache_dtype: str) -> None:
|
||||
llm = vllm_runner(
|
||||
"facebook/opt-125m",
|
||||
dtype=dtype,
|
||||
quantization="ptpc_fp8",
|
||||
enforce_eager=True,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
allow_deprecated_quantization=True,
|
||||
)
|
||||
|
||||
with llm:
|
||||
|
||||
def check_model(model):
|
||||
fc1 = model.model.decoder.layers[0].fc1
|
||||
assert isinstance(fc1.quant_method, PTPCFp8LinearMethod)
|
||||
if kv_cache_dtype == "ptpc_fp8":
|
||||
attn = model.model.decoder.layers[0].self_attn.attn
|
||||
assert isinstance(attn.quant_method, Fp8KVCacheMethod)
|
||||
assert attn._k_scale == 1.0
|
||||
assert attn._v_scale == 1.0
|
||||
|
||||
# For GPUs with hardware support, we keep weights in fp8
|
||||
if current_platform.has_device_capability(94):
|
||||
assert fc1.weight.dtype == current_platform.fp8_dtype()
|
||||
|
||||
llm.apply_model(check_model)
|
||||
|
||||
output = llm.generate_greedy("Hello my name is", max_tokens=4)
|
||||
assert output
|
||||
@@ -12,6 +12,7 @@ import gc
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.utils import large_gpu_mark
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
@@ -32,10 +33,21 @@ def test_duplicated_ignored_sequence_group():
|
||||
assert len(prompts) == len(outputs)
|
||||
|
||||
|
||||
def test_max_tokens_none():
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
pytest.param(
|
||||
"distilbert/distilgpt2",
|
||||
marks=[
|
||||
*([large_gpu_mark(min_gb=80)] if current_platform.is_rocm() else []),
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_max_tokens_none(model):
|
||||
sampling_params = SamplingParams(temperature=0.01, top_p=0.1, max_tokens=None)
|
||||
llm = LLM(
|
||||
model="distilbert/distilgpt2",
|
||||
model=model,
|
||||
max_num_batched_tokens=4096,
|
||||
tensor_parallel_size=1,
|
||||
)
|
||||
|
||||
+9
@@ -4,6 +4,7 @@
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from tests.tool_parsers.utils import (
|
||||
run_tool_extraction,
|
||||
@@ -13,6 +14,14 @@ from vllm.entrypoints.openai.engine.protocol import FunctionCall
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.tool_parsers import ToolParser, ToolParserManager
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def default_tokenizer() -> TokenizerLike:
|
||||
"""Override module-scoped default_tokenizer because gigachat tests
|
||||
mutate the tokenizer via ``add_tokens``."""
|
||||
return AutoTokenizer.from_pretrained("gpt2")
|
||||
|
||||
|
||||
MSG_SEP_TOKEN = "<|message_sep|>\n\n"
|
||||
ROLE_SEP_TOKEN = "<|role_sep|>\n"
|
||||
EOS_TOKEN = "</s>"
|
||||
@@ -0,0 +1,147 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import json
|
||||
import random
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
from vllm.entrypoints.openai.engine.protocol import (
|
||||
DeltaMessage,
|
||||
)
|
||||
from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser
|
||||
|
||||
MODEL = "ibm-granite/granite-4.0-h-tiny"
|
||||
|
||||
|
||||
def create_complex_input(create_string_args: bool):
|
||||
coord_arg: dict | str = {
|
||||
"coordinates": [[23.54, 43.1], [-12.2, 54.3], [4, 5]],
|
||||
"coordinate_type": "latlong",
|
||||
}
|
||||
if create_string_args:
|
||||
# test granite behavior
|
||||
coord_arg = json.dumps(coord_arg)
|
||||
return [
|
||||
{"name": "find_bbox", "arguments": coord_arg},
|
||||
{
|
||||
"name": "get_stock_price",
|
||||
"arguments": {
|
||||
"symbol": "AAPL",
|
||||
"start_date": "2021-01-01",
|
||||
"end_date": "2021-12-31",
|
||||
},
|
||||
},
|
||||
{"name": "find_bbox", "arguments": coord_arg},
|
||||
]
|
||||
|
||||
|
||||
def random_chunks(s: str, min_len: int, max_len: int):
|
||||
chunks = []
|
||||
i = 0
|
||||
n = len(s)
|
||||
|
||||
while i < n:
|
||||
size = random.randint(min_len, max_len)
|
||||
chunks.append(s[i : i + size])
|
||||
i += size
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def tokenizer():
|
||||
return AutoTokenizer.from_pretrained(MODEL)
|
||||
|
||||
|
||||
# create a variety of input chunk sizes
|
||||
@pytest.mark.parametrize(
|
||||
"min_chunk, max_chunk",
|
||||
[
|
||||
(1, 1),
|
||||
(1, 2),
|
||||
(5, 7),
|
||||
(6, 20),
|
||||
],
|
||||
)
|
||||
def test_tool_call_parser_complex(min_chunk: int, max_chunk: int, tokenizer):
|
||||
input_dicts = create_complex_input(True)
|
||||
|
||||
formatted_tcs = [
|
||||
"<tool_call> " + json.dumps(call) + " </tool_call>" for call in input_dicts
|
||||
]
|
||||
|
||||
text_messages = [
|
||||
"Here goes the bbox call: \n",
|
||||
" Now the stock price call: \n ",
|
||||
" Now another bbox call: \n ",
|
||||
" See? I'm a helpful assistant.",
|
||||
]
|
||||
|
||||
test_input = (
|
||||
text_messages[0]
|
||||
+ formatted_tcs[0]
|
||||
+ text_messages[1]
|
||||
+ formatted_tcs[1]
|
||||
+ text_messages[2]
|
||||
+ formatted_tcs[2]
|
||||
+ text_messages[3]
|
||||
)
|
||||
|
||||
any_chat_request = ChatCompletionRequest(
|
||||
seed=42,
|
||||
model=MODEL,
|
||||
messages=[],
|
||||
)
|
||||
|
||||
parser = Granite4ToolParser(tokenizer=tokenizer)
|
||||
|
||||
delta_messages = list[DeltaMessage]()
|
||||
for text in random_chunks(test_input, min_chunk, max_chunk):
|
||||
delta = parser.extract_tool_calls_streaming(
|
||||
previous_text="",
|
||||
current_text="",
|
||||
delta_text=text,
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[],
|
||||
delta_token_ids=[],
|
||||
request=any_chat_request,
|
||||
)
|
||||
if delta is not None:
|
||||
delta_messages.append(delta)
|
||||
|
||||
content = ""
|
||||
tool_calls = list[dict[str, Any]]()
|
||||
|
||||
current_name = "__start__"
|
||||
current_args = ""
|
||||
|
||||
for msg in delta_messages:
|
||||
if msg.content:
|
||||
content += msg.content
|
||||
for tool_call in msg.tool_calls:
|
||||
if delta_func := tool_call.function:
|
||||
if delta_func.name is not None:
|
||||
if current_name == "__start__":
|
||||
current_name = delta_func.name
|
||||
|
||||
if delta_func.name != current_name:
|
||||
tool_calls.append(
|
||||
{
|
||||
"name": current_name,
|
||||
"arguments": json.loads(current_args),
|
||||
}
|
||||
)
|
||||
current_name = delta_func.name
|
||||
current_args = ""
|
||||
|
||||
if delta_func.arguments:
|
||||
current_args += delta_func.arguments
|
||||
|
||||
if current_name != "__start__":
|
||||
tool_calls.append({"name": current_name, "arguments": json.loads(current_args)})
|
||||
|
||||
assert content == "".join(text_messages)
|
||||
assert tool_calls == create_complex_input(False)
|
||||
@@ -0,0 +1,220 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.tool_parsers.abstract_tool_parser import ToolParser
|
||||
from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser
|
||||
from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
|
||||
|
||||
CONFIGS = {
|
||||
"llama": {
|
||||
"tool_parser": Hermes2ProToolParser,
|
||||
},
|
||||
"granite4": {
|
||||
"tool_parser": Granite4ToolParser,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def qwen_tokenizer() -> TokenizerLike:
|
||||
from vllm.tokenizers import get_tokenizer
|
||||
|
||||
return get_tokenizer("Qwen/Qwen3-32B")
|
||||
|
||||
|
||||
@pytest.fixture(params=CONFIGS.keys())
|
||||
def hermes_parser(request, qwen_tokenizer: TokenizerLike) -> ToolParser:
|
||||
config = CONFIGS[request.param]
|
||||
return config["tool_parser"](qwen_tokenizer)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def any_chat_request() -> ChatCompletionRequest:
|
||||
return ChatCompletionRequest(
|
||||
seed=42,
|
||||
model="Qwen/Qwen3-32B",
|
||||
messages=[],
|
||||
)
|
||||
|
||||
|
||||
def test_hermes_parser_streaming_just_forward_text(
|
||||
qwen_tokenizer: TokenizerLike,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """This is some prior text that has nothing to do with tool calling."""
|
||||
tokens = qwen_tokenizer.encode(text)
|
||||
previous_text = ""
|
||||
delta_messages = []
|
||||
for token in tokens:
|
||||
delta_text = qwen_tokenizer.decode([token])
|
||||
current_text = previous_text + delta_text
|
||||
delta = hermes_parser.extract_tool_calls_streaming(
|
||||
previous_text=previous_text,
|
||||
current_text=current_text,
|
||||
delta_text=delta_text,
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[],
|
||||
delta_token_ids=[],
|
||||
request=any_chat_request,
|
||||
)
|
||||
previous_text = current_text
|
||||
delta_messages.append(delta)
|
||||
|
||||
for delta in delta_messages:
|
||||
assert delta is not None
|
||||
assert not delta.tool_calls
|
||||
|
||||
print(delta_messages)
|
||||
assert "".join([delta.content for delta in delta_messages]) == text
|
||||
|
||||
|
||||
def test_hermes_parser_streaming_failure_case_bug_19056(
|
||||
qwen_tokenizer: TokenizerLike,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """<tool_call>
|
||||
{"name": "final_answer", "arguments": {"trigger": true}}
|
||||
</tool_call>"""
|
||||
tokens = qwen_tokenizer.encode(text)
|
||||
previous_text = ""
|
||||
delta_messages = []
|
||||
for token in tokens:
|
||||
text = qwen_tokenizer.decode([token])
|
||||
current_text = previous_text + text
|
||||
delta = hermes_parser.extract_tool_calls_streaming(
|
||||
previous_text=previous_text,
|
||||
current_text=current_text,
|
||||
delta_text=text,
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[],
|
||||
delta_token_ids=[],
|
||||
request=any_chat_request,
|
||||
)
|
||||
previous_text = current_text
|
||||
if delta is not None:
|
||||
delta_messages.append(delta)
|
||||
|
||||
assert delta_messages[0].tool_calls[0].function.name == "final_answer"
|
||||
tool_call_args = "".join(
|
||||
delta.tool_calls[0].function.arguments or "" for delta in delta_messages
|
||||
)
|
||||
assert tool_call_args == '{"trigger": true}'
|
||||
|
||||
|
||||
def test_hermes_parser_streaming(
|
||||
qwen_tokenizer: TokenizerLike,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = '<tool_call>\
|
||||
{"name": "get_current_temperature",\
|
||||
"arguments": {"location":\
|
||||
"San Francisco, California, United States", "unit": "celsius"}}\
|
||||
</tool_call>'
|
||||
|
||||
tokens = qwen_tokenizer.encode(text)
|
||||
previous_text = ""
|
||||
delta_messages = []
|
||||
for token in tokens:
|
||||
text = qwen_tokenizer.decode([token])
|
||||
current_text = previous_text + text
|
||||
delta = hermes_parser.extract_tool_calls_streaming(
|
||||
previous_text=previous_text,
|
||||
current_text=current_text,
|
||||
delta_text=text,
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[],
|
||||
delta_token_ids=[],
|
||||
request=any_chat_request,
|
||||
)
|
||||
previous_text = current_text
|
||||
if delta is not None:
|
||||
delta_messages.append(delta)
|
||||
print(delta_messages)
|
||||
assert delta_messages[0].tool_calls[0].function.name == "get_current_temperature"
|
||||
# load to normalize whitespace
|
||||
tool_call_args = json.loads(
|
||||
"".join(
|
||||
delta.tool_calls[0].function.arguments or "" for delta in delta_messages
|
||||
)
|
||||
)
|
||||
assert tool_call_args == {
|
||||
"location": "San Francisco, California, United States",
|
||||
"unit": "celsius",
|
||||
}
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_no_tool_call(
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """This is not a tool call."""
|
||||
tool_call = hermes_parser.extract_tool_calls(
|
||||
model_output=text,
|
||||
request=any_chat_request,
|
||||
)
|
||||
|
||||
assert tool_call is not None
|
||||
assert not tool_call.tools_called
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_tool_call_between_tags(
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """<tool_call>
|
||||
{"name": "final_answer", "arguments": {"trigger": true}}
|
||||
</tool_call>"""
|
||||
tool_call = hermes_parser.extract_tool_calls(
|
||||
model_output=text,
|
||||
request=any_chat_request,
|
||||
)
|
||||
|
||||
assert tool_call is not None
|
||||
assert tool_call.tools_called
|
||||
assert tool_call.tool_calls[0].function.name == "final_answer"
|
||||
assert tool_call.tool_calls[0].function.arguments == '{"trigger": true}'
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_tool_call_until_eos(
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
if isinstance(hermes_parser, Granite4ToolParser):
|
||||
pytest.skip(reason="The Granite4 tool parser enforces a complete response")
|
||||
|
||||
text = """<tool_call>
|
||||
{"name": "final_answer", "arguments": {"trigger": true}}"""
|
||||
tool_call = hermes_parser.extract_tool_calls(
|
||||
model_output=text,
|
||||
request=any_chat_request,
|
||||
)
|
||||
|
||||
assert tool_call is not None
|
||||
assert tool_call.tools_called
|
||||
assert tool_call.tool_calls[0].function.name == "final_answer"
|
||||
assert tool_call.tool_calls[0].function.arguments == '{"trigger": true}'
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_tool_call_invalid_json(
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
# Missing closing brace to trigger exception
|
||||
text = """<tool_call>
|
||||
{"name": "final_answer", "arguments": {"trigger": true}"""
|
||||
tool_call = hermes_parser.extract_tool_calls(
|
||||
model_output=text,
|
||||
request=any_chat_request,
|
||||
)
|
||||
|
||||
assert tool_call is not None
|
||||
assert not tool_call.tools_called
|
||||
@@ -0,0 +1,451 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for EncoderCudaGraphManager.
|
||||
|
||||
Test organization:
|
||||
No GPU required:
|
||||
- TestFindBudgetGraph — greedy budget selection logic
|
||||
- TestGetCumulativeStats — hit/miss rate statistics
|
||||
GPU required:
|
||||
- TestEncoderCudaGraphCaptureReplay — capture, replay, fallback, counters, chunking
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.worker.gpu.mm.encoder_cudagraph import (
|
||||
EncoderCudaGraphManager,
|
||||
)
|
||||
from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import (
|
||||
EncoderCudaGraphCaptureInputs,
|
||||
EncoderCudaGraphConfig,
|
||||
EncoderCudaGraphReplayBuffers,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_manager_with_budgets(budgets: list[int]) -> EncoderCudaGraphManager:
|
||||
"""Create a minimal EncoderCudaGraphManager with only token_budgets set.
|
||||
|
||||
Skips the parts of __init__ that require a real VllmConfig / model
|
||||
by patching the attributes directly after construction.
|
||||
"""
|
||||
mgr = object.__new__(EncoderCudaGraphManager)
|
||||
mgr.token_budgets = sorted(budgets)
|
||||
mgr.max_batch_size = 16
|
||||
mgr.use_dp = False
|
||||
mgr.budget_graphs = {}
|
||||
mgr.graph_hits = 0
|
||||
mgr.graph_misses = 0
|
||||
mgr.log_stats_interval = 100
|
||||
return mgr
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _generate_budgets
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestGenerateBudgets:
|
||||
"""Auto-generate power-of-2 budgets from min to max."""
|
||||
|
||||
def test_exact_powers_of_2(self):
|
||||
result = EncoderCudaGraphManager._generate_budgets(64, 1024)
|
||||
assert result == [64, 128, 256, 512, 1024]
|
||||
|
||||
def test_max_not_power_of_2(self):
|
||||
result = EncoderCudaGraphManager._generate_budgets(64, 800)
|
||||
assert result == [64, 128, 256, 512, 800]
|
||||
|
||||
def test_min_equals_max(self):
|
||||
result = EncoderCudaGraphManager._generate_budgets(64, 64)
|
||||
assert result == [64]
|
||||
|
||||
def test_large_range(self):
|
||||
result = EncoderCudaGraphManager._generate_budgets(64, 8192)
|
||||
assert result == [64, 128, 256, 512, 1024, 2048, 4096, 8192]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _find_smallest_fitting_budget_given_tokens
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFindBudgetGraph:
|
||||
"""Budget greedy selection: smallest budget >= total_tokens."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"total_tokens,budgets,expected",
|
||||
[
|
||||
# Exact match
|
||||
(2048, [2048, 4096, 8192], 2048),
|
||||
# Below smallest budget — picks smallest
|
||||
(100, [2048, 4096, 8192], 2048),
|
||||
# Zero tokens — picks smallest
|
||||
(0, [2048, 4096, 8192], 2048),
|
||||
# Between budgets — picks next one up
|
||||
(2049, [2048, 4096, 8192], 4096),
|
||||
(4097, [2048, 4096, 8192], 8192),
|
||||
# Exceeds all budgets — returns None (eager fallback)
|
||||
(9000, [2048, 4096, 8192], None),
|
||||
# Single budget, fits
|
||||
(1000, [2048], 2048),
|
||||
# Single budget, does not fit
|
||||
(3000, [2048], None),
|
||||
],
|
||||
)
|
||||
def test_find_budget(self, total_tokens, budgets, expected):
|
||||
mgr = _make_manager_with_budgets(budgets)
|
||||
result = mgr._find_smallest_fitting_budget_given_tokens(total_tokens)
|
||||
assert result == expected
|
||||
|
||||
def test_budgets_are_sorted(self):
|
||||
"""Manager always sorts budgets ascending at init."""
|
||||
mgr = _make_manager_with_budgets([8192, 2048, 4096])
|
||||
assert mgr.token_budgets == [2048, 4096, 8192]
|
||||
# Budget selection still works correctly after sorting
|
||||
assert mgr._find_smallest_fitting_budget_given_tokens(3000) == 4096
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_cumulative_stats
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestGetCumulativeStats:
|
||||
"""Statistics tracking and reporting."""
|
||||
|
||||
def test_initial_stats_are_zero(self):
|
||||
mgr = _make_manager_with_budgets([2048])
|
||||
stats = mgr.get_cumulative_stats()
|
||||
assert stats["graph_hits"] == 0
|
||||
assert stats["graph_misses"] == 0
|
||||
assert stats["hit_rate"] == 0.0
|
||||
|
||||
def test_hit_rate_calculation(self):
|
||||
mgr = _make_manager_with_budgets([2048])
|
||||
mgr.graph_hits = 75
|
||||
mgr.graph_misses = 25
|
||||
stats = mgr.get_cumulative_stats()
|
||||
assert stats["graph_hits"] == 75
|
||||
assert stats["graph_misses"] == 25
|
||||
assert stats["hit_rate"] == pytest.approx(0.75)
|
||||
|
||||
def test_all_hits(self):
|
||||
mgr = _make_manager_with_budgets([2048])
|
||||
mgr.graph_hits = 100
|
||||
mgr.graph_misses = 0
|
||||
assert mgr.get_cumulative_stats()["hit_rate"] == pytest.approx(1.0)
|
||||
|
||||
def test_all_misses(self):
|
||||
mgr = _make_manager_with_budgets([2048])
|
||||
mgr.graph_hits = 0
|
||||
mgr.graph_misses = 50
|
||||
assert mgr.get_cumulative_stats()["hit_rate"] == pytest.approx(0.0)
|
||||
|
||||
def test_stats_report_budget_info(self):
|
||||
budgets = [2048, 4096, 8192]
|
||||
mgr = _make_manager_with_budgets(budgets)
|
||||
stats = mgr.get_cumulative_stats()
|
||||
assert stats["num_budgets"] == 0 # no graphs captured yet
|
||||
assert stats["token_budgets"] == budgets
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GPU fixtures and helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Mock encoder parameters (kept small for fast capture)
|
||||
_SPATIAL_MERGE = 2
|
||||
_HIDDEN = 32
|
||||
_PATCH_SIZE = 4 # H/W per patch in grid_thw units
|
||||
_TEMPORAL_PATCH = 1
|
||||
_IN_CHANNELS = 3
|
||||
# flattened_patch_size = in_channels * temporal_patch * patch_size^2
|
||||
_FLAT = _IN_CHANNELS * _TEMPORAL_PATCH * _PATCH_SIZE * _PATCH_SIZE # 48
|
||||
|
||||
# Test budgets: small to keep capture fast
|
||||
_BUDGETS = [16, 64]
|
||||
_MAX_BATCH = 4
|
||||
|
||||
|
||||
def _count_input_patches(grid_thw_list: list[list[int]]) -> int:
|
||||
return sum(t * h * w for t, h, w in grid_thw_list)
|
||||
|
||||
|
||||
def _count_output_tokens(
|
||||
grid_thw_list: list[list[int]], spatial_merge_size: int
|
||||
) -> int:
|
||||
m = spatial_merge_size
|
||||
return sum(t * (h // m) * (w // m) for t, h, w in grid_thw_list)
|
||||
|
||||
|
||||
class SimpleMockViTModel(torch.nn.Module):
|
||||
"""Minimal ViT model for CUDA graph tests.
|
||||
|
||||
Implements the SupportsEncoderCudaGraph protocol by providing
|
||||
all required methods. The forward pass projects patches and
|
||||
simulates spatial merge by averaging groups of m^2 patches.
|
||||
"""
|
||||
|
||||
supports_encoder_cudagraph = True
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.proj = torch.nn.Linear(_FLAT, _HIDDEN)
|
||||
self.spatial_merge_size = _SPATIAL_MERGE
|
||||
self.out_hidden_size = _HIDDEN
|
||||
|
||||
def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
|
||||
return EncoderCudaGraphConfig(
|
||||
modalities=["image"],
|
||||
input_key="pixel_values",
|
||||
buffer_keys=["dummy_buf"],
|
||||
out_hidden_size=_HIDDEN,
|
||||
)
|
||||
|
||||
def get_encoder_cudagraph_budget_range(
|
||||
self,
|
||||
vllm_config,
|
||||
) -> tuple[int, int]:
|
||||
# For tests: min=4, max=128 (small values for fast capture)
|
||||
return (4, 128)
|
||||
|
||||
def get_encoder_cudagraph_num_items(
|
||||
self,
|
||||
mm_kwargs: dict[str, Any],
|
||||
) -> int:
|
||||
return len(mm_kwargs["image_grid_thw"])
|
||||
|
||||
def get_encoder_cudagraph_per_item_output_tokens(
|
||||
self,
|
||||
mm_kwargs: dict[str, Any],
|
||||
) -> list[int]:
|
||||
m = _SPATIAL_MERGE
|
||||
return [t * (h // m) * (w // m) for t, h, w in mm_kwargs["image_grid_thw"]]
|
||||
|
||||
def get_encoder_cudagraph_per_item_input_sizes(
|
||||
self,
|
||||
mm_kwargs: dict[str, Any],
|
||||
) -> list[int]:
|
||||
return [t * h * w for t, h, w in mm_kwargs["image_grid_thw"]]
|
||||
|
||||
def select_encoder_cudagraph_items(
|
||||
self,
|
||||
mm_kwargs: dict[str, Any],
|
||||
indices: list[int],
|
||||
) -> dict[str, Any]:
|
||||
grid_thw = mm_kwargs["image_grid_thw"]
|
||||
pixel_values = mm_kwargs["pixel_values"]
|
||||
|
||||
if len(indices) == 0:
|
||||
return {
|
||||
"pixel_values": pixel_values[:0],
|
||||
"image_grid_thw": [],
|
||||
}
|
||||
|
||||
patches_per_item = [t * h * w for t, h, w in grid_thw]
|
||||
cum_patches = [0]
|
||||
for p in patches_per_item:
|
||||
cum_patches.append(cum_patches[-1] + p)
|
||||
|
||||
selected_pv = torch.cat(
|
||||
[pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices]
|
||||
)
|
||||
selected_grid = [grid_thw[i] for i in indices]
|
||||
return {
|
||||
"pixel_values": selected_pv,
|
||||
"image_grid_thw": selected_grid,
|
||||
}
|
||||
|
||||
def prepare_encoder_cudagraph_capture_inputs(
|
||||
self,
|
||||
token_budget: int,
|
||||
max_batch_size: int,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
) -> EncoderCudaGraphCaptureInputs:
|
||||
per_image_output = token_budget // max_batch_size
|
||||
grid_config = [
|
||||
[1, _SPATIAL_MERGE, per_image_output * _SPATIAL_MERGE]
|
||||
for _ in range(max_batch_size)
|
||||
]
|
||||
total_patches = _count_input_patches(grid_config)
|
||||
dummy_pixel_values = torch.randn(
|
||||
total_patches, _FLAT, device=device, dtype=dtype
|
||||
)
|
||||
n_out = _count_output_tokens(grid_config, _SPATIAL_MERGE)
|
||||
dummy_buf = torch.zeros(n_out, _HIDDEN, device=device, dtype=dtype)
|
||||
return EncoderCudaGraphCaptureInputs(
|
||||
mm_kwargs={
|
||||
"pixel_values": dummy_pixel_values,
|
||||
"image_grid_thw": grid_config,
|
||||
},
|
||||
buffers={"dummy_buf": dummy_buf},
|
||||
)
|
||||
|
||||
def prepare_encoder_cudagraph_replay_buffers(
|
||||
self,
|
||||
mm_kwargs: dict[str, Any],
|
||||
max_batch_size: int,
|
||||
) -> EncoderCudaGraphReplayBuffers:
|
||||
grid_thw = mm_kwargs["image_grid_thw"]
|
||||
n_out = _count_output_tokens(grid_thw, _SPATIAL_MERGE)
|
||||
p = next(self.parameters())
|
||||
dummy_buf = torch.zeros(n_out, _HIDDEN, device=p.device, dtype=p.dtype)
|
||||
return EncoderCudaGraphReplayBuffers(buffers={"dummy_buf": dummy_buf})
|
||||
|
||||
def encoder_cudagraph_forward(
|
||||
self,
|
||||
mm_kwargs: dict[str, Any],
|
||||
buffers: dict[str, torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
return self._forward(mm_kwargs["pixel_values"])
|
||||
|
||||
def encoder_eager_forward(
|
||||
self,
|
||||
mm_kwargs: dict[str, Any],
|
||||
) -> torch.Tensor:
|
||||
return self._forward(mm_kwargs["pixel_values"])
|
||||
|
||||
def _forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
|
||||
m2 = _SPATIAL_MERGE**2
|
||||
out = self.proj(pixel_values)
|
||||
n_out = out.shape[0] // m2
|
||||
return out[: n_out * m2].view(n_out, m2, _HIDDEN).mean(dim=1)
|
||||
|
||||
|
||||
def _make_manager_for_gpu(
|
||||
model: SimpleMockViTModel,
|
||||
token_budgets: list[int],
|
||||
max_batch_size: int,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
) -> EncoderCudaGraphManager:
|
||||
"""Create EncoderCudaGraphManager bypassing VllmConfig for GPU tests."""
|
||||
mgr = object.__new__(EncoderCudaGraphManager)
|
||||
mgr.token_budgets = sorted(token_budgets)
|
||||
mgr.max_batch_size = max_batch_size
|
||||
mgr.use_dp = False
|
||||
mgr.budget_graphs = {}
|
||||
mgr.graph_hits = 0
|
||||
mgr.graph_misses = 0
|
||||
mgr.log_stats_interval = 100
|
||||
mgr.model = model
|
||||
mgr.config = model.get_encoder_cudagraph_config()
|
||||
mgr.device = device
|
||||
mgr.dtype = dtype
|
||||
return mgr
|
||||
|
||||
|
||||
def _make_pixel_values(
|
||||
grid_thw_list: list[list[int]],
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Random pixel_values matching the total input patch count."""
|
||||
n = _count_input_patches(grid_thw_list)
|
||||
return torch.randn(n, _FLAT, device=device, dtype=dtype)
|
||||
|
||||
|
||||
def _make_mm_kwargs(
|
||||
grid_thw_list: list[list[int]],
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
) -> dict[str, Any]:
|
||||
"""Create mm_kwargs for testing."""
|
||||
return {
|
||||
"pixel_values": _make_pixel_values(grid_thw_list, device, dtype),
|
||||
"image_grid_thw": grid_thw_list,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GPU tests — capture, replay, fallback, counters, chunking
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
|
||||
class TestEncoderCudaGraphCaptureReplay:
|
||||
def setup_method(self):
|
||||
self.device = torch.device("cuda:0")
|
||||
self.dtype = torch.float16
|
||||
self.model = SimpleMockViTModel().to(self.device).half()
|
||||
self.mgr = _make_manager_for_gpu(
|
||||
self.model, _BUDGETS, _MAX_BATCH, self.device, self.dtype
|
||||
)
|
||||
self.mgr.capture()
|
||||
|
||||
# --- capture ---
|
||||
|
||||
def test_capture_creates_one_graph_per_budget(self):
|
||||
assert len(self.mgr.budget_graphs) == len(_BUDGETS)
|
||||
assert set(self.mgr.budget_graphs.keys()) == set(_BUDGETS)
|
||||
|
||||
# --- output shape ---
|
||||
|
||||
def test_execute_returns_one_tensor_per_image(self):
|
||||
grid_thw = [[1, 4, 4], [1, 4, 4]]
|
||||
mm_kwargs = _make_mm_kwargs(grid_thw, self.device, self.dtype)
|
||||
result = self.mgr.execute(mm_kwargs)
|
||||
assert result is not None
|
||||
assert len(result) == 2
|
||||
|
||||
def test_execute_output_tokens_per_image(self):
|
||||
# [1,4,4] → 1*(4//2)*(4//2) = 4 tokens; [1,8,8] → 16 tokens
|
||||
grid_thw = [[1, 4, 4], [1, 8, 8]]
|
||||
mm_kwargs = _make_mm_kwargs(grid_thw, self.device, self.dtype)
|
||||
result = self.mgr.execute(mm_kwargs)
|
||||
assert result is not None
|
||||
assert result[0].shape == (4, _HIDDEN)
|
||||
assert result[1].shape == (16, _HIDDEN)
|
||||
|
||||
# --- budget fallback ---
|
||||
|
||||
def test_eager_fallback_when_tokens_exceed_all_budgets(self):
|
||||
# [1,18,18] → 1*(18//2)*(18//2) = 81 tokens > max budget 64.
|
||||
# Greedy packing handles the fallback internally: the oversized image
|
||||
# gets an eager forward pass and is returned as part of the output list
|
||||
# (execute() no longer returns None for individual image misses).
|
||||
grid_thw = [[1, 18, 18]]
|
||||
mm_kwargs = _make_mm_kwargs(grid_thw, self.device, self.dtype)
|
||||
result = self.mgr.execute(mm_kwargs)
|
||||
assert result is not None
|
||||
assert len(result) == 1
|
||||
# Eager output: SimpleMockViTModel produces n_out = 81 tokens
|
||||
assert result[0].shape == (81, _HIDDEN)
|
||||
assert self.mgr.graph_misses == 1
|
||||
|
||||
# --- counters ---
|
||||
|
||||
def test_hit_counter_increments_by_num_images(self):
|
||||
grid_thw = [[1, 4, 4], [1, 4, 4]]
|
||||
mm_kwargs = _make_mm_kwargs(grid_thw, self.device, self.dtype)
|
||||
self.mgr.execute(mm_kwargs)
|
||||
assert self.mgr.graph_hits == 2
|
||||
|
||||
def test_miss_counter_increments_by_num_images(self):
|
||||
grid_thw = [[1, 18, 18]] # 81 tokens > 64
|
||||
mm_kwargs = _make_mm_kwargs(grid_thw, self.device, self.dtype)
|
||||
self.mgr.execute(mm_kwargs)
|
||||
assert self.mgr.graph_misses == 1
|
||||
|
||||
# --- chunking ---
|
||||
|
||||
def test_chunking_when_images_exceed_max_batch(self):
|
||||
# 8 images > max_batch_size=4 → 2 chunks of 4
|
||||
# each chunk: 4 * 4 = 16 tokens → fits budget 16
|
||||
n_images = _MAX_BATCH * 2
|
||||
grid_thw = [[1, 4, 4]] * n_images
|
||||
mm_kwargs = _make_mm_kwargs(grid_thw, self.device, self.dtype)
|
||||
result = self.mgr.execute(mm_kwargs)
|
||||
assert result is not None
|
||||
assert len(result) == n_images
|
||||
for out in result:
|
||||
assert out.shape == (4, _HIDDEN)
|
||||
@@ -2,11 +2,11 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
|
||||
import vllm.model_executor.layers.batch_invariant as batch_invariant
|
||||
import vllm.envs as envs
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def enable_batch_invariant_mode(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Automatically enable batch invariant kernel overrides for all tests."""
|
||||
monkeypatch.setattr(batch_invariant, "VLLM_BATCH_INVARIANT", True)
|
||||
monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True)
|
||||
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
|
||||
|
||||
@@ -8,14 +8,14 @@ import pytest
|
||||
import torch
|
||||
from utils import (
|
||||
BACKENDS,
|
||||
TEST_MODEL,
|
||||
_extract_step_logprobs,
|
||||
_random_prompt,
|
||||
is_device_capability_below_90,
|
||||
resolve_model_name,
|
||||
skip_unsupported,
|
||||
)
|
||||
|
||||
import vllm.model_executor.layers.batch_invariant as batch_invariant
|
||||
import vllm.envs as envs
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
IS_DEVICE_CAPABILITY_BELOW_90 = is_device_capability_below_90()
|
||||
@@ -57,7 +57,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
|
||||
attention_config = {"backend": backend}
|
||||
# Allow overrides from environment (useful for CI tuning)
|
||||
# "facebook/opt-125m" is too small, doesn't reliably test determinism
|
||||
model = resolve_model_name(backend)
|
||||
model = TEST_MODEL
|
||||
num_trials = int(os.getenv("VLLM_NEEDLE_TRIALS", "5"))
|
||||
max_batch_size = int(os.getenv("VLLM_NEEDLE_BATCH_SIZE", "128"))
|
||||
min_random_prompt = int(os.getenv("VLLM_MIN_PROMPT", "1024"))
|
||||
@@ -169,16 +169,13 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
|
||||
):
|
||||
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
|
||||
random.seed(seed)
|
||||
model_name = resolve_model_name(backend)
|
||||
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
|
||||
|
||||
# For batch invariance, disable custom all-reduce to ensure deterministic
|
||||
# all-reduce operations (custom all-reduce may not be deterministic)
|
||||
from vllm.model_executor.layers.batch_invariant import (
|
||||
vllm_is_batch_invariant,
|
||||
)
|
||||
import vllm.envs as envs
|
||||
|
||||
disable_custom_ar = vllm_is_batch_invariant()
|
||||
disable_custom_ar = envs.VLLM_BATCH_INVARIANT
|
||||
|
||||
if disable_custom_ar:
|
||||
print(f"\n{'=' * 80}")
|
||||
@@ -186,7 +183,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
|
||||
print(f"{'=' * 80}\n")
|
||||
|
||||
llm = LLM(
|
||||
model=model_name,
|
||||
model=TEST_MODEL,
|
||||
tensor_parallel_size=tp_size,
|
||||
max_num_seqs=128,
|
||||
max_model_len=8192,
|
||||
@@ -395,7 +392,7 @@ def test_simple_generation(backend):
|
||||
Simple test that runs the model with a basic prompt and prints the output.
|
||||
Useful for quick smoke testing and debugging.
|
||||
"""
|
||||
model = resolve_model_name(backend)
|
||||
model = TEST_MODEL
|
||||
|
||||
llm = LLM(
|
||||
model=model,
|
||||
@@ -455,10 +452,9 @@ def test_logprobs_without_batch_invariance_should_fail(
|
||||
"""
|
||||
# CRITICAL: Disable batch invariance for this test
|
||||
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0")
|
||||
monkeypatch.setattr(batch_invariant, "VLLM_BATCH_INVARIANT", False)
|
||||
monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", False)
|
||||
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
|
||||
random.seed(seed)
|
||||
model_name = resolve_model_name(backend)
|
||||
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
|
||||
|
||||
print(f"\n{'=' * 80}")
|
||||
@@ -466,7 +462,7 @@ def test_logprobs_without_batch_invariance_should_fail(
|
||||
print(f"{'=' * 80}\n")
|
||||
|
||||
llm = LLM(
|
||||
model=model_name,
|
||||
model=TEST_MODEL,
|
||||
tensor_parallel_size=tp_size,
|
||||
max_num_seqs=32,
|
||||
max_model_len=8192,
|
||||
@@ -674,14 +670,11 @@ def test_decode_logprobs_match_prefill_logprobs(
|
||||
"""
|
||||
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
|
||||
random.seed(seed)
|
||||
model_name = resolve_model_name(backend)
|
||||
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
|
||||
|
||||
from vllm.model_executor.layers.batch_invariant import (
|
||||
vllm_is_batch_invariant,
|
||||
)
|
||||
import vllm.envs as envs
|
||||
|
||||
disable_custom_ar = vllm_is_batch_invariant()
|
||||
disable_custom_ar = envs.VLLM_BATCH_INVARIANT
|
||||
|
||||
if disable_custom_ar:
|
||||
print(f"\n{'=' * 80}")
|
||||
@@ -689,7 +682,7 @@ def test_decode_logprobs_match_prefill_logprobs(
|
||||
print(f"{'=' * 80}\n")
|
||||
|
||||
llm = LLM(
|
||||
model=model_name,
|
||||
model=TEST_MODEL,
|
||||
tensor_parallel_size=tp_size,
|
||||
max_num_seqs=32,
|
||||
max_model_len=8192,
|
||||
|
||||
@@ -17,7 +17,7 @@ from typing import Any
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
from utils import BACKENDS, _random_prompt, resolve_model_name, skip_unsupported
|
||||
from utils import BACKENDS, TEST_MODEL, _random_prompt, skip_unsupported
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
@@ -139,7 +139,6 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
|
||||
backend: str,
|
||||
) -> None:
|
||||
random.seed(int(os.getenv("VLLM_TEST_SEED", "12345")))
|
||||
model_name = resolve_model_name(backend)
|
||||
prompts_all = [_random_prompt(10, 50) for _ in range(32)]
|
||||
|
||||
sp_kwargs: dict[str, Any] = {
|
||||
@@ -159,11 +158,11 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
|
||||
if tp_size:
|
||||
server_args += ["-tp", tp_size]
|
||||
|
||||
with RemoteOpenAIServer(model_name, server_args) as server:
|
||||
with RemoteOpenAIServer(TEST_MODEL, server_args) as server:
|
||||
client = server.get_client()
|
||||
_compare_bs1_vs_bsn_single_process(
|
||||
prompts=prompts_all,
|
||||
sp_kwargs=sp_kwargs,
|
||||
client=client,
|
||||
model_name=model_name,
|
||||
model_name=TEST_MODEL,
|
||||
)
|
||||
|
||||
@@ -7,6 +7,10 @@ import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.transformers_utils.config import get_config
|
||||
from vllm.transformers_utils.model_arch_config_convertor import (
|
||||
ModelArchConfigConvertorBase,
|
||||
)
|
||||
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
|
||||
|
||||
skip_unsupported = pytest.mark.skipif(
|
||||
@@ -16,10 +20,12 @@ skip_unsupported = pytest.mark.skipif(
|
||||
reason="Requires CUDA and >= Ampere (SM80)",
|
||||
)
|
||||
|
||||
DEFAULT_MODEL = "Qwen/Qwen3-1.7B"
|
||||
TEST_MODEL = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL)
|
||||
|
||||
BACKENDS: list[str] = [
|
||||
"FLASH_ATTN",
|
||||
"TRITON_ATTN",
|
||||
"TRITON_MLA",
|
||||
]
|
||||
|
||||
# FlashInfer temporarily disabled due to invariant CTA sizes.
|
||||
@@ -27,19 +33,13 @@ BACKENDS: list[str] = [
|
||||
# if has_flashinfer():
|
||||
# BACKENDS.append("FLASHINFER")
|
||||
|
||||
if flash_attn_supports_mla():
|
||||
BACKENDS.append("FLASH_ATTN_MLA")
|
||||
|
||||
DEFAULT_MODEL = "Qwen/Qwen3-1.7B"
|
||||
MLA_MODEL = "deepseek-ai/DeepSeek-V2-Lite-Chat"
|
||||
|
||||
|
||||
def resolve_model_name(backend: str) -> str:
|
||||
"""Resolve the model name for the given backend."""
|
||||
model = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL)
|
||||
if backend.endswith("MLA") and model == DEFAULT_MODEL:
|
||||
return MLA_MODEL
|
||||
return model
|
||||
# only run MLA backends when the requested test model is itself an MLA model.
|
||||
if os.getenv("VLLM_TEST_MODEL"):
|
||||
config = get_config(TEST_MODEL, trust_remote_code=False)
|
||||
if ModelArchConfigConvertorBase(config, config.get_text_config()).is_deepseek_mla():
|
||||
BACKENDS = ["TRITON_MLA"]
|
||||
if flash_attn_supports_mla():
|
||||
BACKENDS.append("FLASH_ATTN_MLA")
|
||||
|
||||
|
||||
def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:
|
||||
|
||||
@@ -258,8 +258,8 @@ def get_fake_execute_model_fn(original_execute_model_fn: Callable):
|
||||
mamba_kv_cache_dict[
|
||||
num_computed_tokens - num_computed_tokens % BLOCK_SIZE
|
||||
] = (
|
||||
kv_cache[0][0][block_id].clone(),
|
||||
kv_cache[0][1][block_id].clone(),
|
||||
kv_cache[0][block_id].clone(),
|
||||
kv_cache[1][block_id].clone(),
|
||||
)
|
||||
|
||||
last_num_computed_tokens = num_computed_tokens
|
||||
@@ -302,7 +302,7 @@ def get_fake_process_mamba_fn(
|
||||
mamba_layer_name = kv_cache_config.kv_cache_groups[
|
||||
mamba_group_id
|
||||
].layer_names[0]
|
||||
mamba_kv_cache = forward_context[mamba_layer_name].kv_cache[0][-1]
|
||||
mamba_kv_cache = forward_context[mamba_layer_name].kv_cache[-1]
|
||||
mamba_block_table = input_batch.block_table.block_tables[
|
||||
mamba_group_id
|
||||
].block_table.cpu[0]
|
||||
|
||||
@@ -19,9 +19,9 @@ dp_ep_configs=(
|
||||
"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1)
|
||||
)
|
||||
hybrid_ssm_configs=(
|
||||
"ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code"
|
||||
"ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code"
|
||||
# TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models.
|
||||
"ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling"
|
||||
"ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling"
|
||||
)
|
||||
|
||||
# Select config array based on DP_EP env var
|
||||
|
||||
@@ -19,6 +19,7 @@ EXPECTED_VALUES = {
|
||||
"deepseek-ai/DeepSeek-V2-Lite-Chat": 0.65,
|
||||
"google/gemma-3-4b-it": 0.74,
|
||||
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8": 0.84,
|
||||
"ibm-granite/granite-4.0-h-tiny": 0.80,
|
||||
}
|
||||
|
||||
SIMPLE_PROMPT = (
|
||||
|
||||
@@ -12,9 +12,8 @@ from vllm.v1.kv_offload.abstract import (
|
||||
OffloadingEvent,
|
||||
PrepareStoreOutput,
|
||||
)
|
||||
from vllm.v1.kv_offload.arc_manager import ARCOffloadingManager
|
||||
from vllm.v1.kv_offload.backends.cpu import CPUBackend
|
||||
from vllm.v1.kv_offload.lru_manager import LRUOffloadingManager
|
||||
from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
|
||||
from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
|
||||
from vllm.v1.kv_offload.mediums import CPULoadStoreSpec
|
||||
|
||||
|
||||
@@ -79,12 +78,12 @@ def verify_events(
|
||||
assert tuple(stores) == to_hash_sets(expected_stores)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("manager_class", [LRUOffloadingManager, ARCOffloadingManager])
|
||||
def test_already_stored_block_not_evicted_during_prepare_store(manager_class):
|
||||
@pytest.mark.parametrize("eviction_policy", ["lru", "arc"])
|
||||
def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
|
||||
"""
|
||||
Regression test: a block that is already stored must not be evicted
|
||||
by prepare_store() when it needs to make room for new blocks.
|
||||
Applies to both LRUOffloadingManager and ARCOffloadingManager.
|
||||
Applies to both lru and arc policies.
|
||||
|
||||
Scenario:
|
||||
- Store blocks [1, 2] and complete.
|
||||
@@ -96,8 +95,12 @@ def test_already_stored_block_not_evicted_during_prepare_store(manager_class):
|
||||
- After complete_store([2, 3, 4, 5]), block 2 must still be present.
|
||||
"""
|
||||
block_size = 256
|
||||
cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
|
||||
manager = manager_class(cpu_backend, enable_events=True)
|
||||
manager = CPUOffloadingManager(
|
||||
block_size=block_size,
|
||||
num_blocks=4,
|
||||
cache_policy=eviction_policy,
|
||||
enable_events=True,
|
||||
)
|
||||
|
||||
# store [1, 2] and complete
|
||||
manager.prepare_store(to_hashes([1, 2]))
|
||||
@@ -129,12 +132,13 @@ def test_already_stored_block_not_evicted_during_prepare_store(manager_class):
|
||||
|
||||
def test_cpu_manager():
|
||||
"""
|
||||
Tests LRUOffloadingManager with a CPUBackend.
|
||||
Tests CPUOffloadingManager with lru policy.
|
||||
"""
|
||||
# initialize a CPU backend with a capacity of 4 blocks
|
||||
block_size = 256
|
||||
cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
|
||||
cpu_manager = LRUOffloadingManager(cpu_backend, enable_events=True)
|
||||
cpu_manager = CPUOffloadingManager(
|
||||
block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True
|
||||
)
|
||||
|
||||
# prepare store [1, 2]
|
||||
prepare_store_output = cpu_manager.prepare_store(to_hashes([1, 2]))
|
||||
@@ -241,13 +245,15 @@ def test_cpu_manager():
|
||||
|
||||
def test_arc_manager_basic():
|
||||
"""
|
||||
Tests ARCOffloadingManager basic operations with a CPUBackend.
|
||||
Tests CPUOffloadingManager with arc policy.
|
||||
Verifies that ARC handles store, load, and lookup operations correctly.
|
||||
"""
|
||||
# initialize a CPU backend with a capacity of 4 blocks
|
||||
block_size = 256
|
||||
cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
|
||||
arc_manager = ARCOffloadingManager(cpu_backend, enable_events=True)
|
||||
arc_manager = CPUOffloadingManager(
|
||||
block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True
|
||||
)
|
||||
arc_policy = arc_manager._policy
|
||||
assert isinstance(arc_policy, ARCCachePolicy)
|
||||
|
||||
# prepare store [1, 2]
|
||||
prepare_store_output = arc_manager.prepare_store(to_hashes([1, 2]))
|
||||
@@ -278,8 +284,8 @@ def test_arc_manager_basic():
|
||||
assert arc_manager.lookup(to_hashes([1, 2, 3])) == 2
|
||||
|
||||
# blocks should be in T1 (recent)
|
||||
assert len(arc_manager.t1) == 2
|
||||
assert len(arc_manager.t2) == 0
|
||||
assert len(arc_policy.t1) == 2
|
||||
assert len(arc_policy.t2) == 0
|
||||
|
||||
|
||||
def test_arc_manager_t1_to_t2_promotion():
|
||||
@@ -288,23 +294,26 @@ def test_arc_manager_t1_to_t2_promotion():
|
||||
This is a key feature of ARC's adaptive behavior.
|
||||
"""
|
||||
block_size = 256
|
||||
cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
|
||||
arc_manager = ARCOffloadingManager(cpu_backend, enable_events=False)
|
||||
arc_manager = CPUOffloadingManager(
|
||||
block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=False
|
||||
)
|
||||
arc_policy = arc_manager._policy
|
||||
assert isinstance(arc_policy, ARCCachePolicy)
|
||||
|
||||
# store and complete block 1
|
||||
arc_manager.prepare_store(to_hashes([1]))
|
||||
arc_manager.complete_store(to_hashes([1]))
|
||||
|
||||
# block 1 starts in T1 (recent)
|
||||
assert to_hashes([1])[0] in arc_manager.t1
|
||||
assert to_hashes([1])[0] not in arc_manager.t2
|
||||
assert to_hashes([1])[0] in arc_policy.t1
|
||||
assert to_hashes([1])[0] not in arc_policy.t2
|
||||
|
||||
# touch block 1 (simulate second access)
|
||||
arc_manager.touch(to_hashes([1]))
|
||||
|
||||
# block 1 should now be in T2 (frequent)
|
||||
assert to_hashes([1])[0] not in arc_manager.t1
|
||||
assert to_hashes([1])[0] in arc_manager.t2
|
||||
assert to_hashes([1])[0] not in arc_policy.t1
|
||||
assert to_hashes([1])[0] in arc_policy.t2
|
||||
|
||||
|
||||
def test_arc_manager_eviction_with_load():
|
||||
@@ -313,8 +322,9 @@ def test_arc_manager_eviction_with_load():
|
||||
Verifies that blocks being loaded (ref_cnt > 0) cannot be evicted.
|
||||
"""
|
||||
block_size = 256
|
||||
cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
|
||||
arc_manager = ARCOffloadingManager(cpu_backend, enable_events=True)
|
||||
arc_manager = CPUOffloadingManager(
|
||||
block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True
|
||||
)
|
||||
|
||||
# prepare and complete store [1, 2, 3, 4]
|
||||
prepare_store_output = arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
|
||||
@@ -354,28 +364,31 @@ def test_arc_manager_adaptive_target():
|
||||
When a block in B2 is accessed, target_t1_size decreases.
|
||||
"""
|
||||
block_size = 256
|
||||
cpu_backend = CPUBackend(block_size=block_size, num_blocks=2)
|
||||
arc_manager = ARCOffloadingManager(cpu_backend, enable_events=False)
|
||||
arc_manager = CPUOffloadingManager(
|
||||
block_size=block_size, num_blocks=2, cache_policy="arc", enable_events=False
|
||||
)
|
||||
arc_policy = arc_manager._policy
|
||||
assert isinstance(arc_policy, ARCCachePolicy)
|
||||
|
||||
# store blocks 1, 2 (fills cache)
|
||||
arc_manager.prepare_store(to_hashes([1, 2]))
|
||||
arc_manager.complete_store(to_hashes([1, 2]))
|
||||
|
||||
initial_target = arc_manager.target_t1_size
|
||||
initial_target = arc_policy.target_t1_size
|
||||
|
||||
# store block 3, evicting block 1 (moves to B1 ghost list)
|
||||
arc_manager.prepare_store(to_hashes([3]))
|
||||
arc_manager.complete_store(to_hashes([3]))
|
||||
|
||||
# block 1 should be in B1 (ghost list)
|
||||
assert to_hashes([1])[0] in arc_manager.b1
|
||||
assert to_hashes([1])[0] in arc_policy.b1
|
||||
|
||||
# touch block 1 (cache miss, but in B1)
|
||||
# this should increase target_t1_size (favor recency)
|
||||
arc_manager.touch(to_hashes([1]))
|
||||
|
||||
# target should have increased
|
||||
assert arc_manager.target_t1_size > initial_target
|
||||
assert arc_policy.target_t1_size > initial_target
|
||||
|
||||
|
||||
def test_arc_manager_t1_t2_eviction_policy():
|
||||
@@ -384,8 +397,11 @@ def test_arc_manager_t1_t2_eviction_policy():
|
||||
If |T1| >= target_t1_size, evict from T1, otherwise from T2.
|
||||
"""
|
||||
block_size = 256
|
||||
cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
|
||||
arc_manager = ARCOffloadingManager(cpu_backend, enable_events=False)
|
||||
arc_manager = CPUOffloadingManager(
|
||||
block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=False
|
||||
)
|
||||
arc_policy = arc_manager._policy
|
||||
assert isinstance(arc_policy, ARCCachePolicy)
|
||||
|
||||
# store blocks 1, 2, 3, 4
|
||||
arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
|
||||
@@ -395,12 +411,12 @@ def test_arc_manager_t1_t2_eviction_policy():
|
||||
arc_manager.touch(to_hashes([3, 4]))
|
||||
|
||||
# now: T1 = {1, 2}, T2 = {3, 4}
|
||||
assert len(arc_manager.t1) == 2
|
||||
assert len(arc_manager.t2) == 2
|
||||
assert len(arc_policy.t1) == 2
|
||||
assert len(arc_policy.t2) == 2
|
||||
|
||||
# set target_t1_size to prefer evicting from T1
|
||||
# (when |T1| >= target, evict from T1)
|
||||
arc_manager.target_t1_size = 1
|
||||
arc_policy.target_t1_size = 1
|
||||
|
||||
# store block 5, should evict from T1 (block 1, LRU in T1)
|
||||
output = arc_manager.prepare_store(to_hashes([5]))
|
||||
@@ -410,9 +426,9 @@ def test_arc_manager_t1_t2_eviction_policy():
|
||||
arc_manager.complete_store(to_hashes([5]))
|
||||
|
||||
# block 1 should be in B1 (ghost list)
|
||||
assert to_hashes([1])[0] in arc_manager.b1
|
||||
assert to_hashes([1])[0] in arc_policy.b1
|
||||
# block 5 should be in T1
|
||||
assert to_hashes([5])[0] in arc_manager.t1
|
||||
assert to_hashes([5])[0] in arc_policy.t1
|
||||
|
||||
|
||||
def test_arc_manager_ghost_list_bounds():
|
||||
@@ -421,8 +437,11 @@ def test_arc_manager_ghost_list_bounds():
|
||||
They should be capped at cache_capacity.
|
||||
"""
|
||||
block_size = 256
|
||||
cpu_backend = CPUBackend(block_size=block_size, num_blocks=2)
|
||||
arc_manager = ARCOffloadingManager(cpu_backend, enable_events=False)
|
||||
arc_manager = CPUOffloadingManager(
|
||||
block_size=block_size, num_blocks=2, cache_policy="arc", enable_events=False
|
||||
)
|
||||
arc_policy = arc_manager._policy
|
||||
assert isinstance(arc_policy, ARCCachePolicy)
|
||||
|
||||
# fill cache with blocks 1, 2
|
||||
arc_manager.prepare_store(to_hashes([1, 2]))
|
||||
@@ -434,8 +453,8 @@ def test_arc_manager_ghost_list_bounds():
|
||||
arc_manager.complete_store(to_hashes([i]))
|
||||
|
||||
# ghost lists should not exceed cache_capacity
|
||||
assert len(arc_manager.b1) <= arc_manager.cache_capacity
|
||||
assert len(arc_manager.b2) <= arc_manager.cache_capacity
|
||||
assert len(arc_policy.b1) <= arc_policy.cache_capacity
|
||||
assert len(arc_policy.b2) <= arc_policy.cache_capacity
|
||||
|
||||
|
||||
def test_arc_manager_touch_ordering():
|
||||
@@ -444,8 +463,11 @@ def test_arc_manager_touch_ordering():
|
||||
Similar to LRU test but verifies T1/T2 ordering.
|
||||
"""
|
||||
block_size = 256
|
||||
cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
|
||||
arc_manager = ARCOffloadingManager(cpu_backend, enable_events=True)
|
||||
arc_manager = CPUOffloadingManager(
|
||||
block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True
|
||||
)
|
||||
arc_policy = arc_manager._policy
|
||||
assert isinstance(arc_policy, ARCCachePolicy)
|
||||
|
||||
# store blocks 1, 2, 3, 4
|
||||
arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
|
||||
@@ -459,8 +481,8 @@ def test_arc_manager_touch_ordering():
|
||||
arc_manager.touch(to_hashes([1, 3, 4]))
|
||||
|
||||
# T1 = {2}, T2 = {1, 3, 4} (in that order, with 4 most recent)
|
||||
assert len(arc_manager.t1) == 1
|
||||
assert len(arc_manager.t2) == 3
|
||||
assert len(arc_policy.t1) == 1
|
||||
assert len(arc_policy.t2) == 3
|
||||
|
||||
# store block 5, should evict from T1 (block 2, only one in T1)
|
||||
prepare_store_output = arc_manager.prepare_store(to_hashes([5]))
|
||||
@@ -480,8 +502,11 @@ def test_arc_manager_failed_store():
|
||||
Similar to LRU test but for ARC.
|
||||
"""
|
||||
block_size = 256
|
||||
cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
|
||||
arc_manager = ARCOffloadingManager(cpu_backend, enable_events=True)
|
||||
arc_manager = CPUOffloadingManager(
|
||||
block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True
|
||||
)
|
||||
arc_policy = arc_manager._policy
|
||||
assert isinstance(arc_policy, ARCCachePolicy)
|
||||
|
||||
# store blocks 1, 2, 3, 4
|
||||
arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
|
||||
@@ -498,12 +523,12 @@ def test_arc_manager_failed_store():
|
||||
# block 5 should not be in cache
|
||||
assert arc_manager.lookup(to_hashes([5])) == 0
|
||||
# block 5 should not be in T1 or T2
|
||||
assert to_hashes([5])[0] not in arc_manager.t1
|
||||
assert to_hashes([5])[0] not in arc_manager.t2
|
||||
assert to_hashes([5])[0] not in arc_policy.t1
|
||||
assert to_hashes([5])[0] not in arc_policy.t2
|
||||
|
||||
# evicted block should still be gone (in B1 ghost list)
|
||||
evicted_hash = prepare_store_output.block_hashes_evicted[0]
|
||||
assert evicted_hash in arc_manager.b1
|
||||
assert evicted_hash in arc_policy.b1
|
||||
|
||||
|
||||
def test_arc_manager_full_scenario():
|
||||
@@ -512,8 +537,11 @@ def test_arc_manager_full_scenario():
|
||||
Similar to the full LRU test but adapted for ARC behavior.
|
||||
"""
|
||||
block_size = 256
|
||||
cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
|
||||
arc_manager = ARCOffloadingManager(cpu_backend, enable_events=True)
|
||||
arc_manager = CPUOffloadingManager(
|
||||
block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True
|
||||
)
|
||||
arc_policy = arc_manager._policy
|
||||
assert isinstance(arc_policy, ARCCachePolicy)
|
||||
|
||||
# store [1, 2]
|
||||
arc_manager.prepare_store(to_hashes([1, 2]))
|
||||
@@ -529,8 +557,8 @@ def test_arc_manager_full_scenario():
|
||||
arc_manager.touch(to_hashes([2, 3]))
|
||||
|
||||
# T1 has {4, 5}, T2 has {2, 3}
|
||||
assert len(arc_manager.t1) == 2
|
||||
assert len(arc_manager.t2) == 2
|
||||
assert len(arc_policy.t1) == 2
|
||||
assert len(arc_policy.t2) == 2
|
||||
|
||||
# store [6] -> should evict from T1 (4 is oldest in T1)
|
||||
prepare_store_output = arc_manager.prepare_store(to_hashes([6]))
|
||||
@@ -548,11 +576,12 @@ def test_arc_manager_full_scenario():
|
||||
|
||||
def test_filter_reused_manager():
|
||||
"""
|
||||
Tests FilterReusedOffloadingManager with a CPUBackend.
|
||||
Tests FilterReusedOffloadingManager with a CPUOffloadingManager.
|
||||
"""
|
||||
block_size = 256
|
||||
cpu_backend = CPUBackend(block_size=block_size, num_blocks=4)
|
||||
lru_manager = LRUOffloadingManager(cpu_backend, enable_events=True)
|
||||
lru_manager = CPUOffloadingManager(
|
||||
block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True
|
||||
)
|
||||
|
||||
from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
|
||||
|
||||
|
||||
@@ -177,7 +177,7 @@ def test_prepare_next_token_ids():
|
||||
|
||||
next_token_ids_from_padded, valid_sampled_tokens_count = (
|
||||
proposer.prepare_next_token_ids_padded(
|
||||
common_attn_metadata,
|
||||
common_attn_metadata.seq_lens_cpu,
|
||||
sampled_token_ids_tensor,
|
||||
mock_requests,
|
||||
mock_input_batch,
|
||||
|
||||
@@ -187,7 +187,7 @@ def test_prepare_next_token_ids_padded():
|
||||
)
|
||||
|
||||
next_token_ids, valid_sampled_tokens_count = proposer.prepare_next_token_ids_padded(
|
||||
common_attn_metadata,
|
||||
common_attn_metadata.seq_lens_cpu,
|
||||
sampled_token_ids,
|
||||
mock_requests,
|
||||
mock_input_batch,
|
||||
|
||||
@@ -278,3 +278,148 @@ def test_custom_class_serialization_disallowed_without_pickle():
|
||||
with pytest.raises(TypeError):
|
||||
# Attempt to encode the custom class
|
||||
encoder.encode(obj)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestWithTensor:
|
||||
"""Mock request with non-multimodal tensor field like EngineCoreRequest."""
|
||||
|
||||
prompt_embeds: torch.Tensor | None
|
||||
data: str
|
||||
|
||||
|
||||
def test_non_multimodal_tensor_with_ipc():
|
||||
"""Test that non-multimodal tensor fields work correctly with IPC enabled.
|
||||
|
||||
This reproduces the bug where fields like prompt_embeds: torch.Tensor | None
|
||||
would fail to decode when IPC is enabled because _decode_tensor expected a
|
||||
raw tensor tuple but received a msgpack-decoded TensorIpcHandle list.
|
||||
"""
|
||||
import torch.multiprocessing as torch_mp
|
||||
|
||||
from vllm.v1.engine.tensor_ipc import TensorIpcReceiver, TensorIpcSender
|
||||
|
||||
# Create tensor queues for IPC
|
||||
tensor_queues = [torch_mp.Queue()]
|
||||
|
||||
# Create encoder with IPC sender
|
||||
sender = TensorIpcSender(tensor_queues[0])
|
||||
encoder = MsgpackEncoder(oob_tensor_consumer=sender)
|
||||
|
||||
# Create decoder with IPC receiver
|
||||
receiver = TensorIpcReceiver(tensor_queues[0])
|
||||
decoder = MsgpackDecoder(RequestWithTensor, oob_tensor_provider=receiver)
|
||||
|
||||
# Create a request with a non-multimodal tensor
|
||||
original_tensor = torch.randn(5, 10, dtype=torch.float32)
|
||||
request = RequestWithTensor(prompt_embeds=original_tensor, data="test_data")
|
||||
|
||||
# Encode the request - this should send the tensor via IPC
|
||||
encoded = encoder.encode(request)
|
||||
|
||||
# Verify encoding succeeded
|
||||
assert len(encoded) > 0
|
||||
|
||||
# Decode the request - this should retrieve the tensor from IPC queue
|
||||
# Previously this would fail because the decoder tried to unpack the
|
||||
# handle list as raw tensor bytes metadata.
|
||||
decoded = decoder.decode(encoded)
|
||||
|
||||
# Verify the decoded request matches the original
|
||||
assert isinstance(decoded, RequestWithTensor)
|
||||
assert decoded.data == "test_data"
|
||||
assert decoded.prompt_embeds is not None
|
||||
assert torch.allclose(decoded.prompt_embeds, original_tensor), (
|
||||
"Decoded tensor does not match the original tensor."
|
||||
)
|
||||
|
||||
|
||||
def test_non_multimodal_tensor_with_ipc_none_value():
|
||||
"""Test that None values for tensor fields work correctly with IPC enabled."""
|
||||
import torch.multiprocessing as torch_mp
|
||||
|
||||
from vllm.v1.engine.tensor_ipc import TensorIpcReceiver, TensorIpcSender
|
||||
|
||||
# Create tensor queues for IPC
|
||||
tensor_queues = [torch_mp.Queue()]
|
||||
|
||||
# Create encoder with IPC sender
|
||||
sender = TensorIpcSender(tensor_queues[0])
|
||||
encoder = MsgpackEncoder(oob_tensor_consumer=sender)
|
||||
|
||||
# Create decoder with IPC receiver
|
||||
receiver = TensorIpcReceiver(tensor_queues[0])
|
||||
decoder = MsgpackDecoder(RequestWithTensor, oob_tensor_provider=receiver)
|
||||
|
||||
# Create a request with None for the tensor field
|
||||
request = RequestWithTensor(prompt_embeds=None, data="test_data_with_none")
|
||||
|
||||
# Encode and decode the request
|
||||
encoded = encoder.encode(request)
|
||||
decoded = decoder.decode(encoded)
|
||||
|
||||
# Verify the decoded request matches the original
|
||||
assert isinstance(decoded, RequestWithTensor)
|
||||
assert decoded.data == "test_data_with_none"
|
||||
assert decoded.prompt_embeds is None
|
||||
|
||||
|
||||
def test_multiple_senders_single_receiver_ipc():
|
||||
"""Test N senders sharing a queue with a single receiver via msgpack.
|
||||
|
||||
Simulates the real vLLM topology where multiple API server frontends
|
||||
each have their own MsgpackEncoder + TensorIpcSender, all putting
|
||||
tensors onto the same torch.mp queue, and a single engine core
|
||||
decodes them with one MsgpackDecoder + TensorIpcReceiver.
|
||||
"""
|
||||
import torch.multiprocessing as torch_mp
|
||||
|
||||
from vllm.v1.engine.tensor_ipc import TensorIpcReceiver, TensorIpcSender
|
||||
|
||||
num_senders = 3
|
||||
num_messages_per_sender = 2
|
||||
tensor_queue = torch_mp.Queue()
|
||||
|
||||
# Create N independent senders (each gets its own uuid-based sender_id)
|
||||
senders = []
|
||||
encoders = []
|
||||
for _ in range(num_senders):
|
||||
s = TensorIpcSender(tensor_queue)
|
||||
senders.append(s)
|
||||
encoders.append(MsgpackEncoder(oob_tensor_consumer=s))
|
||||
|
||||
# Single receiver
|
||||
receiver = TensorIpcReceiver(tensor_queue)
|
||||
decoder = MsgpackDecoder(RequestWithTensor, oob_tensor_provider=receiver)
|
||||
|
||||
# Encode messages from all senders, interleaving the order
|
||||
# so that tensors from different senders land on the queue interleaved.
|
||||
encoded_payloads: list[tuple[int, int, torch.Tensor, list]] = []
|
||||
for msg_idx in range(num_messages_per_sender):
|
||||
for sender_idx in range(num_senders):
|
||||
tensor = torch.full(
|
||||
(sender_idx + 1, msg_idx + 2),
|
||||
float(sender_idx * 100 + msg_idx),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
req = RequestWithTensor(
|
||||
prompt_embeds=tensor,
|
||||
data=f"s{sender_idx}_m{msg_idx}",
|
||||
)
|
||||
encoded = encoders[sender_idx].encode(req)
|
||||
encoded_payloads.append((sender_idx, msg_idx, tensor, encoded))
|
||||
|
||||
# Decode all messages — the receiver must correctly match each
|
||||
# tensor handle to the right TensorIpcData from the shared queue.
|
||||
for sender_idx, msg_idx, original_tensor, encoded in encoded_payloads:
|
||||
decoded = decoder.decode(encoded)
|
||||
assert isinstance(decoded, RequestWithTensor)
|
||||
assert decoded.data == f"s{sender_idx}_m{msg_idx}"
|
||||
assert decoded.prompt_embeds is not None
|
||||
assert decoded.prompt_embeds.shape == original_tensor.shape, (
|
||||
f"Shape mismatch for sender {sender_idx} msg {msg_idx}: "
|
||||
f"{decoded.prompt_embeds.shape} != {original_tensor.shape}"
|
||||
)
|
||||
assert torch.allclose(decoded.prompt_embeds, original_tensor), (
|
||||
f"Value mismatch for sender {sender_idx} msg {msg_idx}"
|
||||
)
|
||||
|
||||
@@ -0,0 +1,943 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Tests for tensor IPC queue functionality."""
|
||||
|
||||
import contextlib
|
||||
import multiprocessing as mp
|
||||
from dataclasses import dataclass
|
||||
from multiprocessing.synchronize import Barrier as BarrierType
|
||||
from multiprocessing.synchronize import Event as EventType
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.multiprocessing as torch_mp
|
||||
|
||||
from vllm.v1.engine.tensor_ipc import (
|
||||
TensorIpcData,
|
||||
TensorIpcReceiver,
|
||||
TensorIpcSender,
|
||||
)
|
||||
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", autouse=True)
|
||||
def setup_multiprocessing():
|
||||
"""Set multiprocessing start method to 'spawn' for compatibility."""
|
||||
with contextlib.suppress(RuntimeError):
|
||||
# Already set, which is fine
|
||||
torch_mp.set_start_method("spawn", force=True)
|
||||
yield
|
||||
|
||||
|
||||
@dataclass
|
||||
# Use a typed container so the test covers the real vLLM path where tensor IPC
|
||||
# handles are encoded and decoded as fields nested inside larger msgpack payloads.
|
||||
class TensorEnvelope:
|
||||
tensor: torch.Tensor
|
||||
label: str
|
||||
|
||||
|
||||
def encoder_process(
|
||||
tensor_queue: torch_mp.Queue,
|
||||
payload_queue: mp.Queue,
|
||||
result_queue: mp.Queue,
|
||||
tensor_data: dict[str, Any],
|
||||
ready_event: EventType,
|
||||
retrieval_done: EventType,
|
||||
):
|
||||
"""Process that msgpack-encodes and sends tensors via IPC."""
|
||||
try:
|
||||
sender = TensorIpcSender(tensor_queue)
|
||||
encoder = MsgpackEncoder(oob_tensor_consumer=sender)
|
||||
|
||||
if torch.cuda.is_available():
|
||||
device = "cuda:0"
|
||||
tensor = torch.randn(
|
||||
*tensor_data["shape"], dtype=tensor_data["dtype"], device=device
|
||||
)
|
||||
else:
|
||||
# Fall back to CPU for testing
|
||||
device = "cpu"
|
||||
tensor = torch.randn(*tensor_data["shape"], dtype=tensor_data["dtype"])
|
||||
|
||||
message = TensorEnvelope(tensor=tensor, label="cuda-msgpack")
|
||||
encoded = encoder.encode(message)
|
||||
payload_queue.put(encoded, timeout=10.0)
|
||||
|
||||
ready_event.set()
|
||||
|
||||
result_queue.put(
|
||||
{
|
||||
"success": True,
|
||||
"encoded_length": len(encoded),
|
||||
"device": str(device),
|
||||
"tensor_shape": tuple(tensor.shape),
|
||||
}
|
||||
)
|
||||
retrieval_done.wait(timeout=30.0)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
ready_event.set()
|
||||
retrieval_done.set()
|
||||
result_queue.put(
|
||||
{"success": False, "error": str(e), "traceback": traceback.format_exc()}
|
||||
)
|
||||
|
||||
|
||||
def decoder_process(
|
||||
tensor_queue: torch_mp.Queue,
|
||||
payload_queue: mp.Queue,
|
||||
result_queue: mp.Queue,
|
||||
expected_shape: tuple,
|
||||
encoder_ready: EventType,
|
||||
retrieval_done: EventType,
|
||||
):
|
||||
"""Process that msgpack-decodes tensors received via IPC."""
|
||||
try:
|
||||
if not encoder_ready.wait(timeout=10.0):
|
||||
raise TimeoutError("Encoder did not signal ready")
|
||||
|
||||
encoded = payload_queue.get(timeout=5.0)
|
||||
receiver = TensorIpcReceiver(tensor_queue)
|
||||
decoder = MsgpackDecoder(TensorEnvelope, oob_tensor_provider=receiver)
|
||||
decoded = decoder.decode(encoded)
|
||||
|
||||
result_queue.put(
|
||||
{
|
||||
"success": True,
|
||||
"tensor_shape": tuple(decoded.tensor.shape),
|
||||
"device": str(decoded.tensor.device),
|
||||
"label": decoded.label,
|
||||
"matches_expected": tuple(decoded.tensor.shape) == expected_shape,
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
retrieval_done.set()
|
||||
result_queue.put(
|
||||
{"success": False, "error": str(e), "traceback": traceback.format_exc()}
|
||||
)
|
||||
else:
|
||||
retrieval_done.set()
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
def test_cuda_tensor_queue_basic():
|
||||
"""Test CUDA tensor IPC through the msgpack encoder/decoder path."""
|
||||
tensor_queue = torch_mp.Queue()
|
||||
payload_queue: mp.Queue = mp.Queue()
|
||||
result_queue: mp.Queue = mp.Queue()
|
||||
encoder_ready = mp.Event()
|
||||
retrieval_done = mp.Event()
|
||||
|
||||
tensor_shape = (4, 8, 16)
|
||||
tensor_dtype = torch.float32
|
||||
|
||||
encoder_proc = mp.Process(
|
||||
target=encoder_process,
|
||||
args=(
|
||||
tensor_queue,
|
||||
payload_queue,
|
||||
result_queue,
|
||||
{"shape": tensor_shape, "dtype": tensor_dtype},
|
||||
encoder_ready,
|
||||
retrieval_done,
|
||||
),
|
||||
)
|
||||
encoder_proc.start()
|
||||
|
||||
decoder_proc = mp.Process(
|
||||
target=decoder_process,
|
||||
args=(
|
||||
tensor_queue,
|
||||
payload_queue,
|
||||
result_queue,
|
||||
tensor_shape,
|
||||
encoder_ready,
|
||||
retrieval_done,
|
||||
),
|
||||
)
|
||||
decoder_proc.start()
|
||||
|
||||
encoder_result = result_queue.get(timeout=10.0)
|
||||
decoder_result = result_queue.get(timeout=10.0)
|
||||
|
||||
encoder_proc.join(timeout=5.0)
|
||||
decoder_proc.join(timeout=5.0)
|
||||
|
||||
# Verify results
|
||||
assert encoder_result["success"], (
|
||||
f"Encoder failed: {encoder_result.get('error')}\n"
|
||||
f"{encoder_result.get('traceback', '')}"
|
||||
)
|
||||
assert decoder_result["success"], (
|
||||
f"Decoder failed: {decoder_result.get('error')}\n"
|
||||
f"{decoder_result.get('traceback', '')}"
|
||||
)
|
||||
assert decoder_result["matches_expected"], "Tensor shape mismatch"
|
||||
assert "cuda" in decoder_result["device"], "Tensor not on CUDA device"
|
||||
assert decoder_result["label"] == "cuda-msgpack"
|
||||
|
||||
|
||||
def test_cpu_tensor_fallback():
|
||||
"""Test that CPU tensors use standard serialization path."""
|
||||
encoder = MsgpackEncoder()
|
||||
|
||||
# Create a CPU tensor
|
||||
tensor = torch.randn(3, 4, dtype=torch.float32)
|
||||
|
||||
# Encode the tensor (should use standard path, not queue)
|
||||
encoded = encoder.encode({"test_tensor": tensor})
|
||||
|
||||
# Verify encoding succeeded
|
||||
assert len(encoded) > 0
|
||||
assert isinstance(encoded, (list, tuple))
|
||||
|
||||
# Basic check: no queue should be used, so tensor goes through standard path
|
||||
# This is mainly to ensure no exceptions are raised
|
||||
|
||||
|
||||
def test_msgpack_encoder_decoder_with_ipc():
|
||||
"""Test the full msgpack + tensor IPC path in one process."""
|
||||
tensor_queue = torch_mp.Queue()
|
||||
sender = TensorIpcSender(tensor_queue)
|
||||
encoder = MsgpackEncoder(oob_tensor_consumer=sender)
|
||||
receiver = TensorIpcReceiver(tensor_queue)
|
||||
decoder = MsgpackDecoder(TensorEnvelope, oob_tensor_provider=receiver)
|
||||
|
||||
# Use CPU here to exercise the msgpack + sender/receiver integration
|
||||
# without relying on same-process CUDA IPC behavior.
|
||||
tensor = torch.randn(2, 3)
|
||||
|
||||
message = TensorEnvelope(tensor=tensor, label="test")
|
||||
encoded = encoder.encode(message)
|
||||
assert len(encoded) > 0
|
||||
|
||||
decoded = decoder.decode(encoded)
|
||||
assert isinstance(decoded, TensorEnvelope)
|
||||
assert decoded.label == "test"
|
||||
assert torch.allclose(decoded.tensor, tensor)
|
||||
|
||||
|
||||
def test_decoder_buffer_management():
|
||||
"""Test receiver's tensor buffer management when draining queue."""
|
||||
tensor_queue = torch_mp.Queue()
|
||||
|
||||
sender_id = "test_sender"
|
||||
message_id = 1
|
||||
|
||||
# Put multiple tensors in queue using TensorIpcData
|
||||
tensors_data = [
|
||||
(0, torch.randn(2, 3)),
|
||||
(1, torch.randn(4, 5)),
|
||||
(2, torch.randn(6, 7)),
|
||||
]
|
||||
|
||||
for tensor_id, tensor in tensors_data:
|
||||
ipc_data = TensorIpcData(
|
||||
sender_id=sender_id,
|
||||
message_id=message_id,
|
||||
tensor_id=tensor_id,
|
||||
tensor=tensor,
|
||||
)
|
||||
tensor_queue.put(ipc_data)
|
||||
|
||||
# Create receiver directly
|
||||
receiver = TensorIpcReceiver(tensor_queue)
|
||||
|
||||
# Request tensor_id=2 (should buffer tensor_id=0 and tensor_id=1)
|
||||
handle = {"sender_id": sender_id, "message_id": message_id, "tensor_id": 2}
|
||||
|
||||
result = receiver("float32", (6, 7), handle)
|
||||
assert result.shape == (6, 7)
|
||||
|
||||
# Verify buffer has tensor_id 0 and 1
|
||||
sender = receiver._tensor_buffers[sender_id]
|
||||
tensors = sender.tensors.get(message_id, {})
|
||||
assert 0 in tensors
|
||||
assert 1 in tensors
|
||||
|
||||
# Request buffered tensor
|
||||
handle2 = {"sender_id": sender_id, "message_id": message_id, "tensor_id": 0}
|
||||
|
||||
result2 = receiver("float32", (2, 3), handle2)
|
||||
assert result2.shape == (2, 3)
|
||||
# tensor_id 0 should be removed from buffer
|
||||
sender = receiver._tensor_buffers[sender_id]
|
||||
tensors = sender.tensors.get(message_id, {})
|
||||
assert 0 not in tensors
|
||||
|
||||
|
||||
def api_server_worker(
|
||||
server_id: int,
|
||||
tensor_queue: torch_mp.Queue,
|
||||
result_queue: mp.Queue,
|
||||
barrier: BarrierType,
|
||||
retrieval_done: EventType,
|
||||
):
|
||||
"""Worker simulating an API server sending tensors."""
|
||||
try:
|
||||
# Each server sends a unique tensor
|
||||
tensor = torch.ones(server_id + 1, server_id + 2) * server_id
|
||||
sender_id = f"server_{server_id}"
|
||||
|
||||
# Wait for all servers to be ready
|
||||
barrier.wait()
|
||||
|
||||
# Send tensor using TensorIpcData
|
||||
ipc_data = TensorIpcData(
|
||||
sender_id=sender_id,
|
||||
message_id=0,
|
||||
tensor_id=0,
|
||||
tensor=tensor,
|
||||
)
|
||||
tensor_queue.put(ipc_data)
|
||||
|
||||
result_queue.put({"server_id": server_id, "success": True})
|
||||
|
||||
# Keep process alive until main process has retrieved all tensors
|
||||
# This prevents shared memory handles from being invalidated
|
||||
retrieval_done.wait(timeout=30.0)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
result_queue.put(
|
||||
{
|
||||
"server_id": server_id,
|
||||
"success": False,
|
||||
"error": str(e),
|
||||
"traceback": traceback.format_exc(),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_multiple_api_servers_to_engine():
|
||||
"""Test multiple API servers sending to one engine core via multiprocessing."""
|
||||
num_api_servers = 3
|
||||
tensor_queue = torch_mp.Queue()
|
||||
result_queue: mp.Queue = mp.Queue()
|
||||
barrier = mp.Barrier(num_api_servers)
|
||||
retrieval_done = mp.Event()
|
||||
|
||||
# Start multiple API server processes
|
||||
processes = []
|
||||
for server_id in range(num_api_servers):
|
||||
proc = mp.Process(
|
||||
target=api_server_worker,
|
||||
args=(server_id, tensor_queue, result_queue, barrier, retrieval_done),
|
||||
)
|
||||
proc.start()
|
||||
processes.append(proc)
|
||||
|
||||
# Collect results from all servers
|
||||
results = []
|
||||
for _ in range(num_api_servers):
|
||||
result = result_queue.get(timeout=10.0)
|
||||
results.append(result)
|
||||
|
||||
# Verify all servers succeeded
|
||||
for result in results:
|
||||
assert result["success"], (
|
||||
f"Server {result['server_id']} failed: {result.get('error')}"
|
||||
)
|
||||
|
||||
# Verify all tensors are in queue
|
||||
received_tensors = []
|
||||
for _ in range(num_api_servers):
|
||||
ipc_data = tensor_queue.get(timeout=1.0)
|
||||
received_tensors.append((ipc_data.sender_id, ipc_data.tensor))
|
||||
|
||||
assert len(received_tensors) == num_api_servers
|
||||
|
||||
# Verify tensor content (order may vary with multiprocessing)
|
||||
tensor_by_sender = {sid: t for sid, t in received_tensors}
|
||||
for server_id in range(num_api_servers):
|
||||
expected_id = f"server_{server_id}"
|
||||
assert expected_id in tensor_by_sender, (
|
||||
f"Missing tensor from server {server_id}"
|
||||
)
|
||||
expected_tensor = torch.ones(server_id + 1, server_id + 2) * server_id
|
||||
assert torch.allclose(tensor_by_sender[expected_id], expected_tensor)
|
||||
|
||||
# Signal workers that retrieval is complete
|
||||
retrieval_done.set()
|
||||
|
||||
# Wait for all processes to complete
|
||||
for proc in processes:
|
||||
proc.join(timeout=5.0)
|
||||
|
||||
|
||||
def mixed_tensor_encoder_process(
|
||||
tensor_queue: torch_mp.Queue,
|
||||
result_queue: mp.Queue,
|
||||
ready_event: EventType,
|
||||
retrieval_done: EventType,
|
||||
):
|
||||
"""Process that encodes mixed CPU/CUDA tensors."""
|
||||
try:
|
||||
sender = TensorIpcSender(tensor_queue)
|
||||
_encoder = MsgpackEncoder(oob_tensor_consumer=sender)
|
||||
|
||||
# Create only CUDA tensor for IPC (CPU will be serialized)
|
||||
# But actually, let's just send CUDA tensor directly
|
||||
cuda_tensor = torch.randn(4, 5, device="cuda:0")
|
||||
|
||||
# Manually send via IPC to test the mechanism
|
||||
cuda_tensor_shared = cuda_tensor.share_memory_()
|
||||
|
||||
ipc_data = TensorIpcData(
|
||||
sender_id="mixed_encoder",
|
||||
message_id=0,
|
||||
tensor_id=0,
|
||||
tensor=cuda_tensor_shared,
|
||||
)
|
||||
tensor_queue.put(ipc_data, timeout=10.0)
|
||||
|
||||
ready_event.set()
|
||||
|
||||
result_queue.put({"success": True, "sent_cuda": True})
|
||||
|
||||
# Keep process alive until decoder has retrieved the tensor
|
||||
retrieval_done.wait(timeout=30.0)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
ready_event.set()
|
||||
result_queue.put(
|
||||
{"success": False, "error": str(e), "traceback": traceback.format_exc()}
|
||||
)
|
||||
|
||||
|
||||
def mixed_tensor_decoder_process(
|
||||
tensor_queue: torch_mp.Queue,
|
||||
result_queue: mp.Queue,
|
||||
encoder_ready: EventType,
|
||||
retrieval_done: EventType,
|
||||
):
|
||||
"""Process that retrieves mixed tensors from queue."""
|
||||
try:
|
||||
# Wait for encoder to finish
|
||||
if not encoder_ready.wait(timeout=10.0):
|
||||
raise TimeoutError("Encoder did not signal ready")
|
||||
|
||||
# Try to get CUDA tensor from queue
|
||||
ipc_data = tensor_queue.get(timeout=5.0)
|
||||
|
||||
result_queue.put(
|
||||
{
|
||||
"success": True,
|
||||
"is_cuda": ipc_data.tensor.is_cuda,
|
||||
"shape": tuple(ipc_data.tensor.shape),
|
||||
}
|
||||
)
|
||||
|
||||
# Signal that retrieval is complete
|
||||
retrieval_done.set()
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
retrieval_done.set() # Signal even on failure
|
||||
result_queue.put(
|
||||
{"success": False, "error": str(e), "traceback": traceback.format_exc()}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
def test_mixed_cpu_cuda_tensors():
|
||||
"""Test encoding with mixed CPU and CUDA tensors using multiprocessing."""
|
||||
tensor_queue = torch_mp.Queue()
|
||||
result_queue: mp.Queue = mp.Queue()
|
||||
encoder_ready = mp.Event()
|
||||
retrieval_done = mp.Event()
|
||||
|
||||
# Start encoder process
|
||||
encoder_proc = mp.Process(
|
||||
target=mixed_tensor_encoder_process,
|
||||
args=(tensor_queue, result_queue, encoder_ready, retrieval_done),
|
||||
)
|
||||
encoder_proc.start()
|
||||
|
||||
# Start decoder process
|
||||
decoder_proc = mp.Process(
|
||||
target=mixed_tensor_decoder_process,
|
||||
args=(tensor_queue, result_queue, encoder_ready, retrieval_done),
|
||||
)
|
||||
decoder_proc.start()
|
||||
|
||||
# Get results
|
||||
encoder_result = result_queue.get(timeout=10.0)
|
||||
decoder_result = result_queue.get(timeout=10.0)
|
||||
|
||||
encoder_proc.join(timeout=5.0)
|
||||
decoder_proc.join(timeout=5.0)
|
||||
|
||||
# Verify encoder succeeded
|
||||
assert encoder_result["success"], (
|
||||
f"Encoder failed: {encoder_result.get('error')}\n"
|
||||
f"{encoder_result.get('traceback', '')}"
|
||||
)
|
||||
|
||||
# Verify decoder succeeded and got CUDA tensor
|
||||
assert decoder_result["success"], (
|
||||
f"Decoder failed: {decoder_result.get('error')}\n"
|
||||
f"{decoder_result.get('traceback', '')}"
|
||||
)
|
||||
assert decoder_result["is_cuda"], "Retrieved tensor is not on CUDA"
|
||||
assert decoder_result["shape"] == (4, 5), (
|
||||
f"Unexpected shape: {decoder_result['shape']}"
|
||||
)
|
||||
|
||||
|
||||
def cpu_tensor_ipc_encoder_process(
|
||||
tensor_queue: torch_mp.Queue,
|
||||
result_queue: mp.Queue,
|
||||
tensor_shape: tuple,
|
||||
ready_event: EventType,
|
||||
retrieval_done: EventType,
|
||||
):
|
||||
"""Process that encodes and sends CPU tensors via IPC queue."""
|
||||
try:
|
||||
# Create encoder with IPC enabled for all tensors
|
||||
sender = TensorIpcSender(tensor_queue)
|
||||
encoder = MsgpackEncoder(oob_tensor_consumer=sender)
|
||||
|
||||
# Create a CPU tensor
|
||||
tensor = torch.randn(*tensor_shape, dtype=torch.float32)
|
||||
|
||||
# Encode the tensor (should use IPC queue, not standard serialization)
|
||||
encoded = encoder.encode({"test_tensor": tensor})
|
||||
|
||||
# Signal that encoding is complete
|
||||
ready_event.set()
|
||||
|
||||
result_queue.put(
|
||||
{
|
||||
"success": True,
|
||||
"encoded_length": len(encoded),
|
||||
"device": str(tensor.device),
|
||||
"tensor_shape": tuple(tensor.shape),
|
||||
}
|
||||
)
|
||||
|
||||
# Keep process alive until decoder has retrieved the tensor
|
||||
# This is necessary for CPU tensor shared memory to remain valid
|
||||
retrieval_done.wait(timeout=30.0)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
ready_event.set()
|
||||
result_queue.put(
|
||||
{"success": False, "error": str(e), "traceback": traceback.format_exc()}
|
||||
)
|
||||
|
||||
|
||||
def cpu_tensor_ipc_decoder_process(
|
||||
tensor_queue: torch_mp.Queue,
|
||||
result_queue: mp.Queue,
|
||||
expected_shape: tuple,
|
||||
encoder_ready: EventType,
|
||||
retrieval_done: EventType,
|
||||
):
|
||||
"""Process that decodes and receives CPU tensors from IPC queue."""
|
||||
try:
|
||||
# Wait for encoder to finish sending
|
||||
if not encoder_ready.wait(timeout=10.0):
|
||||
raise TimeoutError("Encoder did not signal ready")
|
||||
|
||||
# Get tensor from queue
|
||||
ipc_data = tensor_queue.get(timeout=5.0)
|
||||
|
||||
result_queue.put(
|
||||
{
|
||||
"success": True,
|
||||
"tensor_id": ipc_data.tensor_id,
|
||||
"tensor_shape": tuple(ipc_data.tensor.shape),
|
||||
"device": str(ipc_data.tensor.device),
|
||||
"matches_expected": tuple(ipc_data.tensor.shape) == expected_shape,
|
||||
"is_cpu": ipc_data.tensor.device.type == "cpu",
|
||||
}
|
||||
)
|
||||
|
||||
# Signal that retrieval is complete
|
||||
retrieval_done.set()
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
retrieval_done.set() # Signal even on failure
|
||||
result_queue.put(
|
||||
{"success": False, "error": str(e), "traceback": traceback.format_exc()}
|
||||
)
|
||||
|
||||
|
||||
def test_cpu_tensor_ipc():
|
||||
"""Test CPU tensor sharing via IPC queue when mm_tensor_ipc is enabled."""
|
||||
# Set up single queue and synchronization
|
||||
tensor_queue = torch_mp.Queue()
|
||||
result_queue: mp.Queue = mp.Queue()
|
||||
encoder_ready = mp.Event()
|
||||
retrieval_done = mp.Event()
|
||||
|
||||
tensor_shape = (3, 5, 7)
|
||||
|
||||
# Start encoder process
|
||||
encoder_proc = mp.Process(
|
||||
target=cpu_tensor_ipc_encoder_process,
|
||||
args=(
|
||||
tensor_queue,
|
||||
result_queue,
|
||||
tensor_shape,
|
||||
encoder_ready,
|
||||
retrieval_done,
|
||||
),
|
||||
)
|
||||
encoder_proc.start()
|
||||
|
||||
# Start decoder process
|
||||
decoder_proc = mp.Process(
|
||||
target=cpu_tensor_ipc_decoder_process,
|
||||
args=(
|
||||
tensor_queue,
|
||||
result_queue,
|
||||
tensor_shape,
|
||||
encoder_ready,
|
||||
retrieval_done,
|
||||
),
|
||||
)
|
||||
decoder_proc.start()
|
||||
|
||||
# Wait for processes and collect results
|
||||
encoder_result = result_queue.get(timeout=10.0)
|
||||
decoder_result = result_queue.get(timeout=10.0)
|
||||
|
||||
encoder_proc.join(timeout=5.0)
|
||||
decoder_proc.join(timeout=5.0)
|
||||
|
||||
# Verify results
|
||||
assert encoder_result["success"], (
|
||||
f"Encoder failed: {encoder_result.get('error')}\n"
|
||||
f"{encoder_result.get('traceback', '')}"
|
||||
)
|
||||
assert decoder_result["success"], (
|
||||
f"Decoder failed: {decoder_result.get('error')}\n"
|
||||
f"{decoder_result.get('traceback', '')}"
|
||||
)
|
||||
assert decoder_result["matches_expected"], "Tensor shape mismatch"
|
||||
assert decoder_result["is_cpu"], "Tensor not on CPU device"
|
||||
|
||||
|
||||
def test_ipc_disabled_mode():
|
||||
"""Test that IPC is disabled when no sender is provided."""
|
||||
tensor_queues = [torch_mp.Queue()]
|
||||
|
||||
# Create encoder without IPC sender (IPC disabled)
|
||||
encoder = MsgpackEncoder()
|
||||
|
||||
# Create a CPU tensor
|
||||
cpu_tensor = torch.randn(2, 3, dtype=torch.float32)
|
||||
|
||||
# Encode the tensor (should use standard serialization, not IPC)
|
||||
encoded = encoder.encode({"test_tensor": cpu_tensor})
|
||||
|
||||
# Verify encoding succeeded
|
||||
assert len(encoded) > 0
|
||||
assert isinstance(encoded, (list, tuple))
|
||||
|
||||
# Verify queue is empty (no IPC was used)
|
||||
assert tensor_queues[0].empty(), "Tensor queue should be empty when IPC is disabled"
|
||||
|
||||
# If CUDA is available, test with CUDA tensor too
|
||||
if torch.cuda.is_available():
|
||||
cuda_tensor = torch.randn(4, 5, device="cuda:0")
|
||||
encoded_cuda = encoder.encode({"cuda_tensor": cuda_tensor})
|
||||
assert len(encoded_cuda) > 0
|
||||
assert tensor_queues[0].empty(), (
|
||||
"Tensor queue should be empty for CUDA tensor when IPC is disabled"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MultiTensorMessage:
|
||||
"""Message with multiple tensors to test multi-tensor IPC."""
|
||||
|
||||
t1: torch.Tensor
|
||||
t2: torch.Tensor
|
||||
sender_label: str
|
||||
|
||||
|
||||
def concurrent_sender_process(
|
||||
tensor_queue: torch_mp.Queue,
|
||||
payload_queue: mp.Queue,
|
||||
result_queue: mp.Queue,
|
||||
sender_index: int,
|
||||
num_messages: int,
|
||||
barrier: BarrierType,
|
||||
retrieval_done: EventType,
|
||||
):
|
||||
"""Process that acts as one of N concurrent senders."""
|
||||
try:
|
||||
sender = TensorIpcSender(tensor_queue)
|
||||
encoder = MsgpackEncoder(oob_tensor_consumer=sender)
|
||||
|
||||
# Wait for all senders to be ready before sending
|
||||
barrier.wait(timeout=10.0)
|
||||
|
||||
encoded_payloads = []
|
||||
for msg_idx in range(num_messages):
|
||||
# Each sender creates uniquely-shaped tensors so we can
|
||||
# verify correct routing on the receiver side.
|
||||
t1 = torch.full((sender_index + 1, 3), float(msg_idx), dtype=torch.float32)
|
||||
t2 = torch.full(
|
||||
(2, sender_index + 2), float(msg_idx + 100), dtype=torch.float64
|
||||
)
|
||||
msg = MultiTensorMessage(
|
||||
t1=t1,
|
||||
t2=t2,
|
||||
sender_label=f"sender_{sender_index}_msg_{msg_idx}",
|
||||
)
|
||||
encoded = encoder.encode(msg)
|
||||
encoded_payloads.append(encoded)
|
||||
|
||||
# Send all encoded payloads via the regular (non-tensor) queue
|
||||
for encoded in encoded_payloads:
|
||||
payload_queue.put(encoded, timeout=10.0)
|
||||
|
||||
result_queue.put(
|
||||
{
|
||||
"success": True,
|
||||
"sender_index": sender_index,
|
||||
"num_sent": num_messages,
|
||||
}
|
||||
)
|
||||
|
||||
# Keep alive so shared-memory handles remain valid
|
||||
retrieval_done.wait(timeout=30.0)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
result_queue.put(
|
||||
{
|
||||
"success": False,
|
||||
"sender_index": sender_index,
|
||||
"error": str(e),
|
||||
"traceback": traceback.format_exc(),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_concurrent_senders_single_receiver():
|
||||
"""Test N concurrent senders sharing one queue with a single receiver.
|
||||
|
||||
Each sender encodes multiple messages (each containing two tensors) via
|
||||
its own MsgpackEncoder + TensorIpcSender. A single TensorIpcReceiver
|
||||
on the receiving side must correctly drain-and-buffer interleaved
|
||||
TensorIpcData items from the shared queue and match them back to the
|
||||
right message handles during decode.
|
||||
"""
|
||||
num_senders = 4
|
||||
num_messages_per_sender = 3
|
||||
tensor_queue = torch_mp.Queue()
|
||||
payload_queue: mp.Queue = mp.Queue()
|
||||
result_queue: mp.Queue = mp.Queue()
|
||||
barrier = mp.Barrier(num_senders)
|
||||
retrieval_done = mp.Event()
|
||||
|
||||
# Launch sender processes
|
||||
processes = []
|
||||
for i in range(num_senders):
|
||||
proc = mp.Process(
|
||||
target=concurrent_sender_process,
|
||||
args=(
|
||||
tensor_queue,
|
||||
payload_queue,
|
||||
result_queue,
|
||||
i,
|
||||
num_messages_per_sender,
|
||||
barrier,
|
||||
retrieval_done,
|
||||
),
|
||||
)
|
||||
proc.start()
|
||||
processes.append(proc)
|
||||
|
||||
# Collect send confirmations
|
||||
send_results = []
|
||||
for _ in range(num_senders):
|
||||
send_results.append(result_queue.get(timeout=15.0))
|
||||
for r in send_results:
|
||||
assert r["success"], (
|
||||
f"Sender {r['sender_index']} failed: {r.get('error')}\n"
|
||||
f"{r.get('traceback', '')}"
|
||||
)
|
||||
|
||||
# Now decode all messages from the main process using a single receiver
|
||||
receiver = TensorIpcReceiver(tensor_queue)
|
||||
decoder = MsgpackDecoder(MultiTensorMessage, oob_tensor_provider=receiver)
|
||||
|
||||
decoded_messages: list[MultiTensorMessage] = []
|
||||
total = num_senders * num_messages_per_sender
|
||||
for _ in range(total):
|
||||
encoded = payload_queue.get(timeout=10.0)
|
||||
decoded = decoder.decode(encoded)
|
||||
assert isinstance(decoded, MultiTensorMessage)
|
||||
decoded_messages.append(decoded)
|
||||
|
||||
# Signal senders they can exit
|
||||
retrieval_done.set()
|
||||
|
||||
# Group by sender_label prefix to verify all messages arrived
|
||||
by_sender: dict[int, list[MultiTensorMessage]] = {}
|
||||
for msg in decoded_messages:
|
||||
# label format: "sender_{i}_msg_{j}"
|
||||
parts = msg.sender_label.split("_")
|
||||
sender_idx = int(parts[1])
|
||||
by_sender.setdefault(sender_idx, []).append(msg)
|
||||
|
||||
assert len(by_sender) == num_senders, (
|
||||
f"Expected {num_senders} senders, got {len(by_sender)}"
|
||||
)
|
||||
|
||||
for sender_idx in range(num_senders):
|
||||
msgs = sorted(by_sender[sender_idx], key=lambda m: m.sender_label)
|
||||
assert len(msgs) == num_messages_per_sender, (
|
||||
f"Sender {sender_idx}: expected {num_messages_per_sender} "
|
||||
f"messages, got {len(msgs)}"
|
||||
)
|
||||
for msg_idx, msg in enumerate(msgs):
|
||||
assert msg.sender_label == f"sender_{sender_idx}_msg_{msg_idx}"
|
||||
# Verify tensor shapes match what the sender created
|
||||
assert msg.t1.shape == (sender_idx + 1, 3)
|
||||
assert msg.t2.shape == (2, sender_idx + 2)
|
||||
# Verify tensor values
|
||||
assert torch.allclose(msg.t1, torch.full_like(msg.t1, float(msg_idx)))
|
||||
assert torch.allclose(msg.t2, torch.full_like(msg.t2, float(msg_idx + 100)))
|
||||
|
||||
for proc in processes:
|
||||
proc.join(timeout=5.0)
|
||||
|
||||
|
||||
def test_concurrent_senders_interleaved_buffer():
|
||||
"""Test receiver buffering when tensors from multiple senders interleave.
|
||||
|
||||
Manually enqueue TensorIpcData from two senders in an interleaved order
|
||||
and verify the receiver correctly buffers and retrieves each tensor by
|
||||
its (sender_id, message_id, tensor_id) handle.
|
||||
"""
|
||||
tensor_queue = torch_mp.Queue()
|
||||
|
||||
# Sender A: 2 tensors for message 1
|
||||
a_t0 = torch.randn(2, 3)
|
||||
a_t1 = torch.randn(4, 5)
|
||||
# Sender B: 2 tensors for message 1
|
||||
b_t0 = torch.randn(6, 7)
|
||||
b_t1 = torch.randn(8, 9)
|
||||
|
||||
# Interleave: B_t0, A_t0, B_t1, A_t1
|
||||
for sid, mid, tid, t in [
|
||||
("B", 1, 0, b_t0),
|
||||
("A", 1, 0, a_t0),
|
||||
("B", 1, 1, b_t1),
|
||||
("A", 1, 1, a_t1),
|
||||
]:
|
||||
tensor_queue.put(
|
||||
TensorIpcData(sender_id=sid, message_id=mid, tensor_id=tid, tensor=t)
|
||||
)
|
||||
|
||||
receiver = TensorIpcReceiver(tensor_queue)
|
||||
|
||||
# Request A_t1 first — receiver must drain and buffer B_t0, A_t0, B_t1
|
||||
result = receiver(
|
||||
"float32", a_t1.shape, {"sender_id": "A", "message_id": 1, "tensor_id": 1}
|
||||
)
|
||||
assert torch.equal(result, a_t1)
|
||||
|
||||
# Now request B_t0 from buffer
|
||||
result = receiver(
|
||||
"float32", b_t0.shape, {"sender_id": "B", "message_id": 1, "tensor_id": 0}
|
||||
)
|
||||
assert torch.equal(result, b_t0)
|
||||
|
||||
# Request A_t0 from buffer
|
||||
result = receiver(
|
||||
"float32", a_t0.shape, {"sender_id": "A", "message_id": 1, "tensor_id": 0}
|
||||
)
|
||||
assert torch.equal(result, a_t0)
|
||||
|
||||
# Request B_t1 from buffer
|
||||
result = receiver(
|
||||
"float64", b_t1.shape, {"sender_id": "B", "message_id": 1, "tensor_id": 1}
|
||||
)
|
||||
assert torch.equal(result, b_t1)
|
||||
|
||||
# All buffers should be drained
|
||||
for sid in ("A", "B"):
|
||||
tensors = receiver._tensor_buffers[sid].tensors.get(1, {})
|
||||
assert len(tensors) == 0, f"Sender {sid} buffer not empty: {tensors}"
|
||||
|
||||
|
||||
def test_mixed_cpu_cuda_with_ipc_enabled():
|
||||
"""Test that encoder is configured correctly for IPC with all tensor types."""
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA not available")
|
||||
|
||||
tensor_queue = torch_mp.Queue()
|
||||
|
||||
# Create sender and encoder with IPC enabled
|
||||
sender = TensorIpcSender(tensor_queue)
|
||||
encoder = MsgpackEncoder(oob_tensor_consumer=sender)
|
||||
|
||||
# Verify sender configuration
|
||||
assert encoder.oob_tensor_consumer is not None, "Consumer should be set"
|
||||
|
||||
# Note: Actual IPC transfer only works across processes
|
||||
# (tested in test_cpu_tensor_ipc)
|
||||
# This test just verifies the configuration is correct
|
||||
|
||||
|
||||
def test_tensor_cleanup_after_decode():
|
||||
"""Test that tensors are removed from tracking after successful decode."""
|
||||
# Create a tensor queue
|
||||
tensor_queue = torch_mp.Queue()
|
||||
|
||||
# Create and encode a tensor
|
||||
tensor = torch.randn(5, 5)
|
||||
# Move to shared memory for IPC
|
||||
if not tensor.is_shared():
|
||||
tensor.share_memory_()
|
||||
|
||||
# Manually create a TensorIpcData and put it in the queue
|
||||
sender_id = "test_sender"
|
||||
message_id = 0
|
||||
tensor_id = 0
|
||||
ipc_data = TensorIpcData(
|
||||
sender_id=sender_id,
|
||||
message_id=message_id,
|
||||
tensor_id=tensor_id,
|
||||
tensor=tensor,
|
||||
)
|
||||
tensor_queue.put(ipc_data)
|
||||
|
||||
# Create receiver directly
|
||||
receiver = TensorIpcReceiver(tensor_queue)
|
||||
|
||||
handle = {
|
||||
"sender_id": sender_id,
|
||||
"message_id": message_id,
|
||||
"tensor_id": tensor_id,
|
||||
}
|
||||
|
||||
# Receive the tensor - this should retrieve it from the queue
|
||||
decoded_tensor = receiver(
|
||||
str(tensor.dtype).removeprefix("torch."), tensor.shape, handle
|
||||
)
|
||||
|
||||
# Verify the tensor was decoded
|
||||
assert decoded_tensor.shape == tensor.shape, "Decoded tensor should match shape"
|
||||
|
||||
# Verify the tensor was removed from buffer after decode
|
||||
sender = receiver._tensor_buffers[sender_id]
|
||||
tensors = sender.tensors.get(message_id, {})
|
||||
assert tensor_id not in tensors, "Tensor should be removed from buffer"
|
||||
@@ -378,3 +378,65 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
|
||||
ref_input_batch.refresh_metadata()
|
||||
|
||||
_compare_objs(input_batch, ref_input_batch)
|
||||
|
||||
|
||||
def _construct_pooling_request(req_id_suffix: int):
|
||||
from vllm.pooling_params import PoolingParams
|
||||
|
||||
prompt_token_ids = [
|
||||
np.random.randint(0, VOCAB_SIZE)
|
||||
for _ in range(np.random.randint(10, MAX_PROMPT_SIZE))
|
||||
]
|
||||
return CachedRequestState(
|
||||
req_id=f"pool_req_{req_id_suffix}",
|
||||
prompt_token_ids=prompt_token_ids,
|
||||
sampling_params=None,
|
||||
pooling_params=PoolingParams(task="classify"),
|
||||
mm_features=[],
|
||||
block_ids=([],),
|
||||
generator=None,
|
||||
num_computed_tokens=0,
|
||||
output_token_ids=[],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
def test_pooling_prompt_lens_not_aliased(device: str):
|
||||
"""Verify that prompt_lens in PoolingMetadata does not share memory
|
||||
with the internal num_prompt_tokens pinned buffer. Guards against possible
|
||||
non-determinism in pooling metadata due to mutations to the internal buffer.
|
||||
"""
|
||||
batch_size = 4
|
||||
input_batch = InputBatch(
|
||||
max_num_reqs=batch_size * 2,
|
||||
max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
|
||||
max_num_batched_tokens=batch_size * (MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS),
|
||||
device=torch.device(device),
|
||||
pin_memory=is_pin_memory_available(),
|
||||
vocab_size=VOCAB_SIZE,
|
||||
block_sizes=[16],
|
||||
kernel_block_sizes=[16],
|
||||
is_pooling_model=True,
|
||||
)
|
||||
|
||||
reqs = []
|
||||
# Add requests
|
||||
for i in range(batch_size):
|
||||
req = _construct_pooling_request(i)
|
||||
input_batch.add_request(req)
|
||||
reqs.append(req)
|
||||
input_batch.refresh_metadata()
|
||||
|
||||
# prompt_lens must be a snapshot
|
||||
metadata = input_batch.get_pooling_metadata()
|
||||
prompt_lens_snapshot = metadata.prompt_lens.clone()
|
||||
|
||||
# Mutate the internal buffer (simulates next batch adding new requests)
|
||||
input_batch.num_prompt_tokens_cpu_tensor.fill_(999)
|
||||
|
||||
# prompt_lens must be unaffected by the mutation
|
||||
assert torch.equal(metadata.prompt_lens, prompt_lens_snapshot), (
|
||||
"prompt_lens shares memory with internal pinned buffer; "
|
||||
"mutations to num_prompt_tokens_cpu_tensor corrupted prompt_lens. "
|
||||
f"Expected {prompt_lens_snapshot}, got {metadata.prompt_lens}"
|
||||
)
|
||||
|
||||
@@ -670,8 +670,8 @@ def test_init_kv_cache_without_kv_sharing(default_vllm_config):
|
||||
|
||||
runner.initialize_kv_cache(kv_cache_config)
|
||||
|
||||
layer_0_kv = vllm_ctx[layer_0].kv_cache[0]
|
||||
layer_1_kv = vllm_ctx[layer_1].kv_cache[0]
|
||||
layer_0_kv = vllm_ctx[layer_0].kv_cache
|
||||
layer_1_kv = vllm_ctx[layer_1].kv_cache
|
||||
# check layer 1 kv cache does NOT share memory with layer 0
|
||||
assert id(layer_1_kv) != id(layer_0_kv)
|
||||
|
||||
@@ -740,8 +740,8 @@ def test_init_kv_cache_with_kv_sharing_valid(default_vllm_config):
|
||||
runner.initialize_kv_cache(kv_cache_config)
|
||||
kv_cache_config_after_init = runner.kv_cache_config
|
||||
|
||||
layer_0_kv = vllm_ctx[layer_0].kv_cache[0]
|
||||
layer_1_kv = vllm_ctx[layer_1].kv_cache[0]
|
||||
layer_0_kv = vllm_ctx[layer_0].kv_cache
|
||||
layer_1_kv = vllm_ctx[layer_1].kv_cache
|
||||
# check layer 1 kv cache shares memory with layer 0
|
||||
assert id(layer_1_kv) == id(layer_0_kv)
|
||||
|
||||
@@ -864,9 +864,9 @@ def test_hybrid_attention_mamba_tensor_shapes():
|
||||
np.random.shuffle(ind)
|
||||
blocks0, blocks1 = ind[: (num_blocks // 2)], ind[(num_blocks // 2) :]
|
||||
|
||||
attn_shape = vllm_ctx[layer_0].kv_cache[0].shape
|
||||
conv_shape = vllm_ctx[layer_2].kv_cache[0][0].shape
|
||||
ssm_shape = vllm_ctx[layer_2].kv_cache[0][1].shape
|
||||
attn_shape = vllm_ctx[layer_0].kv_cache.shape
|
||||
conv_shape = vllm_ctx[layer_2].kv_cache[0].shape
|
||||
ssm_shape = vllm_ctx[layer_2].kv_cache[1].shape
|
||||
|
||||
# assert we are using FlashInfer
|
||||
assert attn_shape[0] % num_blocks == 0
|
||||
@@ -905,21 +905,21 @@ def test_hybrid_attention_mamba_tensor_shapes():
|
||||
kernel_blocks_for_attention = kv_blocks_for_attention * block_split_ratio
|
||||
|
||||
for layer in [layer_0, layer_1]:
|
||||
# attention: kv_cache[0][kernel_block_idx, kv_idx, ...]
|
||||
# attention: kv_cache[kernel_block_idx, kv_idx, ...]
|
||||
for i, kernel_block in enumerate(kernel_blocks_for_attention):
|
||||
vllm_ctx[layer].kv_cache[0][kernel_block, :] = attn_blocks_constant[i]
|
||||
vllm_ctx[layer].kv_cache[kernel_block, :] = attn_blocks_constant[i]
|
||||
|
||||
# fill mamba blocks with constants using kernel block indices
|
||||
for layer in [layer_2, layer_3, layer_4, layer_5]:
|
||||
# mamba: kv_cache[0][component][kernel_block_idx, ...]
|
||||
# mamba: kv_cache[component][kernel_block_idx, ...]
|
||||
for i, kv_block in enumerate(kv_blocks_for_mamba):
|
||||
vllm_ctx[layer].kv_cache[0][0][kv_block, :] = conv_blocks_constant[i]
|
||||
vllm_ctx[layer].kv_cache[0][1][kv_block, :] = ssm_blocks_constant[i]
|
||||
vllm_ctx[layer].kv_cache[0][kv_block, :] = conv_blocks_constant[i]
|
||||
vllm_ctx[layer].kv_cache[1][kv_block, :] = ssm_blocks_constant[i]
|
||||
|
||||
# verify attention and mamba contents are correct
|
||||
for layer in [layer_0, layer_1]:
|
||||
for i, kernel_block in enumerate(kernel_blocks_for_attention):
|
||||
actual_kv = vllm_ctx[layer].kv_cache[0][kernel_block, :]
|
||||
actual_kv = vllm_ctx[layer].kv_cache[kernel_block, :]
|
||||
expected = attn_blocks_constant[i]
|
||||
|
||||
# Check K and V separately
|
||||
@@ -928,8 +928,8 @@ def test_hybrid_attention_mamba_tensor_shapes():
|
||||
|
||||
for layer in [layer_2, layer_3, layer_4, layer_5]:
|
||||
for i, kv_block in enumerate(kv_blocks_for_mamba):
|
||||
actual_conv = vllm_ctx[layer].kv_cache[0][0][kv_block, :]
|
||||
actual_ssm = vllm_ctx[layer].kv_cache[0][1][kv_block, :]
|
||||
actual_conv = vllm_ctx[layer].kv_cache[0][kv_block, :]
|
||||
actual_ssm = vllm_ctx[layer].kv_cache[1][kv_block, :]
|
||||
expected_conv = conv_blocks_constant[i]
|
||||
expected_ssm = ssm_blocks_constant[i]
|
||||
|
||||
@@ -938,8 +938,8 @@ def test_hybrid_attention_mamba_tensor_shapes():
|
||||
|
||||
for layer in [layer_2, layer_3, layer_4, layer_5]:
|
||||
for i, kv_block in enumerate(kv_blocks_for_mamba):
|
||||
actual_conv = vllm_ctx[layer].kv_cache[0][0][kv_block, :]
|
||||
actual_ssm = vllm_ctx[layer].kv_cache[0][1][kv_block, :]
|
||||
actual_conv = vllm_ctx[layer].kv_cache[0][kv_block, :]
|
||||
actual_ssm = vllm_ctx[layer].kv_cache[1][kv_block, :]
|
||||
expected_conv = conv_blocks_constant[i]
|
||||
expected_ssm = ssm_blocks_constant[i]
|
||||
assert torch.equal(actual_conv, expected_conv)
|
||||
|
||||
@@ -23,10 +23,10 @@ def test_bind_kv_cache(default_vllm_config):
|
||||
}
|
||||
runner_kv_caches: list[torch.Tensor] = []
|
||||
bind_kv_cache(kv_cache, ctx, runner_kv_caches)
|
||||
assert ctx["layers.0.self_attn"].kv_cache[0] is kv_cache["layers.0.self_attn"]
|
||||
assert ctx["layers.1.self_attn"].kv_cache[0] is kv_cache["layers.1.self_attn"]
|
||||
assert ctx["layers.2.self_attn"].kv_cache[0] is kv_cache["layers.2.self_attn"]
|
||||
assert ctx["layers.3.self_attn"].kv_cache[0] is kv_cache["layers.3.self_attn"]
|
||||
assert ctx["layers.0.self_attn"].kv_cache is kv_cache["layers.0.self_attn"]
|
||||
assert ctx["layers.1.self_attn"].kv_cache is kv_cache["layers.1.self_attn"]
|
||||
assert ctx["layers.2.self_attn"].kv_cache is kv_cache["layers.2.self_attn"]
|
||||
assert ctx["layers.3.self_attn"].kv_cache is kv_cache["layers.3.self_attn"]
|
||||
|
||||
assert runner_kv_caches[0] is kv_cache["layers.0.self_attn"]
|
||||
assert runner_kv_caches[1] is kv_cache["layers.1.self_attn"]
|
||||
@@ -50,8 +50,8 @@ def test_bind_kv_cache_non_attention(default_vllm_config):
|
||||
runner_kv_caches: list[torch.Tensor] = []
|
||||
bind_kv_cache(kv_cache, ctx, runner_kv_caches)
|
||||
|
||||
assert ctx["model.layers.20.attn"].kv_cache[0] is kv_cache["model.layers.20.attn"]
|
||||
assert ctx["model.layers.28.attn"].kv_cache[0] is kv_cache["model.layers.28.attn"]
|
||||
assert ctx["model.layers.20.attn"].kv_cache is kv_cache["model.layers.20.attn"]
|
||||
assert ctx["model.layers.28.attn"].kv_cache is kv_cache["model.layers.28.attn"]
|
||||
|
||||
assert runner_kv_caches[0] is kv_cache["model.layers.20.attn"]
|
||||
assert runner_kv_caches[1] is kv_cache["model.layers.28.attn"]
|
||||
@@ -74,14 +74,14 @@ def test_bind_kv_cache_draft_model(default_vllm_config):
|
||||
runner_kv_caches: list[torch.Tensor] = []
|
||||
bind_kv_cache(kv_cache, ctx, runner_kv_caches)
|
||||
|
||||
assert ctx["model.layers.0.attn"].kv_cache[0] is kv_cache["model.layers.0.attn"]
|
||||
assert ctx["model.layers.1.attn"].kv_cache[0] is kv_cache["model.layers.1.attn"]
|
||||
assert ctx["model.layers.0.attn"].kv_cache is kv_cache["model.layers.0.attn"]
|
||||
assert ctx["model.layers.1.attn"].kv_cache is kv_cache["model.layers.1.attn"]
|
||||
assert (
|
||||
ctx["draft_model.layers.0.attn"].kv_cache[0]
|
||||
ctx["draft_model.layers.0.attn"].kv_cache
|
||||
is kv_cache["draft_model.layers.0.attn"]
|
||||
)
|
||||
assert (
|
||||
ctx["draft_model.layers.1.attn"].kv_cache[0]
|
||||
ctx["draft_model.layers.1.attn"].kv_cache
|
||||
is kv_cache["draft_model.layers.1.attn"]
|
||||
)
|
||||
|
||||
|
||||
@@ -20,8 +20,6 @@ compressed-tensors, nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test, main
|
||||
compressed-tensors, nm-testing/Phi-3-mini-128k-instruct-FP8, main
|
||||
compressed-tensors, neuralmagic/Phi-3-medium-128k-instruct-quantized.w4a16, main
|
||||
#compressed-tensors, mgoin/DeepSeek-Coder-V2-Lite-Instruct-FP8, main
|
||||
compressed-tensors, nm-testing/SparseLlama-3.1-8B-gsm8k-pruned.2of4-FP8-Dynamic-testing, main, 90
|
||||
compressed-tensors, nm-testing/SparseLlama-3.1-8B-gsm8k-pruned.2of4-W8A8-testing, main, 90
|
||||
awq, casperhansen/mixtral-instruct-awq, main
|
||||
awq_marlin, casperhansen/mixtral-instruct-awq, main
|
||||
fp8, neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV, main
|
||||
|
||||
@@ -40,7 +40,6 @@ EXCLUDE = [
|
||||
"vllm/v1/attention/ops",
|
||||
# TODO: Remove these entries after fixing mypy errors.
|
||||
"vllm/benchmarks",
|
||||
"vllm/config",
|
||||
]
|
||||
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user